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
29pub trait PoolMonitor: Send + Sync {
31 fn id(&self) -> PoolID;
33
34 fn space_left(&self) -> usize;
36
37 fn total_size(&self) -> usize;
39
40 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
54fn 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
65pub fn pools_statistics() -> SmallVec<[PoolStats; MAX_POOLS]> {
68 let registry_lock = match POOL_REGISTRY.get() {
70 Some(lock) => lock_unpoison(lock),
71 None => return SmallVec::new(), };
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
85pub 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
91impl<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
275unsafe impl Send for CuSharedMemoryRegion {}
280unsafe 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(®ion));
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
419pub enum CuHandleInner<T: Debug + Send + Sync> {
425 Pooled(ReusableOwned<Box<T>>),
426 Detached(Box<T>), }
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
494pub use crate::config::HandleContent;
498
499#[derive(Debug)]
506struct CuHandleCell<T: Debug + Send + Sync> {
507 touched: AtomicBool,
508 mode: AtomicU8,
509 inner: Mutex<CuHandleInner<T>>,
510}
511
512#[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 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 pub fn new_detached(inner: T) -> Self {
544 Self::new_detached_box(Box::new(inner))
545 }
546
547 pub fn new_detached_box(inner: Box<T>) -> Self {
549 Self::from_inner(CuHandleInner::Detached(inner), HandleContent::default())
550 }
551
552 pub fn new_detached_with_mode(inner: T, mode: HandleContent) -> Self {
557 Self::from_inner(CuHandleInner::Detached(Box::new(inner)), mode)
558 }
559
560 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 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 pub fn strong_count(&self) -> usize {
574 Arc::strong_count(&self.0)
575 }
576
577 pub fn is_unique(&self) -> bool {
579 self.strong_count() == 1
580 }
581
582 pub fn mark_touched(&self) {
588 self.0.touched.store(true, Ordering::Relaxed);
589 }
590
591 pub fn was_touched(&self) -> bool {
594 self.0.touched.load(Ordering::Relaxed)
595 }
596
597 pub fn logging_mode(&self) -> HandleContent {
599 HandleContent::from_u8(self.0.mode.load(Ordering::Relaxed))
600 }
601
602 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 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 pub fn apply_handle_content_policy(&self, mode: HandleContent) {
628 self.0.mode.store(mode as u8, Ordering::Relaxed);
629 }
630}
631
632pub trait HandleContentAware {}
645
646impl<T: Debug + Send + Sync> HandleContentAware for CuHandle<T> {}
647
648pub trait PayloadDefaultLoggingPolicy {
656 fn payload_should_log(&self) -> bool {
657 true
658 }
659}
660
661impl<T: ?Sized> PayloadDefaultLoggingPolicy for T {}
662
663pub 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
803pub trait CuPool<T: ArrayLike>: PoolMonitor {
806 fn acquire(&self) -> Option<CuHandle<T>>;
808
809 fn copy_from<O>(&self, from: &mut CuHandle<O>) -> CuHandle<T>
811 where
812 O: ArrayLike<Element = T::Element>;
813}
814
815pub trait DeviceCuPool<T: ArrayLike>: CuPool<T> {
817 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
828pub struct CuHostMemoryPool<T> {
830 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(); 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
896pub 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 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 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 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 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 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 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 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 impl<E: ElementType + ValidAsZeroBits + DeviceRepr> CuCudaPool<E> {
1083 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 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 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 } to_handle }
1200 }
1201
1202 impl<E> DeviceCuPool<CudaSliceWrapper<E>> for CuCudaPool<E>
1203 where
1204 E: ElementType + ValidAsZeroBits + DeviceRepr,
1205 {
1206 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)]
1235pub 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 let ptr = unsafe { alloc(layout) as *mut E };
1255 if ptr.is_null() {
1256 panic!("Failed to allocate memory");
1257 }
1258 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 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 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 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 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 let v: u64 = 42;
1388 assert!(v.payload_should_log());
1389 }
1390
1391 #[test]
1392 fn test_apply_handle_content_policy_overrides_mode() {
1393 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 let v: u64 = 42;
1416 v.apply_handle_content_policy(HandleContent::TouchedOnly);
1417 }
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 cuda_pool.copy_from(&mut initial_handle)
1499 };
1500
1501 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}