feat: Fiber interruptibility support

Author: 0x676e67

Cargo.toml

@@ -17,7 +17,8 @@ name = "wreq_ruby"
 magnus = { version = "0.8", features = ["bytes"] }
 rb-sys = { version = "0.9.110", default-features = false }
 tokio = { version = "1", features = ["full"] }
-wreq = { version = "6.0.0-rc.23", features = [
+tokio-util = { version = "0.7.17", default-features = false }
+wreq = { version = "6.0.0-rc.24", features = [
     "json",
     "socks",
     "stream",
@@ -50,6 +51,3 @@ incremental = false
 lto = "fat"
 opt-level = 3
 strip = true
-
-[patch.crates-io]
-wreq = { git = "https://github.com/0x676e67/wreq" }

lib/wreq.rb

@@ -18,6 +18,33 @@
 
 unless defined?(Wreq)
   module Wreq
+    # CancellationToken allows Ruby code to cooperatively cancel long-running Rust async tasks.
+    #
+    # This class is a binding to a Rust-side cancellation token, enabling Ruby to signal cancellation
+    # to HTTP requests, streaming, or other operations running in Rust. When `cancel` is called, all
+    # Rust tasks observing this token will be notified and can abort promptly.
+    #
+    # @example
+    #   token = Wreq::CancellationToken.new
+    #   # Pass token to a Wreq request or stream
+    #   token.cancel # Signal cancellation from Ruby
+    #
+    # The actual implementation and state are managed by the Rust extension.
+    class CancellationToken
+      # Create a new cancellation token.
+      #
+      # @return [Wreq::CancellationToken]
+      def self.new
+      end
+
+      # Signal cancellation to all Rust tasks observing this token.
+      #
+      # @return [void]
+      def cancel
+      end
+    end
+
+    # This is a placeholder. The actual value is set by the Rust implementation.
     VERSION = nil
 
     # Send an HTTP request.

src/gvl.rs

@@ -4,55 +4,17 @@
 use std::{ffi::c_void, mem::MaybeUninit, ptr::null_mut};
 
 use rb_sys::rb_thread_call_without_gvl;
-use tokio::sync::watch;
+
+use crate::rt::CancellationToken;
 
 /// Container for safely passing closure and result through C callback.
 struct Args<F, R> {
     func: Option<F>,
     result: MaybeUninit<R>,
 }
 
-/// Cancellation flag for thread interruption support.
-#[derive(Clone)]
-pub struct CancelFlag {
-    rx: watch::Receiver<bool>,
-}
-
-struct CancelSender {
-    tx: watch::Sender<bool>,
-}
-
-impl CancelSender {
-    fn new() -> (Self, CancelFlag) {
-        let (tx, rx) = watch::channel(false);
-        (Self { tx }, CancelFlag { rx })
-    }
-
-    fn cancel(&self) {
-        let _ = self.tx.send(true);
-    }
-}
-
-impl CancelFlag {
-    /// Wait until cancellation is signaled (zero-latency, no polling).
-    pub async fn cancelled(&self) {
-        let mut rx = self.rx.clone();
-        if *rx.borrow_and_update() {
-            return;
-        }
-        loop {
-            if rx.changed().await.is_err() {
-                return;
-            }
-            if *rx.borrow() {
-                return;
-            }
-        }
-    }
-}
-
 struct UnblockData {
-    sender: CancelSender,
+    token: CancellationToken,
 }
 
 unsafe extern "C" fn call_without_gvl<F, R>(arg: *mut c_void) -> *mut c_void
@@ -73,7 +35,7 @@ where
 unsafe extern "C" fn unblock_func(arg: *mut c_void) {
     if !arg.is_null() {
         let data = unsafe { &*(arg as *const UnblockData) };
-        data.sender.cancel();
+        data.token.cancel();
     }
 }
 
@@ -110,31 +72,33 @@ where
 /// This results in all Ruby threads being suspended indefinitely.
 pub fn nogvl_cancellable<F, R>(func: F) -> R
 where
-    F: FnOnce(CancelFlag) -> R,
+    F: FnOnce(CancellationToken) -> R,
     R: Sized,
 {
-    let (sender, flag) = CancelSender::new();
-    let unblock_data = UnblockData { sender };
+    let token = CancellationToken::new();
+    let unblock_data = UnblockData {
+        token: token.clone(),
+    };
 
     struct Wrapper<F, R> {
         func: Option<F>,
-        flag: CancelFlag,
+        token: CancellationToken,
         result: MaybeUninit<R>,
     }
 
     let mut wrapper = Wrapper {
         func: Some(func),
-        flag,
+        token,
         result: MaybeUninit::uninit(),
     };
 
     unsafe extern "C" fn call_with_flag<F, R>(arg: *mut c_void) -> *mut c_void
     where
-        F: FnOnce(CancelFlag) -> R,
+        F: FnOnce(CancellationToken) -> R,
     {
         let wrapper = unsafe { &mut *(arg as *mut Wrapper<F, R>) };
         if let Some(func) = wrapper.func.take() {
-            wrapper.result.write(func(wrapper.flag.clone()));
+            wrapper.result.write(func(wrapper.token.clone()));
         }
         null_mut()
     }

src/lib.rs

@@ -93,5 +93,6 @@ fn init(ruby: &Ruby) -> Result<(), Error> {
     client::include(ruby, &gem_module)?;
     emulation::include(ruby, &gem_module)?;
     error::include(ruby);
+    rt::include(ruby, &gem_module)?;
     Ok(())
 }

src/rt.rs

@@ -1,5 +1,8 @@
 use std::sync::LazyLock;
 
+use magnus::{
+    DataTypeFunctions, Error, Module, Object, RModule, Ruby, TypedData, function, method,
+};
 use tokio::runtime::{Builder, Runtime};
 
 use crate::{error::interrupt_error, gvl};
@@ -17,11 +20,11 @@ pub fn try_block_on<F, T>(future: F) -> F::Output
 where
     F: Future<Output = Result<T, magnus::Error>>,
 {
-    gvl::nogvl_cancellable(|flag| {
+    gvl::nogvl_cancellable(|token| {
         RUNTIME.block_on(async move {
             tokio::select! {
                 biased;
-                _ = flag.cancelled() => Err(interrupt_error()),
+                _ = token.cancelled() => Err(interrupt_error()),
                 result = future => result,
             }
         })
@@ -35,13 +38,57 @@ pub fn maybe_block_on<F, T>(future: F) -> F::Output
 where
     F: Future<Output = Option<T>>,
 {
-    gvl::nogvl_cancellable(|flag| {
+    gvl::nogvl_cancellable(|token| {
         RUNTIME.block_on(async move {
             tokio::select! {
                 biased;
-                _ = flag.cancelled() => None,
+                _ = token.cancelled() => None,
                 result = future => result,
             }
         })
     })
 }
+
+/// A cancellation token for cooperative cancellation of Rust async tasks from Ruby.
+///
+/// This type wraps [`tokio_util::sync::CancellationToken`] and is exposed to Ruby as
+/// `Wreq::CancellationToken`. It allows Ruby code to signal cancellation to long-running or
+/// blocking Rust tasks, enabling graceful interruption.
+///
+/// Typical usage:
+/// - Ruby creates a `Wreq::CancellationToken` and passes it to a Rust-backed async operation.
+/// - Rust code checks or awaits the token to detect cancellation and aborts work if cancelled.
+/// - Calling `cancel` from Ruby triggers cancellation for all tasks observing this token or its
+///   clones.
+///
+/// This is especially useful for interrupting HTTP requests, streaming, or other operations that
+/// may need to be stopped from Ruby.
+#[derive(Clone, DataTypeFunctions, TypedData)]
+#[magnus(class = "Wreq::CancellationToken", free_immediately, size)]
+pub struct CancellationToken(tokio_util::sync::CancellationToken);
+
+impl CancellationToken {
+    /// Create a new [`CancellationToken`].
+    #[inline]
+    pub fn new() -> Self {
+        Self(tokio_util::sync::CancellationToken::new())
+    }
+
+    /// Signal cancellation to all tasks observing this token.
+    #[inline]
+    pub fn cancel(&self) {
+        self.0.cancel()
+    }
+
+    #[inline]
+    async fn cancelled(&self) {
+        self.0.cancelled().await
+    }
+}
+
+pub fn include(ruby: &Ruby, gem_module: &RModule) -> Result<(), Error> {
+    let headers_class = gem_module.define_class("CancellationToken", ruby.class_object())?;
+    headers_class.define_singleton_method("new", function!(CancellationToken::new, 0))?;
+    headers_class.define_method("cancel", method!(CancellationToken::cancel, 0))?;
+    Ok(())
+}

test/interrupt_test.rb

@@ -0,0 +1,64 @@
+require "test_helper"
+
+class ThreadInterruptTest < Minitest::Test
+  HANGING_URL = "http://10.255.255.1:12345/" # Non-routable, hangs forever
+  SLOW_BODY_URL = "https://httpbin.io/drip?duration=5&numbytes=5" # 5s slow body
+
+  def test_connect_phase_interrupt
+    thread = Thread.new do
+      begin
+        Wreq.get(HANGING_URL)
+      rescue => e
+        e
+      end
+    end
+    sleep 2
+    thread.kill
+    result = thread.join(5)
+    assert result, "Thread should be killed and joined"
+  end
+
+  def test_connect_with_timeout_interrupt
+    thread = Thread.new do
+      begin
+        Wreq.get(HANGING_URL, timeout: 60)
+      rescue => e
+        e
+      end
+    end
+    sleep 2
+    thread.kill
+    result = thread.join(5)
+    assert result, "Thread should be killed and joined"
+  end
+
+  def test_body_reading_interrupt
+    thread = Thread.new do
+      resp = Wreq.get(SLOW_BODY_URL)
+      begin
+        resp.text
+      rescue => e
+        e
+      end
+    end
+    sleep 2
+    thread.kill
+    result = thread.join(5)
+    assert result, "Thread should be killed and joined"
+  end
+
+  def test_body_streaming_interrupt
+    thread = Thread.new do
+      resp = Wreq.get(SLOW_BODY_URL)
+      begin
+        resp.chunks { |chunk| chunk }
+      rescue => e
+        e
+      end
+    end
+    sleep 2
+    thread.kill
+    result = thread.join(5)
+    assert result, "Thread should be killed and joined"
+  end
+end