@temporalio/core-bridge 1.5.2 → 1.7.0

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

@@ -10,8 +10,8 @@ use opentelemetry::{
     },
     Context, KeyValue,
 };
-use std::{sync::Arc, time::Duration};
-use temporal_sdk_core_api::telemetry::CoreTelemetry;
+use std::{ops::Deref, sync::Arc, time::Duration};
+use temporal_client::ClientMetricProvider;
 
 /// Used to track context associated with metrics, and record/update them
 ///
@@ -24,6 +24,46 @@ pub(crate) struct MetricsContext {
     instruments: Arc<Instruments>,
 }
 
+/// Wraps OTel's [Meter] to ensure we name our metrics properly, or any other temporal-specific
+/// metrics customizations
+#[derive(derive_more::Constructor)]
+pub struct TemporalMeter<'a> {
+    inner: &'a Meter,
+    metrics_prefix: &'static str,
+}
+
+impl<'a> TemporalMeter<'a> {
+    pub(crate) fn counter(&self, name: &'static str) -> Counter<u64> {
+        self.inner
+            .u64_counter(self.metrics_prefix.to_string() + name)
+            .init()
+    }
+
+    pub(crate) fn histogram(&self, name: &'static str) -> Histogram<u64> {
+        self.inner
+            .u64_histogram(self.metrics_prefix.to_string() + name)
+            .init()
+    }
+}
+
+impl<'a> ClientMetricProvider for TemporalMeter<'a> {
+    fn counter(&self, name: &'static str) -> Counter<u64> {
+        self.counter(name)
+    }
+
+    fn histogram(&self, name: &'static str) -> Histogram<u64> {
+        self.histogram(name)
+    }
+}
+
+impl<'a> Deref for TemporalMeter<'a> {
+    type Target = dyn ClientMetricProvider + 'a;
+
+    fn deref(&self) -> &Self::Target {
+        self as &Self::Target
+    }
+}
+
 struct Instruments {
     wf_completed_counter: Counter<u64>,
     wf_canceled_counter: Counter<u64>,
@@ -54,10 +94,10 @@ impl MetricsContext {
         Self {
             ctx: Default::default(),
             kvs: Default::default(),
-            instruments: Arc::new(Instruments::new_explicit(
+            instruments: Arc::new(Instruments::new_explicit(TemporalMeter::new(
                 &NoopMeterProvider::new().meter("fakemeter"),
                 "fakemetrics",
-            )),
+            ))),
         }
     }
 
@@ -257,42 +297,36 @@ impl Instruments {
             meter
         } else {
             no_op_meter = NoopMeterProvider::default().meter("no_op");
-            &no_op_meter
+            TemporalMeter::new(&no_op_meter, "fakemetrics")
         };
-        Self::new_explicit(meter, telem.metric_prefix)
+        Self::new_explicit(meter)
     }
 
