CI debug testing

Author: Sushisource

temporalio/bridge/src/worker.rs

@@ -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),
@@ -198,10 +210,8 @@ impl<SK: SlotKind + Send + Sync> SlotSupplierTrait for CustomSlotSupplierOfType<
         loop {
             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),),
-                )?;
+                let called = py_obj
+                    .call_method1("reserve_slot", (SlotReserveCtx::from_ctx(SK::kind(), ctx),))?;
                 runtime::THREAD_TASK_LOCAL
                     .with(|tl| pyo3_asyncio::into_future_with_locals(tl.get().unwrap(), called))
             }) {
@@ -217,11 +227,13 @@ impl<SK: SlotKind + Send + Sync> SlotSupplierTrait for CustomSlotSupplierOfType<
             .await;
             match pypermit {
                 Ok(p) => {
+                    dbg!("reserve done", SK::kind());
                     return SlotSupplierPermit::with_user_data(p);
                 }
-                Err(_) => {
+                Err(e) => {
                     // This is a user thrown error, re-raised by the logging wrapper so we can
                     // loop, so do that.
+                    dbg!(e);
                 }
             }
         }
@@ -232,7 +244,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() {
@@ -274,6 +286,7 @@ impl<SK: SlotKind + Send + Sync> SlotSupplierTrait for CustomSlotSupplierOfType<
     }
 
     fn release_slot(&self, ctx: &dyn SlotReleaseContext<SlotKind = Self::SlotKind>) {
+        dbg!("release", SK::kind());
         if let Err(e) = Python::with_gil(|py| {
             let permit = ctx
                 .permit()
@@ -362,6 +375,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

@@ -98,6 +98,7 @@ async def reserve_slot(self, ctx: SlotReserveContext) -> SlotPermit:
         try:
             return await self._supplier.reserve_slot(ctx)
         except asyncio.CancelledError:
+            logger.exception("saw cancelled error")
             raise
         except Exception:
             logger.warning(

tests/worker/test_worker.py

@@ -5,6 +5,7 @@
 import uuid
 from datetime import timedelta
 from typing import Any, Awaitable, Callable, Optional
+import threading
 
 import pytest
 
@@ -351,6 +352,7 @@ def __init__(self, pnum: int):
             self.pnum = pnum
 
     class MySlotSupplier(CustomSlotSupplier):
+        lock = threading.Lock()
         reserves = 0
         releases = 0
         used = 0
@@ -364,7 +366,9 @@ async def reserve_slot(self, ctx: SlotReserveContext) -> SlotPermit:
             self.reserve_asserts(ctx)
             # Verify an async call doesn't bungle things
             await asyncio.sleep(0.01)
-            self.reserves += 1
+            with self.lock:
+                self.reserves += 1
+                print("incremented reserve")
             return MyPermit(self.reserves)
 
         def try_reserve_slot(self, ctx: SlotReserveContext) -> Optional[SlotPermit]:
@@ -376,24 +380,26 @@ def mark_slot_used(self, ctx: SlotMarkUsedContext) -> None:
             assert isinstance(ctx.permit, MyPermit)
             assert ctx.permit.pnum is not None
             assert ctx.slot_info is not None
-            if isinstance(ctx.slot_info, WorkflowSlotInfo):
-                self.seen_used_slot_kinds.add("wf")
-            elif isinstance(ctx.slot_info, ActivitySlotInfo):
-                self.seen_used_slot_kinds.add("a")
-            elif isinstance(ctx.slot_info, LocalActivitySlotInfo):
-                self.seen_used_slot_kinds.add("la")
-            self.used += 1
+            with self.lock:
+                if isinstance(ctx.slot_info, WorkflowSlotInfo):
+                    self.seen_used_slot_kinds.add("wf")
+                elif isinstance(ctx.slot_info, ActivitySlotInfo):
+                    self.seen_used_slot_kinds.add("a")
+                elif isinstance(ctx.slot_info, LocalActivitySlotInfo):
+                    self.seen_used_slot_kinds.add("la")
+                self.used += 1
 
         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
+            with self.lock:
             # Info may be empty, and we should see both empty and not
-            if ctx.slot_info is None:
-                self.seen_release_info_empty = True
-            else:
-                self.seen_release_info_nonempty = True
-            self.releases += 1
+                if ctx.slot_info is None:
+                    self.seen_release_info_empty = True
+                else:
+                    self.seen_release_info_nonempty = True
+                self.releases += 1
 
         def reserve_asserts(self, ctx):
             assert ctx.task_queue is not None
@@ -422,9 +428,6 @@ 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
     # Two workflow tasks, one activity
     assert ss.used == 3