temporalio 0.0.1 → 0.0.2

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

@@ -13,12 +13,11 @@ extern crate core;
 
 mod abstractions;
 pub mod ephemeral_server;
-mod log_export;
 mod pollers;
 mod protosext;
 pub mod replay;
 pub(crate) mod retry_logic;
-pub(crate) mod telemetry;
+pub mod telemetry;
 mod worker;
 
 #[cfg(test)]
@@ -33,10 +32,6 @@ pub use pollers::{
     Client, ClientOptions, ClientOptionsBuilder, ClientTlsConfig, RetryClient, RetryConfig,
     TlsConfig, WorkflowClientTrait,
 };
-pub use telemetry::{
-    fetch_global_buffered_logs, telemetry_init, Logger, MetricTemporality, MetricsExporter,
-    OtelCollectorOptions, TelemetryOptions, TelemetryOptionsBuilder, TraceExporter,
-};
 pub use temporal_sdk_core_api as api;
 pub use temporal_sdk_core_protos as protos;
 pub use temporal_sdk_core_protos::TaskToken;
@@ -44,33 +39,39 @@ pub use url::Url;
 pub use worker::{Worker, WorkerConfig, WorkerConfigBuilder};
 
 use crate::{
-    replay::mock_client_from_history,
-    telemetry::metrics::{MetricsContext, METRIC_METER},
+    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},
-    CoreLog, Worker as WorkerTrait,
+    telemetry::{CoreTelemetry, TelemetryOptions},
+    Worker as WorkerTrait,
 };
-use temporal_sdk_core_protos::{coresdk::ActivityHeartbeat, temporal::api::history::v1::History};
-
-lazy_static::lazy_static! {
-    /// A process-wide unique string, which will be different on every startup
-    static ref PROCCESS_UNIQ_ID: String = {
-        uuid::Uuid::new_v4().simple().to_string()
-    };
-}
+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>(worker_config: WorkerConfig, client: CT) -> Worker
+pub fn init_worker<CT>(
+    runtime: &CoreRuntime,
+    worker_config: WorkerConfig,
+    client: CT,
+) -> Result<Worker, anyhow::Error>
 where
     CT: Into<sealed::AnyClient>,
 {
@@ -96,18 +97,24 @@ where
         worker_config.use_worker_versioning,
     ));
 
-    let metrics = MetricsContext::top_level(worker_config.namespace.clone())
+    let metrics = MetricsContext::top_level(worker_config.namespace.clone(), &runtime.telemetry)
         .with_task_q(worker_config.task_queue.clone());
