refactor(scheduling-utils): improve safety of pointers (#9157)

scheduling-utils/src/pubkeys_ptr.rs

@@ -3,13 +3,12 @@ use {
     std::ptr::NonNull,
 };
 
-pub struct PubkeysPtr<'a> {
+pub struct PubkeysPtr {
     ptr: NonNull<Pubkey>,
     count: usize,
-    allocator: &'a Allocator,
 }
 
-impl<'a> PubkeysPtr<'a> {
+impl PubkeysPtr {
     /// Constructs the pointer from a [`SharablePubkeys`].
     ///
     /// # Safety
@@ -19,7 +18,7 @@ impl<'a> PubkeysPtr<'a> {
     /// - `sharable_pubkeys.num_pubkeys` must be accurate and not overrun the allocation.
     pub unsafe fn from_sharable_pubkeys(
         sharable_pubkeys: &SharablePubkeys,
-        allocator: &'a Allocator,
+        allocator: &Allocator,
     ) -> Self {
         assert_ne!(sharable_pubkeys.num_pubkeys, 0);
         let ptr = allocator.ptr_from_offset(sharable_pubkeys.offset).cast();
@@ -27,7 +26,6 @@ impl<'a> PubkeysPtr<'a> {
         Self {
             ptr,
             count: sharable_pubkeys.num_pubkeys as usize,
-            allocator,
         }
     }
 
@@ -39,9 +37,13 @@ impl<'a> PubkeysPtr<'a> {
     }
 
     /// Frees the underlying allocation.
-    pub fn free(self) {
+    ///
+    /// # Safety
+    ///
+    /// - `Self` must be exclusively owned.
+    pub unsafe fn free(self, allocator: &Allocator) {
         // SAFETY
-        // - Constructor invariants guarantee that we exclusively own this pointer.
-        unsafe { self.allocator.free(self.ptr.cast()) };
+        // - Caller guarantees that we exclusively own this pointer.
+        unsafe { allocator.free(self.ptr.cast()) };
     }
 }

scheduling-utils/src/responses_region.rs

@@ -95,13 +95,12 @@ unsafe fn from_iterator<T: Sized>(
     Some(region)
 }
 
-pub struct CheckResponsesPtr<'a> {
+pub struct CheckResponsesPtr {
     ptr: NonNull<CheckResponse>,
     count: usize,
-    allocator: &'a Allocator,
 }
 
-impl<'a> CheckResponsesPtr<'a> {
+impl CheckResponsesPtr {
     /// Constructs the pointer from a [`TransactionResponseRegion`].
     ///
     /// # Safety
@@ -109,10 +108,9 @@ impl<'a> CheckResponsesPtr<'a> {
     /// - The provided [`TransactionResponseRegion`] must be of type
     ///   [`worker_message_types::CHECK_RESPONSE`].
     /// - The allocation pointed to by this region must not have previously been freed.
-    /// - Pointer must be exclusive so that calling [`Self::free`] is safe.
     pub unsafe fn from_transaction_response_region(
         transaction_response_region: &TransactionResponseRegion,
-        allocator: &'a Allocator,
+        allocator: &Allocator,
     ) -> Self {
         debug_assert!(transaction_response_region.tag == worker_message_types::CHECK_RESPONSE);
 
@@ -121,7 +119,6 @@ impl<'a> CheckResponsesPtr<'a> {
                 .ptr_from_offset(transaction_response_region.transaction_responses_offset)
                 .cast(),
             count: transaction_response_region.num_transaction_responses as usize,
-            allocator,
         }
     }
 
@@ -141,18 +138,21 @@ impl<'a> CheckResponsesPtr<'a> {
     }
 
     /// Free the batch's allocation.
-    pub fn free(self) {
-        unsafe { self.allocator.free(self.ptr.cast()) }
+    ///
+    /// # Safety
+    ///
+    /// - `Self` must be exclusively owned.
+    pub unsafe fn free(self, allocator: &Allocator) {
+        unsafe { allocator.free(self.ptr.cast()) }
     }
 }
 
-pub struct ExecutionResponsesPtr<'a> {
+pub struct ExecutionResponsesPtr {
     ptr: NonNull<ExecutionResponse>,
     count: usize,
-    allocator: &'a Allocator,
 }
 
-impl<'a> ExecutionResponsesPtr<'a> {
+impl ExecutionResponsesPtr {
     /// Constructs the pointer from a [`TransactionResponseRegion`].
     ///
     /// # Safety
@@ -160,10 +160,9 @@ impl<'a> ExecutionResponsesPtr<'a> {
     /// - The provided [`TransactionResponseRegion`] must be of type
     ///   [`worker_message_types::EXECUTION_RESPONSE`].
     /// - The allocation pointed to by this region must not have previously been freed.
-    /// - Pointer must be exclusive so that calling [`Self::free`] is safe.
     pub unsafe fn from_transaction_response_region(
         transaction_response_region: &TransactionResponseRegion,
-        allocator: &'a Allocator,
+        allocator: &Allocator,
     ) -> Self {
         debug_assert!(transaction_response_region.tag == worker_message_types::EXECUTION_RESPONSE);
 
@@ -172,7 +171,6 @@ impl<'a> ExecutionResponsesPtr<'a> {
                 .ptr_from_offset(transaction_response_region.transaction_responses_offset)
                 .cast(),
             count: transaction_response_region.num_transaction_responses as usize,
-            allocator,
         }
     }
 
@@ -192,7 +190,11 @@ impl<'a> ExecutionResponsesPtr<'a> {
     }
 
     /// Free the batch's allocation.
-    pub fn free(self) {
-        unsafe { self.allocator.free(self.ptr.cast()) }
+    ///
+    /// # Safety
+    ///
+    /// - `Self` must be exclusively owned.
+    pub unsafe fn free(self, allocator: &Allocator) {
+        unsafe { allocator.free(self.ptr.cast()) }
     }
 }

scheduling-utils/src/transaction_ptr.rs

@@ -131,12 +131,17 @@ impl<'a, M> TransactionPtrBatch<'a, M> {
         })
     }
 
-    /// Free all transactions in the batch, then free the batch itself.
-    pub fn free(self) {
-        for (transaction_ptr, _) in self.iter() {
-            unsafe { transaction_ptr.free(self.allocator) }
-        }
-
+    /// Free the transaction batch container.
+    ///
+    /// # Safety
+    ///
+    /// - [`SharableTransactionBatchRegion`] must be exclusively owned by this pointer.
+    ///
+    /// # Note
+    ///
+    /// This will not free the underlying transactions as their lifetimes may be differ from that of
+    /// the batch.
+    pub unsafe fn free(self) {
         unsafe { self.allocator.free(self.tx_ptr.cast()) }
     }
 }