temporalio 0.0.0 → 0.0.2

data/bridge/sdk-core/core/src/lib.rs

@@ -0,0 +1,282 @@
+#![warn(missing_docs)] // error if there are missing docs
+#![allow(clippy::upper_case_acronyms)]
+
+//! This crate provides a basis for creating new Temporal SDKs without completely starting from
+//! scratch
+
+#[cfg(test)]
+#[macro_use]
+pub extern crate assert_matches;
+#[macro_use]
+extern crate tracing;
+extern crate core;
+
+mod abstractions;
+pub mod ephemeral_server;
+mod pollers;
+mod protosext;
+pub mod replay;
+pub(crate) mod retry_logic;
+pub mod telemetry;
+mod worker;
+
+#[cfg(test)]
+mod core_tests;
+#[cfg(test)]
+#[macro_use]
+mod test_help;
+
+pub(crate) use temporal_sdk_core_api::errors;
+
+pub use pollers::{
+    Client, ClientOptions, ClientOptionsBuilder, ClientTlsConfig, RetryClient, RetryConfig,
+    TlsConfig, WorkflowClientTrait,
+};
+pub use temporal_sdk_core_api as api;
+pub use temporal_sdk_core_protos as protos;
+pub use temporal_sdk_core_protos::TaskToken;
+pub use url::Url;
+pub use worker::{Worker, WorkerConfig, WorkerConfigBuilder};
+
+use crate::{
+    replay::{mock_client_from_histories, Historator, HistoryForReplay},
+    telemetry::{
+        metrics::MetricsContext, remove_trace_subscriber_for_current_thread,
+        set_trace_subscriber_for_current_thread, telemetry_init, TelemetryInstance,
+    },
+    worker::client::WorkerClientBag,
+};
+use futures::Stream;
+use std::sync::Arc;
+use temporal_client::{ConfiguredClient, TemporalServiceClientWithMetrics};
+use temporal_sdk_core_api::{
+    errors::{CompleteActivityError, PollActivityError, PollWfError},
+    telemetry::{CoreTelemetry, TelemetryOptions},
+    Worker as WorkerTrait,
+};
+use temporal_sdk_core_protos::coresdk::ActivityHeartbeat;
+
+/// Initialize a worker bound to a task queue.
+///
+/// You will need to have already initialized a [CoreRuntime] which will be used for this worker.
+/// After the worker is initialized, you should use [CoreRuntime::tokio_handle] to run the worker's
+/// async functions.
+///
+/// Lang implementations may pass in a [temporal_client::ConfiguredClient] directly (or a
+/// [RetryClient] wrapping one, or a handful of other variants of the same idea). When they do so,
+/// this function will always overwrite the client retry configuration, force the client to use the
+/// namespace defined in the worker config, and set the client identity appropriately. IE: Use
+/// [ClientOptions::connect_no_namespace], not [ClientOptions::connect].
+pub fn init_worker<CT>(
+    runtime: &CoreRuntime,
+    worker_config: WorkerConfig,
+    client: CT,
+) -> Result<Worker, anyhow::Error>
+where
+    CT: Into<sealed::AnyClient>,
+{
+    let client = {
+        let ll = client.into().into_inner();
+        let mut client = Client::new(*ll, worker_config.namespace.clone());
+        client.set_worker_build_id(worker_config.worker_build_id.clone());
+        if let Some(ref id_override) = worker_config.client_identity_override {
+            client.options_mut().identity = id_override.clone();
+        }
+        RetryClient::new(client, RetryConfig::default())
+    };
+    if client.namespace() != worker_config.namespace {
+        panic!("Passed in client is not bound to the same namespace as the worker");
+    }
+    let client_ident = client.get_options().identity.clone();
+    let sticky_q = sticky_q_name_for_worker(&client_ident, &worker_config);
+    let client_bag = Arc::new(WorkerClientBag::new(
+        client,
+        worker_config.namespace.clone(),
+        client_ident,
+        worker_config.worker_build_id.clone(),
+        worker_config.use_worker_versioning,
+    ));
+
+    let metrics = MetricsContext::top_level(worker_config.namespace.clone(), &runtime.telemetry)
+        .with_task_q(worker_config.task_queue.clone());
+    Ok(Worker::new(worker_config, sticky_q, client_bag, metrics))
+}
+
+/// Create a worker for replaying a specific history. It will auto-shutdown as soon as the history
+/// has finished being replayed.
+///
+/// You do not necessarily need a [CoreRuntime] for replay workers, but it's advisable to create
+/// one and use it to run the replay worker's async functions the same way you would for a normal
+/// worker.
+pub fn init_replay_worker<I>(
+    mut config: WorkerConfig,
+    histories: I,
+) -> Result<Worker, anyhow::Error>
+where
+    I: Stream<Item = HistoryForReplay> + Send + 'static,
+{
+    info!(
+        task_queue = config.task_queue.as_str(),
+        "Registering replay worker"
+    );
+    config.max_cached_workflows = 1;
+    config.max_concurrent_wft_polls = 1;
+    config.no_remote_activities = true;
+    let historator = Historator::new(histories);
+    let post_activate = historator.get_post_activate_hook();
+    let shutdown_tok = historator.get_shutdown_setter();
+    let client = mock_client_from_histories(historator);
+    let mut worker = Worker::new(config, None, Arc::new(client), MetricsContext::no_op());
+    worker.set_post_activate_hook(post_activate);
+    shutdown_tok(worker.shutdown_token());
+    Ok(worker)
+}
+
+/// Creates a unique sticky queue name for a worker, iff the config allows for 1 or more cached
+/// workflows.
+pub(crate) fn sticky_q_name_for_worker(
+    process_identity: &str,
+    config: &WorkerConfig,
+) -> Option<String> {
+    if config.max_cached_workflows > 0 {
+        Some(format!(
+            "{}-{}-{}",
+            &process_identity,
+            &config.task_queue,
+            uuid::Uuid::new_v4().simple()
+        ))
+    } else {
+        None
+    }
+}
+
+mod sealed {
+    use super::*;
+
+    /// Allows passing different kinds of clients into things that want to be flexible. Motivating
+    /// use-case was worker initialization.
+    ///
+    /// Needs to exist in this crate to avoid blanket impl conflicts.
+    pub struct AnyClient(Box<ConfiguredClient<TemporalServiceClientWithMetrics>>);
+    impl AnyClient {
+        pub(crate) fn into_inner(self) -> Box<ConfiguredClient<TemporalServiceClientWithMetrics>> {
+            self.0
+        }
+    }
+
+    impl From<RetryClient<ConfiguredClient<TemporalServiceClientWithMetrics>>> for AnyClient {
+        fn from(c: RetryClient<ConfiguredClient<TemporalServiceClientWithMetrics>>) -> Self {
+            Self(Box::new(c.into_inner()))
+        }
+    }
+    impl From<RetryClient<Client>> for AnyClient {
+        fn from(c: RetryClient<Client>) -> Self {
+            Self(Box::new(c.into_inner().into_inner()))
+        }
+    }
+    impl From<Arc<RetryClient<Client>>> for AnyClient {
+        fn from(c: Arc<RetryClient<Client>>) -> Self {
+            Self(Box::new(c.get_client().inner().clone()))
+        }
+    }
+    impl From<ConfiguredClient<TemporalServiceClientWithMetrics>> for AnyClient {
+        fn from(c: ConfiguredClient<TemporalServiceClientWithMetrics>) -> Self {
+            Self(Box::new(c))
+        }
+    }
+}
+
+/// Holds shared state/components needed to back instances of workers and clients. More than one
+/// may be instantiated, but typically only one is needed. More than one runtime instance may be
+/// useful if multiple different telemetry settings are required.
+pub struct CoreRuntime {
+    telemetry: TelemetryInstance,
+    runtime: Option<tokio::runtime::Runtime>,
+    runtime_handle: tokio::runtime::Handle,
+}
+
+impl CoreRuntime {
+    /// Create a new core runtime with the provided telemetry options and tokio runtime builder.
+    /// Also initialize telemetry for the thread this is being called on.
+    ///
+    /// Note that this function will call the [tokio::runtime::Builder::enable_all] builder option
+    /// on the Tokio runtime builder, and will call [tokio::runtime::Builder::on_thread_start] to
+    /// ensure telemetry subscribers are set on every tokio thread.
+    ///
+    /// **Important**: You need to call this *before* calling any async functions on workers or
+    /// clients, otherwise the tracing subscribers will not be properly attached.
+    ///
+    /// # Panics
+    /// If a tokio runtime has already been initialized. To re-use an existing runtime, call
+    /// [CoreRuntime::new_assume_tokio].
+    pub fn new(
+        telemetry_options: TelemetryOptions,
+        mut tokio_builder: tokio::runtime::Builder,
+    ) -> Result<Self, anyhow::Error> {
+        let telemetry = telemetry_init(telemetry_options)?;
+        let subscriber = telemetry.trace_subscriber();
+        let runtime = tokio_builder
+            .enable_all()
+            .on_thread_start(move || {
+                set_trace_subscriber_for_current_thread(subscriber.clone());
+            })
+            .build()?;
+        let _rg = runtime.enter();
+        let mut me = Self::new_assume_tokio_initialized_telem(telemetry);
+        me.runtime = Some(runtime);
+        Ok(me)
+    }
+
+    /// Initialize telemetry for the thread this is being called on, assuming a tokio runtime is
+    /// already active and this call exists in its context. See [Self::new] for more.
+    ///
+    /// # Panics
+    /// If there is no currently active Tokio runtime
+    pub fn new_assume_tokio(telemetry_options: TelemetryOptions) -> Result<Self, anyhow::Error> {
+        let telemetry = telemetry_init(telemetry_options)?;
+        Ok(Self::new_assume_tokio_initialized_telem(telemetry))
+    }
+
+    /// Construct a runtime from an already-initialized telemetry instance, assuming a tokio runtime
+    /// is already active and this call exists in its context. See [Self::new] for more.
+    ///
+    /// # Panics
+    /// If there is no currently active Tokio runtime
+    pub fn new_assume_tokio_initialized_telem(telemetry: TelemetryInstance) -> Self {
+        let runtime_handle = tokio::runtime::Handle::current();
+        set_trace_subscriber_for_current_thread(telemetry.trace_subscriber());
+        Self {
+            telemetry,
+            runtime: None,
+            runtime_handle,
+        }
+    }
+
+    /// Get a handle to the tokio runtime used by this Core runtime.
+    pub fn tokio_handle(&self) -> tokio::runtime::Handle {
+        self.runtime_handle.clone()
+    }
+
+    /// Returns the metric meter used for recording metrics, if they were enabled.
+    pub fn metric_meter(&self) -> Option<&opentelemetry::metrics::Meter> {
+        self.telemetry.get_metric_meter()
+    }
+
+    /// 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) -> Arc<dyn tracing::Subscriber + Send + Sync> {
+        self.telemetry.trace_subscriber()
+    }
+
+    /// Return a reference to the owned [TelemetryInstance]
+    pub fn telemetry(&self) -> &TelemetryInstance {
+        &self.telemetry
+    }
+}
+
+impl Drop for CoreRuntime {
+    fn drop(&mut self) {
+        remove_trace_subscriber_for_current_thread();
+    }
+}

