Logging wrapping / error testing

Author: Sushisource

temporalio/bridge/src/worker.rs

@@ -1,4 +1,5 @@
 use anyhow::Context;
+use log::error;
 use prost::Message;
 use pyo3::exceptions::{PyException, PyRuntimeError, PyValueError};
 use pyo3::prelude::*;
@@ -12,7 +13,7 @@ use temporal_sdk_core::api::errors::{PollActivityError, PollWfError};
 use temporal_sdk_core::replay::{HistoryForReplay, ReplayWorkerInput};
 use temporal_sdk_core_api::errors::WorkflowErrorType;
 use temporal_sdk_core_api::worker::{
-    SlotInfo, SlotInfoTrait, SlotKind, SlotMarkUsedContext, SlotReleaseContext,
+    SlotInfo, SlotInfoTrait, SlotKind, SlotKindType, SlotMarkUsedContext, SlotReleaseContext,
     SlotReservationContext, SlotSupplier as SlotSupplierTrait, SlotSupplierPermit,
 };
 use temporal_sdk_core_api::Worker;
@@ -87,7 +88,7 @@ pub struct ResourceBasedSlotSupplier {
 #[pyclass]
 pub struct SlotReserveCtx {
     #[pyo3(get)]
-    pub slot_type: String, // TODO: Real type
+    pub slot_type: String,
     #[pyo3(get)]
     pub task_queue: String,
     #[pyo3(get)]
@@ -99,9 +100,13 @@ pub struct SlotReserveCtx {
 }
 
 impl SlotReserveCtx {
-    fn from_ctx(slot_type: String, ctx: &dyn SlotReservationContext) -> Self {
+    fn from_ctx(slot_type: SlotKindType, ctx: &dyn SlotReservationContext) -> Self {
         SlotReserveCtx {
-            slot_type,
+            slot_type: match slot_type {
+                SlotKindType::Workflow => "workflow".to_string(),
+                SlotKindType::Activity => "activity".to_string(),
+                SlotKindType::LocalActivity => "local-activity".to_string(),
+            },
             task_queue: ctx.task_queue().to_string(),
             worker_identity: ctx.worker_identity().to_string(),
             worker_build_id: ctx.worker_build_id().to_string(),
@@ -191,26 +196,32 @@ impl<SK: SlotKind + Send + Sync> SlotSupplierTrait for CustomSlotSupplierOfType<
 
     async fn reserve_slot(&self, ctx: &dyn SlotReservationContext) -> SlotSupplierPermit {
         loop {
-            let pypermit = Python::with_gil(|py| {
+            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().to_string(),
-                        ctx,
-                    ),),
+                    (SlotReserveCtx::from_ctx(Self::SlotKind::kind(), ctx),),
                 )?;
                 runtime::THREAD_TASK_LOCAL
                     .with(|tl| pyo3_asyncio::into_future_with_locals(tl.get().unwrap(), called))
-            })
-            .expect("TODO")
+            }) {
+                Ok(f) => f,
+                Err(e) => {
+                    error!(
+                        "Unexpected error in custom slot supplier `reserve_slot`: {}",
+                        e
+                    );
+                    continue;
+                }
+            }
             .await;
             match pypermit {
                 Ok(p) => {
                     return SlotSupplierPermit::with_user_data(p);
                 }
-                Err(e) => {
-                    dbg!("Error in reserve_slot", e);
+                Err(_) => {
+                    // This is a user thrown error, re-raised by the logging wrapper so we can
+                    // loop, so do that.
                 }
             }
         }
@@ -221,22 +232,25 @@ 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().to_string(),
-                    ctx,
-                ),),
+                (SlotReserveCtx::from_ctx(Self::SlotKind::kind(), ctx),),
             )?;
 
             if pa.is_none() {
                 return Ok(None);
             }
             PyResult::Ok(Some(SlotSupplierPermit::with_user_data(pa.into_py(py))))
         })
