Calling async reserve working

Sushisource · Sushisource · commit b641b9e75e50 · 2024-11-13T17:17:41.000-08:00
diff --git a/temporalio/bridge/Cargo.lock b/temporalio/bridge/Cargo.lock
diff --git a/temporalio/bridge/Cargo.toml b/temporalio/bridge/Cargo.toml
@@ -9,6 +9,7 @@ crate-type = ["cdylib"]
 
 [dependencies]
 anyhow = "1.0"
+async-trait = "0.1"
 futures = "0.3"
 log = "0.4"
 once_cell = "1.16"
@@ -17,10 +18,14 @@ prost-types = "0.13"
 pyo3 = { version = "0.20", features = ["extension-module", "abi3-py38", "anyhow"] }
 pyo3-asyncio = { version = "0.20", features = ["tokio-runtime"] }
 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-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-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" }
 tokio = "1.26"
 tokio-stream = "0.1"
 tonic = "0.12"
diff --git a/temporalio/bridge/src/lib.rs b/temporalio/bridge/src/lib.rs
@@ -45,6 +45,7 @@ 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_function(wrap_pyfunction!(new_worker, m)?)?;
     m.add_function(wrap_pyfunction!(new_replay_worker, m)?)?;
     Ok(())
diff --git a/temporalio/bridge/src/runtime.rs b/temporalio/bridge/src/runtime.rs
@@ -13,7 +13,7 @@ use temporal_sdk_core::telemetry::{
     build_otlp_metric_exporter, start_prometheus_metric_exporter, CoreLogStreamConsumer,
     MetricsCallBuffer,
 };
-use temporal_sdk_core::CoreRuntime;
+use temporal_sdk_core::{CoreRuntime, TokioRuntimeBuilder};
 use temporal_sdk_core_api::telemetry::metrics::{CoreMeter, MetricCallBufferer};
 use temporal_sdk_core_api::telemetry::{
     CoreLog, Logger, MetricTemporality, OtelCollectorOptionsBuilder,
@@ -88,6 +88,7 @@ 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();
 
@@ -118,12 +119,33 @@ 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);
+        let locals = pyo3_asyncio::TaskLocals::with_running_loop(py)?.copy_context(py)?;
+        PyResult::Ok(locals)
+    })
+    .expect("Works");
+
     // Create core runtime which starts tokio multi-thread runtime
     let mut core = CoreRuntime::new(
         telemetry_build
             .build()
             .map_err(|err| PyValueError::new_err(format!("Invalid telemetry config: {}", err)))?,
-        tokio::runtime::Builder::new_multi_thread(),
+        TokioRuntimeBuilder {
+            inner: tokio::runtime::Builder::new_multi_thread(),
+            lang_on_thread_start: Some(move || {
+                // Set task locals for each thread
+                Python::with_gil(|py| {
+                    THREAD_TASK_LOCAL.with(|r| {
+                        std::cell::OnceCell::set(r, task_locals.clone()).expect("NOT ALREADY SET");
+                    });
+                    PyResult::Ok(())
+                })
+                .expect("Setting event loop works");
+            }),
+        },
     )
     .map_err(|err| PyRuntimeError::new_err(format!("Failed initializing telemetry: {}", err)))?;
 
