1use std::{fmt, marker::PhantomData, mem::ManuallyDrop, ptr::NonNull, sync::Arc};
9
10use cudf::{
11 CurrentResourceRouteLease, DeviceId, DeviceMemoryManager,
12 interop::contract::{ExecutionBinding, InteropError, InteropErrorKind},
13 runtime::{BindingExtensionGuard, BindingExtensionInitError, BindingExtensionLease},
14};
15use cugraph_sys::{cugraph_resource_handle_get_comm_size, cugraph_resource_handle_get_rank, cugraph_resource_handle_t};
16
17use crate::{
18 cuda::{DeviceGuard, DeviceResourceKind, drop_device_resource},
19 error::{Error, ErrorKind, Result},
20};
21
22enum ResourceHandleState {
23 #[cfg(test)]
24 LocalOwned {
25 raw: *mut cugraph_resource_handle_t,
26 device: i32,
27 },
28 RuntimeBound {
29 lease: ResourceHandleLease,
30 },
31 ExternalBorrowed {
32 raw: *mut cugraph_resource_handle_t,
33 device: i32,
34 },
35 ExternalOwned {
36 raw: *mut cugraph_resource_handle_t,
37 device: i32,
38 },
39}
40
41impl ResourceHandleState {
42 fn standalone_raw(&self) -> Option<*mut cugraph_resource_handle_t> {
43 match self {
44 #[cfg(test)]
45 Self::LocalOwned { raw, .. } => Some(*raw),
46 Self::ExternalBorrowed { raw, .. } | Self::ExternalOwned { raw, .. } => Some(*raw),
47 Self::RuntimeBound { .. } => None,
48 }
49 }
50
51 fn device(&self) -> i32 {
52 match self {
53 #[cfg(test)]
54 Self::LocalOwned { device, .. } => *device,
55 Self::ExternalBorrowed { device, .. } | Self::ExternalOwned { device, .. } => *device,
56 Self::RuntimeBound { lease } => lease.device(),
57 }
58 }
59
60 fn runtime(&self) -> Option<&cudf::Runtime> {
61 match self {
62 Self::RuntimeBound { lease } => Some(lease.runtime()),
63 #[cfg(test)]
64 Self::LocalOwned { .. } => None,
65 Self::ExternalBorrowed { .. } | Self::ExternalOwned { .. } => None,
66 }
67 }
68
69 fn runtime_bound_lease(&self) -> Option<ResourceHandleLease> {
70 match self {
71 Self::RuntimeBound { lease } => Some(lease.clone()),
72 #[cfg(test)]
73 Self::LocalOwned { .. } => None,
74 Self::ExternalBorrowed { .. } | Self::ExternalOwned { .. } => None,
75 }
76 }
77
78 fn is_runtime_bound(&self) -> bool {
79 matches!(self, Self::RuntimeBound { .. })
80 }
81
82 fn uses_external_comm_metadata(&self) -> bool {
83 matches!(self, Self::ExternalBorrowed { .. } | Self::ExternalOwned { .. })
84 }
85}
86
87#[repr(C)]
88struct cugraph_rs_runtime_resource_handle_t {
89 _private: [u8; 0],
90}
91
92unsafe extern "C" {
93 fn cugraph_rs_runtime_resource_handle_create(
94 cuda_stream: *mut std::ffi::c_void,
95 device_memory_resource: *mut std::ffi::c_void,
96 memory_resource_owner: *mut std::ffi::c_void,
97 ) -> *mut cugraph_rs_runtime_resource_handle_t;
98 fn cugraph_rs_runtime_resource_handle_get(
99 handle: *mut cugraph_rs_runtime_resource_handle_t,
100 ) -> *mut cugraph_resource_handle_t;
101 #[cfg(test)]
102 fn cugraph_rs_runtime_resource_handle_effective_device_memory(
103 handle: *mut cugraph_rs_runtime_resource_handle_t,
104 out: *mut usize,
105 ) -> i32;
106}
107
108struct RuntimeResourceHandleOwner {
109 raw: NonNull<cugraph_rs_runtime_resource_handle_t>,
110 device: i32,
111 _current_resource_route: CurrentResourceRouteLease,
116}
117
118unsafe impl Send for RuntimeResourceHandleOwner {}
125unsafe impl Sync for RuntimeResourceHandleOwner {}
129
130impl RuntimeResourceHandleOwner {
131 fn create(runtime: &cudf::Runtime, operation: &'static str) -> Result<Self> {
132 let binding =
133 runtime.execution_binding().map_err(|error| map_runtime_execution_binding_error(error, operation))?;
134 let runtime_device = runtime.device().ordinal() as i32;
135 validate_runtime_binding_device(runtime_device, binding, operation)?;
136
137 let _guard = DeviceGuard::enter_owner_device(runtime_device, operation)?;
138 let manager = DeviceMemoryManager::for_device(runtime.device())
139 .map_err(|error| map_current_resource_route_error(error, operation))?;
140 let current_resource_route = manager
141 .route_current_resource(runtime.stream(), runtime.memory_resource())
142 .map_err(|error| map_current_resource_route_error(error, operation))?;
143 let shim = unsafe {
147 cugraph_rs_runtime_resource_handle_create(
148 binding.cuda_stream().as_raw(),
149 binding.memory_resource().as_raw(),
150 runtime.memory_resource().as_borrowed_sys().as_raw().cast(),
151 )
152 };
153 let raw = NonNull::new(shim).ok_or_else(|| {
154 Error::new(ErrorKind::Unexpected, operation)
155 .with_message("failed to create runtime-bound cuGraph resource handle shim")
156 .with_context("device", runtime_device.to_string())
157 })?;
158
159 let resource_handle = unsafe { cugraph_rs_runtime_resource_handle_get(raw.as_ptr()) };
161 if resource_handle.is_null() {
162 let mut owned = raw.as_ptr();
163 drop_device_resource(&mut owned, runtime_device, operation, DeviceResourceKind::RuntimeResourceHandle);
164 return Err(Error::new(ErrorKind::Unexpected, operation)
165 .with_message("runtime-bound handle shim returned a null cuGraph resource handle")
166 .with_context("device", runtime_device.to_string()));
167 }
168
169 Ok(Self { raw, device: runtime_device, _current_resource_route: current_resource_route })
170 }
171
172 fn resource_handle(&self) -> *mut cugraph_resource_handle_t {
173 unsafe { cugraph_rs_runtime_resource_handle_get(self.raw.as_ptr()) }
175 }
176
177 #[cfg(test)]
178 fn effective_device_memory(&self) -> Result<usize> {
179 let mut out = 0;
180 let status = unsafe { cugraph_rs_runtime_resource_handle_effective_device_memory(self.raw.as_ptr(), &mut out) };
183 if status == 0 {
184 Ok(out)
185 } else {
186 Err(Error::new(ErrorKind::InvalidState, "RuntimeResourceHandleOwner::effective_device_memory")
187 .with_message("cuGraph effective device memory provider rejected the runtime-bound handle")
188 .with_context("provider_status", status.to_string()))
189 }
190 }
191}
192
193impl Drop for RuntimeResourceHandleOwner {
194 fn drop(&mut self) {
195 let mut raw = self.raw.as_ptr();
196 drop_device_resource(
197 &mut raw,
198 self.device,
199 "RuntimeResourceHandleOwner::drop",
200 DeviceResourceKind::RuntimeResourceHandle,
201 );
202 }
203}
204
205#[derive(Clone)]
213pub(crate) struct ResourceHandleLease {
214 owner: BindingExtensionLease<RuntimeResourceHandleOwner>,
215 runtime: Arc<cudf::Runtime>,
216}
217
218impl ResourceHandleLease {
219 fn new(owner: BindingExtensionLease<RuntimeResourceHandleOwner>, runtime: Arc<cudf::Runtime>) -> Self {
220 Self { owner, runtime }
221 }
222
223 fn enter(&self, operation: &'static str) -> Result<BindingExtensionGuard<'_, RuntimeResourceHandleOwner>> {
224 self.owner.enter(operation).map_err(|error| map_binding_extension_lifecycle_error(error, operation))
225 }
226
227 pub(crate) fn with_teardown_gate<R>(&self, drop_native: impl FnOnce() -> R) -> R {
228 self.owner.with_teardown_gate(drop_native)
229 }
230
231 fn device(&self) -> i32 {
232 self.runtime.device().ordinal() as i32
233 }
234
235 pub(crate) fn runtime(&self) -> &cudf::Runtime {
236 self.runtime.as_ref()
237 }
238
239 pub(crate) fn same_binding(&self, other: &Self) -> bool {
240 self.runtime.binding_id() == other.runtime.binding_id()
241 }
242
243 pub(crate) fn same_native_owner(&self, other: &Self) -> bool {
244 self.owner.same_generation(&other.owner)
245 }
246
247 #[cfg(test)]
248 fn generation_matches(&self, other: &Self) -> bool {
249 self.owner.same_generation(&other.owner)
250 }
251}
252
253fn map_binding_extension_lifecycle_error(error: cudf::Error, operation: &'static str) -> Error {
254 Error::new(ErrorKind::InvalidState, operation)
255 .with_message("cudf runtime binding-extension lifecycle rejected cuGraph handle access")
256 .with_context("cudf_code", error.code().to_string())
257 .with_context("cudf_operation", error.operation().to_string())
258 .with_context("cudf_message", error.message().to_string())
259 .set_source(error)
260}
261
262fn map_current_resource_route_error(error: cudf::Error, operation: &'static str) -> Error {
263 Error::new(ErrorKind::InvalidState, operation)
264 .with_message("failed to register the runtime memory resource for legacy current-RMM allocations")
265 .with_context("cudf_code", error.code().to_string())
266 .with_context("cudf_operation", error.operation().to_string())
267 .with_context("cudf_message", error.message().to_string())
268 .set_source(error)
269}
270
271fn map_binding_extension_init_error(error: BindingExtensionInitError<Error>, operation: &'static str) -> Error {
272 match error {
273 BindingExtensionInitError::Lifecycle(error) => map_binding_extension_lifecycle_error(error, operation),
274 BindingExtensionInitError::Factory(error) => error,
275 }
276}
277
278pub(crate) fn map_runtime_execution_binding_error(error: InteropError, operation: &'static str) -> Error {
279 let mut mapped = match error.kind() {
280 InteropErrorKind::ExecutionBindingUnavailable => Error::new(ErrorKind::Unexpected, operation),
281 _ => Error::new(ErrorKind::InvalidState, operation),
282 }
283 .with_message("failed to extract cudf runtime execution binding")
284 .with_context("interop_kind", format!("{:?}", error.kind()))
285 .set_source(error.clone());
286
287 for (key, value) in error.context() {
288 mapped = mapped.with_context(key, value.clone());
289 }
290
291 mapped
292}
293
294fn validate_runtime_binding_device(
295 runtime_device: i32,
296 binding: ExecutionBinding<'_>,
297 operation: &'static str,
298) -> Result<()> {
299 if binding.device() == runtime_device {
300 return Ok(());
301 }
302
303 Err(Error::new(ErrorKind::InvalidState, operation)
304 .with_message("runtime execution binding device did not match the runtime owner device")
305 .with_context("runtime_device", runtime_device.to_string())
306 .with_context("binding_device", binding.device().to_string()))
307}
308
309pub struct ResourceHandle {
317 state: ResourceHandleState,
318 _nosend_nosync: PhantomData<*mut ()>,
320}
321
322pub struct EnteredContext<'a> {
328 owner: &'a ResourceHandle,
329 _guard: DeviceGuard,
330 runtime_extension: Option<BindingExtensionGuard<'a, RuntimeResourceHandleOwner>>,
333 _nosend_nosync: PhantomData<*mut ()>,
334}
335
336impl fmt::Debug for EnteredContext<'_> {
337 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
338 f.debug_struct("EnteredContext").field("device", &self.device()).finish_non_exhaustive()
339 }
340}
341
342impl EnteredContext<'_> {
343 pub fn handle_ptr(&self) -> *mut cugraph_resource_handle_t {
345 match &self.runtime_extension {
346 Some(extension) => extension.value().resource_handle(),
347 None => self
348 .owner
349 .state
350 .standalone_raw()
351 .expect("an entered non-runtime handle always has a standalone pointer"),
352 }
353 }
354
355 pub fn device(&self) -> i32 {
357 self.owner.device()
358 }
359
360 pub fn runtime(&self) -> Option<&cudf::Runtime> {
362 self.owner.runtime()
363 }
364
365 pub(crate) fn runtime_bound_lease(&self) -> Option<ResourceHandleLease> {
366 self.owner.runtime_bound_lease()
367 }
368}
369
370impl fmt::Debug for ResourceHandle {
371 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
372 f.debug_struct("ResourceHandle").finish_non_exhaustive()
373 }
374}
375
376impl ResourceHandle {
377 fn current_device_ordinal(operation: &'static str) -> Result<i32> {
378 let device = DeviceId::current().map_err(|error| {
379 Error::new(ErrorKind::Unexpected, operation)
380 .with_message(format!("failed to query current CUDA device: {}", error))
381 .set_source(error)
382 })?;
383 Ok(device.ordinal() as i32)
384 }
385
386 #[cfg(test)]
390 pub(crate) fn new() -> Result<Self> {
391 Self::for_current_device()
392 }
393
394 #[cfg(test)]
396 pub(crate) fn for_current_device() -> Result<Self> {
397 Self::new_on_device(Self::current_device_ordinal("ResourceHandle::for_current_device")?)
398 }
399
400 pub fn for_runtime(runtime: Arc<cudf::Runtime>) -> Result<Self> {
407 let operation = "ResourceHandle::for_runtime";
408 let owner = runtime
409 .get_or_try_init_binding_extension(operation, || {
410 RuntimeResourceHandleOwner::create(runtime.as_ref(), operation)
411 })
412 .map_err(|error| map_binding_extension_init_error(error, operation))?;
413 let lease = ResourceHandleLease::new(owner, runtime);
414 Ok(Self { state: ResourceHandleState::RuntimeBound { lease }, _nosend_nosync: PhantomData })
415 }
416
417 #[cfg(test)]
418 fn new_on_device(device: i32) -> Result<Self> {
419 if device < 0 {
420 return Err(Error::new(ErrorKind::InvalidInput, "ResourceHandle::new_on_device")
421 .with_message(format!("device ordinal must be >= 0, got {}", device)));
422 }
423
424 DeviceId::new(device as usize).map_err(|error| {
425 Error::new(ErrorKind::InvalidInput, "ResourceHandle::new_on_device")
426 .with_message(format!("invalid CUDA device ordinal {}: {}", device, error))
427 .set_source(error)
428 })?;
429
430 let _guard = DeviceGuard::enter_owner_device(device, "ResourceHandle::new_on_device")?;
431 let raw = unsafe { cugraph_sys::cugraph_create_resource_handle(std::ptr::null_mut()) };
434 if raw.is_null() {
435 return Err(Error::new(ErrorKind::Unexpected, "ResourceHandle::new_on_device")
436 .with_message("cugraph_create_resource_handle returned NULL"));
437 }
438 Ok(Self { state: ResourceHandleState::LocalOwned { raw, device }, _nosend_nosync: PhantomData })
439 }
440
441 pub unsafe fn from_raw_resource_handle(raw: *mut cugraph_resource_handle_t) -> Result<Self> {
463 if raw.is_null() {
464 return Err(Error::new(ErrorKind::InvalidInput, "ResourceHandle::from_raw_resource_handle")
465 .with_message("resource handle pointer must not be null"));
466 }
467 let device = Self::current_device_ordinal("ResourceHandle::from_raw_resource_handle")?;
468 let _guard = DeviceGuard::enter_owner_device(device, "ResourceHandle::from_raw_resource_handle")?;
469 Ok(Self { state: ResourceHandleState::ExternalBorrowed { raw, device }, _nosend_nosync: PhantomData })
470 }
471
472 pub unsafe fn adopt_raw_resource_handle(raw: *mut cugraph_resource_handle_t) -> Result<Self> {
496 if raw.is_null() {
497 return Err(Error::new(ErrorKind::InvalidInput, "ResourceHandle::adopt_raw_resource_handle")
498 .with_message("resource handle pointer must not be null"));
499 }
500 let device = Self::current_device_ordinal("ResourceHandle::adopt_raw_resource_handle")?;
501 let _guard = DeviceGuard::enter_owner_device(device, "ResourceHandle::adopt_raw_resource_handle")?;
502 Ok(Self { state: ResourceHandleState::ExternalOwned { raw, device }, _nosend_nosync: PhantomData })
503 }
504
505 #[cfg(test)]
507 pub(crate) fn as_ptr(&self) -> *mut cugraph_resource_handle_t {
508 self.state.standalone_raw().expect("runtime-bound handle pointers require an EnteredContext")
509 }
510
511 pub fn into_raw_resource_handle(self) -> Result<*mut cugraph_resource_handle_t> {
526 if self.state.is_runtime_bound() {
527 return Err(Error::new(ErrorKind::InvalidState, "ResourceHandle::into_raw_resource_handle")
528 .with_message("runtime-bound handles cannot be converted into standalone raw resource handles"));
529 }
530
531 let this = ManuallyDrop::new(self);
532 Ok(this.state.standalone_raw().expect("runtime-bound handles were rejected before raw extraction"))
533 }
534
535 pub fn device(&self) -> i32 {
537 self.state.device()
538 }
539
540 pub fn runtime(&self) -> Option<&cudf::Runtime> {
541 self.state.runtime()
542 }
543
544 pub(crate) fn runtime_bound_lease(&self) -> Option<ResourceHandleLease> {
545 self.state.runtime_bound_lease()
546 }
547
548 pub fn rank(&self) -> i32 {
557 if self.state.uses_external_comm_metadata() {
558 let Ok(_guard) = DeviceGuard::enter_owner_device(self.device(), "ResourceHandle::rank") else {
559 return 0;
560 };
561 unsafe {
564 cugraph_resource_handle_get_rank(
565 self.state.standalone_raw().expect("external handles always have a standalone pointer"),
566 )
567 }
568 } else {
569 0
570 }
571 }
572
573 pub fn comm_size(&self) -> i32 {
582 if self.state.uses_external_comm_metadata() {
583 let Ok(_guard) = DeviceGuard::enter_owner_device(self.device(), "ResourceHandle::comm_size") else {
584 return 1;
585 };
586 unsafe {
589 cugraph_resource_handle_get_comm_size(
590 self.state.standalone_raw().expect("external handles always have a standalone pointer"),
591 )
592 }
593 } else {
594 1
595 }
596 }
597
598 pub fn enter(&self, operation: &'static str) -> Result<EnteredContext<'_>> {
603 let runtime_extension = match &self.state {
604 ResourceHandleState::RuntimeBound { lease } => Some(lease.enter(operation)?),
605 #[cfg(test)]
606 ResourceHandleState::LocalOwned { .. } => None,
607 ResourceHandleState::ExternalBorrowed { .. } | ResourceHandleState::ExternalOwned { .. } => None,
608 };
609 let guard = DeviceGuard::enter_owner_device(self.device(), operation)?;
610 Ok(EnteredContext { owner: self, _guard: guard, runtime_extension, _nosend_nosync: PhantomData })
611 }
612}
613
614impl Drop for ResourceHandle {
615 fn drop(&mut self) {
616 match &mut self.state {
617 #[cfg(test)]
618 ResourceHandleState::LocalOwned { raw, device } => {
619 drop_device_resource(raw, *device, "ResourceHandle::drop", DeviceResourceKind::ResourceHandle);
620 },
621 ResourceHandleState::ExternalOwned { raw, device } => {
622 drop_device_resource(raw, *device, "ResourceHandle::drop", DeviceResourceKind::ResourceHandle);
623 },
624 ResourceHandleState::RuntimeBound { .. } | ResourceHandleState::ExternalBorrowed { .. } => {},
625 }
626 }
627}
628
629#[cfg(test)]
632mod tests {
633 use std::{
634 ptr,
635 sync::atomic::{AtomicBool, Ordering},
636 };
637
638 use cudf::{
639 DeviceMemoryResource, QueryMemoryDomainOptions,
640 testing::{
641 device_safety::{DeviceEntryTestCall as TestCudaCall, DeviceSafetyTestScope as TestCudaHooks},
642 run_named_cases,
643 },
644 };
645 use cugraph_sys::cugraph_create_resource_handle;
646
647 use super::*;
648 use crate::DeviceArray;
649
650 struct ReleaseProbe(Arc<AtomicBool>);
651
652 impl Drop for ReleaseProbe {
653 fn drop(&mut self) {
654 self.0.store(true, Ordering::SeqCst);
655 }
656 }
657
658 fn probed_runtime() -> (Arc<cudf::Runtime>, Arc<AtomicBool>) {
661 let released = Arc::new(AtomicBool::new(false));
662 let runtime = cudf::Runtime::builder()
663 .build()
664 .expect("runtime")
665 .with_retained_lifetime(Arc::new(ReleaseProbe(Arc::clone(&released))))
666 .expect("probed runtime");
667 (Arc::new(runtime), released)
668 }
669
670 fn budgeted_runtime(capacity_bytes: usize) -> (Arc<cudf::Runtime>, DeviceMemoryResource) {
671 let device = DeviceId::current().expect("current device");
672 let upstream = DeviceMemoryResource::cuda_async_on(device, None, Some(0)).expect("CUDA async resource");
673 let budget = upstream.budget_tracking_limit(capacity_bytes).expect("budget tracking resource");
674 let runtime =
675 Arc::new(cudf::Runtime::builder().with_memory_resource(budget.clone()).build().expect("budgeted runtime"));
676 (runtime, budget)
677 }
678
679 fn effective_device_memory(entered: &EnteredContext<'_>) -> usize {
680 entered
681 .runtime_extension
682 .as_ref()
683 .expect("runtime-bound context")
684 .value()
685 .effective_device_memory()
686 .expect("effective device memory")
687 }
688
689 fn state_variant_name(state: &ResourceHandleState) -> &'static str {
690 match state {
691 ResourceHandleState::LocalOwned { .. } => "LocalOwned",
692 ResourceHandleState::RuntimeBound { .. } => "RuntimeBound",
693 ResourceHandleState::ExternalBorrowed { .. } => "ExternalBorrowed",
694 ResourceHandleState::ExternalOwned { .. } => "ExternalOwned",
695 }
696 }
697
698 #[test]
699 fn native_owner_cases_1() {
700 run_named_cases(&[
701 ("handle::tests::test_for_runtime_retains_execution_owner", test_for_runtime_retains_execution_owner),
702 (
703 "handle::tests::runtime_bound_handles_reuse_one_native_owner_for_exact_binding_hot_hits",
704 runtime_bound_handles_reuse_one_native_owner_for_exact_binding_hot_hits,
705 ),
706 (
707 "handle::tests::runtime_bound_handles_do_not_share_native_owner_across_bindings",
708 runtime_bound_handles_do_not_share_native_owner_across_bindings,
709 ),
710 (
711 "handle::tests::runtime_binding_store_and_native_owner_do_not_form_a_cycle",
712 runtime_binding_store_and_native_owner_do_not_form_a_cycle,
713 ),
714 (
715 "handle::tests::test_for_runtime_uses_runtime_shim_instead_of_direct_cugraph_ownership",
716 test_for_runtime_uses_runtime_shim_instead_of_direct_cugraph_ownership,
717 ),
718 (
719 "handle::tests::test_from_raw_resource_handle_uses_current_device",
720 test_from_raw_resource_handle_uses_current_device,
721 ),
722 (
723 "handle::tests::test_adopt_raw_resource_handle_uses_current_device",
724 test_adopt_raw_resource_handle_uses_current_device,
725 ),
726 (
727 "handle::tests::test_from_raw_resource_handle_does_not_free_on_drop",
728 test_from_raw_resource_handle_does_not_free_on_drop,
729 ),
730 ("handle::tests::test_adopt_raw_resource_handle_marks_owned", test_adopt_raw_resource_handle_marks_owned),
731 (
732 "handle::tests::test_entered_context_exposes_runtime_provenance",
733 test_entered_context_exposes_runtime_provenance,
734 ),
735 (
736 "handle::tests::runtime_handle_effective_memory_tracks_aligned_domain_charge",
737 runtime_handle_effective_memory_tracks_aligned_domain_charge,
738 ),
739 (
740 "handle::tests::runtime_handles_keep_distinct_effective_memory_providers",
741 runtime_handles_keep_distinct_effective_memory_providers,
742 ),
743 (
744 "handle::tests::nested_entry_for_same_native_owner_is_rejected_without_deadlock",
745 nested_entry_for_same_native_owner_is_rejected_without_deadlock,
746 ),
747 (
748 "handle::tests::test_into_raw_resource_handle_releases_rust_ownership",
749 test_into_raw_resource_handle_releases_rust_ownership,
750 ),
751 (
752 "handle::tests::test_runtime_bound_into_raw_resource_handle_is_rejected",
753 test_runtime_bound_into_raw_resource_handle_is_rejected,
754 ),
755 ]);
756 }
757
758 fn test_for_runtime_retains_execution_owner() {
759 let runtime = Arc::new(cudf::Runtime::builder().build().expect("runtime"));
760 assert_eq!(Arc::strong_count(&runtime), 1, "test setup should hold one runtime owner");
761
762 let handle = ResourceHandle::for_runtime(Arc::clone(&runtime)).expect("resource handle");
763 assert_eq!(handle.device(), runtime.device().ordinal() as i32);
764 assert_eq!(Arc::strong_count(&runtime), 2, "runtime-bound handle should retain the runtime owner");
765
766 drop(handle);
767 assert_eq!(Arc::strong_count(&runtime), 1, "dropping the handle should release the retained runtime owner");
768 }
769
770 fn runtime_bound_handles_reuse_one_native_owner_for_exact_binding_hot_hits() {
771 let runtime = Arc::new(cudf::Runtime::builder().build().expect("runtime"));
772 let derived = Arc::new(runtime.with_retained_lifetime(Arc::new(())).expect("derived runtime"));
773 assert!(!Arc::ptr_eq(&runtime, &derived));
774
775 let first = ResourceHandle::for_runtime(Arc::clone(&runtime)).expect("first resource handle");
776 let first_lease = first.runtime_bound_lease().expect("first runtime lease");
777 let first_raw = first.enter("hot_hit_first").expect("enter first").handle_ptr();
778
779 for index in 0..100 {
780 let owner = if index % 2 == 0 { Arc::clone(&runtime) } else { Arc::clone(&derived) };
781 let handle = ResourceHandle::for_runtime(owner).expect("hot-hit resource handle");
782 let lease = handle.runtime_bound_lease().expect("hot-hit runtime lease");
783 assert!(first_lease.same_binding(&lease));
784 assert!(first_lease.same_native_owner(&lease));
785 let entered = handle.enter("hot_hit").expect("enter hot hit");
786 assert_eq!(entered.handle_ptr(), first_raw);
787 }
788 }
789
790 fn runtime_bound_handles_do_not_share_native_owner_across_bindings() {
791 let first_runtime = Arc::new(cudf::Runtime::builder().build().expect("first runtime"));
792 let second_runtime = Arc::new(cudf::Runtime::builder().build().expect("independent runtime"));
793 assert_ne!(first_runtime.binding_id(), second_runtime.binding_id());
794
795 let first = ResourceHandle::for_runtime(first_runtime).expect("first resource handle");
796 let second = ResourceHandle::for_runtime(second_runtime).expect("second resource handle");
797 let first_lease = first.runtime_bound_lease().expect("first runtime lease");
798 let second_lease = second.runtime_bound_lease().expect("second runtime lease");
799
800 assert!(!first_lease.same_binding(&second_lease));
801 assert!(!first_lease.same_native_owner(&second_lease));
802 let first_raw = first.enter("different_binding_first").expect("enter first").handle_ptr();
803 let second_raw = second.enter("different_binding_second").expect("enter second").handle_ptr();
804 assert_ne!(first_raw, second_raw);
805 }
806
807 fn runtime_binding_store_and_native_owner_do_not_form_a_cycle() {
808 let (runtime, released) = probed_runtime();
809 let handle = ResourceHandle::for_runtime(Arc::clone(&runtime)).expect("resource handle");
810 drop(runtime);
811 assert!(!released.load(Ordering::SeqCst), "the resource-handle lease retains its execution domain");
812 drop(handle);
813 assert!(
814 released.load(Ordering::SeqCst),
815 "dropping the last native-owner lease must release the execution domain"
816 );
817 }
818
819 #[test]
820 fn runtime_bound_owner_quarantines_before_releasing_runtime_when_device_entry_fails() {
821 let (runtime, released) = probed_runtime();
822 let handle = ResourceHandle::for_runtime(Arc::clone(&runtime)).expect("resource handle");
823 drop(runtime);
824 assert!(!released.load(Ordering::SeqCst), "the runtime-bound lease retains the execution domain");
825
826 let reset = TestCudaHooks::install(vec![TestCudaCall::GetDevice {
827 status: cudf_sys::CUDF_ERROR_FATAL_CUDA,
828 device: 0,
829 }]);
830 let before = reset.health_snapshot(DeviceId::DEFAULT).quarantined_teardown_count();
831
832 drop(handle);
833
834 assert!(reset.health_snapshot(DeviceId::DEFAULT).quarantined_teardown_count() > before);
835 assert!(
836 released.load(Ordering::SeqCst),
837 "quarantining the native owner must not retain or cycle the cudf runtime"
838 );
839 }
840
841 fn test_for_runtime_uses_runtime_shim_instead_of_direct_cugraph_ownership() {
842 let runtime = Arc::new(cudf::Runtime::builder().build().expect("runtime"));
843
844 let handle = ResourceHandle::for_runtime(runtime).expect("resource handle");
845
846 match &handle.state {
847 ResourceHandleState::RuntimeBound { lease } => {
848 let entered = handle.enter("runtime_state_pointer").expect("enter runtime-bound state");
849 assert!(!entered.handle_ptr().is_null(), "entered runtime-bound state must expose a live pointer");
850 assert_eq!(lease.device(), handle.device(), "runtime-bound state must carry the handle device");
851 },
852 state => panic!("expected runtime-bound state variant, got {}", state_variant_name(state)),
853 }
854 }
855
856 #[test]
857 fn test_from_raw_resource_handle_rejects_null() {
858 let err = unsafe { ResourceHandle::from_raw_resource_handle(ptr::null_mut()) }.expect_err("should reject null");
860 assert_eq!(err.kind(), ErrorKind::InvalidInput);
861 assert_eq!(err.operation(), "ResourceHandle::from_raw_resource_handle");
862 }
863
864 fn test_from_raw_resource_handle_uses_current_device() {
865 let owned = ResourceHandle::new().expect("owned handle");
866 let raw = owned.as_ptr();
867 let current = DeviceId::current().expect("current device");
868
869 let wrapped = unsafe { ResourceHandle::from_raw_resource_handle(raw) }.expect("from_raw_resource_handle");
872 assert_eq!(wrapped.device(), current.ordinal() as i32);
873 match &wrapped.state {
874 ResourceHandleState::ExternalBorrowed { raw: state_raw, device, .. } => {
875 assert_eq!(*state_raw, raw);
876 assert_eq!(*device, current.ordinal() as i32);
877 },
878 state => panic!("expected external borrowed state variant, got {}", state_variant_name(state)),
879 }
880 }
881
882 #[test]
883 fn test_adopt_raw_resource_handle_rejects_null() {
884 let err =
886 unsafe { ResourceHandle::adopt_raw_resource_handle(ptr::null_mut()) }.expect_err("should reject null");
887 assert_eq!(err.kind(), ErrorKind::InvalidInput);
888 assert_eq!(err.operation(), "ResourceHandle::adopt_raw_resource_handle");
889 }
890
891 fn test_adopt_raw_resource_handle_uses_current_device() {
892 let raw = unsafe { cugraph_create_resource_handle(ptr::null_mut()) };
894 assert!(!raw.is_null(), "native handle allocation should succeed");
895 let current = DeviceId::current().expect("current device");
896
897 let adopted = unsafe { ResourceHandle::adopt_raw_resource_handle(raw) }.expect("adopt_raw_resource_handle");
899 assert_eq!(adopted.device(), current.ordinal() as i32);
900 }
901
902 fn test_from_raw_resource_handle_does_not_free_on_drop() {
903 let owned = ResourceHandle::new().expect("owned handle");
907 let raw = owned.as_ptr();
908
909 {
910 let borrowed = unsafe { ResourceHandle::from_raw_resource_handle(raw) }.expect("from_raw_resource_handle");
912 assert!(!borrowed.as_ptr().is_null());
913 }
915
916 assert!(!owned.as_ptr().is_null());
918 }
919
920 fn test_adopt_raw_resource_handle_marks_owned() {
921 let raw = unsafe { cugraph_create_resource_handle(ptr::null_mut()) };
923 assert!(!raw.is_null(), "native handle allocation should succeed");
924
925 let adopted = unsafe { ResourceHandle::adopt_raw_resource_handle(raw) }.expect("adopt_raw_resource_handle");
927 assert_eq!(adopted.as_ptr(), raw);
928 match &adopted.state {
929 ResourceHandleState::ExternalOwned { raw: state_raw, .. } => {
930 assert_eq!(*state_raw, raw, "external owned state must retain the adopted raw handle");
931 },
932 state => panic!("expected external owned state variant, got {}", state_variant_name(state)),
933 }
934 }
935
936 #[test]
937 fn test_enter_temporarily_enters_owner_device_and_restores_ambient_device() {
938 let raw = std::ptr::NonNull::<cugraph_resource_handle_t>::dangling().as_ptr();
939 let handle = ResourceHandle {
940 state: ResourceHandleState::ExternalBorrowed { raw, device: 1 },
941 _nosend_nosync: PhantomData,
942 };
943 let _reset = TestCudaHooks::install(vec![
944 TestCudaCall::GetDevice { status: cudf_sys::CUDF_SUCCESS, device: 0 },
945 TestCudaCall::SetDevice { expected_target: 1, status: cudf_sys::CUDF_SUCCESS },
946 TestCudaCall::SetDevice { expected_target: 0, status: cudf_sys::CUDF_SUCCESS },
947 ]);
948
949 {
950 let entered = handle.enter("ResourceHandle::enter").expect("entered context");
951 assert_eq!(entered.device(), 1);
952 assert_eq!(entered.handle_ptr(), raw);
953 assert!(entered.runtime().is_none());
954 }
955 }
956
957 #[test]
958 fn test_restore_failure_rejects_future_work_on_failed_ambient_device() {
959 let raw = std::ptr::NonNull::<cugraph_resource_handle_t>::dangling().as_ptr();
960 let handle = ResourceHandle {
961 state: ResourceHandleState::ExternalBorrowed { raw, device: 1 },
962 _nosend_nosync: PhantomData,
963 };
964 let _reset = TestCudaHooks::install(vec![
965 TestCudaCall::GetDevice { status: cudf_sys::CUDF_SUCCESS, device: 0 },
966 TestCudaCall::SetDevice { expected_target: 1, status: cudf_sys::CUDF_SUCCESS },
967 TestCudaCall::SetDevice { expected_target: 0, status: cudf_sys::CUDF_ERROR_FATAL_CUDA },
968 ]);
969
970 {
971 let _entered = handle.enter("first_enter").expect("entered context");
972 }
973
974 let failed_ambient_handle = ResourceHandle {
975 state: ResourceHandleState::ExternalBorrowed { raw, device: 0 },
976 _nosend_nosync: PhantomData,
977 };
978 let err = failed_ambient_handle
979 .enter("unhealthy_ambient_enter")
980 .expect_err("failed ambient device should reject later work");
981 assert_eq!(err.kind(), ErrorKind::InvalidState);
982 assert_eq!(err.operation(), "unhealthy_ambient_enter");
983 assert!(err.is_persistent());
984 }
985
986 fn test_entered_context_exposes_runtime_provenance() {
987 let runtime = Arc::new(cudf::Runtime::builder().build().expect("runtime"));
988 let handle = ResourceHandle::for_runtime(Arc::clone(&runtime)).expect("resource handle");
989
990 let entered = handle.enter("ResourceHandle::enter").expect("entered context");
991 let session_runtime = entered.runtime().expect("runtime-bound handle should expose runtime");
992 assert_eq!(session_runtime.device().ordinal(), runtime.device().ordinal());
993 }
994
995 fn runtime_handle_effective_memory_tracks_aligned_domain_charge() {
996 const CAPACITY_BYTES: usize = 1 << 20;
997 let (runtime, budget) = budgeted_runtime(CAPACITY_BYTES);
998 let handle = ResourceHandle::for_runtime(runtime).expect("resource handle");
999 let entered = handle.enter("effective_memory_charge_test").expect("entered context");
1000
1001 assert_eq!(effective_device_memory(&entered), CAPACITY_BYTES);
1002 let allocation = DeviceArray::from_slice(&entered, &[7_i32; 257]).expect("device allocation");
1003 let usage = budget.budget_tracking_usage().expect("budget usage");
1004 assert!(usage.current_aligned_bytes > 0);
1005 assert_eq!(
1006 effective_device_memory(&entered),
1007 CAPACITY_BYTES - usize::try_from(usage.current_aligned_bytes).expect("aligned charge fits usize")
1008 );
1009
1010 drop(allocation);
1011 assert_eq!(budget.budget_tracking_usage().expect("released usage").current_aligned_bytes, 0);
1012 assert_eq!(effective_device_memory(&entered), CAPACITY_BYTES);
1013 }
1014
1015 #[test]
1016 fn runtime_handle_effective_memory_accepts_query_domain_member_resource() {
1017 const CAPACITY_BYTES: usize = 1 << 20;
1018 let base = cudf::Runtime::builder().with_device_memory_manager().build().expect("base runtime");
1019 let manager = base.device_memory_manager().expect("device memory manager");
1020 let domain = manager
1021 .query_memory_domain(
1022 81_001,
1023 QueryMemoryDomainOptions::new(CAPACITY_BYTES).with_ceiling_bytes(CAPACITY_BYTES),
1024 )
1025 .expect("query memory domain");
1026 let lease = domain.reserve_stream(base.stream(), 81_101, CAPACITY_BYTES).expect("domain member");
1027 let runtime = Arc::new(
1028 cudf::Runtime::builder()
1029 .with_device(base.device())
1030 .with_stream(base.stream().clone())
1031 .with_memory_resource(lease.memory_resource().clone())
1032 .build()
1033 .expect("member runtime"),
1034 );
1035 let handle = ResourceHandle::for_runtime(runtime).expect("resource handle");
1036 let entered = handle.enter("query_domain_member_effective_memory_test").expect("entered context");
1037
1038 assert_eq!(effective_device_memory(&entered), CAPACITY_BYTES);
1039 let allocation = DeviceArray::from_slice(&entered, &[9_i32; 257]).expect("device allocation");
1040 let usage = domain.usage().expect("domain usage");
1041 assert!(usage.current_aligned_bytes > 0);
1042 assert_eq!(
1043 effective_device_memory(&entered),
1044 CAPACITY_BYTES - usize::try_from(usage.current_aligned_bytes).expect("aligned charge fits usize")
1045 );
1046
1047 drop(allocation);
1048 assert_eq!(domain.usage().expect("released usage").current_aligned_bytes, 0);
1049 assert_eq!(effective_device_memory(&entered), CAPACITY_BYTES);
1050 }
1051
1052 fn runtime_handles_keep_distinct_effective_memory_providers() {
1053 const SMALL_CAPACITY_BYTES: usize = 1 << 20;
1054 const LARGE_CAPACITY_BYTES: usize = 2 << 20;
1055 let (small_runtime, _small_budget) = budgeted_runtime(SMALL_CAPACITY_BYTES);
1056 let (large_runtime, _large_budget) = budgeted_runtime(LARGE_CAPACITY_BYTES);
1057 let small = ResourceHandle::for_runtime(small_runtime).expect("small resource handle");
1058 let large = ResourceHandle::for_runtime(large_runtime).expect("large resource handle");
1059
1060 let small_entered = small.enter("small_effective_memory_test").expect("small entered context");
1061 assert_eq!(effective_device_memory(&small_entered), SMALL_CAPACITY_BYTES);
1062 drop(small_entered);
1063 let large_entered = large.enter("large_effective_memory_test").expect("large entered context");
1064 assert_eq!(effective_device_memory(&large_entered), LARGE_CAPACITY_BYTES);
1065 }
1066
1067 fn nested_entry_for_same_native_owner_is_rejected_without_deadlock() {
1068 let runtime = Arc::new(cudf::Runtime::builder().build().expect("runtime"));
1069 let first = ResourceHandle::for_runtime(Arc::clone(&runtime)).expect("first resource handle");
1070 let second = ResourceHandle::for_runtime(runtime).expect("second resource handle");
1071 let first_lease = first.runtime_bound_lease().expect("first lease");
1072 let second_lease = second.runtime_bound_lease().expect("second lease");
1073 assert!(first_lease.generation_matches(&second_lease));
1074
1075 let first_entered = first.enter("nested_same_owner_first").expect("first entered context");
1076 let error = second.enter("nested_same_owner_second").expect_err("nested same-owner entry must be rejected");
1077 assert_eq!(error.kind(), ErrorKind::InvalidState);
1078 assert_eq!(error.operation(), "nested_same_owner_second");
1079 assert!(error.to_string().contains("nested_binding_extension_execution"));
1080
1081 drop(first_entered);
1082 second.enter("same_owner_after_release").expect("entry succeeds after prior context drops");
1083 }
1084
1085 #[test]
1086 fn concurrent_entry_for_same_native_owner_is_serialized() {
1087 use std::{sync::mpsc, time::Duration};
1088
1089 let runtime = Arc::new(cudf::Runtime::builder().build().expect("runtime"));
1090 let first = ResourceHandle::for_runtime(Arc::clone(&runtime)).expect("first resource handle");
1091 let first_entered = first.enter("concurrent_same_owner_first").expect("first entered context");
1092 let first_raw = first_entered.handle_ptr() as usize;
1093 let (ready_tx, ready_rx) = mpsc::channel();
1094 let (acquired_tx, acquired_rx) = mpsc::channel();
1095
1096 let worker = std::thread::spawn(move || {
1097 let second = ResourceHandle::for_runtime(runtime).expect("second resource handle");
1098 ready_tx.send(()).expect("send ready");
1099 let _entered = second.enter("concurrent_same_owner_second").expect("serialized second entry");
1100 acquired_tx.send(_entered.handle_ptr() as usize).expect("send acquired");
1101 });
1102
1103 ready_rx.recv_timeout(Duration::from_secs(5)).expect("worker ready");
1104 assert!(
1105 acquired_rx.recv_timeout(Duration::from_millis(100)).is_err(),
1106 "second entry must remain blocked while the first context owns the native handle"
1107 );
1108 drop(first_entered);
1109 assert_eq!(acquired_rx.recv_timeout(Duration::from_secs(5)).expect("worker acquired after release"), first_raw);
1110 worker.join().expect("concurrent entry worker");
1111 }
1112
1113 fn test_into_raw_resource_handle_releases_rust_ownership() {
1114 let handle = ResourceHandle::new().expect("resource handle");
1115 let device = handle.device();
1116 let mut raw = handle.into_raw_resource_handle().expect("into_raw_resource_handle");
1117 assert!(!raw.is_null(), "into_raw_resource_handle should return the underlying pointer");
1118
1119 drop_device_resource(
1120 &mut raw,
1121 device,
1122 "test_into_raw_resource_handle_releases_rust_ownership",
1123 DeviceResourceKind::ResourceHandle,
1124 );
1125 }
1126
1127 fn test_runtime_bound_into_raw_resource_handle_is_rejected() {
1128 let runtime = Arc::new(cudf::Runtime::builder().build().expect("runtime"));
1129 let handle = ResourceHandle::for_runtime(runtime).expect("resource handle");
1130
1131 let err = handle.into_raw_resource_handle().expect_err("runtime-bound handle should reject raw extraction");
1132 assert_eq!(err.kind(), ErrorKind::InvalidState);
1133 assert_eq!(err.operation(), "ResourceHandle::into_raw_resource_handle");
1134 }
1135}