@temporalio/core-bridge 1.1.0 → 1.4.0

package/sdk-core/core/src/pollers/poll_buffer.rs

@@ -1,6 +1,6 @@
 use crate::{
     pollers::{self, Poller},
-    worker::client::WorkerClientBag,
+    worker::client::WorkerClient,
 };
 use futures::{prelude::stream::FuturesUnordered, StreamExt};
 use std::{
@@ -199,7 +199,7 @@ impl Poller<PollWorkflowTaskQueueResponse> for WorkflowTaskPoller {
 
 pub type PollWorkflowTaskBuffer = LongPollBuffer<PollWorkflowTaskQueueResponse>;
 pub(crate) fn new_workflow_task_buffer(
-    client: Arc<WorkerClientBag>,
+    client: Arc<dyn WorkerClient>,
     task_queue: String,
     is_sticky: bool,
     concurrent_pollers: usize,
@@ -220,7 +220,7 @@ pub(crate) fn new_workflow_task_buffer(
 
 pub type PollActivityTaskBuffer = LongPollBuffer<PollActivityTaskQueueResponse>;
 pub(crate) fn new_activity_task_buffer(
-    client: Arc<WorkerClientBag>,
+    client: Arc<dyn WorkerClient>,
     task_queue: String,
     concurrent_pollers: usize,
     buffer_size: usize,
@@ -262,7 +262,7 @@ mod tests {
             });
 
         let pb = new_workflow_task_buffer(
-            Arc::new(mock_client.into()),
+            Arc::new(mock_client),
             "someq".to_string(),
             false,
             1,

package/sdk-core/core/src/replay/mod.rs

@@ -2,7 +2,7 @@
 //! to replay canned histories. It should be used by Lang SDKs to provide replay capabilities to
 //! users during testing.
 
-use crate::{worker::client::mocks::mock_manual_workflow_client, WorkerClientBag};
+use crate::worker::client::{mocks::mock_manual_workflow_client, WorkerClient};
 use futures::FutureExt;
 use std::{
     sync::{
@@ -29,7 +29,7 @@ pub use temporal_sdk_core_protos::{
 pub(crate) fn mock_client_from_history(
     history: &History,
     task_queue: impl Into<String>,
-) -> WorkerClientBag {
+) -> impl WorkerClient {
     let mut mg = mock_manual_workflow_client();
 
     let hist_info = HistoryInfo::new_from_history(history, None).unwrap();
@@ -67,5 +67,5 @@ pub(crate) fn mock_client_from_history(
         async move { Ok(RespondWorkflowTaskFailedResponse {}) }.boxed()
     });
 
-    WorkerClientBag::new(Box::new(mg), "fake_namespace".to_string())
+    mg
 }

package/sdk-core/core/src/retry_logic.rs

@@ -81,7 +81,7 @@ const NANOS_PER_SEC: u32 = 1_000_000_000;
 fn try_from_secs_f64(secs: f64) -> Option<Duration> {
     const MAX_NANOS_F64: f64 = ((u64::MAX as u128 + 1) * (NANOS_PER_SEC as u128)) as f64;
     let nanos = secs * (NANOS_PER_SEC as f64);
-    if !nanos.is_finite() || nanos >= MAX_NANOS_F64 || nanos < 0.0 {
+    if !nanos.is_finite() || !(0.0..MAX_NANOS_F64).contains(&nanos) {
         None
     } else {
         Some(Duration::from_secs_f64(secs))
@@ -91,13 +91,14 @@ fn try_from_secs_f64(secs: f64) -> Option<Duration> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::prost_dur;
 
     #[test]
     fn calcs_backoffs_properly() {
         let rp = RetryPolicy {
-            initial_interval: Some(Duration::from_secs(1).into()),
+            initial_interval: Some(prost_dur!(from_secs(1))),
             backoff_coefficient: 2.0,
-            maximum_interval: Some(Duration::from_secs(10).into()),
+            maximum_interval: Some(prost_dur!(from_secs(10))),
             maximum_attempts: 10,
             non_retryable_error_types: vec![],
         };
@@ -132,7 +133,7 @@ mod tests {
     #[test]
     fn max_attempts_zero_retry_forever() {
         let rp = RetryPolicy {
-            initial_interval: Some(Duration::from_secs(1).into()),
+            initial_interval: Some(prost_dur!(from_secs(1))),
             backoff_coefficient: 1.2,
             maximum_interval: None,
             maximum_attempts: 0,
@@ -146,7 +147,7 @@ mod tests {
     #[test]
     fn no_overflows() {
         let rp = RetryPolicy {
-            initial_interval: Some(Duration::from_secs(1).into()),
+            initial_interval: Some(prost_dur!(from_secs(1))),
             backoff_coefficient: 10.,
             maximum_interval: None,
             maximum_attempts: 0,
@@ -160,9 +161,9 @@ mod tests {
     #[test]
     fn no_retry_err_str_match() {
         let rp = RetryPolicy {
-            initial_interval: Some(Duration::from_secs(1).into()),
+            initial_interval: Some(prost_dur!(from_secs(1))),
             backoff_coefficient: 2.0,
-            maximum_interval: Some(Duration::from_secs(10).into()),
+            maximum_interval: Some(prost_dur!(from_secs(10))),
             maximum_attempts: 10,
             non_retryable_error_types: vec!["no retry".to_string()],
         };
@@ -181,9 +182,9 @@ mod tests {
     #[test]
     fn no_non_retryable_application_failure() {
         let rp = RetryPolicy {
-            initial_interval: Some(Duration::from_secs(1).into()),
+            initial_interval: Some(prost_dur!(from_secs(1))),
             backoff_coefficient: 2.0,
-            maximum_interval: Some(Duration::from_secs(10).into()),
+            maximum_interval: Some(prost_dur!(from_secs(10))),
             maximum_attempts: 10,
             non_retryable_error_types: vec![],
         };

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

@@ -1,15 +1,17 @@
 use super::TELEM_SERVICE_NAME;
 use crate::telemetry::GLOBAL_TELEM_DAT;
+use opentelemetry::sdk::metrics::aggregators::Aggregator;
+use opentelemetry::sdk::metrics::sdk_api::{Descriptor, InstrumentKind};
 use opentelemetry::{
     global,
-    metrics::{Counter, Descriptor, InstrumentKind, Meter, ValueRecorder},
+    metrics::{Counter, Histogram, Meter},
     sdk::{
-        export::metrics::{Aggregator, AggregatorSelector},
+        export::metrics::AggregatorSelector,
         metrics::aggregators::{histogram, last_value, sum},
     },
-    KeyValue,
+    Context, KeyValue,
 };
-use std::{borrow::Cow, sync::Arc, time::Duration};
+use std::{sync::Arc, time::Duration};
 
 /// Used to track context associated with metrics, and record/update them
 ///
@@ -17,12 +19,16 @@ use std::{borrow::Cow, sync::Arc, time::Duration};
 /// appropriate k/vs have already been set.
 #[derive(Default, Clone, Debug)]
 pub(crate) struct MetricsContext {
+    ctx: Context,
     kvs: Arc<Vec<KeyValue>>,
 }
 
 impl MetricsContext {
     fn new(kvs: Vec<KeyValue>) -> Self {
-        Self { kvs: Arc::new(kvs) }
+        Self {
+            ctx: Context::current(),
+            kvs: Arc::new(kvs),
+        }
     }
 
     pub(crate) fn top_level(namespace: String) -> Self {
@@ -38,113 +44,116 @@ impl MetricsContext {
     pub(crate) fn with_new_attrs(&self, new_kvs: impl IntoIterator<Item = KeyValue>) -> Self {
         let mut kvs = self.kvs.clone();
         Arc::make_mut(&mut kvs).extend(new_kvs);
-        Self { kvs }
+        Self {
+            ctx: Context::current(),
+            kvs,
+        }
     }
 
     /// A workflow task queue poll succeeded
     pub(crate) fn wf_tq_poll_ok(&self) {
-        WF_TASK_QUEUE_POLL_SUCCEED_COUNTER.add(1, &self.kvs);
+        WF_TASK_QUEUE_POLL_SUCCEED_COUNTER.add(&self.ctx, 1, &self.kvs);
     }
 
     /// A workflow task queue poll timed out / had empty response
     pub(crate) fn wf_tq_poll_empty(&self) {
-        WF_TASK_QUEUE_POLL_EMPTY_COUNTER.add(1, &self.kvs);
+        WF_TASK_QUEUE_POLL_EMPTY_COUNTER.add(&self.ctx, 1, &self.kvs);
     }
 
     /// A workflow task execution failed
     pub(crate) fn wf_task_failed(&self) {
-        WF_TASK_EXECUTION_FAILURE_COUNTER.add(1, &self.kvs);
+        WF_TASK_EXECUTION_FAILURE_COUNTER.add(&self.ctx, 1, &self.kvs);
     }
 
     /// A workflow completed successfully
     pub(crate) fn wf_completed(&self) {
-        WF_COMPLETED_COUNTER.add(1, &self.kvs);
+        WF_COMPLETED_COUNTER.add(&self.ctx, 1, &self.kvs);
     }
 
     /// A workflow ended cancelled
     pub(crate) fn wf_canceled(&self) {
-        WF_CANCELED_COUNTER.add(1, &self.kvs);
+        WF_CANCELED_COUNTER.add(&self.ctx, 1, &self.kvs);
     }
 
     /// A workflow ended failed
     pub(crate) fn wf_failed(&self) {
-        WF_FAILED_COUNTER.add(1, &self.kvs);
+        WF_FAILED_COUNTER.add(&self.ctx, 1, &self.kvs);
     }
 
     /// A workflow continued as new
     pub(crate) fn wf_continued_as_new(&self) {
-        WF_CONT_COUNTER.add(1, &self.kvs);
+        WF_CONT_COUNTER.add(&self.ctx, 1, &self.kvs);
     }
 
     /// Record workflow total execution time in milliseconds
     pub(crate) fn wf_e2e_latency(&self, dur: Duration) {
-        WF_E2E_LATENCY.record(dur.as_millis() as u64, &self.kvs);
+        WF_E2E_LATENCY.record(&self.ctx, dur.as_millis() as u64, &self.kvs);
     }
 
     /// Record workflow task schedule to start time in millis
     pub(crate) fn wf_task_sched_to_start_latency(&self, dur: Duration) {
-        WF_TASK_SCHED_TO_START_LATENCY.record(dur.as_millis() as u64, &self.kvs);
+        WF_TASK_SCHED_TO_START_LATENCY.record(&self.ctx, dur.as_millis() as u64, &self.kvs);
     }
 
     /// Record workflow task execution time in milliseconds
     pub(crate) fn wf_task_latency(&self, dur: Duration) {
-        WF_TASK_EXECUTION_LATENCY.record(dur.as_millis() as u64, &self.kvs);
+        WF_TASK_EXECUTION_LATENCY.record(&self.ctx, dur.as_millis() as u64, &self.kvs);
     }
 
     /// Record time it takes to catch up on replaying a WFT
     pub(crate) fn wf_task_replay_latency(&self, dur: Duration) {
-        WF_TASK_REPLAY_LATENCY.record(dur.as_millis() as u64, &self.kvs);
+        WF_TASK_REPLAY_LATENCY.record(&self.ctx, dur.as_millis() as u64, &self.kvs);
     }
 
     /// An activity long poll timed out
     pub(crate) fn act_poll_timeout(&self) {
-        ACT_POLL_NO_TASK.add(1, &self.kvs);
+        ACT_POLL_NO_TASK.add(&self.ctx, 1, &self.kvs);
     }
 
     /// An activity execution failed
     pub(crate) fn act_execution_failed(&self) {
-        ACT_EXECUTION_FAILED.add(1, &self.kvs);
+        ACT_EXECUTION_FAILED.add(&self.ctx, 1, &self.kvs);
     }
 
     /// Record activity task schedule to start time in millis
     pub(crate) fn act_sched_to_start_latency(&self, dur: Duration) {
-        ACT_SCHED_TO_START_LATENCY.record(dur.as_millis() as u64, &self.kvs);
+        ACT_SCHED_TO_START_LATENCY.record(&self.ctx, dur.as_millis() as u64, &self.kvs);
     }
 
     /// Record time it took to complete activity execution, from the time core generated the
     /// activity task, to the time lang responded with a completion (failure or success).
     pub(crate) fn act_execution_latency(&self, dur: Duration) {
-        ACT_EXEC_LATENCY.record(dur.as_millis() as u64, &self.kvs);
+        ACT_EXEC_LATENCY.record(&self.ctx, dur.as_millis() as u64, &self.kvs);
     }
 
     /// A worker was registered
     pub(crate) fn worker_registered(&self) {
-        WORKER_REGISTERED.add(1, &self.kvs);
+        WORKER_REGISTERED.add(&self.ctx, 1, &self.kvs);
     }
 
     /// Record current number of available task slots. Context should have worker type set.
     pub(crate) fn available_task_slots(&self, num: usize) {
-        TASK_SLOTS_AVAILABLE.record(num as u64, &self.kvs)
+        TASK_SLOTS_AVAILABLE.record(&self.ctx, num as u64, &self.kvs)
     }
 
     /// Record current number of pollers. Context should include poller type / task queue tag.
     pub(crate) fn record_num_pollers(&self, num: usize) {
-        NUM_POLLERS.record(num as u64, &self.kvs);
+        NUM_POLLERS.record(&self.ctx, num as u64, &self.kvs);
     }
 
     /// A workflow task found a cached workflow to run against
     pub(crate) fn sticky_cache_hit(&self) {
-        STICKY_CACHE_HIT.add(1, &self.kvs);
+        STICKY_CACHE_HIT.add(&self.ctx, 1, &self.kvs);
     }
 
     /// A workflow task did not find a cached workflow
     pub(crate) fn sticky_cache_miss(&self) {
-        STICKY_CACHE_MISS.add(1, &self.kvs);
+        STICKY_CACHE_MISS.add(&self.ctx, 1, &self.kvs);
     }
 
     /// Record current cache size (in number of wfs, not bytes)
     pub(crate) fn cache_size(&self, size: u64) {
-        STICKY_CACHE_SIZE.record(size, &self.kvs);
+        STICKY_CACHE_SIZE.record(&self.ctx, size, &self.kvs);
     }
 }
 
@@ -182,8 +191,8 @@ macro_rules! tm {
     };
     (vr_u64, $ident:ident, $name:expr) => {
         lazy_static::lazy_static! {
-            static ref $ident: ValueRecorder<u64> = {
-                METRIC_METER.u64_value_recorder(metric_prefix().to_string() + $name).init()
+            static ref $ident: Histogram<u64> = {
+                METRIC_METER.u64_histogram(metric_prefix().to_string() + $name).init()
             };
         }
     };
@@ -214,22 +223,22 @@ pub(crate) fn activity_type(ty: String) -> KeyValue {
 pub(crate) fn workflow_type(ty: String) -> KeyValue {
     KeyValue::new(KEY_WF_TYPE, ty)
 }
-pub(crate) const fn workflow_worker_type() -> KeyValue {
+pub(crate) fn workflow_worker_type() -> KeyValue {
     KeyValue {
         key: opentelemetry::Key::from_static_str(KEY_WORKER_TYPE),
-        value: opentelemetry::Value::String(Cow::Borrowed("WorkflowWorker")),
+        value: opentelemetry::Value::String("WorkflowWorker".into()),
     }
 }
-pub(crate) const fn activity_worker_type() -> KeyValue {
+pub(crate) fn activity_worker_type() -> KeyValue {
     KeyValue {
         key: opentelemetry::Key::from_static_str(KEY_WORKER_TYPE),
-        value: opentelemetry::Value::String(Cow::Borrowed("ActivityWorker")),
+        value: opentelemetry::Value::String("ActivityWorker".into()),
     }
 }
-pub(crate) const fn local_activity_worker_type() -> KeyValue {
+pub(crate) fn local_activity_worker_type() -> KeyValue {
     KeyValue {
         key: opentelemetry::Key::from_static_str(KEY_WORKER_TYPE),
-        value: opentelemetry::Value::String(Cow::Borrowed("LocalActivityWorker")),
+        value: opentelemetry::Value::String("LocalActivityWorker".into()),
     }
 }
 
@@ -338,12 +347,12 @@ pub struct SDKAggSelector;
 
 impl AggregatorSelector for SDKAggSelector {
     fn aggregator_for(&self, descriptor: &Descriptor) -> Option<Arc<dyn Aggregator + Send + Sync>> {
-        // Observers are always last value
-        if *descriptor.instrument_kind() == InstrumentKind::ValueObserver {
+        // Gauges are always last value
+        if *descriptor.instrument_kind() == InstrumentKind::GaugeObserver {
             return Some(Arc::new(last_value()));
         }
 
-        if *descriptor.instrument_kind() == InstrumentKind::ValueRecorder {
+        if *descriptor.instrument_kind() == InstrumentKind::Histogram {
             let dname = descriptor
                 .name()
                 .strip_prefix(metric_prefix())
@@ -366,7 +375,7 @@ impl AggregatorSelector for SDKAggSelector {
                 ACT_EXEC_LATENCY_NAME => ACT_EXE_MS_BUCKETS,
                 _ => DEFAULT_MS_BUCKETS,
             };
-            return Some(Arc::new(histogram(descriptor, buckets)));
+            return Some(Arc::new(histogram(buckets)));
         }
 
         Some(Arc::new(sum()))

package/sdk-core/core/src/telemetry/mod.rs

@@ -12,8 +12,12 @@ use once_cell::sync::OnceCell;
 use opentelemetry::{
     global,
     metrics::Meter,
-    sdk::{metrics::PushController, trace::Config, Resource},
-    util::tokio_interval_stream,
+    runtime,
+    sdk::{
+        export::metrics::aggregation::{self, Temporality, TemporalitySelector},
+        trace::Config,
+        Resource,
+    },
     KeyValue,
 };
 use opentelemetry_otlp::WithExportConfig;
@@ -69,7 +73,7 @@ pub enum MetricsExporter {
 }
 
 /// Control where logs go
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy)]
 pub enum Logger {
     /// Log directly to console.
     Console,
@@ -101,6 +105,33 @@ pub struct TelemetryOptions {
     /// the prefix is consistent with other SDKs.
     #[builder(default)]
     pub no_temporal_prefix_for_metrics: bool,
+
+    /// Specifies the aggregation temporality for metric export. Defaults to cumulative.
+    #[builder(default = "MetricTemporality::Cumulative")]
+    pub metric_temporality: MetricTemporality,
+}
+
+/// Types of aggregation temporality for metric export.
+/// See: <https://github.com/open-telemetry/opentelemetry-specification/blob/ce50e4634efcba8da445cc23523243cb893905cb/specification/metrics/datamodel.md#temporality>
+#[derive(Debug, Clone, Copy)]
+pub enum MetricTemporality {
+    /// Successive data points repeat the starting timestamp
+    Cumulative,
+    /// Successive data points advance the starting timestamp
+    Delta,
+}
+
+impl MetricTemporality {
+    fn to_selector(self) -> impl TemporalitySelector + Send + Sync + Clone {
+        match self {
+            MetricTemporality::Cumulative => {
+                aggregation::constant_temporality_selector(Temporality::Cumulative)
+            }
+            MetricTemporality::Delta => {
+                aggregation::constant_temporality_selector(Temporality::Delta)
+            }
+        }
+    }
 }
 
 impl TelemetryOptions {
@@ -123,7 +154,6 @@ impl Default for TelemetryOptions {
 /// Things that need to not be dropped while telemetry is ongoing
 #[derive(Default)]
 pub struct GlobalTelemDat {
-    metric_push_controller: Option<PushController>,
     core_export_logger: Option<CoreExportLogger>,
     runtime: Option<tokio::runtime::Runtime>,
     prom_srv: Option<PromServer>,
@@ -202,7 +232,7 @@ pub fn telemetry_init(opts: &TelemetryOptions) -> Result<&'static GlobalTelemDat
                                 .pretty()
                                 .with_source_location(false);
                             let reg = tracing_subscriber::registry()
-                                .with((&opts).try_get_env_filter()?)
+                                .with((opts).try_get_env_filter()?)
                                 .with(
                                     tracing_subscriber::fmt::layer()
                                         .with_target(false)
@@ -221,16 +251,19 @@ pub fn telemetry_init(opts: &TelemetryOptions) -> Result<&'static GlobalTelemDat
             if let Some(ref metrics) = opts.metrics {
                 match metrics {
                     MetricsExporter::Prometheus(addr) => {
-                        let srv = PromServer::new(*addr)?;
+                        let srv = PromServer::new(*addr, opts.metric_temporality.to_selector())?;
                         globaldat.prom_srv = Some(srv);
                     }
                     MetricsExporter::Otel(OtelCollectorOptions { url, headers }) => {
                         runtime.block_on(async {
                             let metrics = opentelemetry_otlp::new_pipeline()
-                                .metrics(|f| runtime.spawn(f), tokio_interval_stream)
-                                .with_aggregator_selector(SDKAggSelector)
+                                .metrics(
+                                    SDKAggSelector,
+                                    opts.metric_temporality.to_selector(),
+                                    runtime::Tokio,
+                                )
                                 .with_period(Duration::from_secs(1))
-                                .with_resource(default_resource_kvs().iter().cloned())
+                                .with_resource(default_resource())
                                 .with_exporter(
                                     // No joke exporter builder literally not cloneable for some insane
                                     // reason
@@ -242,8 +275,7 @@ pub fn telemetry_init(opts: &TelemetryOptions) -> Result<&'static GlobalTelemDat
                                         )),
                                 )
                                 .build()?;
-                            global::set_meter_provider(metrics.provider());
-                            globaldat.metric_push_controller = Some(metrics);
+                            global::set_meter_provider(metrics);
                             Result::<(), anyhow::Error>::Ok(())
                         })?;
                     }
@@ -266,7 +298,7 @@ pub fn telemetry_init(opts: &TelemetryOptions) -> Result<&'static GlobalTelemDat
                                         )),
                                 )
                                 .with_trace_config(tracer_cfg)
-                                .install_batch(opentelemetry::runtime::Tokio)?;
+                                .install_batch(runtime::Tokio)?;
 
                             let opentelemetry = tracing_opentelemetry::layer().with_tracer(tracer);
 
@@ -331,6 +363,7 @@ pub(crate) fn test_telem_console() {
         tracing: None,
         metrics: None,
         no_temporal_prefix_for_metrics: false,
+        metric_temporality: MetricTemporality::Cumulative,
     })
     .unwrap();
 }
@@ -347,6 +380,7 @@ pub(crate) fn test_telem_collector() {
         })),
         metrics: None,
         no_temporal_prefix_for_metrics: false,
+        metric_temporality: MetricTemporality::Cumulative,
     })
     .unwrap();
 }

package/sdk-core/core/src/telemetry/prometheus_server.rs

@@ -1,13 +1,16 @@
-use crate::telemetry::{
-    default_resource,
-    metrics::{SDKAggSelector, DEFAULT_MS_BUCKETS},
-};
+use crate::telemetry::{default_resource, metrics::SDKAggSelector};
 use hyper::{
     header::CONTENT_TYPE,
     service::{make_service_fn, service_fn},
     Body, Method, Request, Response, Server,
 };
-use opentelemetry::metrics::MetricsError;
+use opentelemetry::{
+    metrics::MetricsError,
+    sdk::{
+        export::metrics::aggregation::TemporalitySelector,
+        metrics::{controllers, processors},
+    },
+};
 use opentelemetry_prometheus::{ExporterBuilder, PrometheusExporter};
 use prometheus::{Encoder, TextEncoder};
 use std::{convert::Infallible, net::SocketAddr, sync::Arc};
@@ -19,14 +22,15 @@ pub(super) struct PromServer {
 }
 
 impl PromServer {
-    pub fn new(addr: SocketAddr) -> Result<Self, MetricsError> {
-        let exporter = ExporterBuilder::default()
-            .with_default_histogram_boundaries(DEFAULT_MS_BUCKETS.to_vec())
-            .with_aggregator_selector(SDKAggSelector)
-            .with_host(addr.ip().to_string())
-            .with_port(addr.port())
-            .with_resource(default_resource())
-            .try_init()?;
+    pub fn new(
+        addr: SocketAddr,
+        temporality: impl TemporalitySelector + Send + Sync + 'static,
+    ) -> Result<Self, MetricsError> {
+        let controller =
+            controllers::basic(processors::factory(SDKAggSelector, temporality).with_memory(true))
+                .with_resource(default_resource())
+                .build();
+        let exporter = ExporterBuilder::new(controller).try_init()?;
         Ok(Self {
             exporter: Arc::new(exporter),
             addr,

package/sdk-core/core/src/test_help/mod.rs

@@ -9,7 +9,7 @@ use crate::{
         client::{mocks::mock_workflow_client, MockWorkerClient, WorkerClient},
         new_wft_poller,
     },
-    TaskToken, Worker, WorkerClientBag, WorkerConfig, WorkerConfigBuilder,
+    TaskToken, Worker, WorkerConfig, WorkerConfigBuilder,
 };
 use bimap::BiMap;
 use futures::{future::BoxFuture, stream, stream::BoxStream, FutureExt, Stream, StreamExt};
@@ -136,7 +136,7 @@ pub(crate) fn mock_worker(mocks: MocksHolder) -> Worker {
     Worker::new_with_pollers(
         mocks.inputs.config,
         sticky_q,
-        Arc::new(mocks.client_bag),
+        mocks.client,
         mocks.inputs.wft_stream,
         mocks.inputs.act_poller,
         Default::default(),
@@ -166,7 +166,7 @@ pub struct FakeWfResponses {
 
 // TODO: Should be all-internal to this module
 pub struct MocksHolder {
-    client_bag: WorkerClientBag,
+    client: Arc<dyn WorkerClient>,
     inputs: MockWorkerInputs,
     pub outstanding_task_map: Option<OutstandingWFTMap>,
 }
@@ -218,11 +218,11 @@ impl MockWorkerInputs {
 
 impl MocksHolder {
     pub(crate) fn from_mock_worker(
-        client_bag: WorkerClientBag,
+        client: impl WorkerClient + 'static,
         mock_worker: MockWorkerInputs,
     ) -> Self {
         Self {
-            client_bag,
+            client: Arc::new(client),
             inputs: mock_worker,
             outstanding_task_map: None,
         }
@@ -245,7 +245,7 @@ impl MocksHolder {
             config: test_worker_cfg().build().unwrap(),
         };
         Self {
-            client_bag: client.into(),
+            client: Arc::new(client),
             inputs: mock_worker,
             outstanding_task_map: None,
         }
@@ -267,7 +267,7 @@ impl MocksHolder {
             config: test_worker_cfg().build().unwrap(),
         };
         Self {
-            client_bag: client.into(),
+            client: Arc::new(client),
             inputs: mock_worker,
             outstanding_task_map: None,
         }
@@ -356,6 +356,7 @@ pub(crate) fn single_hist_mock_sg(
     build_mock_pollers(mh)
 }
 
+#[allow(clippy::type_complexity)]
 pub(crate) struct MockPollCfg {
     pub hists: Vec<FakeWfResponses>,
     pub enforce_correct_number_of_polls: bool,
@@ -604,7 +605,7 @@ pub(crate) fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder {
         });
 
     MocksHolder {
-        client_bag: cfg.mock_client.into(),
+        client: Arc::new(cfg.mock_client),
         inputs: mock_worker,
         outstanding_task_map: Some(outstanding_wf_task_tokens),
     }
@@ -863,3 +864,12 @@ macro_rules! advance_fut {
         }
     };
 }
+
+#[macro_export]
+macro_rules! prost_dur {
+    ($dur_call:ident $args:tt) => {
+        std::time::Duration::$dur_call$args
+            .try_into()
+            .expect("test duration fits")
+    };
+}