-    fn new_explicit(meter: &Meter, metric_prefix: &'static str) -> Self {
-        let ctr = |name: &'static str| -> Counter<u64> {
-            meter.u64_counter(metric_prefix.to_string() + name).init()
-        };
-        let hst = |name: &'static str| -> Histogram<u64> {
-            meter.u64_histogram(metric_prefix.to_string() + name).init()
-        };
+    fn new_explicit(meter: TemporalMeter) -> Self {
         Self {
-            wf_completed_counter: ctr("workflow_completed"),
-            wf_canceled_counter: ctr("workflow_canceled"),
-            wf_failed_counter: ctr("workflow_failed"),
-            wf_cont_counter: ctr("workflow_continue_as_new"),
-            wf_e2e_latency: hst(WF_E2E_LATENCY_NAME),
-            wf_task_queue_poll_empty_counter: ctr("workflow_task_queue_poll_empty"),
-            wf_task_queue_poll_succeed_counter: ctr("workflow_task_queue_poll_succeed"),
-            wf_task_execution_failure_counter: ctr("workflow_task_queue_poll_failed"),
-            wf_task_sched_to_start_latency: hst(WF_TASK_SCHED_TO_START_LATENCY_NAME),
-            wf_task_replay_latency: hst(WF_TASK_REPLAY_LATENCY_NAME),
-            wf_task_execution_latency: hst(WF_TASK_EXECUTION_LATENCY_NAME),
-            act_poll_no_task: ctr("activity_poll_no_task"),
-            act_task_received_counter: ctr("activity_task_received"),
-            act_execution_failed: ctr("activity_execution_failed"),
-            act_sched_to_start_latency: hst(ACT_SCHED_TO_START_LATENCY_NAME),
-            act_exec_latency: hst(ACT_EXEC_LATENCY_NAME),
+            wf_completed_counter: meter.counter("workflow_completed"),
+            wf_canceled_counter: meter.counter("workflow_canceled"),
+            wf_failed_counter: meter.counter("workflow_failed"),
+            wf_cont_counter: meter.counter("workflow_continue_as_new"),
+            wf_e2e_latency: meter.histogram(WF_E2E_LATENCY_NAME),
+            wf_task_queue_poll_empty_counter: meter.counter("workflow_task_queue_poll_empty"),
+            wf_task_queue_poll_succeed_counter: meter.counter("workflow_task_queue_poll_succeed"),
+            wf_task_execution_failure_counter: meter.counter("workflow_task_execution_failed"),
+            wf_task_sched_to_start_latency: meter.histogram(WF_TASK_SCHED_TO_START_LATENCY_NAME),
+            wf_task_replay_latency: meter.histogram(WF_TASK_REPLAY_LATENCY_NAME),
+            wf_task_execution_latency: meter.histogram(WF_TASK_EXECUTION_LATENCY_NAME),
+            act_poll_no_task: meter.counter("activity_poll_no_task"),
+            act_task_received_counter: meter.counter("activity_task_received"),
+            act_execution_failed: meter.counter("activity_execution_failed"),
+            act_sched_to_start_latency: meter.histogram(ACT_SCHED_TO_START_LATENCY_NAME),
+            act_exec_latency: meter.histogram(ACT_EXEC_LATENCY_NAME),
             // name kept as worker start for compat with old sdk / what users expect
-            worker_registered: ctr("worker_start"),
-            num_pollers: hst(NUM_POLLERS_NAME),
-            task_slots_available: hst(TASK_SLOTS_AVAILABLE_NAME),
-            sticky_cache_hit: ctr("sticky_cache_hit"),
-            sticky_cache_miss: ctr("sticky_cache_miss"),
-            sticky_cache_size: hst(STICKY_CACHE_SIZE_NAME),
+            worker_registered: meter.counter("worker_start"),
+            num_pollers: meter.histogram(NUM_POLLERS_NAME),
+            task_slots_available: meter.histogram(TASK_SLOTS_AVAILABLE_NAME),
+            sticky_cache_hit: meter.counter("sticky_cache_hit"),
+            sticky_cache_miss: meter.counter("sticky_cache_miss"),
+            sticky_cache_size: meter.histogram(STICKY_CACHE_SIZE_NAME),
         }
     }
 }

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

@@ -25,7 +25,18 @@ use opentelemetry::{
 };
 use opentelemetry_otlp::WithExportConfig;
 use parking_lot::Mutex;
-use std::{cell::RefCell, collections::VecDeque, convert::TryInto, env, sync::Arc, time::Duration};
+use std::{
+    cell::RefCell,
+    collections::{HashMap, VecDeque},
+    convert::TryInto,
+    env,
+    net::SocketAddr,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
 use temporal_sdk_core_api::telemetry::{
     CoreLog, CoreTelemetry, Logger, MetricTemporality, MetricsExporter, OtelCollectorOptions,
     TelemetryOptions, TraceExporter,
@@ -37,12 +48,10 @@ use tracing_subscriber::{layer::SubscriberExt, EnvFilter, Layer};
 const TELEM_SERVICE_NAME: &str = "temporal-core-sdk";
 
 /// 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_levl.
+/// 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!(
-        "{o},temporal_sdk_core={l},temporal_client={l},temporal_sdk={l}",
-        o = other_level,
-        l = core_level
+        "{other_level},temporal_sdk_core={core_level},temporal_client={core_level},temporal_sdk={core_level}"
     )
 }
 
