@temporalio/core-bridge 1.14.2-canary-release-testing.0 → 1.16.0

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

@@ -1,290 +1,12 @@
 //! This module helps with the initialization and management of telemetry. IE: Metrics and tracing.
 //! Logs from core are all traces, which may be exported to the console, in memory, or externally.
 
-mod log_export;
-pub(crate) mod metrics;
-#[cfg(feature = "otel")]
-mod otel;
-#[cfg(feature = "prom")]
-mod prometheus_meter;
-#[cfg(feature = "prom")]
-mod prometheus_server;
-
-// Always export bucket configuration function since it's used by both OTel and Prometheus
-pub use metrics::default_buckets_for;
-
-#[cfg(feature = "otel")]
-pub use metrics::{
-    ACTIVITY_EXEC_LATENCY_HISTOGRAM_NAME, ACTIVITY_SCHED_TO_START_LATENCY_HISTOGRAM_NAME,
-    MetricsCallBuffer, WORKFLOW_E2E_LATENCY_HISTOGRAM_NAME,
-    WORKFLOW_TASK_EXECUTION_LATENCY_HISTOGRAM_NAME, WORKFLOW_TASK_REPLAY_LATENCY_HISTOGRAM_NAME,
-    WORKFLOW_TASK_SCHED_TO_START_LATENCY_HISTOGRAM_NAME,
-};
-#[cfg(feature = "otel")]
-pub use otel::build_otlp_metric_exporter;
-#[cfg(feature = "prom")]
-pub use prometheus_server::start_prometheus_metric_exporter;
-
-pub use log_export::{CoreLogBuffer, CoreLogBufferedConsumer, CoreLogStreamConsumer};
-
-use crate::telemetry::{log_export::CoreLogConsumerLayer, metrics::PrefixedMetricsMeter};
 use itertools::Itertools;