-    Worker::new(worker_config, sticky_q, client_bag, metrics)
+    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. The provided client should be a mock, and this should only be used
-/// for workflow testing purposes.
-pub fn init_replay_worker(
+/// 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,
-    history: &History,
-) -> Result<Worker, anyhow::Error> {
+    histories: I,
+) -> Result<Worker, anyhow::Error>
+where
+    I: Stream<Item = HistoryForReplay> + Send + 'static,
+{
     info!(
         task_queue = config.task_queue.as_str(),
         "Registering replay worker"
@@ -115,15 +122,18 @@ pub fn init_replay_worker(
     config.max_cached_workflows = 1;
     config.max_concurrent_wft_polls = 1;
     config.no_remote_activities = true;
-    // Could possibly just use mocked pollers here?
-    let client = mock_client_from_history(history, config.task_queue.clone());
-    let run_id = history.extract_run_id_from_start()?.to_string();
-    let last_event = history.last_event_id();
-    let mut worker = Worker::new(config, None, Arc::new(client), MetricsContext::default());
-    worker.set_shutdown_on_run_reaches_event(run_id, last_event);
+    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,
@@ -131,7 +141,9 @@ pub(crate) fn sticky_q_name_for_worker(
     if config.max_cached_workflows > 0 {
         Some(format!(
             "{}-{}-{}",
-            &process_identity, &config.task_queue, *PROCCESS_UNIQ_ID
+            &process_identity,
+            &config.task_queue,
+            uuid::Uuid::new_v4().simple()
         ))
     } else {
         None
@@ -173,3 +185,98 @@ mod sealed {
         }
     }
 }
+
+/// 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/protosext/mod.rs

@@ -274,7 +274,7 @@ impl TryFrom<activity_execution_result::Status> for LocalActivityExecutionResult
             Status::Failed(f)
                 if f.failure
                     .as_ref()
-                    .map(|fail| fail.is_timeout())
+                    .map(|fail| fail.is_timeout().is_some())
                     .unwrap_or_default() =>
             {
                 Ok(LocalActivityExecutionResult::TimedOut(f))

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

@@ -2,70 +2,214 @@
 //! to replay canned histories. It should be used by Lang SDKs to provide replay capabilities to
 //! users during testing.
 
-use crate::worker::client::{mocks::mock_manual_workflow_client, WorkerClient};
-use futures::FutureExt;
+use crate::{
+    worker::client::{mocks::mock_manual_workflow_client, WorkerClient},
+    Worker,
+};
+use futures::{FutureExt, Stream, StreamExt};
+use once_cell::sync::OnceCell;
+use parking_lot::Mutex;
 use std::{
-    sync::{
-        atomic::{AtomicBool, Ordering},
-        Arc,
-    },
-    time::Duration,
+    collections::HashMap,
+    pin::Pin,
+    sync::Arc,
+    task::{Context, Poll},
 };
-use temporal_sdk_core_protos::temporal::api::{
-    common::v1::WorkflowExecution,
-    history::v1::History,
-    workflowservice::v1::{
-        RespondWorkflowTaskCompletedResponse, RespondWorkflowTaskFailedResponse,
+use temporal_sdk_core_protos::{
+    coresdk::workflow_activation::remove_from_cache::EvictionReason,
+    temporal::api::{
+        common::v1::WorkflowExecution,
+        history::v1::History,
+        workflowservice::v1::{
+            RespondWorkflowTaskCompletedResponse, RespondWorkflowTaskFailedResponse,
+        },
     },
 };
 pub use temporal_sdk_core_protos::{
     default_wes_attribs, HistoryInfo, TestHistoryBuilder, DEFAULT_WORKFLOW_TYPE,
 };
+use tokio::sync::{mpsc, mpsc::UnboundedSender, Mutex as TokioMutex};
+use tokio_stream::wrappers::UnboundedReceiverStream;
+use tokio_util::sync::CancellationToken;
+
+/// A history which will be used during replay verification. Since histories do not include the
+/// workflow id, it must be manually attached.
+#[derive(Debug, Clone, derive_more::Constructor)]
+pub struct HistoryForReplay {
+    hist: History,
+    workflow_id: String,
+}
+
+/// Allows lang to feed histories into the replayer one at a time. Simply drop the feeder to signal
+/// to the worker that you're done and it should initiate shutdown.
+pub struct HistoryFeeder {
+    tx: mpsc::Sender<HistoryForReplay>,
+}
+/// The stream half of a [HistoryFeeder]
+pub struct HistoryFeederStream {
+    rcvr: mpsc::Receiver<HistoryForReplay>,
+}
+
+impl HistoryFeeder {
+    /// Make a new history feeder, which will store at most `buffer_size` histories before `feed`
+    /// blocks.
+    ///
+    /// Returns a feeder which will be used to feed in histories, and a stream you can pass to
+    /// one of the replay worker init functions.
+    pub fn new(buffer_size: usize) -> (Self, HistoryFeederStream) {
+        let (tx, rcvr) = mpsc::channel(buffer_size);
+        (Self { tx }, HistoryFeederStream { rcvr })
+    }
+    /// Feed a new history into the replayer, blocking if there is not room to accept another
+    /// history.
+    pub async fn feed(&self, history: HistoryForReplay) -> anyhow::Result<()> {
+        self.tx.send(history).await?;
+        Ok(())
+    }
+}
+
+impl Stream for HistoryFeederStream {
+    type Item = HistoryForReplay;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.rcvr.poll_recv(cx)
+    }
+}
 
-/// Create a mock client which can be used by a replay worker to serve up canned history.
-/// It will return the entire history in one workflow task, after that it will return default
-/// responses (with a 10s wait). If a workflow task failure is sent to the mock, it will send
-/// the complete response again.
-pub(crate) fn mock_client_from_history(
-    history: &History,
-    task_queue: impl Into<String>,
-) -> impl WorkerClient {
+/// Create a mock client which can be used by a replay worker to serve up canned histories. It will
+/// return the entire history in one workflow task. If a workflow task failure is sent to the mock,
+/// it will send the complete response again.
+///
+/// Once it runs out of histories to return, it will serve up default responses after a 10s delay
+pub(crate) fn mock_client_from_histories(historator: Historator) -> impl WorkerClient {
     let mut mg = mock_manual_workflow_client();
 
-    let hist_info = HistoryInfo::new_from_history(history, None).unwrap();
-    let wf = WorkflowExecution {
-        workflow_id: "fake_wf_id".to_string(),
-        run_id: hist_info.orig_run_id().to_string(),
-    };
+    let hist_allow_tx = historator.replay_done_tx.clone();
+    let historator = Arc::new(TokioMutex::new(historator));
 
-    let did_send = Arc::new(AtomicBool::new(false));
-    let did_send_clone = did_send.clone();
-    let tq = task_queue.into();
     mg.expect_poll_workflow_task().returning(move |_, _| {
-        let hist_info = hist_info.clone();
-        let wf = wf.clone();
-        let did_send_clone = did_send_clone.clone();
-        let tq = tq.clone();
+        let historator = historator.clone();
         async move {
-            if !did_send_clone.swap(true, Ordering::AcqRel) {
-                let mut resp = hist_info.as_poll_wft_response(tq);
-                resp.workflow_execution = Some(wf.clone());
+            let mut hlock = historator.lock().await;
+            // Always wait for permission before dispatching the next task
+            let _ = hlock.allow_stream.next().await;
+
+            if let Some(history) = hlock.next().await {
+                let hist_info = HistoryInfo::new_from_history(&history.hist, None).unwrap();
+                let mut resp = hist_info.as_poll_wft_response();
+                resp.workflow_execution = Some(WorkflowExecution {
+                    workflow_id: history.workflow_id,
+                    run_id: hist_info.orig_run_id().to_string(),
+                });
                 Ok(resp)
             } else {
-                tokio::time::sleep(Duration::from_secs(10)).await;
+                if let Some(wc) = hlock.worker_closer.get() {
+                    wc.cancel();
+                }
                 Ok(Default::default())
             }
         }
         .boxed()
     });
 
-    mg.expect_complete_workflow_task()
-        .returning(|_| async move { Ok(RespondWorkflowTaskCompletedResponse::default()) }.boxed());
+    mg.expect_complete_workflow_task().returning(move |_| {
+        async move { Ok(RespondWorkflowTaskCompletedResponse::default()) }.boxed()
+    });
     mg.expect_fail_workflow_task().returning(move |_, _, _| {
-        // We'll need to re-send the history if WFT fails
-        did_send.store(false, Ordering::Release);
-        async move { Ok(RespondWorkflowTaskFailedResponse {}) }.boxed()
+        hist_allow_tx.send("Failed".to_string()).unwrap();
+        async move { Ok(RespondWorkflowTaskFailedResponse::default()) }.boxed()
     });
 
     mg
 }
+
+pub(crate) struct Historator {
+    iter: Pin<Box<dyn Stream<Item = HistoryForReplay> + Send>>,
+    allow_stream: UnboundedReceiverStream<String>,
+    worker_closer: Arc<OnceCell<CancellationToken>>,
+    dat: Arc<Mutex<HistoratorDat>>,
+    replay_done_tx: UnboundedSender<String>,
+}
+impl Historator {
+    pub(crate) fn new(histories: impl Stream<Item = HistoryForReplay> + Send + 'static) -> Self {
+        let dat = Arc::new(Mutex::new(HistoratorDat::default()));
+        let (replay_done_tx, replay_done_rx) = mpsc::unbounded_channel();
+        // Need to allow the first history item
+        replay_done_tx.send("fake".to_string()).unwrap();
+        Self {
+            iter: Box::pin(histories.fuse()),
+            allow_stream: UnboundedReceiverStream::new(replay_done_rx),
+            worker_closer: Arc::new(OnceCell::new()),
+            dat,
+            replay_done_tx,
+        }
+    }
+
+    /// Returns a callback that can be used as the post-activation hook for a worker to indicate
+    /// we're ready to replay the next history, or whatever else.
+    pub(crate) fn get_post_activate_hook(&self) -> impl Fn(&Worker, &str, usize) + Send + Sync {
+        let dat = self.dat.clone();
+        let done_tx = self.replay_done_tx.clone();
+        move |worker, activated_run_id, last_processed_event| {
+            // We can't hold the lock while evaluating the hook, or we'd deadlock.
+            let last_event_in_hist = dat
+                .lock()
+                .run_id_to_last_event_num
+                .get(activated_run_id)
+                .cloned();
+            if let Some(le) = last_event_in_hist {
+                if last_processed_event >= le {
+                    worker.request_wf_eviction(
+                        activated_run_id,
+                        "Always evict workflows after replay",
+                        EvictionReason::LangRequested,
+                    );
+                    done_tx.send(activated_run_id.to_string()).unwrap();
+                }
+            }
+        }
+    }
+
+    pub(crate) fn get_shutdown_setter(&self) -> impl FnOnce(CancellationToken) + 'static {
+        let wc = self.worker_closer.clone();
+        move |ct| {
+            wc.set(ct).expect("Shutdown token must only be set once");
+        }
+    }
+}
+
+impl Stream for Historator {
+    type Item = HistoryForReplay;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match self.iter.poll_next_unpin(cx) {
+            Poll::Ready(Some(history)) => {
+                let run_id = history
+                    .hist
+                    .extract_run_id_from_start()
+                    .expect(
+                        "Histories provided for replay must contain run ids in their workflow \
+                                 execution started events",
+                    )
+                    .to_string();
+                let last_event = history.hist.last_event_id();
+                self.dat
+                    .lock()
+                    .run_id_to_last_event_num
+                    .insert(run_id, last_event as usize);
+                Poll::Ready(Some(history))
+            }
+            Poll::Ready(None) => {
+                self.dat.lock().all_dispatched = true;
+                Poll::Ready(None)
+            }
+            o => o,
+        }
+    }
+}
+
+#[derive(Default)]
+struct HistoratorDat {
+    run_id_to_last_event_num: HashMap<String, usize>,
+    all_dispatched: bool,
+}