agave_scheduling_utils/
transaction_ptr.rs

1use {
2    agave_scheduler_bindings::{SharableTransactionBatchRegion, SharableTransactionRegion},
3    agave_transaction_view::transaction_data::TransactionData,
4    core::ptr::NonNull,
5    rts_alloc::Allocator,
6};
7
8pub struct TransactionPtr {
9    ptr: NonNull<u8>,
10    len: usize,
11}
12
13impl TransactionData for TransactionPtr {
14    fn data(&self) -> &[u8] {
15        unsafe { core::slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
16    }
17}
18
19impl TransactionPtr {
20    /// # Safety
21    /// - `sharable_transaction_region` must reference a valid offset and length
22    ///   within the `allocator`.
23    pub unsafe fn from_sharable_transaction_region(
24        sharable_transaction_region: &SharableTransactionRegion,
25        allocator: &Allocator,
26    ) -> Self {
27        let ptr = allocator.ptr_from_offset(sharable_transaction_region.offset);
28        Self {
29            ptr,
30            len: sharable_transaction_region.length as usize,
31        }
32    }
33
34    /// Translate the ptr type into a sharable region.
35    ///
36    /// # Safety
37    /// - `allocator` must be the allocator owning the memory region pointed
38    ///   to by `self`.
39    pub unsafe fn to_sharable_transaction_region(
40        &self,
41        allocator: &Allocator,
42    ) -> SharableTransactionRegion {
43        // SAFETY: The `TransactionPtr` creation `Self::from_sharable_transaction_region`
44        // is already conditioned on the offset being valid, if that safety constraint
45        // was satisfied translation back to offset is safe.
46        let offset = unsafe { allocator.offset(self.ptr) };
47        SharableTransactionRegion {
48            offset,
49            length: self.len as u32,
50        }
51    }
52
53    /// Frees the memory region pointed to in the `allocator`.
54    /// This should only be called by the owner of the memory
55    /// i.e. the external scheduler.
56    ///
57    /// # Safety
58    /// - Data region pointed to by `TransactionPtr` belongs to the `allocator`.
59    /// - Inner `ptr` must not have been previously freed.
60    pub unsafe fn free(self, allocator: &Allocator) {
61        unsafe { allocator.free(self.ptr) }
62    }
63}
64
65/// A batch of transaction pointers that can be iterated over.
66pub struct TransactionPtrBatch<'a> {
67    ptr: NonNull<SharableTransactionRegion>,
68    num_transactions: usize,
69    allocator: &'a Allocator,
70}
71
72impl<'a> TransactionPtrBatch<'a> {
73    /// # Safety
74    /// - [`SharableTransactionBatchRegion`] must reference a valid offset and length
75    ///   within the `allocator`.
76    /// - ALL [`SharableTransactionRegion`]  within the batch must be valid.
77    ///   See [`TransactionPtr::from_sharable_transaction_region`] for details.
78    pub unsafe fn from_sharable_transaction_batch_region(
79        sharable_transaction_batch_region: &SharableTransactionBatchRegion,
80        allocator: &'a Allocator,
81    ) -> Self {
82        let ptr = allocator
83            .ptr_from_offset(sharable_transaction_batch_region.transactions_offset)
84            .cast();
85
86        Self {
87            ptr,
88            num_transactions: usize::from(sharable_transaction_batch_region.num_transactions),
89            allocator,
90        }
91    }
92
93    /// Iterator returning [`TransactionPtr`] for each transaction in the batch.
94    pub fn iter(&'a self) -> impl Iterator<Item = TransactionPtr> + 'a {
95        (0..self.num_transactions)
96            .map(|idx| unsafe { self.ptr.add(idx) })
97            .map(|ptr| unsafe {
98                TransactionPtr::from_sharable_transaction_region(ptr.as_ref(), self.allocator)
99            })
100    }
101
102    /// Free all transactions in the batch, then free the batch itself.
103    pub fn free(self) {
104        for transaction_ptr in self.iter() {
105            unsafe { transaction_ptr.free(self.allocator) }
106        }
107
108        unsafe { self.allocator.free(self.ptr.cast()) }
109    }
110}