-use parking_lot::Mutex;
-use std::{
-    cell::RefCell,
-    collections::VecDeque,
-    env,
-    sync::{
-        Arc,
-        atomic::{AtomicBool, Ordering},
-    },
-};
-use temporalio_common::telemetry::{
-    CoreLog, CoreTelemetry, Logger, TaskQueueLabelStrategy, TelemetryOptions,
-    TelemetryOptionsBuilder,
-    metrics::{CoreMeter, MetricKeyValue, NewAttributes, TemporalMeter},
-};
-use tracing::{Level, Subscriber};
-use tracing_subscriber::{EnvFilter, Layer, layer::SubscriberExt};
-
-const TELEM_SERVICE_NAME: &str = "temporal-core-sdk";
-
-const FORWARD_LOG_BUFFER_SIZE: usize = 2048;
-
-/// Help you construct an [EnvFilter] compatible filter string which will forward all core module
-/// traces at `core_level` and all others (from 3rd party modules, etc) at `other_level`.
-pub fn construct_filter_string(core_level: Level, other_level: Level) -> String {
-    format!(
-        "{other_level},temporalio_sdk_core={core_level},temporalio_client={core_level},temporalio_sdk={core_level}"
-    )
-}
-
-/// Holds initialized tracing/metrics exporters, etc
-pub struct TelemetryInstance {
-    metric_prefix: String,
-    logs_out: Option<Mutex<CoreLogBuffer>>,
-    metrics: Option<Arc<dyn CoreMeter + 'static>>,
-    /// The tracing subscriber which is associated with this telemetry instance. May be `None` if
-    /// the user has not opted into any tracing configuration.
-    trace_subscriber: Option<Arc<dyn Subscriber + Send + Sync>>,
-    attach_service_name: bool,
-    task_queue_label_strategy: TaskQueueLabelStrategy,
-}
-
-impl TelemetryInstance {
-    fn new(
-        trace_subscriber: Option<Arc<dyn Subscriber + Send + Sync>>,
-        logs_out: Option<Mutex<CoreLogBuffer>>,
-        metric_prefix: String,
-        metrics: Option<Arc<dyn CoreMeter + 'static>>,
-        attach_service_name: bool,
-        task_queue_label_strategy: TaskQueueLabelStrategy,
-    ) -> Self {
-        Self {
-            metric_prefix,
-            logs_out,
-            metrics,
-            trace_subscriber,
-            attach_service_name,
-            task_queue_label_strategy,
-        }
-    }
-
-    /// Return the trace subscriber associated with the telemetry options/instance. Can be used
-    /// to manually set the default for a thread or globally using the `tracing` crate, or with
-    /// [set_trace_subscriber_for_current_thread].
-    pub fn trace_subscriber(&self) -> Option<Arc<dyn Subscriber + Send + Sync>> {
-        self.trace_subscriber.clone()
-    }
-
-    /// Some metric meters cannot be initialized until after a tokio runtime has started and after
-    /// other telemetry has initted (ex: prometheus). They can be attached here.
-    pub fn attach_late_init_metrics(&mut self, meter: Arc<dyn CoreMeter + 'static>) {
-        self.metrics = Some(meter);
-    }
-
-    /// Returns our wrapper for metric meters, including the `metric_prefix` from
-    /// [TelemetryOptions]. This should be used to initialize clients or for any other
-    /// temporal-owned metrics. User defined metrics should use [Self::get_metric_meter].
-    pub fn get_temporal_metric_meter(&self) -> Option<TemporalMeter> {
-        self.metrics.clone().map(|m| {
-            let kvs = self.default_kvs();
-            let attribs = NewAttributes::new(kvs);
-            TemporalMeter::new(
-                Arc::new(PrefixedMetricsMeter::new(self.metric_prefix.clone(), m))
-                    as Arc<dyn CoreMeter>,
-                attribs,
-                self.task_queue_label_strategy,
-            )
-        })
-    }
-
-    /// Returns our wrapper for metric meters, including attaching the service name if enabled.
-    pub fn get_metric_meter(&self) -> Option<TemporalMeter> {
-        self.metrics.clone().map(|m| {
-            let kvs = self.default_kvs();
-            let attribs = NewAttributes::new(kvs);
-            TemporalMeter::new(m, attribs, self.task_queue_label_strategy)
-        })
-    }
-
-    fn default_kvs(&self) -> Vec<MetricKeyValue> {
-        if self.attach_service_name {
-            vec![MetricKeyValue::new("service_name", TELEM_SERVICE_NAME)]
-        } else {
-            vec![]
-        }
-    }
-}
-
-thread_local! {
-    static SUB_GUARD: RefCell<Option<tracing::subscriber::DefaultGuard>> =
-        const { RefCell::new(None) };
-}
-/// Set the trace subscriber for the current thread. This must be done in every thread which uses
-/// core stuff, otherwise traces/logs will not be collected on that thread. For example, if using
-/// a multithreaded Tokio runtime, you should ensure that said runtime uses
-/// [on_thread_start](https://docs.rs/tokio/latest/tokio/runtime/struct.Builder.html#method.on_thread_start)
-/// or a similar mechanism to call this for each thread within the runtime.
-pub fn set_trace_subscriber_for_current_thread(sub: impl Subscriber + Send + Sync + 'static) {
-    SUB_GUARD.with(|sg| {
-        if sg.borrow().is_none() {
-            let g = tracing::subscriber::set_default(sub);
-            *sg.borrow_mut() = Some(g);
-        }
-    })
-}
-
-/// Undoes [set_trace_subscriber_for_current_thread]
-pub fn remove_trace_subscriber_for_current_thread() {
-    SUB_GUARD.take();
-}
-
-impl CoreTelemetry for TelemetryInstance {
-    fn fetch_buffered_logs(&self) -> Vec<CoreLog> {
-        if let Some(logs_out) = self.logs_out.as_ref() {
-            logs_out.lock().drain()
-        } else {
-            vec![]
-        }
-    }
-}
-
-/// Initialize tracing subscribers/output and logging export, returning a [TelemetryInstance]
-/// which can be used to register default / global tracing subscribers.
-///
-/// You should only call this once per unique [TelemetryOptions]
-///
-/// See [TelemetryOptions] docs for more on configuration.
-pub fn telemetry_init(opts: TelemetryOptions) -> Result<TelemetryInstance, anyhow::Error> {
-    let mut logs_out = None;
+use std::collections::VecDeque;
 
-    // Tracing subscriber layers =========
-    let mut console_pretty_layer = None;
-    let mut console_compact_layer = None;
-    let mut forward_layer = None;
-    // ===================================
-
-    let tracing_sub = if let Some(ts) = opts.subscriber_override {
-        Some(ts)
-    } else {
-        opts.logging.map(|logger| {
-            match logger {
-                Logger::Console { filter } => {
-                    // This is silly dupe but can't be avoided without boxing.
-                    if env::var("TEMPORAL_CORE_PRETTY_LOGS").is_ok() {
-                        console_pretty_layer = Some(
-                            tracing_subscriber::fmt::layer()
-                                .with_target(false)
-                                .event_format(
-                                    tracing_subscriber::fmt::format()
-                                        .pretty()
-                                        .with_source_location(false),
-                                )
-                                .with_filter(EnvFilter::new(filter)),
-                        )
-                    } else {
-                        console_compact_layer = Some(
-                            tracing_subscriber::fmt::layer()
-                                .with_target(false)
-                                .event_format(
-                                    tracing_subscriber::fmt::format()
-                                        .compact()
-                                        .with_source_location(false),
-                                )
-                                .with_filter(EnvFilter::new(filter)),
-                        )
-                    }
-                }
-                Logger::Forward { filter } => {
-                    let (export_layer, lo) =
-                        CoreLogConsumerLayer::new_buffered(FORWARD_LOG_BUFFER_SIZE);
-                    logs_out = Some(Mutex::new(lo));
-                    forward_layer = Some(export_layer.with_filter(EnvFilter::new(filter)));
-                }
-                Logger::Push { filter, consumer } => {
-                    forward_layer = Some(
-                        CoreLogConsumerLayer::new(consumer).with_filter(EnvFilter::new(filter)),
-                    );
-                }
-            };
-            let reg = tracing_subscriber::registry()
-                .with(console_pretty_layer)
-                .with(console_compact_layer)
-                .with(forward_layer);
-
-            #[cfg(feature = "tokio-console")]
-            let reg = reg.with(console_subscriber::spawn());
-            Arc::new(reg) as Arc<dyn Subscriber + Send + Sync>
-        })
-    };
-
-    Ok(TelemetryInstance::new(
-        tracing_sub,
-        logs_out,
-        opts.metric_prefix,
-        opts.metrics,
-        opts.attach_service_name,
-        opts.task_queue_label_strategy,
-    ))
-}
-
-/// WARNING: Calling can cause panics because of <https://github.com/tokio-rs/tracing/issues/1656>
-/// Lang must not start using until resolved
-///
-/// Initialize telemetry/tracing globally. Useful for testing. Only takes affect when called
-/// the first time. Subsequent calls are ignored.
-pub fn telemetry_init_global(opts: TelemetryOptions) -> Result<(), anyhow::Error> {
-    static INITTED: AtomicBool = AtomicBool::new(false);
-    if INITTED
-        .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
-        .is_ok()
-    {
-        let ti = telemetry_init(opts)?;
-        if let Some(ts) = ti.trace_subscriber() {
-            tracing::subscriber::set_global_default(ts)?;
-        }
-    }
-    Ok(())
-}
+pub(crate) mod metrics;
 
