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};
@@ -202,16 +202,58 @@ impl CustomSlotSupplier {
     }
 }
 
+#[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 {
+        dbg!("Invoking reserve first time");
         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(SK::kind(), ctx),))?;
+                let called = py_obj.call_method1(
+                    "reserve_slot",
+                    (
+                        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))
             }) {

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,15 +90,18 @@ 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:
-            logger.exception("saw cancelled error")
             raise
         except Exception:
             logger.warning(
@@ -161,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}")