feat: typed storage in pool

Author: ensh63

src/core/pool.rs

@@ -5,7 +5,7 @@ use core::ptr::{self, NonNull};
 
 use nginx_sys::{
     NGX_ALIGNMENT, ngx_buf_t, ngx_create_temp_buf, ngx_palloc, ngx_pcalloc, ngx_pfree,
-    ngx_pmemalign, ngx_pnalloc, ngx_pool_cleanup_add, ngx_pool_t,
+    ngx_pmemalign, ngx_pnalloc, ngx_pool_cleanup_add, ngx_pool_cleanup_t, ngx_pool_t,
 };
 
 use crate::allocator::{AllocError, Allocator, dangling_for_layout};
@@ -128,6 +128,11 @@ impl AsMut<ngx_pool_t> for Pool {
     }
 }
 
+// Wrapper to create an unique value type
+struct Item<T: Sized> {
+    value: T,
+}
+
 impl Pool {
     /// Creates a new `Pool` from an `ngx_pool_t` pointer.
     ///
@@ -200,25 +205,62 @@ impl Pool {
         Some(MemoryBuffer::from_ngx_buf(buf))
     }
 
-    /// Adds a cleanup handler for a value in the memory pool.
+    /// Allocates memory for a value and adds a cleanup handler to the memory pool.
     ///
-    /// Returns `Ok(())` if the cleanup handler is successfully added, or `Err(())` if the cleanup
-    /// handler cannot be added.
+    /// The value is created by calling the provided closure `f`. If allocation fails,
+    /// the closure is not called.
+    ///
+    /// Returns `Some(NonNull<T>)` if the allocation and cleanup handler addition are successful,
+    /// or `None` if allocation fails.
     ///
     /// # Safety
     /// This function is marked as unsafe because it involves raw pointer manipulation.
-    unsafe fn add_cleanup_for_value<T>(&self, value: *mut T) -> Result<(), ()> {
-        let cln = unsafe { ngx_pool_cleanup_add(self.0.as_ptr(), 0) };
+    /// The returned pointer must not outlive the pool, must not be freed manually
+    /// (as it has a cleanup handler), and must not be accessed after the pool is destroyed.
+    pub unsafe fn allocate_with_cleanup<T: Sized, F: FnOnce() -> T>(
+        &self,
+        f: F,
+    ) -> Option<NonNull<T>> {
+        let cln = unsafe { ngx_pool_cleanup_add(self.0.as_ptr(), mem::size_of::<T>()) };
         if cln.is_null() {
-            return Err(());
+            return None;
         }
-
         unsafe {
+            // 'data' may be NULL only if `T` is zero-sized. In that case, no real value is stored,
+            // so we can just use the cleanup structure itself as a placeholder.
+            // Note that zero-sized `T` may implement `Drop`, and this implementation will be
+            // called at cleanup time.
+            if (*cln).data.is_null() {
+                (*cln).data = cln as _;
+            };
             (*cln).handler = Some(cleanup_type::<T>);
-            (*cln).data = value as *mut c_void;
+            // `data` points to the memory allocated for the value by `ngx_pool_cleanup_add()`
+            ptr::write((*cln).data as *mut T, f());
+
+            NonNull::new((*cln).data as *mut T)
         }
+    }
 
-        Ok(())
+    /// Searches for a cleanup handler in the pool's cleanup chain.
+    ///
+    /// Returns `Some(NonNull<ngx_pool_cleanup_t>)` if a cleanup handler is found for type `T`,
+    /// or `None` if no matching handler is found.
+    ///
+    /// # Safety
+    /// This function is marked as unsafe because it involves raw pointer manipulation
+    /// and traverses the nginx cleanup chain structure.
+    unsafe fn cleanup_lookup<T: Sized>(&self) -> Option<NonNull<ngx_pool_cleanup_t>> {
+        let mut cln = (unsafe { *self.0.as_ptr() }).cleanup;
+        unsafe {
+            // SAFETY: comparing function pointers is generally unreliable, but in this specific
+            // case we can assume that the same function pointer was used when adding the cleanup
+            // handler.
+            #[allow(unpredictable_function_pointer_comparisons)]
+            while !cln.is_null() && (*cln).handler != Some(cleanup_type::<T>) {
+                cln = (*cln).next;
+            }
+        }
+        NonNull::new(cln)
     }
 
     /// Allocates memory from the pool of the specified size.
@@ -272,18 +314,136 @@ impl Pool {
     ///
     /// Returns a typed pointer to the allocated memory if successful, or a null pointer if
     /// allocation or cleanup handler addition fails.
-    pub fn allocate<T>(&self, value: T) -> *mut T {
+    pub fn allocate<T: Sized>(&self, value: T) -> *mut T {
         unsafe {
-            let p = self.alloc(mem::size_of::<T>()) as *mut T;
-            ptr::write(p, value);
-            if self.add_cleanup_for_value(p).is_err() {
-                ptr::drop_in_place(p);
-                return ptr::null_mut();
-            };
-            p
+            match self.allocate_with_cleanup(|| value) {
+                None => ptr::null_mut(),
+                Some(mut ptr) => ptr.as_mut(),
+            }
+        }
+    }
+
+    /// Gets the unique value of type `T` from the memory pool, or allocates and initializes it
+    /// using the provided function if it does not exist.
+    ///
+    /// This ensures only one value of type `T` exists in the pool. If a value already exists,
+    /// it is returned and the function `f` is not called. If no value exists, a new one is
+    /// created using `f` and stored with a cleanup handler. If allocation fails, `f` is not called.
+    ///
+    /// Returns a mutable reference to the value if successful, or `None` if allocation fails.
+    pub fn get_or_add_unique<T: Sized, F: FnOnce() -> T>(&mut self, f: F) -> Option<&mut T> {
+        unsafe {
+            self.cleanup_lookup::<Item<T>>()
+                .map(|cln| {
+                    let item = cln.as_ref().data as *mut Item<T>;
+                    &mut (*item).value
+                })
+                .or_else(|| {
+                    self.allocate_with_cleanup(|| Item { value: f() })
+                        .map(|mut ptr| &mut ptr.as_mut().value)
+                })
         }
     }
 
+    /// Gets the unique value of type `T` from the memory pool.
+    ///
+    /// This value must have been previously allocated with [`Pool::get_or_add_unique`].
+    ///
+    /// Returns a reference to the value if found, or `None` if not found.
+    pub fn get_unique<T: Sized>(&self) -> Option<&T> {
+        unsafe {
+            self.cleanup_lookup::<Item<T>>().map(|cln| {
+                let item = cln.as_ref().data as *const Item<T>;
+                &(*item).value
+            })
+        }
+    }
+
+    /// Gets a mutable reference to the unique value of type `T` from the memory pool.
+    ///
+    /// This value must have been previously allocated with [`Pool::get_or_add_unique`].
+    ///
+    /// Returns a mutable reference to the value if found, or `None` if not found.
+    pub fn get_unique_mut<T: Sized>(&mut self) -> Option<&mut T> {
+        unsafe {
+            self.cleanup_lookup::<Item<T>>().map(|cln| {
+                let item = cln.as_ref().data as *mut Item<T>;
+                &mut (*item).value
+            })
+        }
+    }
+
+    /// Runs the cleanup handler for a value and removes it.
+    ///
+    /// Returns `Some(())` if the value was successfully removed,
+    /// or `None` if the value was not found.
+    ///
+    /// # Safety
+    /// The caller must ensure that `value` is a valid pointer to a value that has an
+    /// associated cleanup handler in the pool.
+    pub unsafe fn remove<T: Sized>(&mut self, value: *const T) -> Option<()> {
+        // SAFETY: comparing function pointers is generally unreliable, but in this specific
+        // case we can assume that the same function pointer was used when adding the cleanup
+        // handler.
+        #[allow(unpredictable_function_pointer_comparisons)]
+        self.remove_cleanup_handler(|cln| {
+            cln.handler == Some(cleanup_type::<T>) && core::ptr::eq(cln.data, value as _)
+        })
+        .map(|cln| {
+            unsafe { cln.handler.unwrap()(cln.data) };
+        })
+    }
+
+    /// Runs the cleanup handler for a unique value and removes it.
+    ///
+    /// This value must have been previously allocated with [`Pool::get_or_add_unique`].
+    ///
+    /// Returns `Some(())` if the value was successfully removed,
+    /// or `None` if the value was not found.
+    pub fn remove_unique<T: Sized>(&mut self) -> Option<()> {
+        // SAFETY: comparing function pointers is generally unreliable, but in this specific
+        // case we can assume that the same function pointer was used when adding the cleanup
+        // handler.
+        #[allow(unpredictable_function_pointer_comparisons)]
+        self.remove_cleanup_handler(|cln| cln.handler == Some(cleanup_type::<T>)).map(|cln| {
+            unsafe { cln.handler.unwrap()(cln.data) };
+        })
+    }
+
+    /// Internal method to find and remove a cleanup handler matching a predicate.
+    ///
+    /// Returns `Some(&ngx_pool_cleanup_t)` if a matching handler is found and removed,
+    /// or `None` if not found.
+    fn remove_cleanup_handler(
+        &mut self,
+        predicate: impl Fn(&ngx_pool_cleanup_t) -> bool,
+    ) -> Option<&ngx_pool_cleanup_t> {
+        let mut top = ngx_pool_cleanup_t {
+            handler: None,
+            data: core::ptr::null_mut(),
+            next: unsafe { self.0.as_mut().cleanup },
+        };
+
+        let mut prev = &mut top;
+        let top_ptr = prev as *const _;
+
+        while let Some(cln) = unsafe { prev.next.as_mut() } {
+            if predicate(cln) {
+                if core::ptr::eq(prev, top_ptr) {
+                    // If the removed cleanup was the first in the chain, update the pool's
+                    // cleanup pointer
+                    unsafe { self.0.as_mut().cleanup = cln.next };
+                } else {
+                    prev.next = cln.next;
+                }
+                return Some(cln);
+            }
+            prev = cln;
+        }
+
+        None
+    }
+
     /// Resizes a memory allocation in place if possible.
     ///
     /// If resizing is requested for the last allocation in the pool, it may be
@@ -332,6 +492,8 @@ impl Pool {
 /// * `data` - A raw pointer to the value of type `T` to be cleaned up.
 unsafe extern "C" fn cleanup_type<T>(data: *mut c_void) {
     unsafe {
-        ptr::drop_in_place(data as *mut T);
+        if !data.is_null() {
+            ptr::drop_in_place(data as *mut T);
+        }
     }
 }