This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 23b2ec1b [FIX][Rust] Restore the previous stream when with_stream's
closure fails (#818)
23b2ec1b is described below
commit 23b2ec1b33523b556eef615ab2d0d557213bb372
Author: Gabriel Wu <[email protected]>
AuthorDate: Sat Sep 26 23:57:28 2026 +0800
[FIX][Rust] Restore the previous stream when with_stream's closure fails
(#818)
## Why
`with_stream` sets the device stream, runs the closure, and restores the
previous stream:
```rust
let result = f()?;
// restore prev_stream
```
When the closure returns an error, `?` returns before the restore, so
the device keeps the temporary stream and later work on that thread runs
on it. A panic in the closure skips the restore the same way.
## What
- The previous stream is restored however the closure exits.
- On return (value or error) it is restored explicitly, and a failed
restore is still reported. The closure's own error takes precedence.
- While unwinding, a drop guard restores it.
- The doc comment states the guarantee.
---
rust/tvm-ffi/src/device.rs | 53 ++++++++++++++++++++++++++++++++-------
rust/tvm-ffi/tests/test_device.rs | 30 ++++++++++++++++++++++
2 files changed, 74 insertions(+), 9 deletions(-)
diff --git a/rust/tvm-ffi/src/device.rs b/rust/tvm-ffi/src/device.rs
index b99bdb70..c69129ee 100644
--- a/rust/tvm-ffi/src/device.rs
+++ b/rust/tvm-ffi/src/device.rs
@@ -28,6 +28,9 @@ pub fn current_stream(device: &DLDevice) ->
TVMFFIStreamHandle {
}
/// Call `f` with the device stream temporarily set to `stream`.
///
+/// The previous stream is restored however `f` exits: when it returns a
+/// value, when it returns an error, and when it panics.
+///
/// # Safety
///
/// `stream` must be a valid stream handle for the given device, or null.
@@ -45,16 +48,48 @@ pub unsafe fn with_stream<T>(
&mut prev_stream as *mut TVMFFIStreamHandle
))?;
}
- let result = f()?;
- unsafe {
- crate::check_safe_call!(TVMFFIEnvSetStream(
- device.device_type as i32,
- device.device_id,
- prev_stream,
- std::ptr::null_mut()
- ))?;
+ let restore = RestoreStream {
+ device: *device,
+ stream: prev_stream,
+ };
+ let result = f();
+ let restored = restore.restore();
+ let value = result?;
+ restored?;
+ Ok(value)
+}
+
+/// Restores a device's previous stream when dropped, so that an unwinding
+/// panic in `with_stream`'s closure still restores it.
+struct RestoreStream {
+ device: DLDevice,
+ stream: TVMFFIStreamHandle,
+}
+
+impl RestoreStream {
+ fn set(&self) -> i32 {
+ unsafe {
+ TVMFFIEnvSetStream(
+ self.device.device_type as i32,
+ self.device.device_id,
+ self.stream,
+ std::ptr::null_mut(),
+ )
+ }
+ }
+
+ /// Restores the stream now, reporting a failure.
+ fn restore(self) -> Result<()> {
+ let restore = std::mem::ManuallyDrop::new(self);
+ crate::check_safe_call!(restore.set())
+ }
+}
+
+impl Drop for RestoreStream {
+ fn drop(&mut self) {
+ // Only reached while unwinding; a failure cannot be reported there.
+ self.set();
}
- Ok(result)
}
/// AnyCompatible for DLDevice
diff --git a/rust/tvm-ffi/tests/test_device.rs
b/rust/tvm-ffi/tests/test_device.rs
index 6516460b..31b8c42b 100644
--- a/rust/tvm-ffi/tests/test_device.rs
+++ b/rust/tvm-ffi/tests/test_device.rs
@@ -33,6 +33,36 @@ fn test_device_stream() {
assert_eq!(current_stream(&device), std::ptr::null_mut());
}
+#[test]
+fn test_device_stream_restored_after_error() {
+ let device = DLDevice::new(DLDeviceType::kDLCPU, 0);
+ let dummy_stream = 3 as TVMFFIStreamHandle;
+ // SAFETY: dummy_stream is a test value; CPU device accepts any stream
handle.
+ let result: Result<()> = unsafe {
+ with_stream(&device, dummy_stream, || {
+ Err(Error::new(VALUE_ERROR, "closure failed", ""))
+ })
+ };
+ assert_eq!(result.unwrap_err().message(), "closure failed");
+ assert_eq!(current_stream(&device), std::ptr::null_mut());
+}
+
+#[test]
+fn test_device_stream_restored_after_panic() {
+ let device = DLDevice::new(DLDeviceType::kDLCPU, 0);
+ let dummy_stream = 4 as TVMFFIStreamHandle;
+ let unwound = std::panic::catch_unwind(|| {
+ // SAFETY: dummy_stream is a test value; CPU device accepts any stream
handle.
+ unsafe {
+ with_stream(&device, dummy_stream, || -> Result<()> {
+ panic!("closure panicked")
+ })
+ }
+ });
+ assert!(unwound.is_err());
+ assert_eq!(current_stream(&device), std::ptr::null_mut());
+}
+
#[test]
fn test_device_any_conversion() {
let device = DLDevice::new(DLDeviceType::kDLCPU, 0);