@@ -52,6 +61,7 @@ pub struct TelemetryInstance {
     logs_out: Option<Mutex<CoreLogsOut>>,
     metrics: Option<(Box<dyn MeterProvider + Send + Sync + 'static>, Meter)>,
     trace_subscriber: Arc<dyn Subscriber + Send + Sync>,
+    prom_binding: Option<SocketAddr>,
     _keepalive_rx: Receiver<()>,
 }
 
@@ -61,6 +71,7 @@ impl TelemetryInstance {
         logs_out: Option<Mutex<CoreLogsOut>>,
         metric_prefix: &'static str,
         mut meter_provider: Option<Box<dyn MeterProvider + Send + Sync + 'static>>,
+        prom_binding: Option<SocketAddr>,
         keepalive_rx: Receiver<()>,
     ) -> Self {
         let metrics = meter_provider.take().map(|mp| {
@@ -72,6 +83,7 @@ impl TelemetryInstance {
             logs_out,
             metrics,
             trace_subscriber,
+            prom_binding,
             _keepalive_rx: keepalive_rx,
         }
     }
@@ -81,6 +93,18 @@ impl TelemetryInstance {
     pub fn trace_subscriber(&self) -> Arc<dyn Subscriber + Send + Sync> {
         self.trace_subscriber.clone()
     }
+
+    /// Returns the address the Prometheus server is bound to if it is running
+    pub fn prom_port(&self) -> Option<SocketAddr> {
+        self.prom_binding
+    }
+
+    /// Returns our wrapper for OTel metric meters, can be used to, ex: initialize clients
+    pub fn get_metric_meter(&self) -> Option<TemporalMeter> {
+        self.metrics
+            .as_ref()
+            .map(|(_, m)| TemporalMeter::new(m, self.metric_prefix))
+    }
 }
 
 thread_local! {
@@ -121,10 +145,6 @@ impl CoreTelemetry for TelemetryInstance {
             vec![]
         }
     }
-
-    fn get_metric_meter(&self) -> Option<&Meter> {
-        self.metrics.as_ref().map(|(_, m)| m)
-    }
 }
 
 /// Initialize tracing subscribers/output and logging export, returning a [TelemetryInstance]
@@ -151,6 +171,7 @@ pub fn telemetry_init(opts: TelemetryOptions) -> Result<TelemetryInstance, anyho
         // Parts of telem dat ====
         let mut logs_out = None;
         let metric_prefix = metric_prefix(&opts);
+        let mut prom_binding = None;
         // =======================
 
         // Tracing subscriber layers =========
@@ -200,11 +221,15 @@ pub fn telemetry_init(opts: TelemetryOptions) -> Result<TelemetryInstance, anyho
             let aggregator = SDKAggSelector { metric_prefix };
             match metrics {
                 MetricsExporter::Prometheus(addr) => {
-                    let srv = PromServer::new(
-                        *addr,
-                        aggregator,
-                        metric_temporality_to_selector(opts.metric_temporality),
-                    )?;
+                    let srv = runtime.block_on(async {
+                        PromServer::new(
+                            *addr,
+                            aggregator,
+                            metric_temporality_to_selector(opts.metric_temporality),
+                            &opts.global_tags,
+                        )
+                    })?;
+                    prom_binding = Some(srv.bound_addr());
                     let mp = srv.exporter.meter_provider()?;
                     runtime.spawn(async move { srv.run().await });
                     Some(Box::new(mp) as Box<dyn MeterProvider + Send + Sync>)
@@ -221,7 +246,7 @@ pub fn telemetry_init(opts: TelemetryOptions) -> Result<TelemetryInstance, anyho
                             runtime::Tokio,
                         )
                         .with_period(metric_periodicity.unwrap_or_else(|| Duration::from_secs(1)))
-                        .with_resource(default_resource())
+                        .with_resource(default_resource(&opts.global_tags))
                         .with_exporter(
                             opentelemetry_otlp::new_exporter()
                                 .tonic()
@@ -242,7 +267,8 @@ pub fn telemetry_init(opts: TelemetryOptions) -> Result<TelemetryInstance, anyho
             match &tracing.exporter {
                 TraceExporter::Otel(OtelCollectorOptions { url, headers, .. }) => {
                     runtime.block_on(async {
-                        let tracer_cfg = Config::default().with_resource(default_resource());
+                        let tracer_cfg =
+                            Config::default().with_resource(default_resource(&opts.global_tags));
                         let tracer = opentelemetry_otlp::new_pipeline()
                             .tracing()
                             .with_exporter(
@@ -276,6 +302,7 @@ pub fn telemetry_init(opts: TelemetryOptions) -> Result<TelemetryInstance, anyho
             logs_out,
             metric_prefix,
             meter_provider,
+            prom_binding,
             keepalive_rx,
         ))
         .expect("Must be able to send telem instance out of thread");
@@ -295,10 +322,17 @@ pub fn telemetry_init(opts: TelemetryOptions) -> Result<TelemetryInstance, anyho
     }
 }
 
-/// Initialize telemetry/tracing globally. Useful for testing.
+/// 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> {
-    let ti = telemetry_init(opts)?;
-    tracing::subscriber::set_global_default(ti.trace_subscriber())?;
+    static INITTED: AtomicBool = AtomicBool::new(false);
+    if INITTED
+        .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
+        .is_ok()
+    {
+        let ti = telemetry_init(opts)?;
+        tracing::subscriber::set_global_default(ti.trace_subscriber())?;
+    }
     Ok(())
 }
 
@@ -306,8 +340,12 @@ fn default_resource_kvs() -> &'static [KeyValue] {
     static INSTANCE: OnceCell<[KeyValue; 1]> = OnceCell::new();
     INSTANCE.get_or_init(|| [KeyValue::new("service.name", TELEM_SERVICE_NAME)])
 }
-fn default_resource() -> Resource {
-    Resource::new(default_resource_kvs().iter().cloned())
+
+fn default_resource(override_values: &HashMap<String, String>) -> Resource {
+    let override_kvs = override_values
+        .iter()
+        .map(|(k, v)| KeyValue::new(k.clone(), v.clone()));
+    Resource::new(default_resource_kvs().iter().cloned()).merge(&Resource::new(override_kvs))
 }
 
 fn metric_temporality_to_selector(
@@ -360,6 +398,7 @@ pub mod test_initters {
         .unwrap();
     }
 }
+use crate::telemetry::metrics::TemporalMeter;
 #[cfg(test)]
 pub use test_initters::*;
 

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

@@ -1,23 +1,21 @@
 use crate::telemetry::default_resource;
 use hyper::{
     header::CONTENT_TYPE,
+    server::conn::AddrIncoming,
     service::{make_service_fn, service_fn},
     Body, Method, Request, Response, Server,
 };
-use opentelemetry::{
-    metrics::MetricsError,
-    sdk::{
-        export::metrics::{aggregation::TemporalitySelector, AggregatorSelector},
-        metrics::{controllers, processors},
-    },
+use opentelemetry::sdk::{
+    export::metrics::{aggregation::TemporalitySelector, AggregatorSelector},
+    metrics::{controllers, processors},
 };
 use opentelemetry_prometheus::{ExporterBuilder, PrometheusExporter};
 use prometheus::{Encoder, TextEncoder};
-use std::{convert::Infallible, net::SocketAddr, sync::Arc};
+use std::{collections::HashMap, convert::Infallible, net::SocketAddr, sync::Arc, time::Duration};
 
 /// Exposes prometheus metrics for scraping
 pub(super) struct PromServer {
-    addr: SocketAddr,
+    bound_addr: AddrIncoming,
     pub exporter: Arc<PrometheusExporter>,
 }
 
@@ -26,19 +24,23 @@ impl PromServer {
         addr: SocketAddr,
         aggregation: impl AggregatorSelector + Send + Sync + 'static,
         temporality: impl TemporalitySelector + Send + Sync + 'static,
-    ) -> Result<Self, MetricsError> {
+        tags: &HashMap<String, String>,
+    ) -> Result<Self, anyhow::Error> {
         let controller =
             controllers::basic(processors::factory(aggregation, temporality).with_memory(true))
-                .with_resource(default_resource())
+                // Because Prom is pull-based, make this always refresh
+                .with_collect_period(Duration::from_secs(0))
+                .with_resource(default_resource(tags))
                 .build();
         let exporter = ExporterBuilder::new(controller).try_init()?;
+        let bound_addr = AddrIncoming::bind(&addr)?;
         Ok(Self {
             exporter: Arc::new(exporter),
-            addr,
+            bound_addr,
         })
     }
 
-    pub async fn run(&self) -> hyper::Result<()> {
+    pub async fn run(self) -> hyper::Result<()> {
         // Spin up hyper server to serve metrics for scraping. We use hyper since we already depend
         // on it via Tonic.
         let expclone = self.exporter.clone();
@@ -46,9 +48,13 @@ impl PromServer {
             let expclone = expclone.clone();
             async move { Ok::<_, Infallible>(service_fn(move |req| metrics_req(req, expclone.clone()))) }
         });
-        let server = Server::bind(&self.addr).serve(svc);
+        let server = Server::builder(self.bound_addr).serve(svc);
         server.await
     }
+
+    pub fn bound_addr(&self) -> SocketAddr {
+        self.bound_addr.local_addr()
+    }
 }
 
 /// Serves prometheus metrics in the expected format for scraping

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

@@ -14,6 +14,7 @@ use crate::{
     },
     TaskToken, Worker, WorkerConfig, WorkerConfigBuilder,
 };
+use async_trait::async_trait;
 use bimap::BiMap;
 use futures::{future::BoxFuture, stream, stream::BoxStream, FutureExt, Stream, StreamExt};
 use mockall::TimesRange;
@@ -29,7 +30,10 @@ use std::{
     task::{Context, Poll},
     time::Duration,
 };
-use temporal_sdk_core_api::Worker as WorkerTrait;
+use temporal_sdk_core_api::{
+    errors::{PollActivityError, PollWfError},
+    Worker as WorkerTrait,
+};
 use temporal_sdk_core_protos::{
     coresdk::{
         workflow_activation::WorkflowActivation,
@@ -52,7 +56,6 @@ use tokio_stream::wrappers::UnboundedReceiverStream;
 use tokio_util::sync::CancellationToken;
 
 pub const TEST_Q: &str = "q";
-pub static NO_MORE_WORK_ERROR_MSG: &str = "No more work to do";
 
 pub fn test_worker_cfg() -> WorkerConfigBuilder {
     let mut wcb = WorkerConfigBuilder::default();
@@ -148,6 +151,7 @@ pub(crate) fn mock_worker(mocks: MocksHolder) -> Worker {
         mocks.inputs.wft_stream,
         act_poller,
         MetricsContext::no_op(),
+        None,
         CancellationToken::new(),
     )
 }
@@ -301,7 +305,7 @@ where
             }
             .boxed()
         } else {
-            async { Some(Err(tonic::Status::cancelled(NO_MORE_WORK_ERROR_MSG))) }.boxed()
+            async { None }.boxed()
         }
     });
     Box::new(mock_poller) as BoxedPoller<T>
@@ -375,10 +379,12 @@ pub(crate) struct MockPollCfg {
     pub expect_fail_wft_matcher:
         Box<dyn Fn(&TaskToken, &WorkflowTaskFailedCause, &Option<Failure>) -> bool + Send>,
     pub completion_asserts: Option<Box<dyn Fn(&WorkflowTaskCompletion) + Send>>,
+    pub num_expected_completions: Option<TimesRange>,
     /// If being used with the Rust SDK, this is set true. It ensures pollers will not error out
     /// early with no work, since we cannot know the exact number of times polling will happen.
     /// Instead, they will just block forever.
     pub using_rust_sdk: bool,
+    pub make_poll_stream_interminable: bool,
 }
 
 impl MockPollCfg {
@@ -395,7 +401,9 @@ impl MockPollCfg {
             mock_client: mock_workflow_client(),
             expect_fail_wft_matcher: Box::new(|_, _, _| true),
             completion_asserts: None,
+            num_expected_completions: None,
             using_rust_sdk: false,
+            make_poll_stream_interminable: false,
         }
     }
     pub fn from_resp_batches(
@@ -416,7 +424,9 @@ impl MockPollCfg {
             mock_client,
             expect_fail_wft_matcher: Box::new(|_, _, _| true),
             completion_asserts: None,
+            num_expected_completions: None,
             using_rust_sdk: false,
+            make_poll_stream_interminable: false,
         }
     }
 }
