Allow Stream::add_callback to run on context error again

crates/cust/CHANGELOG.md

@@ -6,6 +6,7 @@ Notable changes to this project will be documented in this file.
 
  - Add `memory::memcpy_dtoh` to allow copying from device to host.
  - Add support in `memory` for pitched malloc and 2D memcpy between device and host.
+ - `Stream::add_callback` now internally uses `cuStreamAddCallback` again, since there are no current plans to remove it (https://stackoverflow.com/a/58173486). As a result, the function again takes a device status as a parameter and *does* execute on context error.
 
 ## 0.3.2 - 2/16/22
 

crates/cust/src/stream.rs

@@ -13,7 +13,7 @@
 use crate::error::{CudaResult, DropResult, ToResult};
 use crate::event::Event;
 use crate::function::{BlockSize, Function, GridSize};
-use crate::sys::{self as cuda, CUstream};
+use crate::sys::{self as cuda, cudaError_enum, CUstream};
 use std::ffi::c_void;
 use std::mem;
 use std::panic;
@@ -151,6 +151,9 @@ impl Stream {
     ///
     /// Callbacks must not make any CUDA API calls.
     ///
+    /// The callback will be passed a `CudaResult<()>` indicating the
+    /// current state of the device with `Ok(())` denoting normal operation.
+    ///
     /// # Examples
     ///
     /// ```
@@ -164,22 +167,23 @@ impl Stream {
     ///
     /// // ... queue up some work on the stream
     ///
-    /// stream.add_callback(Box::new(|| {
-    ///     println!("Work is done!");
+    /// stream.add_callback(Box::new(|status| {
+    ///     println!("Device status is {:?}", status);
     /// }));
     ///
     /// // ... queue up some more work on the stream
     /// # Ok(())
     /// # }
     pub fn add_callback<T>(&self, callback: Box<T>) -> CudaResult<()>
     where
-        T: FnOnce() + Send,
+        T: FnOnce(CudaResult<()>) + Send,
     {
         unsafe {
-            cuda::cuLaunchHostFunc(
+            cuda::cuStreamAddCallback(
                 self.inner,
                 Some(callback_wrapper::<T>),
                 Box::into_raw(callback) as *mut c_void,
+                0,
             )
             .to_result()
         }
@@ -339,13 +343,16 @@ impl Drop for Stream {
         }
     }
 }
-unsafe extern "C" fn callback_wrapper<T>(callback: *mut c_void)
-where
-    T: FnOnce() + Send,
+unsafe extern "C" fn callback_wrapper<T>(
+    _stream: CUstream,
+    status: cudaError_enum,
+    callback: *mut c_void,
+) where
+    T: FnOnce(CudaResult<()>) + Send,
 {
     // Stop panics from unwinding across the FFI
     let _ = panic::catch_unwind(|| {
         let callback: Box<T> = Box::from_raw(callback as *mut T);
-        callback();
+        callback(status.to_result());
     });
 }