Eager Workflow Start (#606)

.gitignore

@@ -6,6 +6,7 @@ Cargo.lock
 
 /.idea/
 *.iml
+*~
 
 # Ignore generated protobuf files
 src/protos/*.rs

client/Cargo.toml

@@ -23,6 +23,7 @@ once_cell = "1.13"
 opentelemetry = { workspace = true, features = ["metrics"] }
 parking_lot = "0.12"
 prost-types = "0.11"
+rand = "0.8.3"
 thiserror = "1.0"
 tokio = "1.1"
 tonic = { workspace = true, features = ["tls", "tls-roots"] }

client/src/lib.rs

@@ -6,10 +6,10 @@
 
 #[macro_use]
 extern crate tracing;
-
 mod metrics;
 mod raw;
 mod retry;
+mod worker_registry;
 mod workflow_handle;
 
 pub use crate::retry::{CallType, RetryClient, RETRYABLE_ERROR_CODES};
@@ -23,6 +23,7 @@ pub use temporal_sdk_core_protos::temporal::api::{
     },
 };
 pub use tonic;
+pub use worker_registry::{Slot, SlotManager, SlotProvider, WorkerRegistry};
 pub use workflow_handle::{WorkflowExecutionInfo, WorkflowExecutionResult};
 
 use crate::{
@@ -280,6 +281,7 @@ pub struct ConfiguredClient<C> {
     headers: Arc<RwLock<HashMap<String, String>>>,
     /// Capabilities as read from the `get_system_info` RPC call made on client connection
     capabilities: Option<get_system_info_response::Capabilities>,
+    workers: Arc<RwLock<SlotManager>>,
 }
 
 impl<C> ConfiguredClient<C> {
@@ -299,6 +301,11 @@ impl<C> ConfiguredClient<C> {
     pub fn capabilities(&self) -> Option<&get_system_info_response::Capabilities> {
         self.capabilities.as_ref()
     }
+
+    /// Returns a cloned reference to a registry with workers using this client instance
+    pub fn workers(&self) -> Arc<RwLock<SlotManager>> {
+        self.workers.clone()
+    }
 }
 
 // The configured client is effectively a "smart" (dumb) pointer
@@ -377,6 +384,7 @@ impl ClientOptions {
             client: TemporalServiceClient::new(svc),
             options: Arc::new(self.clone()),
             capabilities: None,
+            workers: Arc::new(RwLock::new(SlotManager::new())),
         };
         match client
             .get_system_info(GetSystemInfoRequest::default())
@@ -974,6 +982,10 @@ pub struct WorkflowOptions {
 
     /// Optionally associate extra search attributes with a workflow
     pub search_attributes: Option<HashMap<String, Payload>>,
+
+    /// Optionally enable Eager Workflow Start, a latency optimization using local workers
+    /// NOTE: Experimental and incompatible with versioning with BuildIDs
+    pub enable_eager_workflow_start: bool,
 }
 
 #[async_trait::async_trait]
@@ -988,7 +1000,7 @@ impl WorkflowClientTrait for Client {
         options: WorkflowOptions,
     ) -> Result<StartWorkflowExecutionResponse> {
         Ok(WorkflowService::start_workflow_execution(
-            &mut self.inner.client.clone(),
+            &mut self.inner.clone(),
             StartWorkflowExecutionRequest {
                 namespace: self.namespace.clone(),
                 input: input.into_payloads(),
@@ -1006,10 +1018,11 @@ impl WorkflowClientTrait for Client {
                 workflow_execution_timeout: options
                     .execution_timeout
                     .and_then(|d| d.try_into().ok()),
-                workflow_run_timeout: options.execution_timeout.and_then(|d| d.try_into().ok()),
+                workflow_run_timeout: options.run_timeout.and_then(|d| d.try_into().ok()),
                 workflow_task_timeout: options.task_timeout.and_then(|d| d.try_into().ok()),
                 search_attributes: options.search_attributes.and_then(|d| d.try_into().ok()),
                 cron_schedule: options.cron_schedule.unwrap_or_default(),
+                request_eager_execution: options.enable_eager_workflow_start,
                 ..Default::default()
             },
         )
@@ -1169,9 +1182,7 @@ impl WorkflowClientTrait for Client {
                 workflow_execution_timeout: workflow_options
                     .execution_timeout
                     .and_then(|d| d.try_into().ok()),
-                workflow_run_timeout: workflow_options
-                    .execution_timeout
-                    .and_then(|d| d.try_into().ok()),
+                workflow_run_timeout: workflow_options.run_timeout.and_then(|d| d.try_into().ok()),
                 workflow_task_timeout: workflow_options
                     .task_timeout
                     .and_then(|d| d.try_into().ok()),

client/src/raw.rs

@@ -5,10 +5,13 @@
 use crate::{
     metrics::{namespace_kv, task_queue_kv},
     raw::sealed::RawClientLike,
+    worker_registry::{Slot, SlotManager},
     Client, ConfiguredClient, InterceptedMetricsSvc, RetryClient, TemporalServiceClient,
     LONG_POLL_TIMEOUT,
 };
 use futures::{future::BoxFuture, FutureExt, TryFutureExt};
+use parking_lot::RwLock;
+use std::sync::Arc;
 use temporal_sdk_core_api::telemetry::metrics::MetricKeyValue;
 use temporal_sdk_core_protos::{
     grpc::health::v1::{health_client::HealthClient, *},
@@ -55,6 +58,9 @@ pub(super) mod sealed {
         /// Return a mutable ref to the health service client instance
         fn health_client_mut(&mut self) -> &mut HealthClient<Self::SvcType>;
 
+        /// Return a registry with workers using this client instance
+        fn get_workers_info(&self) -> Option<Arc<RwLock<SlotManager>>>;
+
         async fn call<F, Req, Resp>(
             &mut self,
             _call_name: &'static str,
@@ -111,6 +117,10 @@ where
         self.get_client_mut().health_client_mut()
     }
 
+    fn get_workers_info(&self) -> Option<Arc<RwLock<SlotManager>>> {
+        self.get_client().get_workers_info()
+    }
+
     async fn call<F, Req, Resp>(
         &mut self,
         call_name: &'static str,
@@ -173,6 +183,10 @@ where
     fn health_client_mut(&mut self) -> &mut HealthClient<Self::SvcType> {
         self.health_svc_mut()
     }
+
+    fn get_workers_info(&self) -> Option<Arc<RwLock<SlotManager>>> {
+        None
+    }
 }
 
 impl<T> RawClientLike for ConfiguredClient<TemporalServiceClient<T>>
@@ -216,6 +230,10 @@ where
     fn health_client_mut(&mut self) -> &mut HealthClient<Self::SvcType> {
         self.client.health_client_mut()
     }
+
+    fn get_workers_info(&self) -> Option<Arc<RwLock<SlotManager>>> {
+        Some(self.workers())
+    }
 }
 
 impl RawClientLike for Client {
@@ -252,6 +270,10 @@ impl RawClientLike for Client {
     fn health_client_mut(&mut self) -> &mut HealthClient<Self::SvcType> {
         self.inner.health_client_mut()
     }
+
+    fn get_workers_info(&self) -> Option<Arc<RwLock<SlotManager>>> {
+        self.inner.get_workers_info()
+    }
 }
 
 /// Helper for cloning a tonic request as long as the inner message may be cloned.
@@ -356,10 +378,27 @@ macro_rules! proxy {
             self.call(stringify!($method), fact, request.into_request())
         }
     };
+    ($client_type:tt, $client_meth:ident, $method:ident, $req:ty, $resp:ty, $closure_before:expr, $closure_after:expr) => {
+        #[doc = concat!("See [", stringify!($client_type), "::", stringify!($method), "]")]
+        fn $method(
+            &mut self,
+            request: impl tonic::IntoRequest<$req>,
+        ) -> BoxFuture<Result<tonic::Response<$resp>, tonic::Status>> {
+            #[allow(unused_mut)]
+            let fact = |c: &mut Self, mut req: tonic::Request<$req>| {
+                let data = type_closure_two_arg(&mut req, c.get_workers_info().unwrap(), $closure_before);
+                let mut c = c.$client_meth().clone();
+                async move {
+                    type_closure_two_arg(c.$method(req).await, data, $closure_after)
+                }.boxed()
+            };
+            self.call(stringify!($method), fact, request.into_request())
+        }
+    };
 }
 macro_rules! proxier {
     ( $trait_name:ident; $impl_list_name:ident; $client_type:tt; $client_meth:ident;
-      $(($method:ident, $req:ty, $resp:ty $(, $closure:expr)?  );)* ) => {
+      $(($method:ident, $req:ty, $resp:ty $(, $closure:expr)? $(, $_dummy:ident, $closure_after:expr)? );)* ) => {
         #[cfg(test)]
         const $impl_list_name: &'static [&'static str] = &[$(stringify!($method)),*];
         /// Trait version of the generated client with modifications to attach appropriate metric
@@ -377,7 +416,7 @@ macro_rules! proxier {
                 as tonic::codegen::Body>::Error: Into<tonic::codegen::StdError> + Send,
         {
             $(
-               proxy!($client_type, $client_meth, $method, $req, $resp $(,$closure)*);
+               proxy!($client_type, $client_meth, $method, $req, $resp $(,$closure)* $(,$closure_after)*);
             )*
         }
     };
@@ -388,6 +427,10 @@ fn type_closure_arg<T, R>(arg: T, f: impl FnOnce(T) -> R) -> R {
     f(arg)
 }
 
+fn type_closure_two_arg<T, R, S>(arg1: R, arg2: T, f: impl FnOnce(R, T) -> S) -> S {
+    f(arg1, arg2)
+}
+
 proxier! {
     WorkflowService; ALL_IMPLEMENTED_WORKFLOW_SERVICE_RPCS; WorkflowServiceClient; workflow_client_mut;
     (
@@ -435,10 +478,37 @@ proxier! {
         start_workflow_execution,
         StartWorkflowExecutionRequest,
         StartWorkflowExecutionResponse,
-        |r| {
+        |r, workers| {
+            let mut slot: Option<Box<dyn Slot>> = None;
             let mut labels = AttachMetricLabels::namespace(r.get_ref().namespace.clone());
             labels.task_q(r.get_ref().task_queue.clone());
             r.extensions_mut().insert(labels);
+            let req_mut = r.get_mut();
+            if req_mut.request_eager_execution {
+                let namespace = req_mut.namespace.clone();
+                let task_queue = req_mut.task_queue.clone().unwrap().name.clone();
+                match workers.read().try_reserve_wft_slot(namespace, task_queue) {
+                    Some(s) => slot = Some(s),
+                    None => req_mut.request_eager_execution = false
+                }
+            }
+            slot
+        },
+        start_workflow_execution, // Disambiguate the macro
+        |resp, slot| {
+            if let Some(mut s) = slot {
+                if resp.is_ok() {
+                    let response_inner = resp.as_ref().unwrap().get_ref().clone();
+                    if let Some(task) = response_inner.eager_workflow_task {
+                        if let Err(e) = s.schedule_wft(task) {
+                            // This is a latency issue, i.e., the client does not need to handle
+                            //  this error, because the WFT will be retried after a timeout.
+                            warn!(details = ?e, "Eager workflow task rejected by worker.");
+                        }
+                    }
+                }
+            }
+            resp
         }
     );
     (