feat: adds deadline/timeout support to IrrevocableContext

src/core/context/context_test.rs

@@ -149,4 +149,95 @@ mod tests {
             }
         }
     }
+
+    // --- ported "extras" feature tests (candidate for main integration) ---
+
+    /// a context created with a timeout fails a longer-running operation with a deadline error
+    #[tokio::test]
+    async fn test_run_respects_deadline() {
+        let ctx =
+            IrrevocableContext::with_timeout(&span_fixture(), "timeout_ctx", Duration::from_millis(10));
+
+        let result = ctx
+            .run(async {
+                sleep(Duration::from_millis(200)).await;
+                Ok::<(), anyhow::Error>(())
+            })
+            .await;
+
+        assert!(result.is_err());
+        assert!(result.unwrap_err().to_string().contains("deadline exceeded"));
+    }
+
+    /// is_deadline_exceeded reflects the deadline crossing
+    #[tokio::test]
+    async fn test_is_deadline_exceeded() {
+        let ctx =
+            IrrevocableContext::with_timeout(&span_fixture(), "timeout_ctx", Duration::from_millis(5));
+        assert!(!ctx.is_deadline_exceeded());
+        assert!(ctx.deadline().is_some());
+        sleep(Duration::from_millis(20)).await;
+        assert!(ctx.is_deadline_exceeded());
+        assert!(ctx.err().is_some());
+    }
+
+    /// with_cancel returns a child plus a closure that cancels it
+    #[tokio::test]
+    async fn test_with_cancel_closure() {
+        let root = IrrevocableContext::new(&span_fixture(), "cancel_root");
+        let (child, cancel) = root.with_cancel("cancel_child");
+
+        assert!(!child.is_cancelled());
+        cancel();
+
+        let child_clone = child.clone();
+        wait_until(move || child_clone.is_cancelled(), Duration::from_millis(100))
+            .await
+            .expect("child should be cancelled after invoking cancel closure");
+    }
+
+    /// run() on an already-expired context returns the deadline error even when the
+    /// operation itself is immediately ready (guards the short-circuit / select! race)
+    #[tokio::test]
+    async fn test_run_short_circuits_when_deadline_already_passed() {
+        let ctx =
+            IrrevocableContext::with_timeout(&span_fixture(), "expired_ctx", Duration::from_millis(5));
+        // let the deadline lapse before we ever call run()
+        sleep(Duration::from_millis(20)).await;
+        assert!(ctx.is_deadline_exceeded());
+
+        // an immediately-ready future must NOT slip through on an expired context
+        let result = ctx.run(async { Ok::<i32, anyhow::Error>(42) }).await;
+
+        assert!(result.is_err());
+        assert!(result.unwrap_err().to_string().contains("deadline exceeded"));
+    }
+
+    /// cancellation that fires *while* run() is in-flight wins the race inside the
+    /// Some(deadline) arm — exercises select!'s `cancelled()` branch (not the short-circuit)
+    #[tokio::test]
+    async fn test_run_cancellation_beats_deadline() {
+        // generous deadline so the timer never fires; the context is live when run() starts
+        let ctx = IrrevocableContext::with_timeout(
+            &span_fixture(),
+            "cancel_vs_deadline",
+            Duration::from_secs(60),
+        );
+        let canceller = ctx.clone();
+
+        // run a long operation, and concurrently cancel shortly after it begins
+        let (result, ()) = tokio::join!(
+            ctx.run(async {
+                sleep(Duration::from_millis(200)).await;
+                Ok::<(), anyhow::Error>(())
+            }),
+            async move {
+                sleep(Duration::from_millis(10)).await;
+                canceller.cancel();
+            }
+        );
+
+        assert!(result.is_err());
+        assert!(result.unwrap_err().to_string().contains("context cancelled"));
+    }
 }

src/core/context/mod.rs

@@ -1,20 +1,22 @@
 //! Cancelable context with irrecoverable error propagation
 //!
-//! This module provides a simplified context implementation focused on:
+//! This module provides a context implementation focused on:
 //! - Cancellation support via tokio's CancellationToken
-//! - Parent-child context hierarchies  
+//! - Parent-child context hierarchies
 //! - Irrecoverable error propagation that terminates the application
+//! - Optional deadline/timeout support
 //!
-//! Unlike the full Go context API, this implementation focuses only on the core
-//! functionality needed: cancellation and error propagation.
+//! The deadline feature is additive: a context created with `new` behaves
+//! exactly like the previous minimal API (no deadline).
 
 #[cfg(test)]
 mod context_test;
 
 use anyhow::Result;
 use std::sync::Arc;
+use tokio::time::Instant;
 use tokio_util::sync::CancellationToken;
-use tracing::Span;
+use tracing::{Instrument, Span};
 
 /// A cancelable context that supports parent-child hierarchies and irrecoverable error propagation.
 ///
