Passing through contexts working

Author: Sushisource

temporalio/bridge/Cargo.toml

@@ -25,7 +25,7 @@ pythonize = "0.20"
 temporal-client = { version = "0.1.0", path = "../../../sdk-core/client" }
 temporal-sdk-core = { version = "0.1.0", path = "../../../sdk-core/core", features = ["ephemeral-server"] }
 temporal-sdk-core-api = { version = "0.1.0", path = "../../../sdk-core/core-api" }
-temporal-sdk-core-protos = { version = "0.1.0", path = "../../../sdk-core/sdk-core-protos" }
+temporal-sdk-core-protos = { version = "0.1.0", path = "../../../sdk-core/sdk-core-protos", features = ["serde_serialize"] }
 tokio = "1.26"
 tokio-stream = "0.1"
 tonic = "0.12"

temporalio/bridge/src/lib.rs

@@ -46,6 +46,12 @@ fn temporal_sdk_bridge(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<worker::WorkerRef>()?;
     m.add_class::<worker::HistoryPusher>()?;
     m.add_class::<worker::CustomSlotSupplier>()?;
+    m.add_class::<worker::SlotReserveCtx>()?;
+    m.add_class::<worker::SlotReleaseCtx>()?;
+    m.add_class::<worker::SlotMarkUsedCtx>()?;
+    m.add_class::<worker::WorkflowSlotInfo>()?;
+    m.add_class::<worker::ActivitySlotInfo>()?;
+    m.add_class::<worker::LocalActivitySlotInfo>()?;
     m.add_function(wrap_pyfunction!(new_worker, m)?)?;
     m.add_function(wrap_pyfunction!(new_replay_worker, m)?)?;
     Ok(())

temporalio/bridge/src/runtime.rs

@@ -88,7 +88,6 @@ const FORWARD_LOG_BUFFER_SIZE: usize = 2048;
 const FORWARD_LOG_MAX_FREQ_MS: u64 = 10;
 
 pub fn init_runtime(telemetry_config: TelemetryConfig) -> PyResult<RuntimeRef> {
-    dbg!("Initting runtime");
     // Have to build/start telemetry config pieces
     let mut telemetry_build = TelemetryOptionsBuilder::default();
 
@@ -120,9 +119,7 @@ pub fn init_runtime(telemetry_config: TelemetryConfig) -> PyResult<RuntimeRef> {
     }
 
     let task_locals = Python::with_gil(|py| {
-        let asyncio = py.import("asyncio")?;
-        let event_loop = asyncio.call_method0("get_event_loop")?;
-        dbg!(&event_loop);
+        // Event loop is assumed to be running at this point
         let locals = pyo3_asyncio::TaskLocals::with_running_loop(py)?.copy_context(py)?;
         PyResult::Ok(locals)
     })
@@ -137,7 +134,7 @@ pub fn init_runtime(telemetry_config: TelemetryConfig) -> PyResult<RuntimeRef> {
             inner: tokio::runtime::Builder::new_multi_thread(),
             lang_on_thread_start: Some(move || {
                 // Set task locals for each thread
-                Python::with_gil(|py| {
+                Python::with_gil(|_| {
                     THREAD_TASK_LOCAL.with(|r| {
                         std::cell::OnceCell::set(r, task_locals.clone()).expect("NOT ALREADY SET");
                     });

temporalio/bridge/src/worker.rs

@@ -1,9 +1,8 @@
 use anyhow::Context;
 use prost::Message;
-use pyo3::exceptions::{PyException, PyRuntimeError, PyTypeError, PyValueError};
+use pyo3::exceptions::{PyException, PyRuntimeError, PyValueError};
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyTuple};
-use pyo3_asyncio::generic::ContextExt;
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::marker::PhantomData;
@@ -13,8 +12,8 @@ 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::{
-    SlotKind, SlotMarkUsedContext, SlotReleaseContext, SlotReservationContext,
-    SlotSupplier as SlotSupplierTrait, SlotSupplierPermit,
+    SlotInfo, SlotInfoTrait, SlotKind, SlotMarkUsedContext, SlotReleaseContext,
+    SlotReservationContext, SlotSupplier as SlotSupplierTrait, SlotSupplierPermit,
 };
 use temporal_sdk_core_api::Worker;
 use temporal_sdk_core_protos::coresdk::workflow_completion::WorkflowActivationCompletion;
@@ -25,7 +24,6 @@ use tokio_stream::wrappers::ReceiverStream;
 
 use crate::client;
 use crate::runtime;
-use crate::runtime::{TokioRuntime, THREAD_TASK_LOCAL};
 
 pyo3::create_exception!(temporal_sdk_bridge, PollShutdownError, PyException);
 
@@ -88,11 +86,16 @@ pub struct ResourceBasedSlotSupplier {
 
 #[pyclass]
 pub struct SlotReserveCtx {
-    slot_type: String, // TODO: Real type
-    task_queue: String,
-    worker_identity: String,
-    worker_build_id: String,
-    is_sticky: bool,
+    #[pyo3(get)]
+    pub slot_type: String, // TODO: Real type
+    #[pyo3(get)]
+    pub task_queue: String,
+    #[pyo3(get)]
+    pub worker_identity: String,
+    #[pyo3(get)]
+    pub worker_build_id: String,
+    #[pyo3(get)]
+    pub is_sticky: bool,
 }
 
 impl SlotReserveCtx {
@@ -108,10 +111,60 @@ impl SlotReserveCtx {
 }
 
 #[pyclass]
-pub struct SlotMarkUsedCtx {}
+pub struct SlotMarkUsedCtx {
+    #[pyo3(get)]
+    slot_info: PyObject,
+    #[pyo3(get)]
+    permit: PyObject,
+}
 
+// NOTE: this is dumb because we already have the generated proto code, we just can't use
+//   it b/c it's not pyclassable. In theory maybe we could compile-flag enable it in the core
+//   protos crate but... that's a lot for just this. Maybe if there are other use cases.
+
+#[pyclass]
+pub struct WorkflowSlotInfo {
+    #[pyo3(get)]
+    pub workflow_type: String,
+    #[pyo3(get)]
+    pub is_sticky: bool,
+}
+#[pyclass]
+pub struct ActivitySlotInfo {
+    #[pyo3(get)]
+    pub activity_type: String,
+}
 #[pyclass]
-pub struct SlotReleaseCtx {}
+pub struct LocalActivitySlotInfo {
+    #[pyo3(get)]
+    pub activity_type: String,
+}
+
+#[pyclass]
+pub struct SlotReleaseCtx {
+    #[pyo3(get)]
+    slot_info: Option<PyObject>,
+    #[pyo3(get)]
+    permit: PyObject,
+}
+
+fn slot_info_to_py_obj(py: Python<'_>, info: SlotInfo) -> PyObject {
+    match info {
+        SlotInfo::Workflow(w) => WorkflowSlotInfo {
+            workflow_type: w.workflow_type.clone(),
+            is_sticky: w.is_sticky,
+        }
+        .into_py(py),
+        SlotInfo::Activity(a) => ActivitySlotInfo {
+            activity_type: a.activity_type.clone(),
+        }
+        .into_py(py),
+        SlotInfo::LocalActivity(a) => LocalActivitySlotInfo {
+            activity_type: a.activity_type.clone(),
+        }
+        .into_py(py),
+    }
+}
 
 #[pyclass]
 #[derive(Clone)]
@@ -132,72 +185,94 @@ impl CustomSlotSupplier {
     }
 }
 
-impl<SK: SlotKind> CustomSlotSupplierOfType<SK> {
-    fn call_method<P: IntoPy<PyObject>, F: FnOnce(Python<'_>, &PyAny) -> FR, FR>(
-        &self,
-        method_name: &str,
-        arg: P,
-        post_closure: F,
-    ) -> FR {
-        Python::with_gil(|py| {
-            let py_obj = self.inner.as_ref(py);
-            let method = py_obj
-                .getattr(method_name)
-                .map_err(|_| {
-                    PyTypeError::new_err(format!(
-                        "CustomSlotSupplier must implement '{}' method",
-                        method_name
-                    ))
-                })
-                .expect("TODO");
-
-            post_closure(py, method.call((arg.into_py(py),), None).expect("TODO"))
-        })
-    }
-}
-
 #[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!("Trying to reserve slot");
-        let pypermit = Python::with_gil(|py| {
+        loop {
+            let pypermit = 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,
+                    ),),
+                )?;
+                runtime::THREAD_TASK_LOCAL
+                    .with(|tl| pyo3_asyncio::into_future_with_locals(tl.get().unwrap(), called))
+            })
+            .expect("TODO")
+            .await;
+            match pypermit {
+                Ok(p) => {
+                    return SlotSupplierPermit::with_user_data(p);
+                }
+                Err(e) => {
+                    dbg!("Error in reserve_slot", e);
+                }
+            }
+        }
+    }
+
+    fn try_reserve_slot(&self, ctx: &dyn SlotReservationContext) -> Option<SlotSupplierPermit> {
+        Python::with_gil(|py| {
             let py_obj = self.inner.as_ref(py);
-            let called = py_obj.call_method1(
-                "reserve_slot",
+            let pa = py_obj.call_method1(
+                "try_reserve_slot",
                 (SlotReserveCtx::from_ctx(
                     Self::SlotKind::kind().to_string(),
                     ctx,
                 ),),
             )?;
-            THREAD_TASK_LOCAL
-                .with(|tl| pyo3_asyncio::into_future_with_locals(tl.get().unwrap(), called))
+
+            if pa.is_none() {
+                return Ok(None);
+            }
+            PyResult::Ok(Some(SlotSupplierPermit::with_user_data(pa.into_py(py))))
         })
         .expect("TODO")
-        .await;
-        SlotSupplierPermit::with_user_data(pypermit)
     }
 
-    fn try_reserve_slot(&self, ctx: &dyn SlotReservationContext) -> Option<SlotSupplierPermit> {
-        self.call_method(
-            "try_reserve_slot",
-            SlotReserveCtx::from_ctx(Self::SlotKind::kind().to_string(), ctx),
-            |py, pa| {
-                if pa.is_none() {
-                    return None;
-                }
-                Some(SlotSupplierPermit::with_user_data(pa.into_py(py)))
-            },
-        )
-    }
-
-    fn mark_slot_used(&self, _ctx: &dyn SlotMarkUsedContext<SlotKind = Self::SlotKind>) {
-        self.call_method("mark_slot_used", SlotMarkUsedCtx {}, |_, _| ())
+    fn mark_slot_used(&self, ctx: &dyn SlotMarkUsedContext<SlotKind = Self::SlotKind>) {
+        Python::with_gil(|py| {
+            let permit = ctx
+                .permit()
+                .user_data::<PyObject>()
+                .cloned()
+                .unwrap_or_else(|| py.None());
+            let py_obj = self.inner.as_ref(py);
+            py_obj.call_method1(
+                "mark_slot_used",
+                (SlotMarkUsedCtx {
+                    slot_info: slot_info_to_py_obj(py, ctx.info().downcast()),
+                    permit,
+                },),
+            )?;
+            PyResult::Ok(())
+        })
+        .expect("TODO");
     }
 
-    fn release_slot(&self, _ctx: &dyn SlotReleaseContext<SlotKind = Self::SlotKind>) {
-        self.call_method("release_slot", SlotReleaseCtx {}, |_, _| ())
+    fn release_slot(&self, ctx: &dyn SlotReleaseContext<SlotKind = Self::SlotKind>) {
+        Python::with_gil(|py| {
+            let permit = ctx
+                .permit()
+                .user_data::<PyObject>()
+                .cloned()
+                .unwrap_or_else(|| py.None());
+            let py_obj = self.inner.as_ref(py);
+            py_obj.call_method1(
+                "release_slot",
+                (SlotReleaseCtx {
+                    slot_info: ctx.info().map(|i| slot_info_to_py_obj(py, i.downcast())),
+                    permit,
+                },),
+            )?;
+            PyResult::Ok(())
+        })
+        .expect("TODO");
     }
 
     fn available_slots(&self) -> Option<usize> {