1 // Copyright 2024, The Android Open Source Project 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 use super::shared::Shared; 16 use core::{ 17 mem::{swap, take}, 18 ops::{Deref, DerefMut}, 19 }; 20 use gbl_async::yield_now; 21 22 /// Provides interfaces for allocating and deallocating buffers. 23 pub trait BufferPool { 24 /// The type that can be dereferenced into a buffer. 25 type Buffer: DerefMut<Target = [u8]>; 26 27 /// Allocates a buffer. 28 /// 29 /// * Returns Some(_) on success. 30 /// * Returns None if buffer is not available. allocate(&mut self) -> Option<Self::Buffer>31 fn allocate(&mut self) -> Option<Self::Buffer>; 32 33 /// Deallocates a buffer. deallocate(&mut self, buf: Self::Buffer)34 fn deallocate(&mut self, buf: Self::Buffer); 35 } 36 37 // Implements for all types of fixed size preallocated buffers, including `&mut [&mut [u8]]`, 38 // `Vec<Vec<u8>`, `&mut [Vec<u8>]` and `Vec<&mut [u8]>`. 39 impl<B, T> BufferPool for T 40 where 41 B: DerefMut<Target = [u8]> + Default, // Can be `&mut [u8]`, `Vec<u8>` 42 T: DerefMut<Target = [B]>, // Can be `&mut [B]`, `Vec<B>` 43 { 44 type Buffer = B; 45 allocate(&mut self) -> Option<B>46 fn allocate(&mut self) -> Option<B> { 47 self.iter_mut().find_map(|v| (v.len() > 0).then_some(take(v))) 48 } 49 deallocate(&mut self, mut buf: B)50 fn deallocate(&mut self, mut buf: B) { 51 swap(&mut buf, self.iter_mut().find(|v| v.len() == 0).unwrap()); 52 } 53 } 54 55 impl<T: BufferPool> Shared<T> { 56 // Try allocate a [ScopedBuffer] allocate(&self) -> Option<ScopedBuffer<T>>57 pub(crate) fn allocate(&self) -> Option<ScopedBuffer<T>> { 58 self.borrow_mut().allocate().map(|v| ScopedBuffer { buf: Some(v), pool: self }) 59 } 60 61 // Allocates a [ScopedBuffer] and waits until succeed. allocate_async(&self) -> ScopedBuffer<T>62 pub(crate) async fn allocate_async(&self) -> ScopedBuffer<T> { 63 loop { 64 match self.allocate() { 65 Some(v) => return v, 66 _ => yield_now().await, 67 } 68 } 69 } 70 } 71 72 /// Represents a scoped buffer allocated by `BufferPool`. 73 pub(crate) struct ScopedBuffer<'a, T: BufferPool> { 74 buf: Option<T::Buffer>, 75 pool: &'a Shared<T>, 76 } 77 78 impl<T: BufferPool> Drop for ScopedBuffer<'_, T> { drop(&mut self)79 fn drop(&mut self) { 80 self.pool.borrow_mut().deallocate(self.buf.take().unwrap()) 81 } 82 } 83 84 impl<T: BufferPool> Deref for ScopedBuffer<'_, T> { 85 type Target = [u8]; 86 deref(&self) -> &Self::Target87 fn deref(&self) -> &Self::Target { 88 self.buf.as_ref().unwrap() 89 } 90 } 91 92 impl<T: BufferPool> DerefMut for ScopedBuffer<'_, T> { deref_mut(&mut self) -> &mut Self::Target93 fn deref_mut(&mut self) -> &mut Self::Target { 94 self.buf.as_mut().unwrap() 95 } 96 } 97