@temporalio/core-bridge 1.12.0 → 1.12.2

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

@@ -1,14 +1,48 @@
+use crate::telemetry::prometheus_meter::Registry;
 use http_body_util::Full;
 use hyper::{Method, Request, Response, body::Bytes, header::CONTENT_TYPE, service::service_fn};
 use hyper_util::{
     rt::{TokioExecutor, TokioIo},
     server::conn::auto,
 };
-use opentelemetry_prometheus::PrometheusExporter;
-use prometheus::{Encoder, Registry, TextEncoder};
-use std::net::{SocketAddr, TcpListener};
+use prometheus::{Encoder, TextEncoder};
+use std::{
+    net::{SocketAddr, TcpListener},
+    sync::Arc,
+};
 use temporal_sdk_core_api::telemetry::PrometheusExporterOptions;
-use tokio::io;
+use tokio::{io, task::AbortHandle};
+
+pub struct StartedPromServer {
+    pub meter: Arc<crate::telemetry::prometheus_meter::CorePrometheusMeter>,
+    pub bound_addr: SocketAddr,
+    pub abort_handle: AbortHandle,
+}
+
+/// Builds and runs a prometheus endpoint which can be scraped by prom instances for metrics export.
+/// Returns the meter that can be used as a [CoreMeter].
+///
+/// Requires a Tokio runtime to exist, and will block briefly while binding the server endpoint.
+pub fn start_prometheus_metric_exporter(
+    opts: PrometheusExporterOptions,
+) -> Result<StartedPromServer, anyhow::Error> {
+    let srv = PromServer::new(&opts)?;
+    let meter = Arc::new(
+        crate::telemetry::prometheus_meter::CorePrometheusMeter::new(
+            srv.registry().clone(),
+            opts.use_seconds_for_durations,
+            opts.unit_suffix,
+            opts.histogram_bucket_overrides,
+        ),
+    );
+    let bound_addr = srv.bound_addr()?;
+    let handle = tokio::spawn(async move { srv.run().await });
+    Ok(StartedPromServer {
+        meter,
+        bound_addr,
+        abort_handle: handle.abort_handle(),
+    })
+}
 
 /// Exposes prometheus metrics for scraping
 pub(super) struct PromServer {
@@ -17,30 +51,16 @@ pub(super) struct PromServer {
 }
 
 impl PromServer {
-    pub(super) fn new(
-        opts: &PrometheusExporterOptions,
-    ) -> Result<(Self, PrometheusExporter), anyhow::Error> {
-        let registry = Registry::new();
-        let exporter = opentelemetry_prometheus::exporter()
-            .without_scope_info()
-            .with_registry(registry.clone());
-        let exporter = if !opts.counters_total_suffix {
-            exporter.without_counter_suffixes()
-        } else {
-            exporter
-        };
-        let exporter = if !opts.unit_suffix {
-            exporter.without_units()
-        } else {
-            exporter
-        };
-        Ok((
-            Self {
-                listener: TcpListener::bind(opts.socket_addr)?,
-                registry,
-            },
-            exporter.build()?,
-        ))
+    pub(super) fn new(opts: &PrometheusExporterOptions) -> Result<Self, anyhow::Error> {
+        let registry = Registry::new(opts.global_tags.clone());
+        Ok(Self {
+            listener: TcpListener::bind(opts.socket_addr)?,
+            registry,
+        })
+    }
+
+    pub(super) fn registry(&self) -> &Registry {
+        &self.registry
     }
 
     pub(super) async fn run(self) -> Result<(), anyhow::Error> {

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

@@ -9,7 +9,8 @@ use crate::{
     worker::{
         TaskPollers,
         client::{
-            MockWorkerClient, WorkerClient, WorkflowTaskCompletion, mocks::mock_workflow_client,
+            LegacyQueryResult, MockWorkerClient, WorkerClient, WorkflowTaskCompletion,
+            mocks::mock_worker_client,
         },
     },
 };
@@ -179,6 +180,7 @@ pub(crate) fn mock_worker(mocks: MocksHolder) -> Worker {
                 .unwrap_or_else(|| mock_poller_from_resps([])),
         },
         None,
+        None,
     )
 }
 
@@ -406,6 +408,9 @@ pub(crate) struct MockPollCfg {
     /// All calls to fail WFTs must match this predicate
     pub(crate) expect_fail_wft_matcher:
         Box<dyn Fn(&TaskToken, &WorkflowTaskFailedCause, &Option<Failure>) -> bool + Send>,
+    /// All calls to legacy query responses must match this predicate
+    pub(crate) expect_legacy_query_matcher:
+        Box<dyn Fn(&TaskToken, &LegacyQueryResult) -> bool + Send>,
     pub(crate) completion_mock_fn: Option<Box<WFTCompletionMockFn>>,
     pub(crate) num_expected_completions: Option<TimesRange>,
     /// If being used with the Rust SDK, this is set true. It ensures pollers will not error out
@@ -426,8 +431,9 @@ impl MockPollCfg {
             enforce_correct_number_of_polls,
             num_expected_fails,
             num_expected_legacy_query_resps: 0,
-            mock_client: mock_workflow_client(),
+            mock_client: mock_worker_client(),
             expect_fail_wft_matcher: Box::new(|_, _, _| true),
+            expect_legacy_query_matcher: Box::new(|_, _| true),
             completion_mock_fn: None,
             num_expected_completions: None,
             using_rust_sdk: false,
@@ -439,14 +445,14 @@ impl MockPollCfg {
     pub(crate) fn from_hist_builder(t: TestHistoryBuilder) -> Self {
         let full_hist_info = t.get_full_history_info().unwrap();
         let tasks = 1..=full_hist_info.wf_task_count();
-        Self::from_resp_batches("fake_wf_id", t, tasks, mock_workflow_client())
+        Self::from_resp_batches("fake_wf_id", t, tasks, mock_worker_client())
     }
 
     pub(crate) fn from_resps(
         t: TestHistoryBuilder,
         resps: impl IntoIterator<Item = impl Into<ResponseType>>,
     ) -> Self {
-        Self::from_resp_batches("fake_wf_id", t, resps, mock_workflow_client())
+        Self::from_resp_batches("fake_wf_id", t, resps, mock_worker_client())
     }
 
     pub(crate) fn from_resp_batches(
@@ -466,6 +472,7 @@ impl MockPollCfg {
             num_expected_legacy_query_resps: 0,
             mock_client,
             expect_fail_wft_matcher: Box::new(|_, _, _| true),
+            expect_legacy_query_matcher: Box::new(|_, _| true),
             completion_mock_fn: None,
             num_expected_completions: None,
             using_rust_sdk: false,
@@ -709,6 +716,7 @@ pub(crate) fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder {
     let outstanding = outstanding_wf_task_tokens.clone();
     cfg.mock_client
         .expect_respond_legacy_query()
+        .withf(cfg.expect_legacy_query_matcher)
         .times::<TimesRange>(cfg.num_expected_legacy_query_resps.into())
         .returning(move |tt, _| {
             outstanding.release_token(&tt);
@@ -948,10 +956,8 @@ pub(crate) async fn poll_and_reply_clears_outstanding_evicts<'a>(
 
             worker.complete_workflow_activation(reply).await.unwrap();
 
-            if do_release {
-                if let Some(omap) = outstanding_map.as_ref() {
-                    omap.release_run(&res.run_id);
-                }
+            if do_release && let Some(omap) = outstanding_map.as_ref() {
+                omap.release_run(&res.run_id);
             }
             // Restart assertions from the beginning if it was an eviction (and workflow execution
             // isn't over)
@@ -1068,10 +1074,19 @@ impl WorkerExt for Worker {
                 );
             },
             async {
-                assert_matches!(
-                    self.poll_workflow_activation().await.unwrap_err(),
-                    PollError::ShutDown
-                );
+                loop {
+                    match self.poll_workflow_activation().await {
+                        Err(PollError::ShutDown) => break,
+                        Ok(a) if a.is_only_eviction() => {
+                            self.complete_workflow_activation(WorkflowActivationCompletion::empty(
+                                a.run_id,
+                            ))
+                            .await
+                            .unwrap();
+                        }
+                        o => panic!("Unexpected activation while draining: {o:?}"),
+                    }
+                }
             }
         );
         self.finalize_shutdown().await;

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

@@ -423,7 +423,7 @@ impl HeartbeatStreamState {
 mod test {
     use super::*;
 
-    use crate::worker::client::mocks::mock_workflow_client;
+    use crate::worker::client::mocks::mock_worker_client;
     use std::time::Duration;
     use temporal_sdk_core_protos::temporal::api::{
         common::v1::Payload, workflowservice::v1::RecordActivityTaskHeartbeatResponse,
@@ -434,7 +434,7 @@ mod test {
     /// every 1/2 of the heartbeat timeout.
     #[tokio::test]
     async fn process_heartbeats_and_shutdown() {
-        let mut mock_client = mock_workflow_client();
+        let mut mock_client = mock_worker_client();
         mock_client
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
@@ -456,7 +456,7 @@ mod test {
 
     #[tokio::test]
     async fn send_heartbeats_less_frequently_than_throttle_interval() {
-        let mut mock_client = mock_workflow_client();
+        let mut mock_client = mock_worker_client();
         mock_client
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
@@ -475,7 +475,7 @@ mod test {
     /// Ensure that heartbeat can be called from a tight loop and correctly throttle
     #[tokio::test]
     async fn process_tight_loop_and_shutdown() {
-        let mut mock_client = mock_workflow_client();
+        let mut mock_client = mock_worker_client();
         mock_client
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
@@ -495,7 +495,7 @@ mod test {
     /// This test reports one heartbeat and waits for the throttle_interval to elapse before sending another
     #[tokio::test]
     async fn report_heartbeat_after_timeout() {
-        let mut mock_client = mock_workflow_client();
+        let mut mock_client = mock_worker_client();
         mock_client
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
@@ -513,7 +513,7 @@ mod test {
 
     #[tokio::test]
     async fn evict_works() {
-        let mut mock_client = mock_workflow_client();
+        let mut mock_client = mock_worker_client();
         mock_client
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
@@ -534,7 +534,7 @@ mod test {
 
     #[tokio::test]
     async fn evict_immediate_after_record() {
-        let mut mock_client = mock_workflow_client();
+        let mut mock_client = mock_worker_client();
         mock_client
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))

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

@@ -726,14 +726,14 @@ mod tests {
         abstractions::tests::fixed_size_permit_dealer,
         pollers::{ActivityTaskOptions, LongPollBuffer},
         prost_dur,
-        worker::client::mocks::mock_workflow_client,
+        worker::client::mocks::mock_worker_client,
     };
     use temporal_sdk_core_api::worker::PollerBehavior;
     use temporal_sdk_core_protos::coresdk::activity_result::ActivityExecutionResult;
 
     #[tokio::test]
     async fn per_worker_ratelimit() {
-        let mut mock_client = mock_workflow_client();
+        let mut mock_client = mock_worker_client();
         mock_client
             .expect_poll_activity_task()
             .times(1)
@@ -812,7 +812,7 @@ mod tests {
 
     #[tokio::test]
     async fn local_timeouts() {
-        let mut mock_client = mock_workflow_client();
+        let mut mock_client = mock_worker_client();
         mock_client
             .expect_poll_activity_task()
             .times(1)
@@ -902,7 +902,7 @@ mod tests {
 
     #[tokio::test]
     async fn local_timeout_heartbeating() {
-        let mut mock_client = mock_workflow_client();
+        let mut mock_client = mock_worker_client();
         mock_client
             .expect_poll_activity_task()
             .times(1)

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

@@ -22,7 +22,7 @@ pub(crate) static DEFAULT_TEST_CAPABILITIES: &Capabilities = &Capabilities {
 
 #[cfg(test)]
 /// Create a mock client primed with basic necessary expectations
-pub(crate) fn mock_workflow_client() -> MockWorkerClient {
+pub(crate) fn mock_worker_client() -> MockWorkerClient {
     let mut r = MockWorkerClient::new();
     r.expect_capabilities()
         .returning(|| Some(*DEFAULT_TEST_CAPABILITIES));
@@ -33,11 +33,13 @@ pub(crate) fn mock_workflow_client() -> MockWorkerClient {
         .returning(|_| Ok(ShutdownWorkerResponse {}));
     r.expect_sdk_name_and_version()
         .returning(|| ("test-core".to_string(), "0.0.0".to_string()));
+    r.expect_get_identity()
+        .returning(|| "test-identity".to_string());
     r
 }
 
 /// Create a mock manual client primed with basic necessary expectations
-pub(crate) fn mock_manual_workflow_client() -> MockManualWorkerClient {
+pub(crate) fn mock_manual_worker_client() -> MockManualWorkerClient {
     let mut r = MockManualWorkerClient::new();
     r.expect_capabilities()
         .returning(|| Some(*DEFAULT_TEST_CAPABILITIES));
@@ -46,6 +48,8 @@ pub(crate) fn mock_manual_workflow_client() -> MockManualWorkerClient {
     r.expect_is_mock().returning(|| true);
     r.expect_sdk_name_and_version()
         .returning(|| ("test-core".to_string(), "0.0.0".to_string()));
+    r.expect_get_identity()
+        .returning(|| "test-identity".to_string());
     r
 }
 
@@ -135,7 +139,7 @@ mockall::mock! {
         fn respond_legacy_query<'a, 'b>(
             &self,
             task_token: TaskToken,
-            query_result: QueryResult,
+        query_result: LegacyQueryResult,
         ) -> impl Future<Output = Result<RespondQueryTaskCompletedResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 
@@ -146,10 +150,13 @@ mockall::mock! {
         fn shutdown_worker<'a, 'b>(&self, sticky_task_queue: String) -> impl Future<Output = Result<ShutdownWorkerResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 
+        fn record_worker_heartbeat<'a, 'b>(&self, heartbeat: WorkerHeartbeat) -> impl Future<Output = Result<RecordWorkerHeartbeatResponse>> + Send + 'b where 'a: 'b, Self: 'b;
+
         fn replace_client(&self, new_client: RetryClient<Client>);
         fn capabilities(&self) -> Option<Capabilities>;
         fn workers(&self) -> Arc<SlotManager>;
         fn is_mock(&self) -> bool;
         fn sdk_name_and_version(&self) -> (String, String);
+        fn get_identity(&self) -> String;
     }
 }

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

@@ -1,16 +1,21 @@
 //! Worker-specific client needs
 
 pub(crate) mod mocks;
+use crate::abstractions::dbg_panic;
+use crate::protosext::legacy_query_failure;
+use crate::worker::heartbeat::HeartbeatFn;
 use parking_lot::RwLock;
+use std::sync::OnceLock;
 use std::{sync::Arc, time::Duration};
 use temporal_client::{
     Client, IsWorkerTaskLongPoll, Namespace, NamespacedClient, NoRetryOnMatching, RetryClient,
     SlotManager, WorkflowService,
 };
 use temporal_sdk_core_api::worker::WorkerVersioningStrategy;
+use temporal_sdk_core_protos::temporal::api::worker::v1::WorkerHeartbeat;
 use temporal_sdk_core_protos::{
     TaskToken,
-    coresdk::workflow_commands::QueryResult,
+    coresdk::{workflow_commands::QueryResult, workflow_completion},
     temporal::api::{
         command::v1::Command,
         common::v1::{
@@ -34,12 +39,18 @@ use tonic::IntoRequest;
 
 type Result<T, E = tonic::Status> = std::result::Result<T, E>;
 
+pub enum LegacyQueryResult {
+    Succeeded(QueryResult),
+    Failed(workflow_completion::Failure),
+}
+
 /// Contains everything a worker needs to interact with the server
 pub(crate) struct WorkerClientBag {
     replaceable_client: RwLock<RetryClient<Client>>,
     namespace: String,
     identity: String,
     worker_versioning_strategy: WorkerVersioningStrategy,
+    heartbeat_data: Arc<OnceLock<HeartbeatFn>>,
 }
 
 impl WorkerClientBag {
@@ -48,12 +59,14 @@ impl WorkerClientBag {
         namespace: String,
         identity: String,
         worker_versioning_strategy: WorkerVersioningStrategy,
+        heartbeat_data: Arc<OnceLock<HeartbeatFn>>,
     ) -> Self {
         Self {
             replaceable_client: RwLock::new(client),
             namespace,
             identity,
             worker_versioning_strategy,
+            heartbeat_data,
         }
     }
 
@@ -114,6 +127,15 @@ impl WorkerClientBag {
             None
         }
     }
+
+    fn capture_heartbeat(&self) -> Option<WorkerHeartbeat> {
+        if let Some(hb) = self.heartbeat_data.get() {
+            hb()
+        } else {
+            dbg_panic!("Heartbeat function never set");
+            None
+        }
+    }
 }
 
 /// This trait contains everything workers need to interact with Temporal, and hence provides a
@@ -197,12 +219,17 @@ pub trait WorkerClient: Sync + Send {
     async fn respond_legacy_query(
         &self,
         task_token: TaskToken,
-        query_result: QueryResult,
+        query_result: LegacyQueryResult,
     ) -> Result<RespondQueryTaskCompletedResponse>;
     /// Describe the namespace
     async fn describe_namespace(&self) -> Result<DescribeNamespaceResponse>;
     /// Shutdown the worker
     async fn shutdown_worker(&self, sticky_task_queue: String) -> Result<ShutdownWorkerResponse>;
+    /// Record a worker heartbeat
+    async fn record_worker_heartbeat(
+        &self,
+        heartbeat: WorkerHeartbeat,
+    ) -> Result<RecordWorkerHeartbeatResponse>;
 
     /// Replace the underlying client
     fn replace_client(&self, new_client: RetryClient<Client>);
@@ -214,6 +241,8 @@ pub trait WorkerClient: Sync + Send {
     fn is_mock(&self) -> bool;
     /// Return name and version of the SDK
     fn sdk_name_and_version(&self) -> (String, String);
+    /// Get worker identity
+    fn get_identity(&self) -> String;
 }
 
 /// Configuration options shared by workflow, activity, and Nexus polling calls
@@ -267,6 +296,7 @@ impl WorkerClient for WorkerClientBag {
             binary_checksum: self.binary_checksum(),
             worker_version_capabilities: self.worker_version_capabilities(),
             deployment_options: self.deployment_options(),
+            worker_heartbeat: None,
         }
         .into_request();
         request.extensions_mut().insert(IsWorkerTaskLongPoll);
@@ -303,6 +333,7 @@ impl WorkerClient for WorkerClientBag {
             }),
             worker_version_capabilities: self.worker_version_capabilities(),
             deployment_options: self.deployment_options(),
+            worker_heartbeat: None,
         }
         .into_request();
         request.extensions_mut().insert(IsWorkerTaskLongPoll);
@@ -335,6 +366,7 @@ impl WorkerClient for WorkerClientBag {
             identity: self.identity.clone(),
             worker_version_capabilities: self.worker_version_capabilities(),
             deployment_options: self.deployment_options(),
+            worker_heartbeat: self.capture_heartbeat().into_iter().collect(),
         }
         .into_request();
         request.extensions_mut().insert(IsWorkerTaskLongPoll);
@@ -578,9 +610,20 @@ impl WorkerClient for WorkerClientBag {
     async fn respond_legacy_query(
         &self,
         task_token: TaskToken,
-        query_result: QueryResult,
+        query_result: LegacyQueryResult,
     ) -> Result<RespondQueryTaskCompletedResponse> {
+        let mut failure = None;
+        let (query_result, cause) = match query_result {
+            LegacyQueryResult::Succeeded(s) => (s, WorkflowTaskFailedCause::Unspecified),
+            #[allow(deprecated)]
+            LegacyQueryResult::Failed(f) => {
+                let cause = f.force_cause();
+                failure = f.failure.clone();
+                (legacy_query_failure(f), cause)
+            }
+        };
         let (_, completed_type, query_result, error_message) = query_result.into_components();
+
         Ok(self
             .cloned_client()
             .respond_query_task_completed(RespondQueryTaskCompletedRequest {
@@ -589,8 +632,8 @@ impl WorkerClient for WorkerClientBag {
                 query_result,
                 error_message,
                 namespace: self.namespace.clone(),
-                // TODO: https://github.com/temporalio/sdk-core/issues/867
-                failure: None,
+                failure,
+                cause: cause.into(),
             })
             .await?
             .into_inner())
@@ -612,6 +655,7 @@ impl WorkerClient for WorkerClientBag {
             identity: self.identity.clone(),
             sticky_task_queue,
             reason: "graceful shutdown".to_string(),
+            worker_heartbeat: self.capture_heartbeat(),
         };
 
         Ok(
@@ -626,6 +670,21 @@ impl WorkerClient for WorkerClientBag {
         *replaceable_client = new_client;
     }
 
+    async fn record_worker_heartbeat(
+        &self,
+        heartbeat: WorkerHeartbeat,
+    ) -> Result<RecordWorkerHeartbeatResponse> {
+        Ok(self
+            .cloned_client()
+            .record_worker_heartbeat(RecordWorkerHeartbeatRequest {
+                namespace: self.namespace.clone(),
+                identity: self.identity.clone(),
+                worker_heartbeat: vec![heartbeat],
+            })
+            .await?
+            .into_inner())
+    }
+
     fn capabilities(&self) -> Option<Capabilities> {
         let client = self.replaceable_client.read();
         client.get_client().inner().capabilities().cloned()
@@ -645,6 +704,10 @@ impl WorkerClient for WorkerClientBag {
         let opts = lock.get_client().inner().options();
         (opts.client_name.clone(), opts.client_version.clone())
     }
+
+    fn get_identity(&self) -> String {
+        self.identity.clone()
+    }
 }
 
 impl NamespacedClient for WorkerClientBag {