-/// WARNING: Calling can cause panics because of <https://github.com/tokio-rs/tracing/issues/1656>
-/// Lang must not start using until resolved
-///
-/// Initialize the fallback global handler. All lang SDKs should call this somewhere, once, at
-/// startup, as it initializes a fallback handler for any dependencies (looking at you, otel) that
-/// don't provide good ways to customize their tracing usage. It sets a WARN-level global filter
-/// that uses the default console logger.
-pub fn telemetry_init_fallback() -> Result<(), anyhow::Error> {
-    telemetry_init_global(
-        TelemetryOptionsBuilder::default()
-            .logging(Logger::Console {
-                filter: construct_filter_string(Level::DEBUG, Level::WARN),
-            })
-            .build()?,
-    )?;
-    Ok(())
-}
+pub use metrics::MetricsCallBuffer;
 
 /// A trait for using [Display] on the contents of vecs, etc, which don't implement it.
 ///

package/sdk-core/crates/sdk-core/src/test_help/integ_helpers.rs

@@ -8,14 +8,16 @@ pub use crate::{
 };
 
 use crate::{
-    TaskToken, Worker, WorkerConfig, WorkerConfigBuilder,
+    PollError, PollerBehavior, TaskToken, Worker, WorkerConfig, WorkerConfigBuilder,
+    WorkerVersioningStrategy,
     pollers::{BoxedPoller, MockManualPoller},
     protosext::ValidPollWFTQResponse,
     replay::TestHistoryBuilder,
     sticky_q_name_for_worker,
     worker::{
-        TaskPollers,
+        TaskPollers, WorkerTelemetry,
         client::{LegacyQueryResult, MockWorkerClient, WorkerClient, WorkflowTaskCompletion},
+        worker_config_builder,
     },
 };
 use assert_matches::assert_matches;
