Skip to main content

cu29_runtime/
pool.rs

1use arrayvec::ArrayString;
2use bincode::de::Decoder;
3use bincode::enc::Encoder;
4use bincode::error::{DecodeError, EncodeError};
5use bincode::{Decode, Encode};
6use cu29_traits::CuResult;
7use hashbrown::HashMap;
8use object_pool::{Pool, ReusableOwned};
9use serde::de::{self, MapAccess, SeqAccess, Visitor};
10use serde::{Deserialize, Deserializer, Serialize, Serializer};
11use smallvec::SmallVec;
12use std::alloc::{Layout, alloc, dealloc};
13use std::cell::Cell;
14use std::cell::UnsafeCell;
15use std::fmt::Debug;
16use std::fs::OpenOptions;
17use std::marker::PhantomData;
18use std::mem::{align_of, size_of};
19use std::ops::{Deref, DerefMut};
20use std::path::{Path, PathBuf};
21use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
22use std::sync::{Arc, Mutex, MutexGuard, OnceLock};
23
24use memmap2::{MmapMut, MmapOptions};
25use tempfile::NamedTempFile;
26
27type PoolID = ArrayString<64>;
28
29/// Trait for a Pool to exposed to be monitored by the monitoring API.
30pub trait PoolMonitor: Send + Sync {
31    /// A unique and descriptive identifier for the pool.
32    fn id(&self) -> PoolID;
33
34    /// Number of buffer slots left in the pool.
35    fn space_left(&self) -> usize;
36
37    /// Total size of the pool in number of buffers.
38    fn total_size(&self) -> usize;
39
40    /// Size of one buffer
41    fn buffer_size(&self) -> usize;
42}
43
44static POOL_REGISTRY: OnceLock<Mutex<HashMap<String, Arc<dyn PoolMonitor>>>> = OnceLock::new();
45const MAX_POOLS: usize = 16;
46
47fn lock_unpoison<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
48    match mutex.lock() {
49        Ok(guard) => guard,
50        Err(poison) => poison.into_inner(),
51    }
52}
53
54// Register a pool to the global registry.
55fn register_pool(pool: Arc<dyn PoolMonitor>) {
56    POOL_REGISTRY
57        .get_or_init(|| Mutex::new(HashMap::new()))
58        .lock()
59        .unwrap_or_else(|poison| poison.into_inner())
60        .insert(pool.id().to_string(), pool);
61}
62
63type PoolStats = (PoolID, usize, usize, usize);
64
65/// Get the list of pools and their statistics.
66/// We use SmallVec here to avoid heap allocations while the stack is running.
67pub fn pools_statistics() -> SmallVec<[PoolStats; MAX_POOLS]> {
68    // Safely get the registry, returning empty stats if not initialized.
69    let registry_lock = match POOL_REGISTRY.get() {
70        Some(lock) => lock_unpoison(lock),
71        None => return SmallVec::new(), // Return empty if registry is not initialized
72    };
73    let mut result = SmallVec::with_capacity(MAX_POOLS);
74    for pool in registry_lock.values() {
75        result.push((
76            pool.id(),
77            pool.space_left(),
78            pool.total_size(),
79            pool.buffer_size(),
80        ));
81    }
82    result
83}
84
85/// Basic Type that can be used in a buffer in a CuPool.
86pub trait ElementType: Default + Sized + Copy + Debug + Unpin + Send + Sync {
87    fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError>;
88    fn decode<D: Decoder<Context = ()>>(decoder: &mut D) -> Result<Self, DecodeError>;
89}
90
91/// Blanket implementation for all types that are Sized, Copy, Encode, Decode and Debug.
92impl<T> ElementType for T
93where
94    T: Default + Sized + Copy + Debug + Unpin + Send + Sync,
95    T: Encode,
96    T: Decode<()>,
97{
98    fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
99        self.encode(encoder)
100    }
101
102    fn decode<D: Decoder<Context = ()>>(decoder: &mut D) -> Result<Self, DecodeError> {
103        Self::decode(decoder)
104    }
105}
106
107pub trait ArrayLike: Deref<Target = [Self::Element]> + DerefMut + Debug + Sync + Send {
108    type Element: ElementType;
109}
110
111thread_local! {
112    static SHARED_HANDLE_SERIALIZATION_ENABLED: Cell<bool> = const { Cell::new(false) };
113}
114
115pub struct SharedHandleSerializationGuard {
116    previous: bool,
117}
118
119impl Drop for SharedHandleSerializationGuard {
120    fn drop(&mut self) {
121        SHARED_HANDLE_SERIALIZATION_ENABLED.with(|enabled| enabled.set(self.previous));
122    }
123}
124
125pub fn enable_shared_handle_serialization() -> SharedHandleSerializationGuard {
126    let previous = SHARED_HANDLE_SERIALIZATION_ENABLED.with(|enabled| {
127        let previous = enabled.get();
128        enabled.set(true);
129        previous
130    });
131    SharedHandleSerializationGuard { previous }
132}
133
134fn shared_handle_serialization_enabled() -> bool {
135    SHARED_HANDLE_SERIALIZATION_ENABLED.with(Cell::get)
136}
137
138#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
139#[serde(rename_all = "snake_case")]
140pub enum CuSharedMemoryElementType {
141    U8,
142    U16,
143    U32,
144    U64,
145    I8,
146    I16,
147    I32,
148    I64,
149    F32,
150    F64,
151}
152
153impl CuSharedMemoryElementType {
154    pub fn of<E: ElementType + 'static>() -> Option<Self> {
155        let type_id = core::any::TypeId::of::<E>();
156        if type_id == core::any::TypeId::of::<u8>() {
157            Some(Self::U8)
158        } else if type_id == core::any::TypeId::of::<u16>() {
159            Some(Self::U16)
160        } else if type_id == core::any::TypeId::of::<u32>() {
161            Some(Self::U32)
162        } else if type_id == core::any::TypeId::of::<u64>() {
163            Some(Self::U64)
164        } else if type_id == core::any::TypeId::of::<i8>() {
165            Some(Self::I8)
166        } else if type_id == core::any::TypeId::of::<i16>() {
167            Some(Self::I16)
168        } else if type_id == core::any::TypeId::of::<i32>() {
169            Some(Self::I32)
170        } else if type_id == core::any::TypeId::of::<i64>() {
171            Some(Self::I64)
172        } else if type_id == core::any::TypeId::of::<f32>() {
173            Some(Self::F32)
174        } else if type_id == core::any::TypeId::of::<f64>() {
175            Some(Self::F64)
176        } else {
177            None
178        }
179    }
180}
181
182#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
183pub struct CuSharedMemoryHandleDescriptor {
184    #[serde(rename = "__cu_shm_handle__")]
185    pub marker: bool,
186    pub path: String,
187    pub offset_bytes: usize,
188    pub len_elements: usize,
189    pub element_type: CuSharedMemoryElementType,
190}
191
192impl CuSharedMemoryHandleDescriptor {
193    fn new(
194        path: String,
195        offset_bytes: usize,
196        len_elements: usize,
197        element_type: CuSharedMemoryElementType,
198    ) -> Self {
199        Self {
200            marker: true,
201            path,
202            offset_bytes,
203            len_elements,
204            element_type,
205        }
206    }
207}
208
209struct CuSharedMemoryRegion {
210    path: PathBuf,
211    mmap: UnsafeCell<MmapMut>,
212    _backing_file: Option<NamedTempFile>,
213}
214
215impl CuSharedMemoryRegion {
216    fn create(byte_len: usize) -> CuResult<Arc<Self>> {
217        let file = NamedTempFile::new()
218            .map_err(|e| cu29_traits::CuError::new_with_cause("create shared memory file", e))?;
219        file.as_file()
220            .set_len(byte_len as u64)
221            .map_err(|e| cu29_traits::CuError::new_with_cause("size shared memory file", e))?;
222        let mmap = unsafe {
223            MmapOptions::new()
224                .len(byte_len)
225                .map_mut(file.as_file())
226                .map_err(|e| cu29_traits::CuError::new_with_cause("map shared memory file", e))?
227        };
228        let region = Arc::new(Self {
229            path: file.path().to_path_buf(),
230            mmap: UnsafeCell::new(mmap),
231            _backing_file: Some(file),
232        });
233        cache_shared_region(region.clone());
234        Ok(region)
235    }
236
237    fn open(path: &Path) -> CuResult<Arc<Self>> {
238        if let Some(region) = cached_shared_region(path) {
239            return Ok(region);
240        }
241
242        let file = OpenOptions::new()
243            .read(true)
244            .write(true)
245            .open(path)
246            .map_err(|e| cu29_traits::CuError::new_with_cause("open shared memory file", e))?;
247        let len = file
248            .metadata()
249            .map_err(|e| cu29_traits::CuError::new_with_cause("stat shared memory file", e))?
250            .len() as usize;
251        let mmap = unsafe {
252            MmapOptions::new()
253                .len(len)
254                .map_mut(&file)
255                .map_err(|e| cu29_traits::CuError::new_with_cause("map shared memory file", e))?
256        };
257        let region = Arc::new(Self {
258            path: path.to_path_buf(),
259            mmap: UnsafeCell::new(mmap),
260            _backing_file: None,
261        });
262        cache_shared_region(region.clone());
263        Ok(region)
264    }
265}
266
267impl Debug for CuSharedMemoryRegion {
268    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
269        f.debug_struct("CuSharedMemoryRegion")
270            .field("path", &self.path)
271            .finish_non_exhaustive()
272    }
273}
274
275// SAFETY:
276// Access to the mapped bytes is mediated through Copper handles and pool slot
277// leasing, so cross-thread aliasing follows the same external synchronization as
278// other mutable payload buffers.
279unsafe impl Send for CuSharedMemoryRegion {}
280// SAFETY:
281// See `Send` rationale above.
282unsafe impl Sync for CuSharedMemoryRegion {}
283
284fn shared_region_cache() -> &'static Mutex<HashMap<PathBuf, std::sync::Weak<CuSharedMemoryRegion>>>
285{
286    static CACHE: OnceLock<Mutex<HashMap<PathBuf, std::sync::Weak<CuSharedMemoryRegion>>>> =
287        OnceLock::new();
288    CACHE.get_or_init(|| Mutex::new(HashMap::new()))
289}
290
291fn cache_shared_region(region: Arc<CuSharedMemoryRegion>) {
292    lock_unpoison(shared_region_cache()).insert(region.path.clone(), Arc::downgrade(&region));
293}
294
295fn cached_shared_region(path: &Path) -> Option<Arc<CuSharedMemoryRegion>> {
296    lock_unpoison(shared_region_cache())
297        .get(path)
298        .and_then(std::sync::Weak::upgrade)
299}
300
301fn shared_slot_stride<E: ElementType>(len_elements: usize) -> usize {
302    let raw_bytes = len_elements
303        .checked_mul(size_of::<E>())
304        .expect("shared memory slot size overflow");
305    let alignment = align_of::<E>().max(1);
306    raw_bytes.div_ceil(alignment) * alignment
307}
308
309#[derive(Debug)]
310pub struct CuSharedMemoryBuffer<E: ElementType> {
311    region: Arc<CuSharedMemoryRegion>,
312    offset_bytes: usize,
313    len_elements: usize,
314    _marker: PhantomData<E>,
315}
316
317impl<E: ElementType + 'static> CuSharedMemoryBuffer<E> {
318    fn from_region(
319        region: Arc<CuSharedMemoryRegion>,
320        offset_bytes: usize,
321        len_elements: usize,
322    ) -> Self {
323        Self {
324            region,
325            offset_bytes,
326            len_elements,
327            _marker: PhantomData,
328        }
329    }
330
331    pub fn from_vec_detached(data: Vec<E>) -> CuResult<Self> {
332        let len_elements = data.len();
333        let slot_stride = shared_slot_stride::<E>(len_elements.max(1));
334        let region = CuSharedMemoryRegion::create(slot_stride)?;
335        let mut buffer = Self::from_region(region, 0, len_elements);
336        if !data.is_empty() {
337            buffer.copy_from_slice(&data);
338        }
339        Ok(buffer)
340    }
341
342    pub fn from_descriptor(descriptor: &CuSharedMemoryHandleDescriptor) -> CuResult<Self> {
343        let expected = CuSharedMemoryElementType::of::<E>()
344            .ok_or_else(|| cu29_traits::CuError::from("unsupported shared memory element type"))?;
345        if descriptor.element_type != expected {
346            return Err(cu29_traits::CuError::from(
347                "shared memory descriptor element type mismatch",
348            ));
349        }
350        let region = CuSharedMemoryRegion::open(Path::new(&descriptor.path))?;
351        Ok(Self::from_region(
352            region,
353            descriptor.offset_bytes,
354            descriptor.len_elements,
355        ))
356    }
357
358    pub fn descriptor(&self) -> Option<CuSharedMemoryHandleDescriptor>
359    where
360        E: 'static,
361    {
362        CuSharedMemoryElementType::of::<E>().map(|element_type| {
363            CuSharedMemoryHandleDescriptor::new(
364                self.region.path.display().to_string(),
365                self.offset_bytes,
366                self.len_elements,
367                element_type,
368            )
369        })
370    }
371}
372
373impl<E: ElementType> Deref for CuSharedMemoryBuffer<E> {
374    type Target = [E];
375
376    fn deref(&self) -> &Self::Target {
377        let ptr = unsafe { (*self.region.mmap.get()).as_ptr().add(self.offset_bytes) as *const E };
378        unsafe { std::slice::from_raw_parts(ptr, self.len_elements) }
379    }
380}
381
382impl<E: ElementType> DerefMut for CuSharedMemoryBuffer<E> {
383    fn deref_mut(&mut self) -> &mut Self::Target {
384        let ptr = unsafe {
385            (*self.region.mmap.get())
386                .as_mut_ptr()
387                .add(self.offset_bytes) as *mut E
388        };
389        unsafe { std::slice::from_raw_parts_mut(ptr, self.len_elements) }
390    }
391}
392
393impl<E: ElementType> ArrayLike for CuSharedMemoryBuffer<E> {
394    type Element = E;
395}
396
397impl<E: ElementType> Encode for CuSharedMemoryBuffer<E> {
398    fn encode<Enc: Encoder>(&self, encoder: &mut Enc) -> Result<(), EncodeError> {
399        let len = self.len_elements as u64;
400        Encode::encode(&len, encoder)?;
401        for value in self.deref() {
402            value.encode(encoder)?;
403        }
404        Ok(())
405    }
406}
407
408impl<E: ElementType + 'static> Decode<()> for CuSharedMemoryBuffer<E> {
409    fn decode<D: Decoder<Context = ()>>(decoder: &mut D) -> Result<Self, DecodeError> {
410        let len = <u64 as Decode<()>>::decode(decoder)? as usize;
411        let mut vec = Vec::with_capacity(len);
412        for _ in 0..len {
413            vec.push(E::decode(decoder)?);
414        }
415        Self::from_vec_detached(vec).map_err(|e| DecodeError::OtherString(e.to_string()))
416    }
417}
418
419/// A handle to a pooled or detached object.
420///
421/// For onboard usages, large payloads should typically be pooled. The detached form exists for
422/// offline/deserialization flows and for payloads that are intentionally heap-backed instead of
423/// pool-backed.
424pub enum CuHandleInner<T: Debug + Send + Sync> {
425    Pooled(ReusableOwned<Box<T>>),
426    Detached(Box<T>), // Should only be used in offline cases (e.g. deserialization)
427}
428
429impl<T> Debug for CuHandleInner<T>
430where
431    T: Debug + Send + Sync,
432{
433    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
434        match self {
435            CuHandleInner::Pooled(r) => {
436                write!(f, "Pooled: {:?}", r.deref().deref())
437            }
438            CuHandleInner::Detached(r) => write!(f, "Detached: {r:?}"),
439        }
440    }
441}
442
443impl<T> CuHandleInner<T>
444where
445    T: Debug + Send + Sync,
446{
447    fn inner_ref(&self) -> &T {
448        match self {
449            CuHandleInner::Pooled(pooled) => pooled.deref().as_ref(),
450            CuHandleInner::Detached(detached) => detached.deref(),
451        }
452    }
453
454    fn inner_mut(&mut self) -> &mut T {
455        match self {
456            CuHandleInner::Pooled(pooled) => pooled.deref_mut().as_mut(),
457            CuHandleInner::Detached(detached) => detached.deref_mut(),
458        }
459    }
460}
461
462impl<T> AsRef<T> for CuHandleInner<T>
463where
464    T: Debug + Send + Sync,
465{
466    fn as_ref(&self) -> &T {
467        self.inner_ref()
468    }
469}
470
471impl<T> AsMut<T> for CuHandleInner<T>
472where
473    T: Debug + Send + Sync,
474{
475    fn as_mut(&mut self) -> &mut T {
476        self.inner_mut()
477    }
478}
479
480impl<T: ArrayLike> Deref for CuHandleInner<T> {
481    type Target = [T::Element];
482
483    fn deref(&self) -> &Self::Target {
484        self.inner_ref().deref()
485    }
486}
487
488impl<T: ArrayLike> DerefMut for CuHandleInner<T> {
489    fn deref_mut(&mut self) -> &mut Self::Target {
490        self.inner_mut().deref_mut()
491    }
492}
493
494// `HandleContent` is defined in `config.rs` so it lives in both the library and the
495// `cu29-rendercfg` bin (which includes config.rs standalone). Re-export it here for
496// pool consumers that don't otherwise reach for the config module.
497pub use crate::config::HandleContent;
498
499/// Backing storage for a [`CuHandle`]: the payload mutex plus the per-handle touched flag
500/// and logging mode. Shared across handle clones via [`Arc`].
501///
502/// `mode` is stored as an [`AtomicU8`] so the runtime can override it once after the
503/// source's `process()` returns (when the configured `handle_content` policy is known
504/// but the source itself didn't construct the handle with that policy in mind).
505#[derive(Debug)]
506struct CuHandleCell<T: Debug + Send + Sync> {
507    touched: AtomicBool,
508    mode: AtomicU8,
509    inner: Mutex<CuHandleInner<T>>,
510}
511
512/// A shareable handle to a pooled or detached object.
513///
514/// When `T: ArrayLike`, the handle also participates in Copper's buffer pool APIs.
515#[derive(Debug)]
516pub struct CuHandle<T: Debug + Send + Sync>(Arc<CuHandleCell<T>>);
517
518impl<T: Debug + Send + Sync> Clone for CuHandle<T> {
519    fn clone(&self) -> Self {
520        Self(self.0.clone())
521    }
522}
523
524impl<T: Debug + Send + Sync> Deref for CuHandle<T> {
525    type Target = Mutex<CuHandleInner<T>>;
526
527    fn deref(&self) -> &Self::Target {
528        &self.0.inner
529    }
530}
531
532impl<T: Debug + Send + Sync> CuHandle<T> {
533    /// Wrap a raw [`CuHandleInner`] into a fresh handle with the given logging mode.
534    fn from_inner(inner: CuHandleInner<T>, mode: HandleContent) -> Self {
535        CuHandle(Arc::new(CuHandleCell {
536            touched: AtomicBool::new(false),
537            mode: AtomicU8::new(mode as u8),
538            inner: Mutex::new(inner),
539        }))
540    }
541
542    /// Create a new CuHandle not part of a Pool (not for onboard usages, use pools instead)
543    pub fn new_detached(inner: T) -> Self {
544        Self::new_detached_box(Box::new(inner))
545    }
546
547    /// Create a detached handle from an already heap-allocated object.
548    pub fn new_detached_box(inner: Box<T>) -> Self {
549        Self::from_inner(CuHandleInner::Detached(inner), HandleContent::default())
550    }
551
552    /// Create a detached handle with a non-default logging mode.
553    ///
554    /// Mostly useful for tests and for sources that want to mint detached handles
555    /// (instead of pool-acquired ones) under a specific [`HandleContent`] policy.
556    pub fn new_detached_with_mode(inner: T, mode: HandleContent) -> Self {
557        Self::from_inner(CuHandleInner::Detached(Box::new(inner)), mode)
558    }
559
560    /// Safely access the inner value, applying a closure to it.
561    pub fn with_inner<R>(&self, f: impl FnOnce(&CuHandleInner<T>) -> R) -> R {
562        let lock = lock_unpoison(&self.0.inner);
563        f(&*lock)
564    }
565
566    /// Mutably access the inner value, applying a closure to it.
567    pub fn with_inner_mut<R>(&self, f: impl FnOnce(&mut CuHandleInner<T>) -> R) -> R {
568        let mut lock = lock_unpoison(&self.0.inner);
569        f(&mut *lock)
570    }
571
572    /// Returns the number of handles sharing this payload storage.
573    pub fn strong_count(&self) -> usize {
574        Arc::strong_count(&self.0)
575    }
576
577    /// Returns true when this is the only handle to this payload storage.
578    pub fn is_unique(&self) -> bool {
579        self.strong_count() == 1
580    }
581
582    /// Mark this handle as read by a downstream consumer.
583    ///
584    /// When the source is configured with [`HandleContent::TouchedOnly`], the unified-log
585    /// encoder will only write the payload bytes if at least one consumer called this.
586    /// Cheap (one relaxed atomic store); safe to call multiple times.
587    pub fn mark_touched(&self) {
588        self.0.touched.store(true, Ordering::Relaxed);
589    }
590
591    /// Returns true if [`mark_touched`](Self::mark_touched) has ever been called on this
592    /// handle (or any clone of it).
593    pub fn was_touched(&self) -> bool {
594        self.0.touched.load(Ordering::Relaxed)
595    }
596
597    /// Logging mode currently in effect for this handle.
598    pub fn logging_mode(&self) -> HandleContent {
599        HandleContent::from_u8(self.0.mode.load(Ordering::Relaxed))
600    }
601
602    /// Convenience: [`with_inner`](Self::with_inner) plus [`mark_touched`](Self::mark_touched)
603    /// in one call, for the common consumer-side access pattern.
604    pub fn with_touched_inner<R>(&self, f: impl FnOnce(&CuHandleInner<T>) -> R) -> R {
605        self.mark_touched();
606        self.with_inner(f)
607    }
608
609    /// Decides whether the unified-log encoder should write this handle's payload bytes
610    /// for the current frame.
611    ///
612    /// This is the inherent "specific" arm of the autoref-specialization pattern; the
613    /// encoder resolves to this method for handle payloads (or composite payloads that
614    /// forward to it) and to [`PayloadDefaultLoggingPolicy::payload_should_log`] for
615    /// every other payload type.
616    pub fn payload_should_log(&self) -> bool {
617        match self.logging_mode() {
618            HandleContent::All => true,
619            HandleContent::None => false,
620            HandleContent::TouchedOnly => self.was_touched(),
621        }
622    }
623
624    /// Apply a source's configured [`HandleContent`] policy to this handle. Inherent
625    /// "specific" arm of the autoref-specialization pattern; visible to every clone.
626    /// `Relaxed` is sufficient — synchronization piggy-backs on the copperlist handoff.
627    pub fn apply_handle_content_policy(&self, mode: HandleContent) {
628        self.0.mode.store(mode as u8, Ordering::Relaxed);
629    }
630}
631
632/// Opt-in marker required when `NodeLogging.handle_content` is non-default. Codegen
633/// emits a compile-time bound check against this trait, so an unmarked payload
634/// surfaces as a clear compile error instead of a silent runtime no-op.
635///
636/// ```ignore
637/// impl cu29::pool::HandleContentAware for MyPayload {}
638/// ```
639///
640/// The impl is the gate; for the policy to actually fire, the payload must also
641/// forward [`apply_handle_content_policy`](CuHandle::apply_handle_content_policy)
642/// and [`payload_should_log`](CuHandle::payload_should_log) to an inner [`CuHandle`]
643/// (see `CuImage`). [`CuHandle`] itself satisfies both halves.
644pub trait HandleContentAware {}
645
646impl<T: Debug + Send + Sync> HandleContentAware for CuHandle<T> {}
647
648/// Default arm of the autoref-specialization pattern used by the unified-log encoder.
649///
650/// Blanket-impl'd for every type, so any payload that doesn't define its own inherent
651/// `payload_should_log` method inherits this default-true implementation. Types like
652/// [`CuHandle`] (and composite payloads wrapping one) provide an inherent method with
653/// the same name, which wins method resolution because inherent methods are tried
654/// before trait methods at the same candidate self-type.
655pub trait PayloadDefaultLoggingPolicy {
656    fn payload_should_log(&self) -> bool {
657        true
658    }
659}
660
661impl<T: ?Sized> PayloadDefaultLoggingPolicy for T {}
662
663/// Default arm of the autoref-specialization pattern used by the runtime to push a
664/// source's configured [`HandleContent`] policy into the [`CuHandle`]s that live inside
665/// a payload. Default is a no-op; [`CuHandle`] and composite payloads provide inherent
666/// overrides that propagate the mode.
667pub trait PayloadDefaultHandlePolicyApply {
668    fn apply_handle_content_policy(&self, _mode: HandleContent) {}
669}
670
671impl<T: ?Sized> PayloadDefaultHandlePolicyApply for T {}
672
673impl<U> Serialize for CuHandle<Vec<U>>
674where
675    U: ElementType + Serialize + 'static,
676{
677    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
678        let inner = lock_unpoison(&self.0.inner);
679        inner.inner_ref().serialize(serializer)
680    }
681}
682
683impl<'de, U> Deserialize<'de> for CuHandle<Vec<U>>
684where
685    U: ElementType + Deserialize<'de> + 'static,
686{
687    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
688        Vec::<U>::deserialize(deserializer).map(CuHandle::new_detached)
689    }
690}
691
692impl<U> Serialize for CuHandle<CuSharedMemoryBuffer<U>>
693where
694    U: ElementType + Serialize + 'static,
695{
696    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
697        let inner = lock_unpoison(&self.0.inner);
698        let buffer = inner.inner_ref();
699
700        if shared_handle_serialization_enabled()
701            && let Some(descriptor) = buffer.descriptor()
702        {
703            return descriptor.serialize(serializer);
704        }
705
706        buffer.deref().serialize(serializer)
707    }
708}
709
710impl<'de, U> Deserialize<'de> for CuHandle<CuSharedMemoryBuffer<U>>
711where
712    U: ElementType + Deserialize<'de> + 'static,
713{
714    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
715        enum Repr<U> {
716            Descriptor(CuSharedMemoryHandleDescriptor),
717            Data(Vec<U>),
718        }
719
720        impl<'de, U> Deserialize<'de> for Repr<U>
721        where
722            U: ElementType + Deserialize<'de>,
723        {
724            fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
725                struct ReprVisitor<U>(PhantomData<U>);
726
727                impl<'de, U> Visitor<'de> for ReprVisitor<U>
728                where
729                    U: ElementType + Deserialize<'de>,
730                {
731                    type Value = Repr<U>;
732
733                    fn expecting(
734                        &self,
735                        formatter: &mut std::fmt::Formatter<'_>,
736                    ) -> std::fmt::Result {
737                        formatter
738                            .write_str("a shared-memory handle descriptor or an element sequence")
739                    }
740
741                    fn visit_seq<A: SeqAccess<'de>>(self, seq: A) -> Result<Self::Value, A::Error> {
742                        let data =
743                            Vec::<U>::deserialize(de::value::SeqAccessDeserializer::new(seq))?;
744                        Ok(Repr::Data(data))
745                    }
746
747                    fn visit_map<A: MapAccess<'de>>(self, map: A) -> Result<Self::Value, A::Error> {
748                        let descriptor = CuSharedMemoryHandleDescriptor::deserialize(
749                            de::value::MapAccessDeserializer::new(map),
750                        )?;
751                        Ok(Repr::Descriptor(descriptor))
752                    }
753                }
754
755                deserializer.deserialize_any(ReprVisitor(PhantomData))
756            }
757        }
758
759        match Repr::<U>::deserialize(deserializer)? {
760            Repr::Descriptor(descriptor) => CuSharedMemoryBuffer::from_descriptor(&descriptor)
761                .map(CuHandle::new_detached)
762                .map_err(de::Error::custom),
763            Repr::Data(data) => CuSharedMemoryBuffer::from_vec_detached(data)
764                .map(CuHandle::new_detached)
765                .map_err(de::Error::custom),
766        }
767    }
768}
769
770impl<T: ArrayLike + Encode> Encode for CuHandle<T>
771where
772    <T as ArrayLike>::Element: 'static,
773{
774    fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
775        let inner = lock_unpoison(&self.0.inner);
776        crate::monitoring::record_payload_handle_bytes(
777            inner.inner_ref().len() * size_of::<T::Element>(),
778        );
779        inner.inner_ref().encode(encoder)
780    }
781}
782
783impl<T: Debug + Send + Sync> Default for CuHandle<T> {
784    fn default() -> Self {
785        panic!("Cannot create a default CuHandle")
786    }
787}
788
789impl<U: ElementType + Decode<()> + 'static> Decode<()> for CuHandle<Vec<U>> {
790    fn decode<D: Decoder<Context = ()>>(decoder: &mut D) -> Result<Self, DecodeError> {
791        let vec: Vec<U> = Vec::decode(decoder)?;
792        Ok(CuHandle::new_detached(vec))
793    }
794}
795
796impl<U: ElementType + Decode<()> + 'static> Decode<()> for CuHandle<CuSharedMemoryBuffer<U>> {
797    fn decode<D: Decoder<Context = ()>>(decoder: &mut D) -> Result<Self, DecodeError> {
798        let buffer = CuSharedMemoryBuffer::<U>::decode(decoder)?;
799        Ok(CuHandle::new_detached(buffer))
800    }
801}
802
803/// A CuPool is a pool of buffers that can be shared between different parts of the code.
804/// Handles can be stored locally in the tasks and shared between them.
805pub trait CuPool<T: ArrayLike>: PoolMonitor {
806    /// Acquire a buffer from the pool.
807    fn acquire(&self) -> Option<CuHandle<T>>;
808
809    /// Copy data from a handle to a new handle from the pool.
810    fn copy_from<O>(&self, from: &mut CuHandle<O>) -> CuHandle<T>
811    where
812        O: ArrayLike<Element = T::Element>;
813}
814
815/// A device memory pool can copy data from a device to a host memory pool on top.
816pub trait DeviceCuPool<T: ArrayLike>: CuPool<T> {
817    /// Takes a handle to a device buffer and copies it into a host buffer pool.
818    /// It returns a new handle from the host pool with the data from the device handle given.
819    fn copy_to_host_pool<O>(
820        &self,
821        from_device_handle: &CuHandle<T>,
822        to_host_handle: &mut CuHandle<O>,
823    ) -> CuResult<()>
824    where
825        O: ArrayLike<Element = T::Element>;
826}
827
828/// A pool of host memory buffers.
829pub struct CuHostMemoryPool<T> {
830    /// Underlying pool of host buffers.
831    // Being an Arc is a requirement of try_pull_owned() so buffers can refer back to the pool.
832    id: PoolID,
833    pool: Arc<Pool<Box<T>>>,
834    size: usize,
835    buffer_size: usize,
836}
837
838impl<T: ArrayLike + 'static> CuHostMemoryPool<T> {
839    pub fn new<F>(id: &str, size: usize, buffer_initializer: F) -> CuResult<Arc<Self>>
840    where
841        F: Fn() -> T,
842    {
843        let pool = Arc::new(Pool::new(size, move || Box::new(buffer_initializer())));
844        let buffer_size = pool.try_pull().unwrap().len() * size_of::<T::Element>();
845
846        let og = Self {
847            id: PoolID::from(id).map_err(|_| "Failed to create PoolID")?,
848            pool,
849            size,
850            buffer_size,
851        };
852        let og = Arc::new(og);
853        register_pool(og.clone());
854        Ok(og)
855    }
856}
857
858impl<T: ArrayLike> PoolMonitor for CuHostMemoryPool<T> {
859    fn id(&self) -> PoolID {
860        self.id
861    }
862
863    fn space_left(&self) -> usize {
864        self.pool.len()
865    }
866
867    fn total_size(&self) -> usize {
868        self.size
869    }
870
871    fn buffer_size(&self) -> usize {
872        self.buffer_size
873    }
874}
875
876impl<T: ArrayLike> CuPool<T> for CuHostMemoryPool<T> {
877    fn acquire(&self) -> Option<CuHandle<T>> {
878        let owned_object = self.pool.try_pull_owned(); // Use the owned version
879
880        owned_object.map(|reusable| {
881            CuHandle::from_inner(CuHandleInner::Pooled(reusable), HandleContent::default())
882        })
883    }
884
885    fn copy_from<O: ArrayLike<Element = T::Element>>(&self, from: &mut CuHandle<O>) -> CuHandle<T> {
886        let to_handle = self.acquire().expect("No available buffers in the pool");
887        {
888            let from_lock = lock_unpoison(&from.0.inner);
889            let mut to_lock = lock_unpoison(&to_handle.0.inner);
890            to_lock.inner_mut().copy_from_slice(from_lock.inner_ref());
891        }
892        to_handle
893    }
894}
895
896/// A pool of fixed-size shared-memory buffers that can be leased to a child
897/// process without copying the underlying bytes.
898pub struct CuSharedMemoryPool<E: ElementType> {
899    id: PoolID,
900    pool: Arc<Pool<Box<CuSharedMemoryBuffer<E>>>>,
901    size: usize,
902    buffer_size: usize,
903}
904
905impl<E: ElementType + 'static> CuSharedMemoryPool<E> {
906    pub fn new(id: &str, size: usize, elements_per_buffer: usize) -> CuResult<Arc<Self>> {
907        let slot_stride = shared_slot_stride::<E>(elements_per_buffer.max(1));
908        let region = CuSharedMemoryRegion::create(
909            slot_stride
910                .checked_mul(size)
911                .ok_or_else(|| cu29_traits::CuError::from("shared memory pool size overflow"))?,
912        )?;
913        let next_slot = Arc::new(std::sync::atomic::AtomicUsize::new(0));
914        let initializer_region = region.clone();
915        let initializer_next_slot = next_slot.clone();
916        let pool = Arc::new(Pool::new(size, move || {
917            let slot = initializer_next_slot.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
918            assert!(slot < size, "shared memory pool slot index overflow");
919            Box::new(CuSharedMemoryBuffer::from_region(
920                initializer_region.clone(),
921                slot * slot_stride,
922                elements_per_buffer,
923            ))
924        }));
925
926        let pool = Arc::new(Self {
927            id: PoolID::from(id).map_err(|_| "Failed to create PoolID")?,
928            pool,
929            size,
930            buffer_size: elements_per_buffer * size_of::<E>(),
931        });
932        register_pool(pool.clone());
933        Ok(pool)
934    }
935}
936
937impl<E: ElementType> PoolMonitor for CuSharedMemoryPool<E> {
938    fn id(&self) -> PoolID {
939        self.id
940    }
941
942    fn space_left(&self) -> usize {
943        self.pool.len()
944    }
945
946    fn total_size(&self) -> usize {
947        self.size
948    }
949
950    fn buffer_size(&self) -> usize {
951        self.buffer_size
952    }
953}
954
955impl<E: ElementType> CuPool<CuSharedMemoryBuffer<E>> for CuSharedMemoryPool<E> {
956    fn acquire(&self) -> Option<CuHandle<CuSharedMemoryBuffer<E>>> {
957        self.pool.try_pull_owned().map(|reusable| {
958            CuHandle::from_inner(CuHandleInner::Pooled(reusable), HandleContent::default())
959        })
960    }
961
962    fn copy_from<O>(&self, from: &mut CuHandle<O>) -> CuHandle<CuSharedMemoryBuffer<E>>
963    where
964        O: ArrayLike<Element = E>,
965    {
966        let to_handle = self.acquire().expect("No available buffers in the pool");
967        {
968            let from_lock = lock_unpoison(&from.0.inner);
969            let mut to_lock = lock_unpoison(&to_handle.0.inner);
970            to_lock.inner_mut().copy_from_slice(from_lock.inner_ref());
971        }
972        to_handle
973    }
974}
975
976impl<E: ElementType + 'static> ArrayLike for Vec<E> {
977    type Element = E;
978}
979
980#[cfg(all(feature = "cuda", not(target_os = "macos")))]
981mod cuda {
982    use super::*;
983    use cu29_traits::CuError;
984    use cudarc::driver::{
985        CudaContext, CudaSlice, CudaStream, DeviceRepr, HostSlice, SyncOnDrop, ValidAsZeroBits,
986    };
987    use std::sync::Arc;
988
989    #[derive(Debug)]
990    pub struct CudaSliceWrapper<E>(CudaSlice<E>);
991
992    impl<E> Deref for CudaSliceWrapper<E>
993    where
994        E: ElementType,
995    {
996        type Target = [E];
997
998        fn deref(&self) -> &Self::Target {
999            // Implement logic to return a slice
1000            panic!("You need to copy data to host memory pool before accessing it.");
1001        }
1002    }
1003
1004    impl<E> DerefMut for CudaSliceWrapper<E>
1005    where
1006        E: ElementType,
1007    {
1008        fn deref_mut(&mut self) -> &mut Self::Target {
1009            panic!("You need to copy data to host memory pool before accessing it.");
1010        }
1011    }
1012
1013    impl<E: ElementType> ArrayLike for CudaSliceWrapper<E> {
1014        type Element = E;
1015    }
1016
1017    impl<E> CudaSliceWrapper<E> {
1018        pub fn as_cuda_slice(&self) -> &CudaSlice<E> {
1019            &self.0
1020        }
1021
1022        pub fn as_cuda_slice_mut(&mut self) -> &mut CudaSlice<E> {
1023            &mut self.0
1024        }
1025    }
1026
1027    // Create a wrapper type to bridge between ArrayLike and HostSlice
1028    pub struct HostSliceWrapper<'a, T: ArrayLike> {
1029        inner: &'a T,
1030    }
1031
1032    impl<T: ArrayLike> HostSlice<T::Element> for HostSliceWrapper<'_, T> {
1033        fn len(&self) -> usize {
1034            self.inner.len()
1035        }
1036
1037        // SAFETY: HostSlice requires the returned slice to remain valid for 'b.
1038        unsafe fn stream_synced_slice<'b>(
1039            &'b self,
1040            stream: &'b CudaStream,
1041        ) -> (&'b [T::Element], SyncOnDrop<'b>) {
1042            (self.inner.deref(), SyncOnDrop::sync_stream(stream))
1043        }
1044
1045        // SAFETY: This wrapper cannot provide mutable access; callers must not rely on this.
1046        unsafe fn stream_synced_mut_slice<'b>(
1047            &'b mut self,
1048            _stream: &'b CudaStream,
1049        ) -> (&'b mut [T::Element], SyncOnDrop<'b>) {
1050            panic!("Cannot get mutable reference from immutable wrapper")
1051        }
1052    }
1053
1054    // Mutable wrapper
1055    pub struct HostSliceMutWrapper<'a, T: ArrayLike> {
1056        inner: &'a mut T,
1057    }
1058
1059    impl<T: ArrayLike> HostSlice<T::Element> for HostSliceMutWrapper<'_, T> {
1060        fn len(&self) -> usize {
1061            self.inner.len()
1062        }
1063
1064        // SAFETY: HostSlice requires the returned slice to remain valid for 'b.
1065        unsafe fn stream_synced_slice<'b>(
1066            &'b self,
1067            stream: &'b CudaStream,
1068        ) -> (&'b [T::Element], SyncOnDrop<'b>) {
1069            (self.inner.deref(), SyncOnDrop::sync_stream(stream))
1070        }
1071
1072        // SAFETY: HostSlice requires the returned slice to remain valid for 'b.
1073        unsafe fn stream_synced_mut_slice<'b>(
1074            &'b mut self,
1075            stream: &'b CudaStream,
1076        ) -> (&'b mut [T::Element], SyncOnDrop<'b>) {
1077            (self.inner.deref_mut(), SyncOnDrop::sync_stream(stream))
1078        }
1079    }
1080
1081    // Add helper methods to the CuCudaPool implementation
1082    impl<E: ElementType + ValidAsZeroBits + DeviceRepr> CuCudaPool<E> {
1083        // Helper method to get a HostSliceWrapper from a CuHandleInner
1084        fn get_host_slice_wrapper<O: ArrayLike<Element = E>>(
1085            handle_inner: &CuHandleInner<O>,
1086        ) -> HostSliceWrapper<'_, O> {
1087            HostSliceWrapper {
1088                inner: handle_inner.inner_ref(),
1089            }
1090        }
1091
1092        // Helper method to get a HostSliceMutWrapper from a CuHandleInner
1093        fn get_host_slice_mut_wrapper<O: ArrayLike<Element = E>>(
1094            handle_inner: &mut CuHandleInner<O>,
1095        ) -> HostSliceMutWrapper<'_, O> {
1096            HostSliceMutWrapper {
1097                inner: handle_inner.inner_mut(),
1098            }
1099        }
1100    }
1101    /// A pool of CUDA memory buffers.
1102    pub struct CuCudaPool<E>
1103    where
1104        E: ElementType + ValidAsZeroBits + DeviceRepr + Unpin,
1105    {
1106        id: PoolID,
1107        stream: Arc<CudaStream>,
1108        pool: Arc<Pool<Box<CudaSliceWrapper<E>>>>,
1109        nb_buffers: usize,
1110        nb_element_per_buffer: usize,
1111    }
1112
1113    impl<E: ElementType + ValidAsZeroBits + DeviceRepr> CuCudaPool<E> {
1114        #[allow(dead_code)]
1115        pub fn new(
1116            id: &'static str,
1117            ctx: Arc<CudaContext>,
1118            nb_buffers: usize,
1119            nb_element_per_buffer: usize,
1120        ) -> CuResult<Self> {
1121            let stream = ctx.default_stream();
1122            let pool = (0..nb_buffers)
1123                .map(|_| {
1124                    stream
1125                        .alloc_zeros(nb_element_per_buffer)
1126                        .map(CudaSliceWrapper)
1127                        .map(Box::new)
1128                        .map_err(|_| "Failed to allocate device memory")
1129                })
1130                .collect::<Result<Vec<_>, _>>()?;
1131
1132            Ok(Self {
1133                id: PoolID::from(id).map_err(|_| "Failed to create PoolID")?,
1134                stream,
1135                pool: Arc::new(Pool::from_vec(pool)),
1136                nb_buffers,
1137                nb_element_per_buffer,
1138            })
1139        }
1140    }
1141
1142    impl<E> PoolMonitor for CuCudaPool<E>
1143    where
1144        E: DeviceRepr + ElementType + ValidAsZeroBits,
1145    {
1146        fn id(&self) -> PoolID {
1147            self.id
1148        }
1149
1150        fn space_left(&self) -> usize {
1151            self.pool.len()
1152        }
1153
1154        fn total_size(&self) -> usize {
1155            self.nb_buffers
1156        }
1157
1158        fn buffer_size(&self) -> usize {
1159            self.nb_element_per_buffer * size_of::<E>()
1160        }
1161    }
1162
1163    impl<E> CuPool<CudaSliceWrapper<E>> for CuCudaPool<E>
1164    where
1165        E: DeviceRepr + ElementType + ValidAsZeroBits,
1166    {
1167        fn acquire(&self) -> Option<CuHandle<CudaSliceWrapper<E>>> {
1168            self.pool
1169                .try_pull_owned()
1170                .map(|x| CuHandle::from_inner(CuHandleInner::Pooled(x), HandleContent::default()))
1171        }
1172
1173        fn copy_from<O>(&self, from_handle: &mut CuHandle<O>) -> CuHandle<CudaSliceWrapper<E>>
1174        where
1175            O: ArrayLike<Element = E>,
1176        {
1177            let to_handle = self.acquire().expect("No available buffers in the pool");
1178
1179            {
1180                let from_lock = lock_unpoison(&from_handle.0.inner);
1181                let mut to_lock = lock_unpoison(&to_handle.0.inner);
1182
1183                match &mut *to_lock {
1184                    CuHandleInner::Detached(to) => {
1185                        let wrapper = Self::get_host_slice_wrapper(&*from_lock);
1186                        self.stream
1187                            .memcpy_htod(&wrapper, to.deref_mut().as_cuda_slice_mut())
1188                            .expect("Failed to copy data to device");
1189                    }
1190                    CuHandleInner::Pooled(to) => {
1191                        let wrapper = Self::get_host_slice_wrapper(&*from_lock);
1192                        self.stream
1193                            .memcpy_htod(&wrapper, to.deref_mut().as_mut().as_cuda_slice_mut())
1194                            .expect("Failed to copy data to device");
1195                    }
1196                }
1197            } // locks are dropped here
1198            to_handle // now we can safely return to_handle
1199        }
1200    }
1201
1202    impl<E> DeviceCuPool<CudaSliceWrapper<E>> for CuCudaPool<E>
1203    where
1204        E: ElementType + ValidAsZeroBits + DeviceRepr,
1205    {
1206        /// Copy from device to host
1207        fn copy_to_host_pool<O>(
1208            &self,
1209            device_handle: &CuHandle<CudaSliceWrapper<E>>,
1210            host_handle: &mut CuHandle<O>,
1211        ) -> Result<(), CuError>
1212        where
1213            O: ArrayLike<Element = E>,
1214        {
1215            let device_lock = device_handle.lock().map_err(|e| {
1216                CuError::from("Device handle mutex poisoned").add_cause(&e.to_string())
1217            })?;
1218            let mut host_lock = host_handle.lock().map_err(|e| {
1219                CuError::from("Host handle mutex poisoned").add_cause(&e.to_string())
1220            })?;
1221            let src = match &*device_lock {
1222                CuHandleInner::Pooled(source) => source.deref().as_ref().as_cuda_slice(),
1223                CuHandleInner::Detached(source) => source.deref().as_cuda_slice(),
1224            };
1225            let mut wrapper = Self::get_host_slice_mut_wrapper(&mut *host_lock);
1226            self.stream.memcpy_dtoh(src, &mut wrapper).map_err(|e| {
1227                CuError::from("Failed to copy data from device to host").add_cause(&e.to_string())
1228            })?;
1229            Ok(())
1230        }
1231    }
1232}
1233
1234#[derive(Debug)]
1235/// A buffer that is aligned to a specific size with the Element of type E.
1236pub struct AlignedBuffer<E: ElementType> {
1237    ptr: *mut E,
1238    size: usize,
1239    layout: Layout,
1240}
1241
1242impl<E: ElementType> AlignedBuffer<E> {
1243    pub fn new(num_elements: usize, alignment: usize) -> Self {
1244        assert!(
1245            num_elements > 0 && size_of::<E>() > 0,
1246            "AlignedBuffer requires a non-zero element count and non-zero-sized element type"
1247        );
1248        let alignment = alignment.max(align_of::<E>());
1249        let alloc_size = num_elements
1250            .checked_mul(size_of::<E>())
1251            .expect("AlignedBuffer allocation size overflow");
1252        let layout = Layout::from_size_align(alloc_size, alignment).unwrap();
1253        // SAFETY: layout describes a valid, non-zero allocation request.
1254        let ptr = unsafe { alloc(layout) as *mut E };
1255        if ptr.is_null() {
1256            panic!("Failed to allocate memory");
1257        }
1258        // SAFETY: ptr is valid for writes of `num_elements` elements.
1259        unsafe {
1260            for i in 0..num_elements {
1261                std::ptr::write(ptr.add(i), E::default());
1262            }
1263        }
1264        Self {
1265            ptr,
1266            size: num_elements,
1267            layout,
1268        }
1269    }
1270}
1271
1272impl<E: ElementType> Deref for AlignedBuffer<E> {
1273    type Target = [E];
1274
1275    fn deref(&self) -> &Self::Target {
1276        // SAFETY: `new` initializes all elements and keeps the pointer aligned.
1277        unsafe { std::slice::from_raw_parts(self.ptr, self.size) }
1278    }
1279}
1280
1281impl<E: ElementType> DerefMut for AlignedBuffer<E> {
1282    fn deref_mut(&mut self) -> &mut Self::Target {
1283        // SAFETY: `new` initializes all elements and keeps the pointer aligned.
1284        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.size) }
1285    }
1286}
1287
1288impl<E: ElementType> Drop for AlignedBuffer<E> {
1289    fn drop(&mut self) {
1290        // SAFETY: `ptr` was allocated with `layout` in `new`.
1291        unsafe { dealloc(self.ptr as *mut u8, self.layout) }
1292    }
1293}
1294
1295#[cfg(test)]
1296mod tests {
1297    use super::*;
1298
1299    #[test]
1300    fn test_handle_touched_flag_defaults_false() {
1301        let h: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![1, 2, 3]);
1302        assert!(!h.was_touched());
1303        assert_eq!(h.logging_mode(), HandleContent::All);
1304    }
1305
1306    #[test]
1307    fn test_handle_mark_touched_propagates_across_clones() {
1308        let h: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![1, 2, 3]);
1309        let clone = h.clone();
1310        assert!(!h.was_touched());
1311        assert!(!clone.was_touched());
1312
1313        clone.mark_touched();
1314        // Shared Arc<CuHandleCell>: any clone sees the flag flip.
1315        assert!(h.was_touched());
1316        assert!(clone.was_touched());
1317    }
1318
1319    #[test]
1320    fn test_handle_strong_count_tracks_clones() {
1321        let h: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![1, 2, 3]);
1322        assert_eq!(h.strong_count(), 1);
1323        assert!(h.is_unique());
1324
1325        let clone = h.clone();
1326        assert_eq!(h.strong_count(), 2);
1327        assert_eq!(clone.strong_count(), 2);
1328        assert!(!h.is_unique());
1329        assert!(!clone.is_unique());
1330
1331        drop(clone);
1332        assert_eq!(h.strong_count(), 1);
1333        assert!(h.is_unique());
1334    }
1335
1336    #[test]
1337    fn test_with_touched_inner_marks_and_reads() {
1338        let h: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![10, 20]);
1339        let first = h.with_touched_inner(|inner| inner.as_ref()[0]);
1340        assert_eq!(first, 10);
1341        assert!(h.was_touched());
1342    }
1343
1344    #[test]
1345    fn test_with_inner_does_not_mark_touched() {
1346        let h: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![10, 20]);
1347        let _ = h.with_inner(|inner| inner.as_ref()[0]);
1348        assert!(
1349            !h.was_touched(),
1350            "with_inner must not flip the touched flag"
1351        );
1352    }
1353
1354    #[test]
1355    fn test_payload_should_log_mode_all() {
1356        let h: CuHandle<Vec<u8>> = CuHandle::new_detached_with_mode(vec![1], HandleContent::All);
1357        assert!(h.payload_should_log());
1358        h.mark_touched();
1359        assert!(h.payload_should_log());
1360    }
1361
1362    #[test]
1363    fn test_payload_should_log_mode_none() {
1364        let h: CuHandle<Vec<u8>> = CuHandle::new_detached_with_mode(vec![1], HandleContent::None);
1365        assert!(!h.payload_should_log());
1366        h.mark_touched();
1367        assert!(
1368            !h.payload_should_log(),
1369            "HandleContent::None must never log payload, even when touched"
1370        );
1371    }
1372
1373    #[test]
1374    fn test_payload_should_log_mode_touched_only() {
1375        let h: CuHandle<Vec<u8>> =
1376            CuHandle::new_detached_with_mode(vec![1], HandleContent::TouchedOnly);
1377        assert!(!h.payload_should_log(), "untouched + TouchedOnly => skip");
1378        h.mark_touched();
1379        assert!(h.payload_should_log(), "touched + TouchedOnly => log");
1380    }
1381
1382    #[test]
1383    fn test_default_policy_returns_true_for_non_handle_types() {
1384        use crate::pool::PayloadDefaultLoggingPolicy as _;
1385        // The autoref-specialization fallback applies to any type that doesn't define
1386        // its own inherent payload_should_log. A plain integer is a fine stand-in.
1387        let v: u64 = 42;
1388        assert!(v.payload_should_log());
1389    }
1390
1391    #[test]
1392    fn test_apply_handle_content_policy_overrides_mode() {
1393        // Handle minted with the default All policy; runtime hook flips it to
1394        // TouchedOnly before downstream consumers see it. Encoder decision tracks.
1395        let h: CuHandle<Vec<u8>> = CuHandle::new_detached(vec![1, 2, 3]);
1396        assert_eq!(h.logging_mode(), HandleContent::All);
1397        assert!(h.payload_should_log());
1398
1399        h.apply_handle_content_policy(HandleContent::TouchedOnly);
1400        assert_eq!(h.logging_mode(), HandleContent::TouchedOnly);
1401        assert!(
1402            !h.payload_should_log(),
1403            "after switch to TouchedOnly an untouched handle must skip"
1404        );
1405
1406        h.mark_touched();
1407        assert!(h.payload_should_log(), "touched + TouchedOnly => log");
1408    }
1409
1410    #[test]
1411    fn test_default_apply_policy_is_noop_for_non_handle_types() {
1412        use crate::pool::PayloadDefaultHandlePolicyApply as _;
1413        // No-op for any payload that doesn't wrap a handle — compile-time check that
1414        // codegen can call this on every slot without knowing the payload shape.
1415        let v: u64 = 42;
1416        v.apply_handle_content_policy(HandleContent::TouchedOnly);
1417        // No observable change; the test just proves the call compiles and returns.
1418    }
1419
1420    #[test]
1421    fn test_pool() {
1422        use std::cell::RefCell;
1423        let objs = RefCell::new(vec![vec![1], vec![2], vec![3]]);
1424        let holding = objs.borrow().clone();
1425        let objs_as_slices = holding.iter().map(|x| x.as_slice()).collect::<Vec<_>>();
1426        let pool = CuHostMemoryPool::new("mytestcudapool", 3, || objs.borrow_mut().pop().unwrap())
1427            .unwrap();
1428
1429        let obj1 = pool.acquire().unwrap();
1430        {
1431            let obj2 = pool.acquire().unwrap();
1432            assert!(objs_as_slices.contains(&obj1.lock().unwrap().deref().deref()));
1433            assert!(objs_as_slices.contains(&obj2.lock().unwrap().deref().deref()));
1434            assert_eq!(pool.space_left(), 1);
1435        }
1436        assert_eq!(pool.space_left(), 2);
1437
1438        let obj3 = pool.acquire().unwrap();
1439        assert!(objs_as_slices.contains(&obj3.lock().unwrap().deref().deref()));
1440
1441        assert_eq!(pool.space_left(), 1);
1442
1443        let _obj4 = pool.acquire().unwrap();
1444        assert_eq!(pool.space_left(), 0);
1445
1446        let obj5 = pool.acquire();
1447        assert!(obj5.is_none());
1448    }
1449
1450    #[cfg(all(feature = "cuda", has_nvidia_gpu))]
1451    #[test]
1452    fn test_cuda_pool() {
1453        use crate::pool::cuda::CuCudaPool;
1454        use cudarc::driver::CudaContext;
1455        let ctx = CudaContext::new(0).unwrap();
1456        let pool = CuCudaPool::<f32>::new("mytestcudapool", ctx, 3, 1).unwrap();
1457
1458        let _obj1 = pool.acquire().unwrap();
1459
1460        {
1461            let _obj2 = pool.acquire().unwrap();
1462            assert_eq!(pool.space_left(), 1);
1463        }
1464        assert_eq!(pool.space_left(), 2);
1465
1466        let _obj3 = pool.acquire().unwrap();
1467
1468        assert_eq!(pool.space_left(), 1);
1469
1470        let _obj4 = pool.acquire().unwrap();
1471        assert_eq!(pool.space_left(), 0);
1472
1473        let obj5 = pool.acquire();
1474        assert!(obj5.is_none());
1475    }
1476
1477    #[cfg(all(feature = "cuda", has_nvidia_gpu))]
1478    #[test]
1479    fn test_copy_roundtrip() {
1480        use crate::pool::cuda::CuCudaPool;
1481        use cudarc::driver::CudaContext;
1482        let ctx = CudaContext::new(0).unwrap();
1483        let host_pool = CuHostMemoryPool::new("mytesthostpool", 3, || vec![0.0; 1]).unwrap();
1484        let cuda_pool = CuCudaPool::<f32>::new("mytestcudapool", ctx, 3, 1).unwrap();
1485
1486        let cuda_handle = {
1487            let mut initial_handle = host_pool.acquire().unwrap();
1488            {
1489                let mut inner_initial_handle = initial_handle.lock().unwrap();
1490                if let CuHandleInner::Pooled(ref mut pooled) = *inner_initial_handle {
1491                    pooled[0] = 42.0;
1492                } else {
1493                    panic!();
1494                }
1495            }
1496
1497            // send that to the GPU
1498            cuda_pool.copy_from(&mut initial_handle)
1499        };
1500
1501        // get it back to the host
1502        let mut final_handle = host_pool.acquire().unwrap();
1503        cuda_pool
1504            .copy_to_host_pool(&cuda_handle, &mut final_handle)
1505            .unwrap();
1506
1507        let value = final_handle.lock().unwrap().deref().deref()[0];
1508        assert_eq!(value, 42.0);
1509    }
1510}