@@ -584,16 +594,20 @@ pub(crate) fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder {
     );
 
     let outstanding = outstanding_wf_task_tokens.clone();
-    cfg.mock_client
-        .expect_complete_workflow_task()
-        .returning(move |comp| {
-            if let Some(ass) = cfg.completion_asserts.as_ref() {
-                // tee hee
-                ass(&comp)
-            }
-            outstanding.release_token(&comp.task_token);
-            Ok(RespondWorkflowTaskCompletedResponse::default())
-        });
+    let expect_completes = cfg.mock_client.expect_complete_workflow_task();
+    if let Some(range) = cfg.num_expected_completions {
+        expect_completes.times(range);
+    } else if cfg.completion_asserts.is_some() {
+        expect_completes.times(1..);
+    }
+    expect_completes.returning(move |comp| {
+        if let Some(ass) = cfg.completion_asserts.as_ref() {
+            // tee hee
+            ass(&comp)
+        }
+        outstanding.release_token(&comp.task_token);
+        Ok(RespondWorkflowTaskCompletedResponse::default())
+    });
     let outstanding = outstanding_wf_task_tokens.clone();
     cfg.mock_client
         .expect_fail_workflow_task()
@@ -612,11 +626,15 @@ pub(crate) fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder {
             Ok(Default::default())
         });
 