@@ -364,12 +386,17 @@ impl TryFrom<MetricsConfig> for Arc<dyn CoreMeter> {
 // altered to support spawning based on current Tokio runtime instead of a
 // single static one
 
-struct TokioRuntime;
+pub(crate) struct TokioRuntime;
 
 tokio::task_local! {
     static TASK_LOCALS: once_cell::unsync::OnceCell<pyo3_asyncio::TaskLocals>;
 }
 
+thread_local! {
+    pub(crate) static THREAD_TASK_LOCAL: std::cell::OnceCell<pyo3_asyncio::TaskLocals> =
+        std::cell::OnceCell::new();
+}
+
 impl pyo3_asyncio::generic::Runtime for TokioRuntime {
     type JoinError = tokio::task::JoinError;
     type JoinHandle = tokio::task::JoinHandle<()>;
@@ -378,9 +405,7 @@ impl pyo3_asyncio::generic::Runtime for TokioRuntime {
     where
         F: Future<Output = ()> + Send + 'static,
     {
-        tokio::runtime::Handle::current().spawn(async move {
-            fut.await;
-        })
+        tokio::runtime::Handle::current().spawn(fut)
     }
 }
 
diff --git a/temporalio/bridge/src/worker.rs b/temporalio/bridge/src/worker.rs
@@ -1,16 +1,21 @@
 use anyhow::Context;
 use prost::Message;
-use pyo3::exceptions::{PyException, PyRuntimeError, PyValueError};
+use pyo3::exceptions::{PyException, PyRuntimeError, PyTypeError, 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;
 use std::sync::Arc;
 use std::time::Duration;
 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;
+use temporal_sdk_core_api::worker::{
+    SlotKind, SlotMarkUsedContext, SlotReleaseContext, SlotReservationContext,
+    SlotSupplier as SlotSupplierTrait, SlotSupplierPermit,
+};
 use temporal_sdk_core_api::Worker;
 use temporal_sdk_core_protos::coresdk::workflow_completion::WorkflowActivationCompletion;
 use temporal_sdk_core_protos::coresdk::{ActivityHeartbeat, ActivityTaskCompletion};
@@ -20,6 +25,7 @@ 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);
 
@@ -63,6 +69,7 @@ pub struct TunerHolder {
 pub enum SlotSupplier {
     FixedSize(FixedSizeSlotSupplier),
     ResourceBased(ResourceBasedSlotSupplier),
+    Custom(CustomSlotSupplier),
 }
 
 #[derive(FromPyObject)]
@@ -79,6 +86,125 @@ pub struct ResourceBasedSlotSupplier {
     tuner_config: ResourceBasedTunerConfig,
 }
 
+#[pyclass]
+pub struct SlotReserveCtx {
+    slot_type: String, // TODO: Real type
+    task_queue: String,
+    worker_identity: String,
+    worker_build_id: String,
+    is_sticky: bool,
+}
+
+impl SlotReserveCtx {
+    fn from_ctx(slot_type: String, ctx: &dyn SlotReservationContext) -> Self {
+        SlotReserveCtx {
+            slot_type,
+            task_queue: ctx.task_queue().to_string(),
+            worker_identity: ctx.worker_identity().to_string(),
+            worker_build_id: ctx.worker_build_id().to_string(),
+            is_sticky: ctx.is_sticky(),
+        }
+    }
+}
+
+#[pyclass]
+pub struct SlotMarkUsedCtx {}
+
+#[pyclass]
+pub struct SlotReleaseCtx {}
+
+#[pyclass]
+#[derive(Clone)]
+pub struct CustomSlotSupplier {
+    inner: PyObject,
+}
+
+struct CustomSlotSupplierOfType<SK: SlotKind> {
+    inner: PyObject,
+    _phantom: PhantomData<SK>,
+}
+
+#[pymethods]
+impl CustomSlotSupplier {
+    #[new]
+    fn new(inner: PyObject) -> Self {
+        CustomSlotSupplier { inner }
+    }
+}
+
+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| {
+            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,
+                ),),
+            )?;
+            THREAD_TASK_LOCAL
+                .with(|tl| pyo3_asyncio::into_future_with_locals(tl.get().unwrap(), called))
+        })
+        .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 release_slot(&self, _ctx: &dyn SlotReleaseContext<SlotKind = Self::SlotKind>) {
+        self.call_method("release_slot", SlotReleaseCtx {}, |_, _| ())
+    }
+
+    fn available_slots(&self) -> Option<usize> {
+        None
+    }
+}
+
 #[derive(FromPyObject, Clone, Copy, PartialEq)]
 pub struct ResourceBasedTunerConfig {
     target_memory_usage: f64,
@@ -369,7 +495,9 @@ impl TryFrom<TunerHolder> for temporal_sdk_core::TunerHolder {
     }
 }
 
-impl<SK: SlotKind> TryFrom<SlotSupplier> for temporal_sdk_core::SlotSupplierOptions<SK> {
+impl<SK: SlotKind + Send + Sync + 'static> TryFrom<SlotSupplier>
+    for temporal_sdk_core::SlotSupplierOptions<SK>
+{
     type Error = PyErr;
 
     fn try_from(supplier: SlotSupplier) -> PyResult<temporal_sdk_core::SlotSupplierOptions<SK>> {
@@ -386,6 +514,12 @@ impl<SK: SlotKind> TryFrom<SlotSupplier> for temporal_sdk_core::SlotSupplierOpti
                     ),
                 )
             }
+            SlotSupplier::Custom(cs) => temporal_sdk_core::SlotSupplierOptions::Custom(Arc::new(
+                CustomSlotSupplierOfType::<SK> {
+                    inner: cs.inner,
+                    _phantom: PhantomData,
+                },
+            )),
         })
     }
 }
diff --git a/temporalio/worker/_tuning.py b/temporalio/worker/_tuning.py
@@ -6,6 +6,7 @@
 from typing_extensions import TypeAlias
 
 import temporalio.bridge.worker
+import temporalio.bridge.temporal_sdk_bridge
 from temporalio.bridge.worker import (
     CustomSlotSupplier,
     SlotMarkUsedContext,
@@ -108,6 +109,8 @@ def _to_bridge_slot_supplier(
                 slot_supplier.tuner_config.target_cpu_usage,
             ),
         )
+    elif isinstance(slot_supplier, CustomSlotSupplier):
+        return temporalio.bridge.temporal_sdk_bridge.CustomSlotSupplier(slot_supplier)
     else:
         raise TypeError(f"Unknown slot supplier type: {slot_supplier}")
 
diff --git a/tests/worker/test_worker.py b/tests/worker/test_worker.py
@@ -344,22 +344,25 @@ async def test_warns_when_workers_too_lot(client: Client, env: WorkflowEnvironme
 async def test_custom_slot_supplier(client: Client, env: WorkflowEnvironment):
     class MySlotSupplier(CustomSlotSupplier):
         async def reserve_slot(self, ctx: SlotReserveContext) -> SlotPermit:
+            print("Reserving slot")
             return SlotPermit()
 
         def try_reserve_slot(self, ctx: SlotReserveContext) -> Optional[SlotPermit]:
-            pass
+            print("Try Reserving slot")
+            return None
 
         def mark_slot_used(self, ctx: SlotMarkUsedContext) -> None:
-            pass
+            print("Marking slot used")
 
         def release_slot(self, ctx: SlotReleaseContext) -> None:
-            pass
+            print("Releasing slot")
 
     ss = MySlotSupplier()
 
     tuner = WorkerTuner.create_composite(
         workflow_supplier=ss, activity_supplier=ss, local_activity_supplier=ss
     )
+    print("!!!!!!!!1 About to create new worker, event loop is", asyncio.get_event_loop())
     async with new_worker(
         client,
         WaitOnSignalWorkflow,