@@ -26,6 +28,8 @@ pub struct IrrevocableContext {
 
 struct ContextInner {
     token: CancellationToken,
+    /// Absolute deadline after which the context is considered expired, if any.
+    deadline: Option<Instant>,
     parent: Option<IrrevocableContext>,
     span: Span,
 }
@@ -38,18 +42,34 @@ impl IrrevocableContext {
         Self {
             inner: Arc::new(ContextInner {
                 token: CancellationToken::new(),
+                deadline: None,
                 parent: None,
                 span,
             }),
         }
     }
 
-    /// Create a child context that inherits cancellation from the parent
+    /// Create a new root context that expires after `timeout` elapses.
+    pub fn with_timeout(parent_span: &Span, tag: &str, timeout: std::time::Duration) -> Self {
+        let span = tracing::span!(parent: parent_span, tracing::Level::TRACE, "irrevocable_context_timeout", tag = tag);
+
+        Self {
+            inner: Arc::new(ContextInner {
+                token: CancellationToken::new(),
+                deadline: Some(Instant::now() + timeout),
+                parent: None,
+                span,
+            }),
+        }
+    }
+
+    /// Create a child context that inherits cancellation and deadline from the parent.
     pub fn child(&self, tag: &str) -> Self {
         let span = tracing::span!(parent: &self.inner.span, tracing::Level::TRACE, "irrevocable_context_child", tag = tag);
         Self {
             inner: Arc::new(ContextInner {
                 token: self.inner.token.child_token(),
+                deadline: self.inner.deadline,
                 parent: Some(self.clone()),
                 span,
             }),
@@ -74,21 +94,65 @@ impl IrrevocableContext {
         self.inner.token.cancelled().await;
     }
 
-    /// Run an operation with cancellation support
-    /// If the context is cancelled before the operation completes, it returns an error.
-    /// otherwise, it returns the operation's result.
+    /// Returns the absolute deadline for this context, if one was set.
+    pub fn deadline(&self) -> Option<Instant> {
+        self.inner.deadline
+    }
+
+    /// Returns true if the context has a deadline that has already passed.
+    pub fn is_deadline_exceeded(&self) -> bool {
+        self.inner.deadline.is_some_and(|d| Instant::now() >= d)
+    }
+
+    /// Returns an error describing why the context is unusable (cancelled or past deadline), if so.
+    pub fn err(&self) -> Option<anyhow::Error> {
+        if self.is_cancelled() {
+            Some(anyhow::anyhow!("context cancelled"))
+        } else if self.is_deadline_exceeded() {
+            Some(anyhow::anyhow!("context deadline exceeded"))
+        } else {
+            None
+        }
+    }
+
+    /// Run an operation with cancellation and deadline support.
+    /// Returns an error if the context is cancelled or its deadline elapses before the operation completes.
     pub async fn run<F, T>(&self, future: F) -> Result<T>
     where
         F: std::future::Future<Output = Result<T>>,
     {
-        let _enter = self.inner.span.enter();
+        // Attach the context span via `.instrument()` rather than holding an
+        // `enter()` guard across the `.await`: the guard is `!Send` and would
+        // stay entered while the future is suspended, leaking the span onto
+        // whatever task resumes it. `.instrument()` enters the span only while
+        // this future is actively polled.
+        async move {
+            // Short-circuit if the context is already unusable, so an already-cancelled
+            // or already-expired context never races an immediately-ready future (which
+            // `select!` would otherwise resolve pseudo-randomly).
+            if let Some(err) = self.err() {
+                return Err(err);
+            }
 
-        tokio::select! {
-            result = future => result,
-            _ = self.cancelled() => {
-                Err(anyhow::anyhow!("context cancelled"))
+            match self.inner.deadline {
+                Some(deadline) => {
+                    let remaining = deadline.saturating_duration_since(Instant::now());
+                    tokio::select! {
+                        result = future => result,
+                        _ = tokio::time::sleep(remaining) => Err(anyhow::anyhow!("context deadline exceeded")),
+                        _ = self.cancelled() => Err(anyhow::anyhow!("context cancelled")),
+                    }
+                }
+                None => {
+                    tokio::select! {
+                        result = future => result,
+                        _ = self.cancelled() => Err(anyhow::anyhow!("context cancelled")),
+                    }
+                }
             }
         }
+        .instrument(self.inner.span.clone())
+        .await
     }
 
     /// Propagate an irrecoverable error up the context chain.
@@ -120,13 +184,25 @@ impl IrrevocableContext {
             Err(err) => self.throw_irrecoverable(err),
         }
     }
+
+    /// Returns a cancellable child context together with a function that cancels it.
+    /// Mirrors Go's `context.WithCancel`.
+    ///
+    /// The returned closure is `Send + Sync + Clone`, so it can be moved across
+    /// threads/tasks or stored — the intended way to hand out cancel authority.
+    pub fn with_cancel(&self, tag: &str) -> (Self, impl Fn() + Send + Sync + Clone) {
+        let child = self.child(tag);
+        let token = child.inner.token.clone();
+        (child, move || token.cancel())
+    }
 }
 
 // Custom Debug implementation for better visibility into the context state
 impl std::fmt::Debug for IrrevocableContext {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("IrrevocableContext")
             .field("is_cancelled", &self.is_cancelled())
+            .field("deadline", &self.inner.deadline)
             .field("has_parent", &self.inner.parent.is_some())
             .finish()
     }