data/bridge/sdk-core/core/src/pollers/mod.rs

@@ -0,0 +1,54 @@
+mod poll_buffer;
+
+pub(crate) use poll_buffer::{
+    new_activity_task_buffer, new_workflow_task_buffer, WorkflowTaskPoller,
+};
+pub use temporal_client::{
+    Client, ClientOptions, ClientOptionsBuilder, ClientTlsConfig, RetryClient, RetryConfig,
+    TlsConfig, WorkflowClientTrait,
+};
+use temporal_sdk_core_protos::temporal::api::workflowservice::v1::{
+    PollActivityTaskQueueResponse, PollWorkflowTaskQueueResponse,
+};
+
+#[cfg(test)]
+use futures::Future;
+
+pub type Result<T, E = tonic::Status> = std::result::Result<T, E>;
+
+/// A trait for things that poll the server. Hides complexity of concurrent polling or polling
+/// on sticky/nonsticky queues simultaneously.
+#[cfg_attr(test, mockall::automock)]
+#[cfg_attr(test, allow(unused))]
+#[async_trait::async_trait]
+pub trait Poller<PollResult>
+where
+    PollResult: Send + Sync + 'static,
+{
+    async fn poll(&self) -> Option<Result<PollResult>>;
+    fn notify_shutdown(&self);
+    async fn shutdown(self);
+    /// Need a separate shutdown to be able to consume boxes :(
+    async fn shutdown_box(self: Box<Self>);
+}
+pub type BoxedPoller<T> = Box<dyn Poller<T> + Send + Sync + 'static>;
+pub type BoxedWFPoller = BoxedPoller<PollWorkflowTaskQueueResponse>;
+pub type BoxedActPoller = BoxedPoller<PollActivityTaskQueueResponse>;
+
+#[cfg(test)]
+mockall::mock! {
+    pub ManualPoller<T: Send + Sync + 'static> {}
+    #[allow(unused)]
+    impl<T: Send + Sync + 'static> Poller<T> for ManualPoller<T> {
+        fn poll<'a, 'b>(&self)
+          -> impl Future<Output = Option<Result<T>>> + Send + 'b
+            where 'a: 'b, Self: 'b;
+        fn notify_shutdown(&self);
+        fn shutdown<'a>(self)
+          -> impl Future<Output = ()> + Send + 'a
+            where Self: 'a;
+        fn shutdown_box<'a>(self: Box<Self>)
+          -> impl Future<Output = ()> + Send + 'a
+            where Self: 'a;
+    }
+}

