@temporalio/core-bridge 1.5.2 → 1.7.0

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

@@ -1,7 +1,7 @@
 use crate::{AttachMetricLabels, LONG_POLL_METHOD_NAMES};
 use futures::{future::BoxFuture, FutureExt};
 use opentelemetry::{
-    metrics::{Counter, Histogram, Meter},
+    metrics::{Counter, Histogram},
     KeyValue,
 };
 use std::{
@@ -30,18 +30,27 @@ pub struct MetricsContext {
     long_svc_request_latency: Histogram<u64>,
 }
 
+/// Things that can provide metrics for the client implement this. Trait exists to avoid having
+/// to make a whole new lower-level crate just for a tiny shared wrapper around OTel meters.
+pub trait ClientMetricProvider: Send + Sync {
+    /// Construct a counter metric
+    fn counter(&self, name: &'static str) -> Counter<u64>;
+    /// Construct a histogram metric
+    fn histogram(&self, name: &'static str) -> Histogram<u64>;
+}
+
 impl MetricsContext {
-    pub(crate) fn new(kvs: Vec<KeyValue>, meter: &Meter) -> Self {
+    pub(crate) fn new(kvs: Vec<KeyValue>, metric_provider: &dyn ClientMetricProvider) -> Self {
         Self {
             ctx: opentelemetry::Context::current(),
             kvs: Arc::new(kvs),
             poll_is_long: false,
-            svc_request: meter.u64_counter("request").init(),
-            svc_request_failed: meter.u64_counter("request_failure").init(),
-            long_svc_request: meter.u64_counter("long_request").init(),
-            long_svc_request_failed: meter.u64_counter("long_request_failure").init(),
-            svc_request_latency: meter.u64_histogram("request_latency").init(),
-            long_svc_request_latency: meter.u64_histogram("long_request_latency").init(),
+            svc_request: metric_provider.counter("request"),
+            svc_request_failed: metric_provider.counter("request_failure"),
+            long_svc_request: metric_provider.counter("long_request"),
+            long_svc_request_failed: metric_provider.counter("long_request_failure"),
+            svc_request_latency: metric_provider.histogram("request_latency"),
+            long_svc_request_latency: metric_provider.histogram("long_request_latency"),
         }
     }
 

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

@@ -760,9 +760,9 @@ proxier! {
         }
     );
     (
-        update_workflow,
-        UpdateWorkflowRequest,
-        UpdateWorkflowResponse,
+        update_workflow_execution,
+        UpdateWorkflowExecutionRequest,
+        UpdateWorkflowExecutionResponse,
         |r| {
             let labels = AttachMetricLabels::namespace(r.get_ref().namespace.clone());
             r.extensions_mut().insert(labels);
@@ -898,7 +898,7 @@ mod tests {
         let no_underscores: HashSet<_> = impl_list.iter().map(|x| x.replace('_', "")).collect();
         for method in methods {
             if !no_underscores.contains(&method.to_lowercase()) {
-                panic!("RPC method {} is not implemented by raw client", method)
+                panic!("RPC method {method} is not implemented by raw client")
             }
         }
     }

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

@@ -1,6 +1,6 @@
 use crate::{
-    ClientOptions, ListClosedFilters, ListOpenFilters, Namespace, Result, RetryConfig,
-    StartTimeFilter, WorkflowClientTrait, WorkflowOptions,
+    ClientOptions, ListClosedFilters, ListOpenFilters, Namespace, RegisterNamespaceOptions, Result,
+    RetryConfig, SignalWithStartOptions, StartTimeFilter, WorkflowClientTrait, WorkflowOptions,
 };
 use backoff::{backoff::Backoff, exponential::ExponentialBackoff, Clock, SystemClock};
 use futures_retry::{ErrorHandler, FutureRetry, RetryPolicy};
@@ -8,7 +8,7 @@ use std::{fmt::Debug, future::Future, sync::Arc, time::Duration};
 use temporal_sdk_core_protos::{
     coresdk::workflow_commands::QueryResult,
     temporal::api::{
-        common::v1::{Header, Payload, Payloads},
+        common::v1::{Payload, Payloads},
         enums::v1::WorkflowTaskFailedCause,
         failure::v1::Failure,
         query::v1::WorkflowQuery,
@@ -362,28 +362,14 @@ where
 
     async fn signal_with_start_workflow_execution(
         &self,
-        input: Option<Payloads>,
-        task_queue: String,
-        workflow_id: String,
-        workflow_type: String,
-        request_id: Option<String>,
-        options: WorkflowOptions,
-        signal_name: String,
-        signal_input: Option<Payloads>,
-        signal_header: Option<Header>,
+        options: SignalWithStartOptions,
+        workflow_options: WorkflowOptions,
     ) -> Result<SignalWithStartWorkflowExecutionResponse> {
         retry_call!(
             self,
             signal_with_start_workflow_execution,
-            input.clone(),
-            task_queue.clone(),
-            workflow_id.clone(),
-            workflow_type.clone(),
-            request_id.clone(),
             options.clone(),
-            signal_name.clone(),
-            signal_input.clone(),
-            signal_header.clone()
+            workflow_options.clone()
         )
     }
 
@@ -473,6 +459,13 @@ where
         )
     }
 
+    async fn register_namespace(
+        &self,
+        options: RegisterNamespaceOptions,
+    ) -> Result<RegisterNamespaceResponse> {
+        retry_call!(self, register_namespace, options.clone())
+    }
+
     async fn list_namespaces(&self) -> Result<ListNamespacesResponse> {
         retry_call!(self, list_namespaces,)
     }
@@ -530,6 +523,25 @@ where
         )
     }
 
+    async fn list_archived_workflow_executions(
+        &self,
+        page_size: i32,
+        next_page_token: Vec<u8>,
+        query: String,
+    ) -> Result<ListArchivedWorkflowExecutionsResponse> {
+        retry_call!(
+            self,
+            list_archived_workflow_executions,
+            page_size,
+            next_page_token.clone(),
+            query.clone()
+        )
+    }
+
+    async fn get_search_attributes(&self) -> Result<GetSearchAttributesResponse> {
+        retry_call!(self, get_search_attributes)
+    }
+
     fn get_options(&self) -> &ClientOptions {
         self.client.get_options()
     }

package/sdk-core/core/Cargo.toml

@@ -12,12 +12,17 @@ categories = ["development-tools"]
 
 [lib]
 
+[features]
+# Do not enable this feature when building production SDKs. If we ever want a user in the field to
+# record WF input data, we can build them a custom SDK or they can build - it adds significant extra
+# code size in the form of [de]serializers.
+save_wf_inputs = ["rmp-serde", "temporal-sdk-core-protos/serde_serialize"]
+
 [dependencies]
 anyhow = "1.0"
 arc-swap = "1.3"
-async-channel = "1.6"
 async-trait = "0.1"
-base64 = "0.13"
+base64 = "0.21"
 crossbeam = "0.8"
 dashmap = "5.0"
 derive_builder = "0.12"
@@ -31,8 +36,7 @@ http = "0.2"
 hyper = "0.14"
 itertools = "0.10"
 lazy_static = "1.4"
-log = "0.4"
-lru = "0.8"
+lru = "0.10"
 mockall = "0.11"
 nix = "0.26"
 once_cell = "1.5"
@@ -40,29 +44,31 @@ opentelemetry = { version = "0.18", features = ["rt-tokio"] }
 opentelemetry-otlp = { version = "0.11", features = ["tokio", "metrics"] }
 opentelemetry-prometheus = "0.11"
 parking_lot = { version = "0.12", features = ["send_guard"] }
+pin-project = "1.0"
 prometheus = "0.13"
 prost = "0.11"
-prost-types = "0.11"
+prost-types = { version = "0.4", package = "prost-wkt-types" }
 rand = "0.8.3"
 reqwest = { version = "0.11", features = ["json", "stream", "rustls-tls", "tokio-rustls"], default-features = false }
 ringbuf = "0.3"
+rmp-serde = { version = "1.1", optional = true }
 serde = "1.0"
 serde_json = "1.0"
 siphasher = "0.3"
 slotmap = "1.0"
 tar = "0.4"
 thiserror = "1.0"
-tokio = { version = "1.1", features = ["rt", "rt-multi-thread", "parking_lot", "time", "fs", "process"] }
+tokio = { version = "1.26", features = ["rt", "rt-multi-thread", "parking_lot", "time", "fs", "process"] }
 tokio-util = { version = "0.7", features = ["io", "io-util"] }
 tokio-stream = "0.1"
 tonic = { version = "0.8", features = ["tls", "tls-roots"] }
-tracing = { version = "0.1", features = ["log-always"] }
+tracing = "0.1"
 tracing-futures = "0.2"
 tracing-opentelemetry = "0.18"
 tracing-subscriber = { version = "0.3", features = ["parking_lot", "env-filter", "registry"] }
 url = "2.2"
 uuid = { version = "1.1", features = ["v4"] }
-zip = "0.6"
+zip = "0.6.3"
 
 # 1st party local deps
 [dependencies.temporal-sdk-core-api]
@@ -85,8 +91,9 @@ version = "0.1"
 [dev-dependencies]
 assert_matches = "1.4"
 bimap = "0.6.1"
+clap = { version = "4.0", features = ["derive"] }
 criterion = "0.4"
-rstest = "0.16"
+rstest = "0.17"
 temporal-sdk-core-test-utils = { path = "../test-utils" }
 temporal-sdk = { path = "../sdk" }
 
@@ -101,9 +108,10 @@ path = "../tests/main.rs"
 test = false
 
 [[test]]
-name = "load_tests"
-path = "../tests/load_tests.rs"
+name = "heavy_tests"
+path = "../tests/heavy_tests.rs"
 test = false
+required-features = ["save_wf_inputs"]
 
 [[bench]]
 name = "workflow_replay"
@@ -113,4 +121,9 @@ harness = false
 # the dev-dependencies, which we want.
 [[example]]
 name = "integ_runner"
-path = "../tests/runner.rs"
+path = "../tests/runner.rs"
+
+[[example]]
+name = "wf_input_replay"
+path = "../tests/wf_input_replay.rs"
+required-features = ["save_wf_inputs"]

package/sdk-core/core/src/abstractions/take_cell.rs

@@ -0,0 +1,28 @@
+use parking_lot::Mutex;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+/// Implements something a bit like a `OnceCell`, but starts already initialized and allows you
+/// to take everything out of it only once in a thread-safe way. This isn't optimized for super
+/// fast-path usage.
+pub struct TakeCell<T> {
+    taken: AtomicBool,
+    data: Mutex<Option<T>>,
+}
+
+impl<T> TakeCell<T> {
+    pub fn new(val: T) -> Self {
+        Self {
+            taken: AtomicBool::new(false),
+            data: Mutex::new(Some(val)),
+        }
+    }
+
+    /// If the cell has not already been taken from, takes the value and returns it
+    pub fn take_once(&self) -> Option<T> {
+        if self.taken.load(Ordering::Acquire) {
+            return None;
+        }
+        self.taken.store(true, Ordering::Release);
+        self.data.lock().take()
+    }
+}

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

@@ -1,18 +1,30 @@
 //! This module contains very generic helpers that can be used codebase-wide
 
+pub mod take_cell;
+
 use crate::MetricsContext;
+use derive_more::DebugCustom;
 use futures::{stream, Stream, StreamExt};
 use std::{
     fmt::{Debug, Formatter},
-    sync::Arc,
+    sync::{
+        atomic::{AtomicBool, AtomicUsize, Ordering},
+        Arc,
+    },
 };
 use tokio::sync::{AcquireError, OwnedSemaphorePermit, Semaphore, TryAcquireError};
+use tokio_util::sync::CancellationToken;
 
 /// Wraps a [Semaphore] with a function call that is fed the available permits any time a permit is
 /// acquired or restored through the provided methods
 #[derive(Clone)]
 pub(crate) struct MeteredSemaphore {
     sem: Arc<Semaphore>,
+    /// The number of permit owners who have acquired a permit from the semaphore, but are not yet
+    /// meaningfully using that permit. This is useful for giving a more semantically accurate count
+    /// of used task slots, since we typically wait for a permit first before polling, but that slot
+    /// isn't used in the sense the user expects until we actually also get the corresponding task.
+    unused_claimants: Arc<AtomicUsize>,
     metrics_ctx: MetricsContext,
     record_fn: fn(&MetricsContext, usize),
 }
@@ -25,6 +37,7 @@ impl MeteredSemaphore {
     ) -> Self {
         Self {
             sem: Arc::new(Semaphore::new(inital_permits)),
+            unused_claimants: Arc::new(AtomicUsize::new(0)),
             metrics_ctx,
             record_fn,
         }
@@ -36,42 +49,154 @@ impl MeteredSemaphore {
 
     pub async fn acquire_owned(&self) -> Result<OwnedMeteredSemPermit, AcquireError> {
         let res = self.sem.clone().acquire_owned().await?;
-        self.record();
-        Ok(OwnedMeteredSemPermit {
-            inner: res,
-            record_fn: self.record_drop_owned(),
-        })
+        Ok(self.build_owned(res))
     }
 
     pub fn try_acquire_owned(&self) -> Result<OwnedMeteredSemPermit, TryAcquireError> {
         let res = self.sem.clone().try_acquire_owned()?;
+        Ok(self.build_owned(res))
+    }
+
+    fn build_owned(&self, res: OwnedSemaphorePermit) -> OwnedMeteredSemPermit {
+        self.unused_claimants.fetch_add(1, Ordering::Release);
         self.record();
-        Ok(OwnedMeteredSemPermit {
+        OwnedMeteredSemPermit {
             inner: res,
-            record_fn: self.record_drop_owned(),
-        })
+            unused_claimants: Some(self.unused_claimants.clone()),
+            record_fn: self.record_owned(),
+        }
     }
 
     fn record(&self) {
-        (self.record_fn)(&self.metrics_ctx, self.sem.available_permits());
+        (self.record_fn)(
+            &self.metrics_ctx,
+            self.sem.available_permits() + self.unused_claimants.load(Ordering::Acquire),
+        );
     }
 
-    fn record_drop_owned(&self) -> Box<dyn Fn() + Send + Sync> {
+    fn record_owned(&self) -> Box<dyn Fn(bool) + Send + Sync> {
         let rcf = self.record_fn;
         let mets = self.metrics_ctx.clone();
         let sem = self.sem.clone();
-        Box::new(move || rcf(&mets, sem.available_permits() + 1))
+        let uc = self.unused_claimants.clone();
+        // When being called from the drop impl, the semaphore permit isn't actually dropped yet,
+        // so account for that.
+        Box::new(move |add_one: bool| {
+            let extra = usize::from(add_one);
+            rcf(
+                &mets,
+                sem.available_permits() + uc.load(Ordering::Acquire) + extra,
+            )
+        })
+    }
+}
+
+/// A version of [MeteredSemaphore] that can be closed and supports waiting for close to complete.
+/// Once closed, no permits will be handed out.
+/// Close completes when all permits have been returned.
+pub(crate) struct ClosableMeteredSemaphore {
+    inner: Arc<MeteredSemaphore>,
+    outstanding_permits: AtomicUsize,
+    close_requested: AtomicBool,
+    close_complete_token: CancellationToken,
+}
+
+impl ClosableMeteredSemaphore {
+    pub fn new_arc(sem: Arc<MeteredSemaphore>) -> Arc<Self> {
+        Arc::new(Self {
+            inner: sem,
+            outstanding_permits: Default::default(),
+            close_requested: AtomicBool::new(false),
+            close_complete_token: CancellationToken::new(),
+        })
+    }
+}
+
+impl ClosableMeteredSemaphore {
+    #[cfg(test)]
+    pub fn available_permits(&self) -> usize {
+        self.inner.available_permits()
+    }
+
+    /// Request to close the semaphore and prevent new permits from being acquired.
+    pub fn close(&self) {
+        self.close_requested.store(true, Ordering::Release);
+        if self.outstanding_permits.load(Ordering::Acquire) == 0 {
+            self.close_complete_token.cancel();
+        }
+    }
+
+    /// Returns after close has been requested and all outstanding permits have been returned.
+    pub async fn close_complete(&self) {
+        self.close_complete_token.cancelled().await;
+    }
+
+    /// Acquire a permit if one is available and close was not requested.
+    pub fn try_acquire_owned(
+        self: &Arc<Self>,
+    ) -> Result<TrackedOwnedMeteredSemPermit, TryAcquireError> {
+        if self.close_requested.load(Ordering::Acquire) {
+            return Err(TryAcquireError::Closed);
+        }
+        self.outstanding_permits.fetch_add(1, Ordering::Release);
+        let res = self.inner.try_acquire_owned();
+        if res.is_err() {
+            self.outstanding_permits.fetch_sub(1, Ordering::Release);
+        }
+        res.map(|permit| TrackedOwnedMeteredSemPermit {
+            inner: Some(permit),
+            on_drop: self.on_permit_dropped(),
+        })
+    }
+
+    fn on_permit_dropped(self: &Arc<Self>) -> Box<dyn Fn() + Send + Sync> {
+        let sem = self.clone();
+        Box::new(move || {
+            sem.outstanding_permits.fetch_sub(1, Ordering::Release);
+            if sem.close_requested.load(Ordering::Acquire)
+                && sem.outstanding_permits.load(Ordering::Acquire) == 0
+            {
+                sem.close_complete_token.cancel();
+            }
+        })
+    }
+}
+
+/// Tracks an OwnedMeteredSemPermit and calls on_drop when dropped.
+#[derive(DebugCustom)]
+#[debug(fmt = "Tracked({inner:?})")]
+pub(crate) struct TrackedOwnedMeteredSemPermit {
+    inner: Option<OwnedMeteredSemPermit>,
+    on_drop: Box<dyn Fn() + Send + Sync>,
+}
+impl From<TrackedOwnedMeteredSemPermit> for OwnedMeteredSemPermit {
+    fn from(mut value: TrackedOwnedMeteredSemPermit) -> Self {
+        value
+            .inner
+            .take()
+            .expect("Inner permit should be available")
+    }
+}
+impl Drop for TrackedOwnedMeteredSemPermit {
+    fn drop(&mut self) {
+        (self.on_drop)();
     }
 }
 
 /// Wraps an [OwnedSemaphorePermit] to update metrics when it's dropped
 pub(crate) struct OwnedMeteredSemPermit {
     inner: OwnedSemaphorePermit,
-    record_fn: Box<dyn Fn() + Send + Sync>,
+    /// See [MeteredSemaphore::unused_claimants]. If present when dropping, used to decrement the
+    /// count.
+    unused_claimants: Option<Arc<AtomicUsize>>,
+    record_fn: Box<dyn Fn(bool) + Send + Sync>,
 }
 impl Drop for OwnedMeteredSemPermit {
     fn drop(&mut self) {
-        (self.record_fn)()
+        if let Some(uc) = self.unused_claimants.take() {
+            uc.fetch_sub(1, Ordering::Release);
+        }
+        (self.record_fn)(true)
     }
 }
 impl Debug for OwnedMeteredSemPermit {
@@ -79,6 +204,32 @@ impl Debug for OwnedMeteredSemPermit {
         self.inner.fmt(f)
     }
 }
+impl OwnedMeteredSemPermit {
+    /// Should be called once this permit is actually being "used" for the work it was meant to
+    /// permit.
+    pub(crate) fn into_used(mut self) -> UsedMeteredSemPermit {
+        if let Some(uc) = self.unused_claimants.take() {
+            uc.fetch_sub(1, Ordering::Release);
+            (self.record_fn)(false)
+        }
+        UsedMeteredSemPermit(self)
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct UsedMeteredSemPermit(OwnedMeteredSemPermit);
+impl UsedMeteredSemPermit {
+    #[cfg(feature = "save_wf_inputs")]
+    pub(crate) fn fake_deserialized() -> Self {
+        let sem = Arc::new(Semaphore::new(1));
+        let inner = sem.try_acquire_owned().unwrap();
+        Self(OwnedMeteredSemPermit {
+            inner,
+            unused_claimants: None,
+            record_fn: Box::new(|_| {}),
+        })
+    }
+}
 
 /// From the input stream, create a new stream which only pulls from the input stream when allowed.
 /// When allowed is determined by the passed in `proceeder` emitting an item. The input stream is
@@ -163,4 +314,43 @@ mod tests {
         allow_tx.send(()).unwrap();
         assert_eq!(when_allowed.poll_next_unpin(&mut cx), Poll::Ready(None));
     }
+
+    #[tokio::test]
+    async fn closable_semaphore_permit_drop_returns_permit() {
+        let inner = MeteredSemaphore::new(2, MetricsContext::no_op(), |_, _| {});
+        let sem = ClosableMeteredSemaphore::new_arc(Arc::new(inner));
+        let perm = sem.try_acquire_owned().unwrap();
+        let permits = sem.outstanding_permits.load(Ordering::Acquire);
+        assert_eq!(permits, 1);
+        drop(perm);
+        let permits = sem.outstanding_permits.load(Ordering::Acquire);
+        assert_eq!(permits, 0);
+    }
+
+    #[tokio::test]
+    async fn closable_semaphore_permit_drop_after_close_resolves_close_complete() {
+        let inner = MeteredSemaphore::new(2, MetricsContext::no_op(), |_, _| {});
+        let sem = ClosableMeteredSemaphore::new_arc(Arc::new(inner));
+        let perm = sem.try_acquire_owned().unwrap();
+        sem.close();
+        drop(perm);
+        sem.close_complete().await;
+    }
+
+    #[tokio::test]
+    async fn closable_semaphore_close_complete_ready_if_unused() {
+        let inner = MeteredSemaphore::new(2, MetricsContext::no_op(), |_, _| {});
+        let sem = ClosableMeteredSemaphore::new_arc(Arc::new(inner));
+        sem.close();
+        sem.close_complete().await;
+    }
+
+    #[tokio::test]
+    async fn closable_semaphore_does_not_hand_out_permits_after_closed() {
+        let inner = MeteredSemaphore::new(2, MetricsContext::no_op(), |_, _| {});
+        let sem = ClosableMeteredSemaphore::new_arc(Arc::new(inner));
+        sem.close();
+        let perm = sem.try_acquire_owned().unwrap_err();
+        assert_matches!(perm, TryAcquireError::Closed);
+    }
 }