-    MocksHolder {
+    let mut mh = MocksHolder {
         client: Arc::new(cfg.mock_client),
         inputs: mock_worker,
         outstanding_task_map: Some(outstanding_wf_task_tokens),
+    };
+    if cfg.make_poll_stream_interminable {
+        mh.make_wft_stream_interminable();
     }
+    mh
 }
 
 pub struct QueueResponse<T> {
@@ -835,6 +853,7 @@ pub(crate) fn gen_assert_and_fail(asserter: &dyn Fn(&WorkflowActivation)) -> Ass
                 message: "Intentional test failure".to_string(),
                 ..Default::default()
             }),
+            ..Default::default()
         }
         .into(),
     )
@@ -880,3 +899,43 @@ macro_rules! prost_dur {
             .expect("test duration fits")
     };
 }
+
+#[async_trait]
+pub(crate) trait WorkerExt {
+    /// Initiate shutdown, drain the pollers, and wait for shutdown to complete.
+    async fn drain_pollers_and_shutdown(self);
+    /// Initiate shutdown, drain the *activity* poller, and wait for shutdown to complete.
+    /// Takes a ref because of that one test that needs it.
+    async fn drain_activity_poller_and_shutdown(&self);
+}
+
+#[async_trait]
+impl WorkerExt for Worker {
+    async fn drain_pollers_and_shutdown(self) {
+        self.initiate_shutdown();
+        tokio::join!(
+            async {
+                assert_matches!(
+                    self.poll_activity_task().await.unwrap_err(),
+                    PollActivityError::ShutDown
+                );
+            },
+            async {
+                assert_matches!(
+                    self.poll_workflow_activation().await.unwrap_err(),
+                    PollWfError::ShutDown
+                );
+            }
+        );
+        self.finalize_shutdown().await;
+    }
+
+    async fn drain_activity_poller_and_shutdown(&self) {
+        self.initiate_shutdown();
+        assert_matches!(
+            self.poll_activity_task().await.unwrap_err(),
+            PollActivityError::ShutDown
+        );
+        self.shutdown().await;
+    }
+}