@temporalio/core-bridge 1.8.5 → 1.9.0-rc.0

package/sdk-core/client/src/metrics.rs

@@ -1,69 +1,81 @@
 use crate::{AttachMetricLabels, LONG_POLL_METHOD_NAMES};
 use futures::{future::BoxFuture, FutureExt};
-use opentelemetry::{
-    metrics::{Counter, Histogram},
-    KeyValue,
-};
 use std::{
     sync::Arc,
     task::{Context, Poll},
     time::{Duration, Instant},
 };
+use temporal_sdk_core_api::telemetry::metrics::{
+    CoreMeter, Counter, Histogram, MetricAttributes, MetricKeyValue, MetricParameters,
+    TemporalMeter,
+};
 use tonic::{body::BoxBody, transport::Channel};
 use tower::Service;
 
 /// Used to track context associated with metrics, and record/update them
 // Possible improvement: make generic over some type tag so that methods are only exposed if the
 // appropriate k/vs have already been set.
-#[derive(Clone, Debug)]
+#[derive(Clone, derive_more::DebugCustom)]
+#[debug(fmt = "MetricsContext {{ attribs: {kvs:?}, poll_is_long: {poll_is_long} }}")]
 pub struct MetricsContext {
-    ctx: opentelemetry::Context,
-    kvs: Arc<Vec<KeyValue>>,
+    meter: Arc<dyn CoreMeter>,
+    kvs: MetricAttributes,
     poll_is_long: bool,
 
-    svc_request: Counter<u64>,
-    svc_request_failed: Counter<u64>,
-    long_svc_request: Counter<u64>,
-    long_svc_request_failed: Counter<u64>,
-
-    svc_request_latency: Histogram<u64>,
-    long_svc_request_latency: Histogram<u64>,
-}
+    svc_request: Arc<dyn Counter>,
+    svc_request_failed: Arc<dyn Counter>,
+    long_svc_request: Arc<dyn Counter>,
+    long_svc_request_failed: Arc<dyn Counter>,
 