-        .expect("TODO")
+        .unwrap_or_else(|e| {
+            error!(
+                "Uncaught error in custom slot supplier `try_reserve_slot`: {}",
+                e
+            );
+            None
+        })
     }
 
     fn mark_slot_used(&self, ctx: &dyn SlotMarkUsedContext<SlotKind = Self::SlotKind>) {
-        Python::with_gil(|py| {
+        if let Err(e) = Python::with_gil(|py| {
             let permit = ctx
                 .permit()
                 .user_data::<PyObject>()
@@ -251,12 +265,16 @@ impl<SK: SlotKind + Send + Sync> SlotSupplierTrait for CustomSlotSupplierOfType<
                 },),
             )?;
             PyResult::Ok(())
-        })
-        .expect("TODO");
+        }) {
+            error!(
+                "Uncaught error in custom slot supplier `mark_slot_used`: {}",
+                e
+            );
+        }
     }
 
     fn release_slot(&self, ctx: &dyn SlotReleaseContext<SlotKind = Self::SlotKind>) {
-        Python::with_gil(|py| {
+        if let Err(e) = Python::with_gil(|py| {
             let permit = ctx
                 .permit()
                 .user_data::<PyObject>()
@@ -271,8 +289,12 @@ impl<SK: SlotKind + Send + Sync> SlotSupplierTrait for CustomSlotSupplierOfType<
                 },),
             )?;
             PyResult::Ok(())
-        })
-        .expect("TODO");
+        }) {
+            error!(
+                "Uncaught error in custom slot supplier `release_slot`: {}",
+                e
+            );
+        }
     }
 
     fn available_slots(&self) -> Option<usize> {

temporalio/bridge/worker.py

@@ -12,6 +12,7 @@
     Awaitable,
     Callable,
     List,
+    Literal,
     Optional,
     Protocol,
     Sequence,
@@ -34,6 +35,7 @@
 import temporalio.bridge.temporal_sdk_bridge
 import temporalio.converter
 import temporalio.exceptions
+from temporalio.bridge.temporal_sdk_bridge import PollShutdownError
 
 
 @dataclass
@@ -99,8 +101,8 @@ class SlotPermit:
 class SlotReserveContext(Protocol):
     """Context for reserving a slot from a :py:class:`CustomSlotSupplier`."""
 
-    slot_type: str  # TODO real type
-    """The type of slot trying to be reserved."""
+    slot_type: Literal["workflow", "activity", "local-activity"]
+    """The type of slot trying to be reserved. Always one of "workflow", "activity", or "local-activity"."""
     task_queue: str
     """The name of the task queue for which this reservation request is associated."""
     worker_identity: str
@@ -159,14 +161,13 @@ class SlotReleaseContext:
 class CustomSlotSupplier(ABC):
     """This class can be implemented to provide custom slot supplier behavior."""
 
-    # TODO: AbortError equivalent
     @abstractmethod
     async def reserve_slot(self, ctx: SlotReserveContext) -> SlotPermit:
         """This function is called before polling for new tasks. Your implementation must block until a
         slot is available then return a permit to use that slot.
 
-        The only acceptable exception to throw is AbortError, any other exceptions thrown will be
-        logged and ignored.
+        The only acceptable exception to throw is :py:class:`asyncio.CancelledError`, as invocations of this method may
+        be cancelled. Any other exceptions thrown will be logged and ignored.
 
         Args:
             ctx: The context for slot reservation.

temporalio/worker/_tuning.py

@@ -1,3 +1,5 @@
+import asyncio
+import logging
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from datetime import timedelta
@@ -21,6 +23,8 @@
 
 _DEFAULT_RESOURCE_ACTIVITY_MAX = 500
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass(frozen=True)
 class FixedSizeSlotSupplier:
@@ -87,6 +91,48 @@ class ResourceBasedSlotSupplier:
 ]
 
 
+class _ErrorLoggingSlotSupplier(CustomSlotSupplier):
+    def __init__(self, supplier: CustomSlotSupplier):
+        self._supplier = supplier
+
+    async def reserve_slot(self, ctx: SlotReserveContext) -> SlotPermit:
+        try:
+            return await self._supplier.reserve_slot(ctx)
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            logger.warning(
+                "Error in custom slot supplier `reserve_slot`", exc_info=True
+            )
+            # Error needs to be re-thrown here so the rust code will loop
+            raise
+
+    def try_reserve_slot(self, ctx: SlotReserveContext) -> Optional[SlotPermit]:
+        try:
+            return self._supplier.try_reserve_slot(ctx)
+        except Exception:
+            logger.warning(
+                "Error in custom slot supplier `try_reserve_slot`", exc_info=True
+            )
+            return None
+
+    def release_slot(self, ctx: SlotReleaseContext) -> None:
+        try:
+            self._supplier.release_slot(ctx)
+        except Exception:
+            logger.warning(
+                "Error in custom slot supplier `release_slot`", exc_info=True
+            )
+
+    def mark_slot_used(self, ctx: SlotMarkUsedContext) -> None:
+        try:
+            self._supplier.mark_slot_used(ctx)
+        except Exception:
+            logger.warning(
+                "Error in custom slot supplier `mark_slot_used`", exc_info=True
+            )
+
+
 def _to_bridge_slot_supplier(
     slot_supplier: SlotSupplier, kind: Literal["workflow", "activity", "local_activity"]
 ) -> temporalio.bridge.worker.SlotSupplier:
@@ -114,7 +160,9 @@ def _to_bridge_slot_supplier(
             ),
         )
     elif isinstance(slot_supplier, CustomSlotSupplier):
-        return temporalio.bridge.temporal_sdk_bridge.CustomSlotSupplier(slot_supplier)
+        return temporalio.bridge.temporal_sdk_bridge.CustomSlotSupplier(
+            _ErrorLoggingSlotSupplier(slot_supplier)
+        )
     else:
         raise TypeError(f"Unknown slot supplier type: {slot_supplier}")
 

tests/worker/test_worker.py

@@ -426,12 +426,89 @@ def reserve_asserts(self, ctx):
     # Two workflow tasks, one activity
     assert ss.used == 3
     assert ss.seen_sticky_kinds == {True, False}
-    assert ss.seen_slot_kinds == {"Workflow", "Activity", "LocalActivity"}
+    assert ss.seen_slot_kinds == {"workflow", "activity", "local-activity"}
     assert ss.seen_used_slot_kinds == {"wf", "a"}
     assert ss.seen_release_info_empty
     assert ss.seen_release_info_nonempty
 
 
+@workflow.defn
+class SimpleWorkflow:
+    @workflow.run
+    async def run(self) -> str:
+        return "hi"
+
+
+async def test_throwing_slot_supplier(client: Client, env: WorkflowEnvironment):
+    """Ensures a (mostly) broken slot supplier doesn't hose everything up"""
+
+    class ThrowingSlotSupplier(CustomSlotSupplier):
+        marked_used = False
+
+        async def reserve_slot(self, ctx: SlotReserveContext) -> SlotPermit:
+            # Hand out workflow tasks until one is used
+            if ctx.slot_type == "workflow" and not self.marked_used:
+                return SlotPermit()
+            raise ValueError("I always throw")
+
+        def try_reserve_slot(self, ctx: SlotReserveContext) -> Optional[SlotPermit]:
+            raise ValueError("I always throw")
+
+        def mark_slot_used(self, ctx: SlotMarkUsedContext) -> None:
+            raise ValueError("I always throw")
+
+        def release_slot(self, ctx: SlotReleaseContext) -> None:
+            raise ValueError("I always throw")
+
+    ss = ThrowingSlotSupplier()
+
+    tuner = WorkerTuner.create_composite(
+        workflow_supplier=ss, activity_supplier=ss, local_activity_supplier=ss
+    )
+    async with new_worker(
+        client,
+        SimpleWorkflow,
+        activities=[say_hello],
+        tuner=tuner,
+    ) as w:
+        wf1 = await client.start_workflow(
+            SimpleWorkflow.run,
+            id=f"throwing-slot-supplier-{uuid.uuid4()}",
+            task_queue=w.task_queue,
+        )
+        await wf1.result()
+
+
+async def test_blocking_slot_supplier(client: Client, env: WorkflowEnvironment):
+    class BlockingSlotSupplier(CustomSlotSupplier):
+        marked_used = False
+
+        async def reserve_slot(self, ctx: SlotReserveContext) -> SlotPermit:
+            await asyncio.get_event_loop().create_future()
+            raise ValueError("Should be unreachable")
+
+        def try_reserve_slot(self, ctx: SlotReserveContext) -> Optional[SlotPermit]:
+            return None
+
+        def mark_slot_used(self, ctx: SlotMarkUsedContext) -> None:
+            return None
+
+        def release_slot(self, ctx: SlotReleaseContext) -> None:
+            return None
+
+    ss = BlockingSlotSupplier()
+
+    tuner = WorkerTuner.create_composite(
+        workflow_supplier=ss, activity_supplier=ss, local_activity_supplier=ss
+    )
+    async with new_worker(
+        client,
+        SimpleWorkflow,
+        activities=[say_hello],
+        tuner=tuner,
+    ) as _w:
+        await asyncio.sleep(1)
+
 def create_worker(
     client: Client,
     on_fatal_error: Optional[Callable[[BaseException], Awaitable[None]]] = None,