data/bridge/sdk-core/core/src/pollers/poll_buffer.rs

@@ -0,0 +1,297 @@
+use crate::{
+    pollers::{self, Poller},
+    worker::client::WorkerClient,
+};
+use futures::{prelude::stream::FuturesUnordered, StreamExt};
+use std::{
+    fmt::Debug,
+    future::Future,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+};
+use temporal_sdk_core_protos::temporal::api::workflowservice::v1::{
+    PollActivityTaskQueueResponse, PollWorkflowTaskQueueResponse,
+};
+use tokio::{
+    sync::{
+        mpsc::{channel, Receiver},
+        Mutex, Semaphore,
+    },
+    task::JoinHandle,
+};
+use tokio_util::sync::CancellationToken;
+
+pub struct LongPollBuffer<T> {
+    buffered_polls: Mutex<Receiver<pollers::Result<T>>>,
+    shutdown: CancellationToken,
+    /// This semaphore exists to ensure that we only poll server as many times as core actually
+    /// *asked* it to be polled - otherwise we might spin and buffer polls constantly. This also
+    /// means unit tests can continue to function in a predictable manner when calling mocks.
+    polls_requested: Arc<Semaphore>,
+    join_handles: FuturesUnordered<JoinHandle<()>>,
+    /// Called every time the number of pollers is changed
+    num_pollers_changed: Option<Box<dyn Fn(usize) + Send + Sync>>,
+    active_pollers: Arc<AtomicUsize>,
+}
+
+struct ActiveCounter<'a>(&'a AtomicUsize);
+impl<'a> ActiveCounter<'a> {
+    fn new(a: &'a AtomicUsize) -> Self {
+        a.fetch_add(1, Ordering::Relaxed);
+        Self(a)
+    }
+}
+impl Drop for ActiveCounter<'_> {
+    fn drop(&mut self) {
+        self.0.fetch_sub(1, Ordering::Relaxed);
+    }
+}
+
+impl<T> LongPollBuffer<T>
+where
+    T: Send + Debug + 'static,
+{
+    pub fn new<FT>(
+        poll_fn: impl Fn() -> FT + Send + Sync + 'static,
+        max_pollers: usize,
+        buffer_size: usize,
+        shutdown: CancellationToken,
+    ) -> Self
+    where
+        FT: Future<Output = pollers::Result<T>> + Send,
+    {
+        let (tx, rx) = channel(buffer_size);
+        let polls_requested = Arc::new(Semaphore::new(0));
+        let active_pollers = Arc::new(AtomicUsize::new(0));
+        let join_handles = FuturesUnordered::new();
+        let pf = Arc::new(poll_fn);
+        for _ in 0..max_pollers {
+            let tx = tx.clone();
+            let pf = pf.clone();
+            let shutdown = shutdown.clone();
+            let polls_requested = polls_requested.clone();
+            let ap = active_pollers.clone();
+            let jh = tokio::spawn(async move {
+                loop {
+                    if shutdown.is_cancelled() {
+                        break;
+                    }
+                    let sp = tokio::select! {
+                        sp = polls_requested.acquire() => sp.expect("Polls semaphore not dropped"),
+                        _ = shutdown.cancelled() => continue,
+                    };
+                    let _active_guard = ActiveCounter::new(ap.as_ref());
+                    let r = tokio::select! {
+                        r = pf() => r,
+                        _ = shutdown.cancelled() => continue,
+                    };
+                    sp.forget();
+                    let _ = tx.send(r).await;
+                }
+            });
+            join_handles.push(jh);
+        }
+        Self {
+            buffered_polls: Mutex::new(rx),
+            shutdown,
+            polls_requested,
+            join_handles,
+            num_pollers_changed: None,
+            active_pollers,
+        }
+    }
+
+    /// Set a function that will be called every time the number of pollers changes.
+    /// TODO: Currently a bit weird, will make more sense once we implement dynamic poller scaling.
+    pub fn set_num_pollers_handler(&mut self, handler: impl Fn(usize) + Send + Sync + 'static) {
+        self.num_pollers_changed = Some(Box::new(handler));
+    }
+}
+
+#[async_trait::async_trait]
+impl<T> Poller<T> for LongPollBuffer<T>
+where
+    T: Send + Sync + Debug + 'static,
+{
+    /// Poll the buffer. Adds one permit to the polling pool - the point of this being that the
+    /// buffer may support many concurrent pollers, but there is no reason to have them poll unless
+    /// enough polls have actually been requested. Calling this function adds a permit that any
+    /// concurrent poller may fulfill.
+    ///
+    /// EX: If this function is only ever called serially and always `await`ed, there will be no
+    /// concurrent polling. If it is called many times and the futures are awaited concurrently,
+    /// then polling will happen concurrently.
+    ///
+    /// Returns `None` if the poll buffer has been shut down
+    #[instrument(name = "long_poll", level = "trace", skip(self))]
+    async fn poll(&self) -> Option<pollers::Result<T>> {
+        self.polls_requested.add_permits(1);
+        if let Some(fun) = self.num_pollers_changed.as_ref() {
+            fun(self.active_pollers.load(Ordering::Relaxed));
+        }
+
+        let mut locked = self.buffered_polls.lock().await;
+        let res = (*locked).recv().await;
+
+        if let Some(fun) = self.num_pollers_changed.as_ref() {
+            fun(self.active_pollers.load(Ordering::Relaxed));
+        }
+
+        res
+    }
+
+    fn notify_shutdown(&self) {
+        self.shutdown.cancel();
+    }
+
+    async fn shutdown(mut self) {
+        self.notify_shutdown();
+        while self.join_handles.next().await.is_some() {}
+    }
+
+    async fn shutdown_box(self: Box<Self>) {
+        let this = *self;
+        this.shutdown().await;
+    }
+}
+
+/// A poller capable of polling on a sticky and a nonsticky queue simultaneously for workflow tasks.
+#[derive(derive_more::Constructor)]
+pub struct WorkflowTaskPoller {
+    normal_poller: PollWorkflowTaskBuffer,
+    sticky_poller: Option<PollWorkflowTaskBuffer>,
+}
+
+#[async_trait::async_trait]
+impl Poller<PollWorkflowTaskQueueResponse> for WorkflowTaskPoller {
+    async fn poll(&self) -> Option<pollers::Result<PollWorkflowTaskQueueResponse>> {
+        if let Some(sq) = self.sticky_poller.as_ref() {
+            tokio::select! {
+                r = self.normal_poller.poll() => r,
+                r = sq.poll() => r,
+            }
+        } else {
+            self.normal_poller.poll().await
+        }
+    }
+
+    fn notify_shutdown(&self) {
+        self.normal_poller.notify_shutdown();
+        if let Some(sq) = self.sticky_poller.as_ref() {
+            sq.notify_shutdown();
+        }
+    }
+
+    async fn shutdown(mut self) {
+        self.normal_poller.shutdown().await;
+        if let Some(sq) = self.sticky_poller {
+            sq.shutdown().await;
+        }
+    }
+
+    async fn shutdown_box(self: Box<Self>) {
+        let this = *self;
+        this.shutdown().await;
+    }
+}
+
+pub type PollWorkflowTaskBuffer = LongPollBuffer<PollWorkflowTaskQueueResponse>;
+pub(crate) fn new_workflow_task_buffer(
+    client: Arc<dyn WorkerClient>,
+    task_queue: String,
+    is_sticky: bool,
+    concurrent_pollers: usize,
+    buffer_size: usize,
+    shutdown: CancellationToken,
+) -> PollWorkflowTaskBuffer {
+    LongPollBuffer::new(
+        move || {
+            let client = client.clone();
+            let task_queue = task_queue.clone();
+            async move { client.poll_workflow_task(task_queue, is_sticky).await }
+        },
+        concurrent_pollers,
+        buffer_size,
+        shutdown,
+    )
+}
+
+pub type PollActivityTaskBuffer = LongPollBuffer<PollActivityTaskQueueResponse>;
+pub(crate) fn new_activity_task_buffer(
+    client: Arc<dyn WorkerClient>,
+    task_queue: String,
+    concurrent_pollers: usize,
+    buffer_size: usize,
+    max_tps: Option<f64>,
+    shutdown: CancellationToken,
+) -> PollActivityTaskBuffer {
+    LongPollBuffer::new(
+        move || {
+            let client = client.clone();
+            let task_queue = task_queue.clone();
+            async move { client.poll_activity_task(task_queue, max_tps).await }
+        },
+        concurrent_pollers,
+        buffer_size,
+        shutdown,
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::worker::client::mocks::mock_manual_workflow_client;
+    use futures::FutureExt;
+    use std::time::Duration;
+    use tokio::{select, sync::mpsc::channel};
+
+    #[tokio::test]
+    async fn only_polls_once_with_1_poller() {
+        let mut mock_client = mock_manual_workflow_client();
+        mock_client
+            .expect_poll_workflow_task()
+            .times(2)
+            .returning(move |_, _| {
+                async {
+                    tokio::time::sleep(Duration::from_millis(100)).await;
+                    Ok(Default::default())
+                }
+                .boxed()
+            });
+
+        let pb = new_workflow_task_buffer(
+            Arc::new(mock_client),
+            "someq".to_string(),
+            false,
+            1,
+            1,
+            CancellationToken::new(),
+        );
+
+        // Poll a bunch of times, "interrupting" it each time, we should only actually have polled
+        // once since the poll takes a while
+        let (interrupter_tx, mut interrupter_rx) = channel(50);
+        for _ in 0..10 {
+            interrupter_tx.send(()).await.unwrap();
+        }
+
+        // We should never get anything out since we interrupted 100% of polls
+        let mut last_val = false;
+        for _ in 0..10 {
+            select! {
+                _ = interrupter_rx.recv() => {
+                    last_val = true;
+                }
+                _ = pb.poll() => {
+                }
+            }
+        }
+        assert!(last_val);
+        // Now we grab the buffered poll response, the poll task will go again but we don't grab it,
+        // therefore we will have only polled twice.
+        pb.poll().await.unwrap().unwrap();
+        pb.shutdown().await;
+    }
+}