-/// Things that can provide metrics for the client implement this. Trait exists to avoid having
-/// to make a whole new lower-level crate just for a tiny shared wrapper around OTel meters.
-pub trait ClientMetricProvider: Send + Sync {
-    /// Construct a counter metric
-    fn counter(&self, name: &'static str) -> Counter<u64>;
-    /// Construct a histogram metric
-    fn histogram(&self, name: &'static str) -> Histogram<u64>;
+    svc_request_latency: Arc<dyn Histogram>,
+    long_svc_request_latency: Arc<dyn Histogram>,
 }
 
 impl MetricsContext {
-    pub(crate) fn new(kvs: Vec<KeyValue>, metric_provider: &dyn ClientMetricProvider) -> Self {
+    pub(crate) fn new(tm: TemporalMeter) -> Self {
+        let meter = tm.inner;
         Self {
-            ctx: opentelemetry::Context::current(),
-            kvs: Arc::new(kvs),
+            kvs: meter.new_attributes(tm.default_attribs),
             poll_is_long: false,
-            svc_request: metric_provider.counter("request"),
-            svc_request_failed: metric_provider.counter("request_failure"),
-            long_svc_request: metric_provider.counter("long_request"),
-            long_svc_request_failed: metric_provider.counter("long_request_failure"),
-            svc_request_latency: metric_provider.histogram("request_latency"),
-            long_svc_request_latency: metric_provider.histogram("long_request_latency"),
+            svc_request: meter.counter(MetricParameters {
+                name: "request".into(),
+                description: "Count of client request successes by rpc name".into(),
+                unit: "".into(),
+            }),
+            svc_request_failed: meter.counter(MetricParameters {
+                name: "request_failure".into(),
+                description: "Count of client request failures by rpc name".into(),
+                unit: "".into(),
+            }),
+            long_svc_request: meter.counter(MetricParameters {
+                name: "long_request".into(),
+                description: "Count of long-poll request successes by rpc name".into(),
+                unit: "".into(),
+            }),
+            long_svc_request_failed: meter.counter(MetricParameters {
+                name: "long_request_failure".into(),
+                description: "Count of long-poll request failures by rpc name".into(),
+                unit: "".into(),
+            }),
+            svc_request_latency: meter.histogram(MetricParameters {
+                name: "request_latency".into(),
+                unit: "ms".into(),
+                description: "Histogram of client request latencies".into(),
+            }),
+            long_svc_request_latency: meter.histogram(MetricParameters {
+                name: "long_request_latency".into(),
+                unit: "ms".into(),
+                description: "Histogram of client long-poll request latencies".into(),
+            }),
+            meter,
         }
     }
 
-    /// Extend an existing metrics context with new attributes, returning a new one
-    pub(crate) fn with_new_attrs(&self, new_kvs: impl IntoIterator<Item = KeyValue>) -> Self {
-        let mut r = self.clone();
-        r.add_new_attrs(new_kvs);
-        r
-    }
-
-    /// Add new attributes to the context, mutating it
-    pub(crate) fn add_new_attrs(&mut self, new_kvs: impl IntoIterator<Item = KeyValue>) {
-        Arc::make_mut(&mut self.kvs).extend(new_kvs);
+    /// Mutate this metrics context with new attributes
+    pub(crate) fn with_new_attrs(&mut self, new_kvs: impl IntoIterator<Item = MetricKeyValue>) {
+        self.kvs = self
+            .meter
+            .extend_attributes(self.kvs.clone(), new_kvs.into());
     }
 
     pub(crate) fn set_is_long_poll(&mut self) {
@@ -73,18 +85,18 @@ impl MetricsContext {
     /// A request to the temporal service was made
     pub(crate) fn svc_request(&self) {
         if self.poll_is_long {
-            self.long_svc_request.add(&self.ctx, 1, &self.kvs);
+            self.long_svc_request.add(1, &self.kvs);
         } else {
-            self.svc_request.add(&self.ctx, 1, &self.kvs);
+            self.svc_request.add(1, &self.kvs);
         }
     }
 
     /// A request to the temporal service failed
     pub(crate) fn svc_request_failed(&self) {
         if self.poll_is_long {
-            self.long_svc_request_failed.add(&self.ctx, 1, &self.kvs);
+            self.long_svc_request_failed.add(1, &self.kvs);
         } else {
-            self.svc_request_failed.add(&self.ctx, 1, &self.kvs);
+            self.svc_request_failed.add(1, &self.kvs);
         }
     }
 
@@ -92,10 +104,10 @@ impl MetricsContext {
     pub(crate) fn record_svc_req_latency(&self, dur: Duration) {
         if self.poll_is_long {
             self.long_svc_request_latency
-                .record(&self.ctx, dur.as_millis() as u64, &self.kvs);
+                .record(dur.as_millis() as u64, &self.kvs);
         } else {
             self.svc_request_latency
-                .record(&self.ctx, dur.as_millis() as u64, &self.kvs);
+                .record(dur.as_millis() as u64, &self.kvs);
         }
     }
 }
@@ -104,16 +116,16 @@ const KEY_NAMESPACE: &str = "namespace";
 const KEY_SVC_METHOD: &str = "operation";
 const KEY_TASK_QUEUE: &str = "task_queue";
 
-pub(crate) fn namespace_kv(ns: String) -> KeyValue {
-    KeyValue::new(KEY_NAMESPACE, ns)
+pub(crate) fn namespace_kv(ns: String) -> MetricKeyValue {
+    MetricKeyValue::new(KEY_NAMESPACE, ns)
 }
 
-pub(crate) fn task_queue_kv(tq: String) -> KeyValue {
-    KeyValue::new(KEY_TASK_QUEUE, tq)
+pub(crate) fn task_queue_kv(tq: String) -> MetricKeyValue {
+    MetricKeyValue::new(KEY_TASK_QUEUE, tq)
 }
 
-pub(crate) fn svc_operation(op: String) -> KeyValue {
-    KeyValue::new(KEY_SVC_METHOD, op)
+pub(crate) fn svc_operation(op: String) -> MetricKeyValue {
+    MetricKeyValue::new(KEY_SVC_METHOD, op)
 }
 
 /// Implements metrics functionality for gRPC (really, any http) calls
@@ -137,19 +149,18 @@ impl Service<http::Request<BoxBody>> for GrpcMetricSvc {
         let metrics = self
             .metrics
             .clone()
-            .map(|m| {
+            .map(|mut m| {
                 // Attach labels from client wrapper
                 if let Some(other_labels) = req.extensions_mut().remove::<AttachMetricLabels>() {
                     m.with_new_attrs(other_labels.labels)
-                } else {
-                    m
                 }
+                m
             })
             .and_then(|mut metrics| {
                 // Attach method name label if possible
                 req.uri().to_string().rsplit_once('/').map(|split_tup| {
                     let method_name = split_tup.1;
-                    metrics.add_new_attrs([svc_operation(method_name.to_string())]);
+                    metrics.with_new_attrs([svc_operation(method_name.to_string())]);
                     if LONG_POLL_METHOD_NAMES.contains(&method_name) {
                         metrics.set_is_long_poll();
                     }

package/sdk-core/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 std::sync::Arc;
+use temporal_sdk_core_api::telemetry::metrics::MetricKeyValue;
 use temporal_sdk_core_protos::{
     grpc::health::v1::{health_client::HealthClient, *},
     temporal::api::{
@@ -25,22 +28,37 @@ use tonic::{
 pub(super) mod sealed {
     use super::*;
 
-    /// Something that has a workflow service client
+    /// Something that has access to the raw grpc services
     #[async_trait::async_trait]
     pub trait RawClientLike: Send {
         type SvcType: Send + Sync + Clone + 'static;
 
-        /// Return the workflow service client instance
-        fn workflow_client(&mut self) -> &mut WorkflowServiceClient<Self::SvcType>;
+        /// Return a ref to the workflow service client instance
+        fn workflow_client(&self) -> &WorkflowServiceClient<Self::SvcType>;
 
-        /// Return the operator service client instance
-        fn operator_client(&mut self) -> &mut OperatorServiceClient<Self::SvcType>;
+        /// Return a mutable ref to the workflow service client instance
+        fn workflow_client_mut(&mut self) -> &mut WorkflowServiceClient<Self::SvcType>;
 
-        /// Return the test service client instance
-        fn test_client(&mut self) -> &mut TestServiceClient<Self::SvcType>;
+        /// Return a ref to the operator service client instance
+        fn operator_client(&self) -> &OperatorServiceClient<Self::SvcType>;
 
-        /// Return the health service client instance
-        fn health_client(&mut self) -> &mut HealthClient<Self::SvcType>;
+        /// Return a mutable ref to the operator service client instance
+        fn operator_client_mut(&mut self) -> &mut OperatorServiceClient<Self::SvcType>;
+
+        /// Return a ref to the test service client instance
+        fn test_client(&self) -> &TestServiceClient<Self::SvcType>;
+
+        /// Return a mutable ref to the test service client instance
+        fn test_client_mut(&mut self) -> &mut TestServiceClient<Self::SvcType>;
+
+        /// Return a ref to the health service client instance
+        fn health_client(&self) -> &HealthClient<Self::SvcType>;
+
+        /// 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<SlotManager>>;
 
         async fn call<F, Req, Resp>(
             &mut self,
@@ -58,7 +76,6 @@ pub(super) mod sealed {
     }
 }
 
-// Here we implement retry on anything that is already RawClientLike
 #[async_trait::async_trait]
 impl<RC, T> RawClientLike for RetryClient<RC>
 where
@@ -67,20 +84,40 @@ where
 {
     type SvcType = T;
 
-    fn workflow_client(&mut self) -> &mut WorkflowServiceClient<Self::SvcType> {
-        self.get_client_mut().workflow_client()
+    fn workflow_client(&self) -> &WorkflowServiceClient<Self::SvcType> {
+        self.get_client().workflow_client()
+    }
+
+    fn workflow_client_mut(&mut self) -> &mut WorkflowServiceClient<Self::SvcType> {
+        self.get_client_mut().workflow_client_mut()
+    }
+
+    fn operator_client(&self) -> &OperatorServiceClient<Self::SvcType> {
+        self.get_client().operator_client()
+    }
+
+    fn operator_client_mut(&mut self) -> &mut OperatorServiceClient<Self::SvcType> {
+        self.get_client_mut().operator_client_mut()
+    }
+
+    fn test_client(&self) -> &TestServiceClient<Self::SvcType> {
+        self.get_client().test_client()
+    }
+
+    fn test_client_mut(&mut self) -> &mut TestServiceClient<Self::SvcType> {
+        self.get_client_mut().test_client_mut()
     }
 
-    fn operator_client(&mut self) -> &mut OperatorServiceClient<Self::SvcType> {
-        self.get_client_mut().operator_client()
+    fn health_client(&self) -> &HealthClient<Self::SvcType> {
+        self.get_client().health_client()
     }
 
-    fn test_client(&mut self) -> &mut TestServiceClient<Self::SvcType> {
-        self.get_client_mut().test_client()
+    fn health_client_mut(&mut self) -> &mut HealthClient<Self::SvcType> {
+        self.get_client_mut().health_client_mut()
     }
 
-    fn health_client(&mut self) -> &mut HealthClient<Self::SvcType> {
-        self.get_client_mut().health_client()
+    fn get_workers_info(&self) -> Option<Arc<SlotManager>> {
+        self.get_client().get_workers_info()
     }
 
     async fn call<F, Req, Resp>(
@@ -114,21 +151,41 @@ where
 {
     type SvcType = T;
 
-    fn workflow_client(&mut self) -> &mut WorkflowServiceClient<Self::SvcType> {
+    fn workflow_client(&self) -> &WorkflowServiceClient<Self::SvcType> {
+        self.workflow_svc()
+    }
+
+    fn workflow_client_mut(&mut self) -> &mut WorkflowServiceClient<Self::SvcType> {
         self.workflow_svc_mut()
     }
 
-    fn operator_client(&mut self) -> &mut OperatorServiceClient<Self::SvcType> {
+    fn operator_client(&self) -> &OperatorServiceClient<Self::SvcType> {
+        self.operator_svc()
+    }
+
+    fn operator_client_mut(&mut self) -> &mut OperatorServiceClient<Self::SvcType> {
         self.operator_svc_mut()
     }
 
-    fn test_client(&mut self) -> &mut TestServiceClient<Self::SvcType> {
+    fn test_client(&self) -> &TestServiceClient<Self::SvcType> {
+        self.test_svc()
+    }
+
+    fn test_client_mut(&mut self) -> &mut TestServiceClient<Self::SvcType> {
         self.test_svc_mut()
     }
 
-    fn health_client(&mut self) -> &mut HealthClient<Self::SvcType> {
+    fn health_client(&self) -> &HealthClient<Self::SvcType> {
+        self.health_svc()
+    }
+
+    fn health_client_mut(&mut self) -> &mut HealthClient<Self::SvcType> {
         self.health_svc_mut()
     }
+
+    fn get_workers_info(&self) -> Option<Arc<SlotManager>> {
+        None
+    }
 }
 
 impl<T> RawClientLike for ConfiguredClient<TemporalServiceClient<T>>
@@ -141,41 +198,81 @@ where
 {
     type SvcType = T;
 
-    fn workflow_client(&mut self) -> &mut WorkflowServiceClient<Self::SvcType> {
+    fn workflow_client(&self) -> &WorkflowServiceClient<Self::SvcType> {
         self.client.workflow_client()
     }
 
-    fn operator_client(&mut self) -> &mut OperatorServiceClient<Self::SvcType> {
+    fn workflow_client_mut(&mut self) -> &mut WorkflowServiceClient<Self::SvcType> {
+        self.client.workflow_client_mut()
+    }
+
+    fn operator_client(&self) -> &OperatorServiceClient<Self::SvcType> {
         self.client.operator_client()
     }
 
-    fn test_client(&mut self) -> &mut TestServiceClient<Self::SvcType> {
+    fn operator_client_mut(&mut self) -> &mut OperatorServiceClient<Self::SvcType> {
+        self.client.operator_client_mut()
+    }
+
+    fn test_client(&self) -> &TestServiceClient<Self::SvcType> {
         self.client.test_client()
     }
 
-    fn health_client(&mut self) -> &mut HealthClient<Self::SvcType> {
+    fn test_client_mut(&mut self) -> &mut TestServiceClient<Self::SvcType> {
+        self.client.test_client_mut()
+    }
+
+    fn health_client(&self) -> &HealthClient<Self::SvcType> {
         self.client.health_client()
     }
+
+    fn health_client_mut(&mut self) -> &mut HealthClient<Self::SvcType> {
+        self.client.health_client_mut()
+    }
+
+    fn get_workers_info(&self) -> Option<Arc<SlotManager>> {
+        Some(self.workers())
+    }
 }
 
 impl RawClientLike for Client {
     type SvcType = InterceptedMetricsSvc;
 
-    fn workflow_client(&mut self) -> &mut WorkflowServiceClient<Self::SvcType> {
+    fn workflow_client(&self) -> &WorkflowServiceClient<Self::SvcType> {
         self.inner.workflow_client()
     }
 
-    fn operator_client(&mut self) -> &mut OperatorServiceClient<Self::SvcType> {
+    fn workflow_client_mut(&mut self) -> &mut WorkflowServiceClient<Self::SvcType> {
+        self.inner.workflow_client_mut()
+    }
+
+    fn operator_client(&self) -> &OperatorServiceClient<Self::SvcType> {
         self.inner.operator_client()
     }
 
-    fn test_client(&mut self) -> &mut TestServiceClient<Self::SvcType> {
+    fn operator_client_mut(&mut self) -> &mut OperatorServiceClient<Self::SvcType> {
+        self.inner.operator_client_mut()
+    }
+
+    fn test_client(&self) -> &TestServiceClient<Self::SvcType> {
         self.inner.test_client()
     }
 
-    fn health_client(&mut self) -> &mut HealthClient<Self::SvcType> {
+    fn test_client_mut(&mut self) -> &mut TestServiceClient<Self::SvcType> {
+        self.inner.test_client_mut()
+    }
+
+    fn health_client(&self) -> &HealthClient<Self::SvcType> {
         self.inner.health_client()
     }
+
+    fn health_client_mut(&mut self) -> &mut HealthClient<Self::SvcType> {
+        self.inner.health_client_mut()
+    }
+
+    fn get_workers_info(&self) -> Option<Arc<SlotManager>> {
+        self.inner.get_workers_info()
+    }
 }
 
 /// Helper for cloning a tonic request as long as the inner message may be cloned.
@@ -199,10 +296,10 @@ fn req_cloner<T: Clone>(cloneme: &Request<T>) -> Request<T> {
 
 #[derive(Debug)]
 pub(super) struct AttachMetricLabels {
-    pub(super) labels: Vec<opentelemetry::KeyValue>,
+    pub(super) labels: Vec<MetricKeyValue>,
 }
 impl AttachMetricLabels {
-    pub fn new(kvs: impl Into<Vec<opentelemetry::KeyValue>>) -> Self {
+    pub fn new(kvs: impl Into<Vec<MetricKeyValue>>) -> Self {
         Self { labels: kvs.into() }
     }
     pub fn namespace(ns: impl Into<String>) -> Self {
@@ -280,10 +377,29 @@ 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 $(, $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
@@ -301,7 +417,8 @@ 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)*)*);
             )*
         }
     };
@@ -312,8 +429,12 @@ 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;
+    WorkflowService; ALL_IMPLEMENTED_WORKFLOW_SERVICE_RPCS; WorkflowServiceClient; workflow_client_mut;
     (
         register_namespace,
         RegisterNamespaceRequest,
@@ -359,10 +480,35 @@ proxier! {
         start_workflow_execution,
         StartWorkflowExecutionRequest,
         StartWorkflowExecutionResponse,
-        |r| {
+        |r, workers| {
+            let mut slot: Option<Box<dyn Slot + Send>> = 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.try_reserve_wft_slot(namespace, task_queue) {
+                    Some(s) => slot = Some(s),
+                    None => req_mut.request_eager_execution = false
+                }
+            }
+            slot
+        },
+        |resp, slot| {
+            if let Some(mut s) = slot {
+                if let Ok(response) = resp.as_ref() {
+                    if let Some(task) = response.get_ref().clone().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
         }
     );
     (
@@ -825,7 +971,7 @@ proxier! {
 }
 
 proxier! {
-    OperatorService; ALL_IMPLEMENTED_OPERATOR_SERVICE_RPCS; OperatorServiceClient; operator_client;
+    OperatorService; ALL_IMPLEMENTED_OPERATOR_SERVICE_RPCS; OperatorServiceClient; operator_client_mut;
     (add_search_attributes, AddSearchAttributesRequest, AddSearchAttributesResponse);
     (remove_search_attributes, RemoveSearchAttributesRequest, RemoveSearchAttributesResponse);
     (list_search_attributes, ListSearchAttributesRequest, ListSearchAttributesResponse);
@@ -841,7 +987,7 @@ proxier! {
 }
 
 proxier! {
-    TestService; ALL_IMPLEMENTED_TEST_SERVICE_RPCS; TestServiceClient; test_client;
+    TestService; ALL_IMPLEMENTED_TEST_SERVICE_RPCS; TestServiceClient; test_client_mut;
     (lock_time_skipping, LockTimeSkippingRequest, LockTimeSkippingResponse);
     (unlock_time_skipping, UnlockTimeSkippingRequest, UnlockTimeSkippingResponse);
     (sleep, SleepRequest, SleepResponse);
@@ -851,7 +997,7 @@ proxier! {
 }
 
 proxier! {
-    HealthService; ALL_IMPLEMENTED_HEALTH_SERVICE_RPCS; HealthClient; health_client;
+    HealthService; ALL_IMPLEMENTED_HEALTH_SERVICE_RPCS; HealthClient; health_client_mut;
     (check, HealthCheckRequest, HealthCheckResponse);
     (watch, HealthCheckRequest, tonic::codec::Streaming<HealthCheckResponse>);
 }
@@ -874,7 +1020,7 @@ mod tests {
 
         let list_ns_req = ListNamespacesRequest::default();
         let fact = |c: &mut RetryClient<_>, req| {
-            let mut c = c.workflow_client().clone();
+            let mut c = c.workflow_client_mut().clone();
             async move { c.list_namespaces(req).await }.boxed()
         };
         retry_client
@@ -885,7 +1031,7 @@ mod tests {
         // Operator svc method
         let del_ns_req = DeleteNamespaceRequest::default();
         let fact = |c: &mut RetryClient<_>, req| {
-            let mut c = c.operator_client().clone();
+            let mut c = c.operator_client_mut().clone();
             async move { c.delete_namespace(req).await }.boxed()
         };
         retry_client
@@ -924,27 +1070,27 @@ mod tests {
     fn verify_all_workflow_service_methods_implemented() {
         // This is less work than trying to hook into the codegen process
         let proto_def =
-            include_str!("../../protos/api_upstream/temporal/api/workflowservice/v1/service.proto");
+            include_str!("../../sdk-core-protos/protos/api_upstream/temporal/api/workflowservice/v1/service.proto");
         verify_methods(proto_def, ALL_IMPLEMENTED_WORKFLOW_SERVICE_RPCS);
     }
 
     #[test]
     fn verify_all_operator_service_methods_implemented() {
         let proto_def =
-            include_str!("../../protos/api_upstream/temporal/api/operatorservice/v1/service.proto");
+            include_str!("../../sdk-core-protos/protos/api_upstream/temporal/api/operatorservice/v1/service.proto");
         verify_methods(proto_def, ALL_IMPLEMENTED_OPERATOR_SERVICE_RPCS);
     }
 
     #[test]
     fn verify_all_test_service_methods_implemented() {
         let proto_def =
-            include_str!("../../protos/testsrv_upstream/temporal/api/testservice/v1/service.proto");
+            include_str!("../../sdk-core-protos/protos/testsrv_upstream/temporal/api/testservice/v1/service.proto");
         verify_methods(proto_def, ALL_IMPLEMENTED_TEST_SERVICE_RPCS);
     }
 
     #[test]
     fn verify_all_health_service_methods_implemented() {
-        let proto_def = include_str!("../../protos/grpc/health/v1/health.proto");
+        let proto_def = include_str!("../../sdk-core-protos/protos/grpc/health/v1/health.proto");
         verify_methods(proto_def, ALL_IMPLEMENTED_HEALTH_SERVICE_RPCS);
     }
 }

package/sdk-core/client/src/retry.rs

@@ -11,6 +11,7 @@ use temporal_sdk_core_protos::{
         common::v1::{Payload, Payloads},
         failure::v1::Failure,
         query::v1::WorkflowQuery,
+        update,
         workflowservice::v1::*,
     },
     TaskToken,
@@ -526,6 +527,25 @@ where
         retry_call!(self, get_search_attributes)
     }
 
+    async fn update_workflow_execution(
+        &self,
+        workflow_id: String,
+        run_id: String,
+        name: String,
+        wait_policy: update::v1::WaitPolicy,
+        args: Option<Payloads>,
+    ) -> Result<UpdateWorkflowExecutionResponse> {
+        retry_call!(
+            self,
+            update_workflow_execution,
+            workflow_id.clone(),
+            run_id.clone(),
+            name.clone(),
+            wait_policy.clone(),
+            args.clone()
+        )
+    }
+
     fn get_options(&self) -> &ClientOptions {
         self.client.get_options()
     }