@@ -36,9 +38,8 @@ use std::{
     task::{Context, Poll},
     time::Duration,
 };
+pub use temporalio_common::telemetry::metrics::TemporalMeter;
 use temporalio_common::{
-    Worker as WorkerTrait,
-    errors::PollError,
     protos::{
         coresdk::{
             workflow_activation::{WorkflowActivationJob, workflow_activation_job},
@@ -58,7 +59,7 @@ use temporalio_common::{
         },
         utilities::pack_any,
     },
-    worker::{PollerBehavior, WorkerTaskTypes, WorkerVersioningStrategy},
+    worker::WorkerTaskTypes,
 };
 use tokio::sync::{Notify, mpsc::unbounded_channel};
 use tokio_stream::wrappers::UnboundedReceiverStream;
@@ -68,7 +69,7 @@ use uuid::Uuid;
 pub const NAMESPACE: &str = "default";
 
 /// Initiate shutdown, drain the pollers (handling evictions), and wait for shutdown to complete.
-pub async fn drain_pollers_and_shutdown(worker: &dyn WorkerTrait) {
+pub async fn drain_pollers_and_shutdown(worker: &Worker) {
     worker.initiate_shutdown();
     tokio::join!(
         async {
@@ -97,9 +98,20 @@ pub async fn drain_pollers_and_shutdown(worker: &dyn WorkerTrait) {
     worker.shutdown().await;
 }
 
-pub fn test_worker_cfg() -> WorkerConfigBuilder {
-    let mut wcb = WorkerConfigBuilder::default();
-    wcb.namespace(NAMESPACE)
+#[allow(clippy::type_complexity)]
+pub fn test_worker_cfg() -> WorkerConfigBuilder<
+    worker_config_builder::SetWorkflowTaskPollerBehavior<
+        worker_config_builder::SetTaskTypes<
+            worker_config_builder::SetIgnoreEvictsOnShutdown<
+                worker_config_builder::SetVersioningStrategy<
+                    worker_config_builder::SetTaskQueue<worker_config_builder::SetNamespace>,
+                >,
+            >,
+        >,
+    >,
+> {
+    WorkerConfig::builder()
+        .namespace(NAMESPACE)
         .task_queue(Uuid::new_v4().to_string())
         .versioning_strategy(WorkerVersioningStrategy::None {
             build_id: "test_bin_id".to_string(),
@@ -107,8 +119,7 @@ pub fn test_worker_cfg() -> WorkerConfigBuilder {
         .ignore_evicts_on_shutdown(true)
         .task_types(WorkerTaskTypes::all())
         // Serial polling since it makes mocking much easier.
-        .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(1_usize));
-    wcb
+        .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(1_usize))
 }
 
 /// When constructing responses for mocks, indicates how a given response should be built
@@ -204,7 +215,7 @@ pub fn mock_worker(mocks: MocksHolder) -> Worker {
             act_poller,
             nexus_poller,
         },
-        None,
+        mocks.worker_telemetry,
         None,
         false,
     )
@@ -222,6 +233,7 @@ pub struct MocksHolder {
     client: Arc<dyn WorkerClient>,
     inputs: MockWorkerInputs,
     pub outstanding_task_map: Option<OutstandingWFTMap>,
+    worker_telemetry: Option<WorkerTelemetry>,
 }
 
 impl MocksHolder {
@@ -255,6 +267,10 @@ impl MocksHolder {
         self.inputs.nexus_poller = Some(nexus_poller);
     }
 
+    pub fn set_temporal_meter(&mut self, meter: TemporalMeter) {
+        self.worker_telemetry = Some(WorkerTelemetry::from_meter(meter));
+    }
+
     /// Can be used for tests that need to avoid auto-shutdown due to running out of mock responses
     pub fn make_wft_stream_interminable(&mut self) {
         if let Some(old_stream) = self.inputs.wft_stream.take() {
@@ -298,6 +314,7 @@ impl MocksHolder {
             client: Arc::new(client),
             inputs: mock_worker,
             outstanding_task_map: None,
+            worker_telemetry: None,
         }
     }
 
@@ -321,6 +338,7 @@ impl MocksHolder {
             client: Arc::new(client),
             inputs: mock_worker,
             outstanding_task_map: None,
+            worker_telemetry: None,
         }
     }
 
@@ -344,6 +362,7 @@ impl MocksHolder {
             client: Arc::new(client),
             inputs: mock_worker,
             outstanding_task_map: None,
+            worker_telemetry: None,
         }
     }
 
@@ -381,6 +400,7 @@ impl MocksHolder {
             client: Arc::new(client),
             inputs: mock_worker,
             outstanding_task_map: None,
+            worker_telemetry: None,
         }
     }
 
@@ -406,6 +426,7 @@ impl MocksHolder {
             client: Arc::new(client),
             inputs: mock_worker,
             outstanding_task_map: None,
+            worker_telemetry: None,
         }
     }
 }
@@ -816,6 +837,7 @@ pub fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder {
         client: Arc::new(cfg.mock_client),
         inputs: mock_worker,
         outstanding_task_map: Some(outstanding_wf_task_tokens),
+        worker_telemetry: None,
     };
     if let Some(activity_responses) = cfg.activity_responses {
         let act_poller = mock_poller_from_resps(activity_responses);
@@ -1015,10 +1037,7 @@ pub trait WorkerTestHelpers {
 }
 
 #[async_trait::async_trait]
-impl<T> WorkerTestHelpers for T
-where
-    T: WorkerTrait + ?Sized,
-{
+impl WorkerTestHelpers for Worker {
     async fn complete_execution(&self, run_id: &str) {
         self.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
             run_id.to_string(),

package/sdk-core/crates/sdk-core/src/test_help/unit_helpers.rs

@@ -2,12 +2,9 @@
 
 use futures_util::{StreamExt, stream::FuturesUnordered};
 use std::{collections::HashSet, future::Future};
-use temporalio_common::{
-    Worker as CoreWorker,
-    protos::coresdk::{
-        workflow_activation::workflow_activation_job,
-        workflow_completion::{WorkflowActivationCompletion, workflow_activation_completion},
-    },
+use temporalio_common::protos::coresdk::{
+    workflow_activation::workflow_activation_job,
+    workflow_completion::{WorkflowActivationCompletion, workflow_activation_completion},
 };
 
 /// Given a desired number of concurrent executions and a provided function that produces a future,

package/sdk-core/crates/sdk-core/src/worker/activities/activity_heartbeat_manager.rs

@@ -595,6 +595,7 @@ mod test {
                 details: vec![Payload {
                     metadata: Default::default(),
                     data: vec![payload_data],
+                    external_payloads: Default::default(),
                 }],
             },
             // Mimic the same delay we would apply in activity task manager

package/sdk-core/crates/sdk-core/src/worker/activities/local_activities.rs

@@ -3,7 +3,7 @@ use crate::{
     abstractions::{MeteredPermitDealer, OwnedMeteredSemPermit, UsedMeteredSemPermit, dbg_panic},
     protosext::ValidScheduleLA,
     telemetry::metrics::{activity_type, should_record_failure_metric, workflow_type},
-    worker::workflow::HeartbeatTimeoutMsg,
+    worker::{LocalActivitySlotKind, workflow::HeartbeatTimeoutMsg},
 };
 use futures_util::{
     Stream, StreamExt, future, future::AbortRegistration, stream, stream::BoxStream,
@@ -17,20 +17,17 @@ use std::{
     task::{Context, Poll},
     time::{Duration, Instant, SystemTime},
 };
-use temporalio_common::{
-    protos::{
-        coresdk::{
-            LocalActivitySlotInfo,
-            activity_result::{Cancellation, Failure as ActFail, Success},
-            activity_task::{ActivityCancelReason, ActivityTask, Start, activity_task},
-        },
-        temporal::api::{
-            common::v1::WorkflowExecution,
-            enums::v1::TimeoutType,
-            failure::v1::{Failure as APIFailure, TimeoutFailureInfo, failure},
-        },
+use temporalio_common::protos::{
+    coresdk::{
+        LocalActivitySlotInfo,
+        activity_result::{Cancellation, Failure as ActFail, Success},
+        activity_task::{ActivityCancelReason, ActivityTask, Start, activity_task},
+    },
+    temporal::api::{
+        common::v1::WorkflowExecution,
+        enums::v1::TimeoutType,
+        failure::v1::{Failure as APIFailure, TimeoutFailureInfo, failure},
     },
-    worker::LocalActivitySlotKind,
 };
 use tokio::{
     sync::{

package/sdk-core/crates/sdk-core/src/worker/activities.rs

@@ -7,7 +7,7 @@ pub(crate) use local_activities::{
 };
 
 use crate::{
-    PollError, TaskToken,
+    TaskToken,
     abstractions::{
         ClosableMeteredPermitDealer, MeteredPermitDealer, TrackedOwnedMeteredSemPermit,
         UsedMeteredSemPermit,
@@ -17,6 +17,7 @@ use crate::{
         MetricsContext, activity_type, eager, should_record_failure_metric, workflow_type,
     },
     worker::{
+        ActivitySlotKind, PollError,
         activities::activity_heartbeat_manager::ActivityHeartbeatError, client::WorkerClient,
     },
 };
@@ -35,21 +36,16 @@ use std::{
     },
     time::{Duration, Instant, SystemTime},
 };
-use temporalio_common::{
-    protos::{
-        coresdk::{
-            ActivityHeartbeat, ActivitySlotInfo,
-            activity_result::{self as ar, activity_execution_result as aer},
-            activity_task::{ActivityCancelReason, ActivityCancellationDetails, ActivityTask},
-        },
-        temporal::api::{
-            failure::v1::{
-                ApplicationFailureInfo, CanceledFailureInfo, Failure, failure::FailureInfo,
-            },
-            workflowservice::v1::PollActivityTaskQueueResponse,
-        },
+use temporalio_common::protos::{
+    coresdk::{
+        ActivityHeartbeat, ActivitySlotInfo,
+        activity_result::{self as ar, activity_execution_result as aer},
+        activity_task::{ActivityCancelReason, ActivityCancellationDetails, ActivityTask},
+    },
+    temporal::api::{
+        failure::v1::{ApplicationFailureInfo, CanceledFailureInfo, Failure, failure::FailureInfo},
+        workflowservice::v1::PollActivityTaskQueueResponse,
     },
-    worker::ActivitySlotKind,
 };
 use tokio::{
     join,
@@ -737,12 +733,10 @@ mod tests {
         abstractions::tests::fixed_size_permit_dealer,
         pollers::{ActivityTaskOptions, LongPollBuffer},
         prost_dur,
-        worker::client::mocks::mock_worker_client,
+        worker::{PollerBehavior, client::mocks::mock_worker_client},
     };
     use crossbeam_utils::atomic::AtomicCell;
-    use temporalio_common::{
-        protos::coresdk::activity_result::ActivityExecutionResult, worker::PollerBehavior,
-    };
+    use temporalio_common::protos::coresdk::activity_result::ActivityExecutionResult;
 
     #[tokio::test]
     async fn per_worker_ratelimit() {
@@ -787,6 +781,7 @@ mod tests {
                 max_tps: None,
             },
             Arc::new(AtomicCell::new(None)),
+            Arc::new(AtomicBool::new(false)),
         );
         let atm = WorkerActivityTasks::new(
             sem.clone(),
@@ -879,6 +874,7 @@ mod tests {
                 max_tps: None,
             },
             Arc::new(AtomicCell::new(None)),
+            Arc::new(AtomicBool::new(false)),
         );
         let atm = WorkerActivityTasks::new(
             sem.clone(),
@@ -953,6 +949,7 @@ mod tests {
                 max_tps: None,
             },
             Arc::new(AtomicCell::new(None)),
+            Arc::new(AtomicBool::new(false)),
         );
         let atm = WorkerActivityTasks::new(
             sem.clone(),

package/sdk-core/crates/sdk-core/src/worker/client/mocks.rs

@@ -1,5 +1,5 @@
 use super::*;
-use futures_util::Future;
+use futures_util::{Future, FutureExt};
 use std::sync::{Arc, LazyLock};
 use temporalio_client::worker::ClientWorkerSet;
 
@@ -30,12 +30,13 @@ pub fn mock_worker_client() -> MockWorkerClient {
         .returning(|| DEFAULT_WORKERS_REGISTRY.clone());
     r.expect_is_mock().returning(|| true);
     r.expect_shutdown_worker()
-        .returning(|_, _| Ok(ShutdownWorkerResponse {}));
+        .returning(|_, _, _, _| Ok(ShutdownWorkerResponse {}));
     r.expect_sdk_name_and_version()
         .returning(|| ("test-core".to_string(), "0.0.0".to_string()));
     r.expect_identity()
         .returning(|| "test-identity".to_string());
     r.expect_worker_grouping_key().returning(Uuid::new_v4);
+    r.expect_worker_instance_key().returning(Uuid::new_v4);
     r.expect_set_heartbeat_client_fields().returning(|hb| {
         hb.sdk_name = "test-core".to_string();
         hb.sdk_version = "0.0.0".to_string();
@@ -53,10 +54,14 @@ pub(crate) fn mock_manual_worker_client() -> MockManualWorkerClient {
     r.expect_workers()
         .returning(|| DEFAULT_WORKERS_REGISTRY.clone());
     r.expect_is_mock().returning(|| true);
+    r.expect_shutdown_worker()
+        .returning(|_, _, _, _| async { Ok(ShutdownWorkerResponse {}) }.boxed());
     r.expect_sdk_name_and_version()
         .returning(|| ("test-core".to_string(), "0.0.0".to_string()));
     r.expect_identity()
         .returning(|| "test-identity".to_string());
+    r.expect_worker_grouping_key().returning(Uuid::new_v4);
+    r.expect_worker_instance_key().returning(Uuid::new_v4);
     r
 }
 
@@ -124,7 +129,7 @@ mockall::mock! {
         fn fail_nexus_task<'a, 'b>(
             &self,
             task_token: TaskToken,
-            error: nexus::v1::HandlerError,
+            error: NexusTaskFailure,
         ) -> impl Future<Output = Result<RespondNexusTaskFailedResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 
@@ -154,7 +159,7 @@ mockall::mock! {
           impl Future<Output = Result<DescribeNamespaceResponse>> + Send + 'b
           where 'a: 'b, Self: 'b;
 
-        fn shutdown_worker<'a, 'b>(&self, sticky_task_queue: String, worker_heartbeat: Option<WorkerHeartbeat>) -> impl Future<Output = Result<ShutdownWorkerResponse>> + Send + 'b
+        fn shutdown_worker<'a, 'b>(&self, sticky_task_queue: String, task_queue: String, task_queue_types: Vec<TaskQueueType>, worker_heartbeat: Option<WorkerHeartbeat>) -> impl Future<Output = Result<ShutdownWorkerResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 
         fn record_worker_heartbeat<'a, 'b>(
@@ -163,13 +168,14 @@ mockall::mock! {
             heartbeat: Vec<WorkerHeartbeat>
         ) -> impl Future<Output = Result<RecordWorkerHeartbeatResponse>> + Send + 'b where 'a: 'b, Self: 'b;
 
-        fn replace_client(&self, new_client: Client);
+        fn replace_connection(&self, new_connection: Connection);
         fn capabilities(&self) -> Option<Capabilities>;
         fn workers(&self) -> Arc<ClientWorkerSet>;
         fn is_mock(&self) -> bool;
         fn sdk_name_and_version(&self) -> (String, String);
         fn identity(&self) -> String;
         fn worker_grouping_key(&self) -> Uuid;
+        fn worker_instance_key(&self) -> Uuid;
         fn set_heartbeat_client_fields(&self, heartbeat: &mut WorkerHeartbeat);
     }
 }