@temporalio/core-bridge 1.6.0 → 1.7.1

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

@@ -27,9 +27,10 @@ use opentelemetry_otlp::WithExportConfig;
 use parking_lot::Mutex;
 use std::{
     cell::RefCell,
-    collections::VecDeque,
+    collections::{HashMap, VecDeque},
     convert::TryInto,
     env,
+    net::SocketAddr,
     sync::{
         atomic::{AtomicBool, Ordering},
         Arc,
@@ -60,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<()>,
 }
 
@@ -69,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| {
@@ -80,6 +83,7 @@ impl TelemetryInstance {
             logs_out,
             metrics,
             trace_subscriber,
+            prom_binding,
             _keepalive_rx: keepalive_rx,
         }
     }
@@ -89,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! {
@@ -129,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]
@@ -159,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 =========
@@ -208,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>)
@@ -229,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()
@@ -250,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(
@@ -284,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");
@@ -321,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(
@@ -375,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,6 +379,7 @@ 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.
@@ -396,6 +401,7 @@ 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,
         }
@@ -418,6 +424,7 @@ 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,
         }
@@ -587,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()
@@ -842,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(),
     )
@@ -887,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;
+    }
+}