Cancellation Tokens

Author: Sushisource

src/Temporalio/Bridge/CustomSlotSupplier.cs

@@ -46,13 +46,15 @@ private unsafe void Reserve(Interop.SlotReserveCtx ctx, void* sender)
 
         private void SafeReserve(Interop.SlotReserveCtx ctx, IntPtr sender)
         {
+            var cancelTokenSrc = new System.Threading.CancellationTokenSource();
             var reserveTask = Task.Run(async () =>
             {
                 while (true)
                 {
                     try
                     {
-                        var permit = await userSupplier.ReserveSlotAsync(new(ctx)).ConfigureAwait(false);
+                        var permit = await userSupplier.ReserveSlotAsync(
+                            new(ctx), cancelTokenSrc.Token).ConfigureAwait(false);
                         var usedPermitId = AddPermitToMap(permit);
                         unsafe
                         {

src/Temporalio/Bridge/include/temporal-sdk-bridge.h

@@ -389,10 +389,13 @@ typedef struct SlotReserveCtx {
   struct ByteArrayRef worker_identity;
   struct ByteArrayRef worker_build_id;
   bool is_sticky;
+  void *token_src;
 } SlotReserveCtx;
 
 typedef void (*CustomReserveSlotCallback)(struct SlotReserveCtx ctx, void *sender);
 
+typedef void (*CustomCancelReserveCallback)(void *token_source);
+
 /**
  * Must return C#-tracked id for the permit. A zero value means no permit was reserved.
  */
@@ -448,6 +451,7 @@ typedef void (*CustomReleaseSlotCallback)(struct SlotReleaseCtx ctx);
 
 typedef struct CustomSlotSupplierCallbacks {
   CustomReserveSlotCallback reserve;
+  CustomCancelReserveCallback cancel_reserve;
   CustomTryReserveSlotCallback try_reserve;
   CustomMarkSlotUsedCallback mark_used;
   CustomReleaseSlotCallback release;
@@ -694,6 +698,8 @@ struct WorkerReplayPushResult worker_replay_push(struct Worker *worker,
 
 void complete_async_reserve(void *sender, uintptr_t permit_id);
 
+void set_reserve_cancel_target(struct SlotReserveCtx *ctx, void *token_ptr);
+
 #ifdef __cplusplus
 } // extern "C"
 #endif // __cplusplus

src/Temporalio/Bridge/src/worker.rs

@@ -96,6 +96,7 @@ unsafe impl<SK> Sync for CustomSlotSupplier<SK> {}
 
 type CustomReserveSlotCallback =
     unsafe extern "C" fn(ctx: SlotReserveCtx, sender: *mut libc::c_void);
+type CustomCancelReserveCallback = unsafe extern "C" fn(token_source: *mut libc::c_void);
 /// Must return C#-tracked id for the permit. A zero value means no permit was reserved.
 type CustomTryReserveSlotCallback = unsafe extern "C" fn(ctx: SlotReserveCtx) -> usize;
 type CustomMarkSlotUsedCallback = unsafe extern "C" fn(ctx: SlotMarkUsedCtx);
@@ -108,6 +109,7 @@ pub struct CustomSlotSupplierCallbacksImpl(*const CustomSlotSupplierCallbacks);
 #[repr(C)]
 pub struct CustomSlotSupplierCallbacks {
     reserve: CustomReserveSlotCallback,
+    cancel_reserve: CustomCancelReserveCallback,
     try_reserve: CustomTryReserveSlotCallback,
     mark_used: CustomMarkSlotUsedCallback,
     release: CustomReleaseSlotCallback,
@@ -139,6 +141,8 @@ pub struct SlotReserveCtx {
     worker_identity: ByteArrayRef,
     worker_build_id: ByteArrayRef,
     is_sticky: bool,
+    // The C# side will store a pointer here to the cancellation token source
+    token_src: *mut libc::c_void,
 }
 
 #[repr(C)]
@@ -169,6 +173,21 @@ pub struct SlotReleaseCtx {
     slot_permit: usize,
 }
 
+struct CancelReserveGuard {
+    token_src: *mut libc::c_void,
+    callback: CustomCancelReserveCallback,
+}
+impl Drop for CancelReserveGuard {
+    fn drop(&mut self) {
+        if !self.token_src.is_null() {
+            unsafe {
+                (self.callback)(self.token_src);
+            }
+        }
+    }
+}
+unsafe impl Send for CancelReserveGuard {}
+
 #[async_trait::async_trait]
 impl<SK: SlotKind + Send + Sync> temporal_sdk_core_api::worker::SlotSupplier
     for CustomSlotSupplier<SK>
@@ -180,10 +199,13 @@ impl<SK: SlotKind + Send + Sync> temporal_sdk_core_api::worker::SlotSupplier
         let ctx = Self::convert_reserve_ctx(ctx);
         let tx = Box::into_raw(Box::new(tx)) as *mut libc::c_void;
         unsafe {
+            let _drop_guard = CancelReserveGuard {
+                token_src: ctx.token_src,
+                callback: (*self.inner.0).cancel_reserve,
+            };
             ((*self.inner.0).reserve)(ctx, tx);
+            rx.await.expect("reserve channel is not closed")
         }
-        let r = rx.await.expect("reserve channel is not closed");
-        r
     }
 
     fn try_reserve_slot(&self, ctx: &dyn SlotReservationContext) -> Option<SlotSupplierPermit> {
@@ -244,6 +266,7 @@ impl<SK: SlotKind + Send + Sync> CustomSlotSupplier<SK> {
             worker_identity: ctx.worker_identity().into(),
             worker_build_id: ctx.worker_build_id().into(),
             is_sticky: ctx.is_sticky(),
+            token_src: std::ptr::null_mut(),
         }
     }
 
@@ -738,6 +761,16 @@ pub extern "C" fn complete_async_reserve(sender: *mut libc::c_void, permit_id: u
     }
 }
 
+#[no_mangle]
+pub unsafe extern "C" fn set_reserve_cancel_target(
+    ctx: *mut SlotReserveCtx,
+    token_ptr: *mut libc::c_void,
+) {
+    if let Some(ctx) = ctx.as_mut() {
+        ctx.token_src = token_ptr;
+    }
+}
+
 impl TryFrom<&WorkerOptions> for temporal_sdk_core::WorkerConfig {
     type Error = anyhow::Error;
 

src/Temporalio/Worker/Tuning/ICustomSlotSupplier.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Temporalio.Worker.Tuning
@@ -22,9 +23,11 @@ public interface ICustomSlotSupplier : ISlotSupplier
         /// This method will be called concurrently from multiple threads, so it must be thread-safe.
         /// </remarks>
         /// <param name="ctx">The context for slot reservation.</param>
+        /// <param name="cancellationToken">A cancellation token that the SDK may use
+        ///  to cancel the operation.</param>
         /// <returns>A permit to use the slot which may be populated with your own data.</returns>
         /// <exception cref="OperationCanceledException">Cancellation requested.</exception>
-        public Task<ISlotPermit> ReserveSlotAsync(SlotReserveContext ctx);
+        public Task<ISlotPermit> ReserveSlotAsync(SlotReserveContext ctx, CancellationToken cancellationToken);
 
         /// <summary>
         /// This function is called when trying to reserve slots for "eager" workflow and activity tasks.

tests/Temporalio.Tests/Worker/WorkerTuningTests.cs

@@ -156,7 +156,7 @@ private class MySlotSupplier : ICustomSlotSupplier
 
         public HashSet<bool> SeenReleaseInfoPresence { get; } = new();
 
-        public async Task<ISlotPermit> ReserveSlotAsync(SlotReserveContext ctx)
+        public async Task<ISlotPermit> ReserveSlotAsync(SlotReserveContext ctx, CancellationToken cancellationToken)
         {
             // Do something async to make sure that works
             await Task.Delay(10);
@@ -243,7 +243,7 @@ await Env.Client.ExecuteWorkflowAsync(
 
     private class ThrowingSlotSupplier : ICustomSlotSupplier
     {
-        public Task<ISlotPermit> ReserveSlotAsync(SlotReserveContext ctx)
+        public Task<ISlotPermit> ReserveSlotAsync(SlotReserveContext ctx, CancellationToken cancellationToken)
         {
             // Let the workflow complete, but other reservations fail
             if (ctx.SlotType == SlotType.Workflow)