Actually cancel reserve slot tasks

Author: Sushisource

temporalio/bridge/src/worker.rs

@@ -7,7 +7,7 @@ use pyo3::types::{PyBytes, PyTuple};
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::marker::PhantomData;
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
 use std::time::Duration;
 use temporal_sdk_core::api::errors::{PollActivityError, PollWfError};
 use temporal_sdk_core::replay::{HistoryForReplay, ReplayWorkerInput};
@@ -64,6 +64,18 @@ pub struct TunerHolder {
     local_activity_slot_supplier: SlotSupplier,
 }
 
+// pub fn set_task_locals_on_tuner<'a>(py: Python<'a>, tuner: &TunerHolder) -> PyResult<()> {
+//     // TODO: All suppliers
+//     if let SlotSupplier::Custom(ref cs) = tuner.workflow_slot_supplier {
+//         Python::with_gil(|py| {
+//             let py_obj = cs.inner.as_ref(py);
+//             py_obj.call_method0("set_task_locals")?;
+//             Ok(())
+//         })?;
+//     };
+//     Ok(())
+// }
+
 #[derive(FromPyObject)]
 pub enum SlotSupplier {
     FixedSize(FixedSizeSlotSupplier),
@@ -190,17 +202,60 @@ impl CustomSlotSupplier {
     }
 }
 
+// Shouldn't really need this callback nonsense, it should be possible to do this from the pyo3
+// asyncio library, but we'd have to vendor the whole thing to make the right improvements. When
+// pyo3 is upgraded and we are using
+
+#[pyclass]
+struct CreatedTaskForSlotCallback {
+    stored_task: Arc<OnceLock<PyObject>>,
+}
+
+#[pymethods]
+impl CreatedTaskForSlotCallback {
+    fn __call__(&self, task: PyObject) -> PyResult<()> {
+        self.stored_task.set(task).expect("must only be set once");
+        Ok(())
+    }
+}
+
+struct TaskCanceller {
+    stored_task: Arc<OnceLock<PyObject>>,
+}
+
+impl TaskCanceller {
+    fn new(stored_task: Arc<OnceLock<PyObject>>) -> Self {
+        TaskCanceller { stored_task }
+    }
+}
+
+impl Drop for TaskCanceller {
+    fn drop(&mut self) {
+        if let Some(task) = self.stored_task.get() {
+            Python::with_gil(|py| {
+                task.call_method0(py, "cancel")
+                    .expect("Failed to cancel task");
+            });
+        }
+    }
+}
+
 #[async_trait::async_trait]
 impl<SK: SlotKind + Send + Sync> SlotSupplierTrait for CustomSlotSupplierOfType<SK> {
     type SlotKind = SK;
 
     async fn reserve_slot(&self, ctx: &dyn SlotReservationContext) -> SlotSupplierPermit {
         loop {
+            let stored_task = Arc::new(OnceLock::new());
+            let _task_canceller = TaskCanceller::new(stored_task.clone());
             let pypermit = match Python::with_gil(|py| {
                 let py_obj = self.inner.as_ref(py);
                 let called = py_obj.call_method1(
                     "reserve_slot",
-                    (SlotReserveCtx::from_ctx(Self::SlotKind::kind(), ctx),),
+                    (
+                        SlotReserveCtx::from_ctx(SK::kind(), ctx),
+                        CreatedTaskForSlotCallback { stored_task },
+                    ),
                 )?;
                 runtime::THREAD_TASK_LOCAL
                     .with(|tl| pyo3_asyncio::into_future_with_locals(tl.get().unwrap(), called))
@@ -232,7 +287,7 @@ impl<SK: SlotKind + Send + Sync> SlotSupplierTrait for CustomSlotSupplierOfType<
             let py_obj = self.inner.as_ref(py);
             let pa = py_obj.call_method1(
                 "try_reserve_slot",
-                (SlotReserveCtx::from_ctx(Self::SlotKind::kind(), ctx),),
+                (SlotReserveCtx::from_ctx(SK::kind(), ctx),),
             )?;
 
             if pa.is_none() {
@@ -362,6 +417,8 @@ pub fn new_replay_worker<'a>(
 impl WorkerRef {
     fn validate<'p>(&self, py: Python<'p>) -> PyResult<&'p PyAny> {
         let worker = self.worker.as_ref().unwrap().clone();
+        // Set custom slot supplier task locals so they can run futures
+        // match worker.get_config().tuner {}
         self.runtime.future_into_py(py, async move {
             worker
                 .validate()

temporalio/worker/_tuning.py

@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import Literal, Optional, Union
+from typing import Any, Callable, Literal, Optional, Union
 
 from typing_extensions import TypeAlias
 
@@ -90,13 +90,17 @@ class ResourceBasedSlotSupplier:
 ]
 
 
-class _ErrorLoggingSlotSupplier(CustomSlotSupplier):
+class _BridgeSlotSupplierWrapper:
     def __init__(self, supplier: CustomSlotSupplier):
         self._supplier = supplier
 
-    async def reserve_slot(self, ctx: SlotReserveContext) -> SlotPermit:
+    async def reserve_slot(
+        self, ctx: SlotReserveContext, reserve_cb: Callable[[Any], None]
+    ) -> SlotPermit:
         try:
-            return await self._supplier.reserve_slot(ctx)
+            reserve_fut = asyncio.create_task(self._supplier.reserve_slot(ctx))
+            reserve_cb(reserve_fut)
+            return await reserve_fut
         except asyncio.CancelledError:
             raise
         except Exception:
@@ -160,7 +164,7 @@ def _to_bridge_slot_supplier(
         )
     elif isinstance(slot_supplier, CustomSlotSupplier):
         return temporalio.bridge.worker.BridgeCustomSlotSupplier(
-            _ErrorLoggingSlotSupplier(slot_supplier)
+            _BridgeSlotSupplierWrapper(slot_supplier)
         )
     else:
         raise TypeError(f"Unknown slot supplier type: {slot_supplier}")

tests/worker/test_worker.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import concurrent.futures
+import threading
 import uuid
 from datetime import timedelta
 from typing import Any, Awaitable, Callable, Optional
@@ -353,6 +354,7 @@ def __init__(self, pnum: int):
     class MySlotSupplier(CustomSlotSupplier):
         reserves = 0
         releases = 0
+        highest_seen_reserve_on_release = 0
         used = 0
         seen_sticky_kinds = set()
         seen_slot_kinds = set()
@@ -388,6 +390,9 @@ def release_slot(self, ctx: SlotReleaseContext) -> None:
             assert ctx.permit is not None
             assert isinstance(ctx.permit, MyPermit)
             assert ctx.permit.pnum is not None
+            self.highest_seen_reserve_on_release = max(
+                ctx.permit.pnum, self.highest_seen_reserve_on_release
+            )
             # Info may be empty, and we should see both empty and not
             if ctx.slot_info is None:
                 self.seen_release_info_empty = True
@@ -422,10 +427,11 @@ def reserve_asserts(self, ctx):
         await wf1.signal(WaitOnSignalWorkflow.my_signal, "finish")
         await wf1.result()
 
-    async def releases() -> int:
-        return ss.releases
-
-    assert ss.reserves == ss.releases
+    # We can't use reserve number directly because there is a technically possible race
+    # where the python reserve function appears to complete, but Rust doesn't see that.
+    # This isn't solvable without redoing a chunk of pyo3-asyncio. So we only check
+    # that the permits passed to release line up.
+    assert ss.highest_seen_reserve_on_release == ss.releases
     # Two workflow tasks, one activity
     assert ss.used == 3
     assert ss.seen_sticky_kinds == {True, False}