@temporalio/core-bridge 0.20.2 → 0.22.0

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

@@ -9,7 +9,7 @@
 //! An example of running an activity worker:
 //! ```no_run
 //! use std::{sync::Arc, str::FromStr};
-//! use temporal_sdk::{sdk_client_options, Worker};
+//! use temporal_sdk::{sdk_client_options, Worker, ActContext};
 //! use temporal_sdk_core::{init_worker, Url};
 //! use temporal_sdk_core_api::worker::{WorkerConfig, WorkerConfigBuilder};
 //!
@@ -23,7 +23,7 @@
 //!     let mut worker = Worker::new_from_core(Arc::new(core_worker), "task_queue");
 //!     worker.register_activity(
 //!         "echo_activity",
-//!         |echo_me: String| async move { Ok(echo_me) },
+//!         |_ctx: ActContext, echo_me: String| async move { Ok(echo_me) },
 //!     );
 //!     worker.run().await?;
 //!     Ok(())
@@ -33,12 +33,15 @@
 #[macro_use]
 extern crate tracing;
 
+mod activity_context;
 mod conversions;
 pub mod interceptors;
 mod payload_converter;
 mod workflow_context;
 mod workflow_future;
 
+pub use activity_context::ActContext;
+
 pub use workflow_context::{
     ActivityOptions, CancellableFuture, ChildWorkflow, ChildWorkflowOptions, LocalActivityOptions,
     Signal, SignalData, SignalWorkflowOptions, WfContext,
@@ -49,9 +52,10 @@ use crate::{
     workflow_context::{ChildWfCommon, PendingChildWorkflow},
 };
 use anyhow::{anyhow, bail};
-use futures::{future::BoxFuture, stream::FuturesUnordered, FutureExt, StreamExt};
+use futures::{future::BoxFuture, FutureExt, StreamExt, TryFutureExt, TryStreamExt};
 use once_cell::sync::OnceCell;
 use std::{
+    cell::RefCell,
     collections::HashMap,
     fmt::{Debug, Display, Formatter},
     future::Future,
@@ -82,12 +86,12 @@ use temporal_sdk_core_protos::{
 };
 use tokio::{
     sync::{
-        mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender},
-        oneshot, watch,
-        watch::Receiver,
+        mpsc::{unbounded_channel, UnboundedSender},
+        oneshot,
     },
     task::JoinError,
 };
+use tokio_stream::wrappers::UnboundedReceiverStream;
 use tokio_util::sync::CancellationToken;
 
 const VERSION: &str = env!("CARGO_PKG_VERSION");
@@ -120,13 +124,19 @@ struct CommonWorker {
 }
 
 struct WorkflowHalf {
-    /// Maps run id to the driver
-    workflows: HashMap<String, UnboundedSender<WorkflowActivation>>,
+    /// Maps run id to cached workflow state
+    workflows: RefCell<HashMap<String, WorkflowData>>,
     /// Maps workflow type to the function for executing workflow runs with that ID
-    workflow_fns: HashMap<String, WorkflowFunction>,
-    /// Handles for each spawned workflow run are inserted here to be cleaned up when all runs
-    /// are finished
-    join_handles: FuturesUnordered<BoxFuture<'static, Result<WorkflowResult<()>, JoinError>>>,
+    workflow_fns: RefCell<HashMap<String, WorkflowFunction>>,
+}
+struct WorkflowData {
+    /// Channel used to send the workflow activations
+    activation_chan: UnboundedSender<WorkflowActivation>,
+}
+
+struct WorkflowFutureHandle<F: Future<Output = Result<WorkflowResult<()>, JoinError>>> {
+    join_handle: F,
+    run_id: String,
 }
 
 struct ActivityHalf {
@@ -136,8 +146,6 @@ struct ActivityHalf {
 }
 
 impl Worker {
-    // pub fn new(cfg: WorkerConfig) -> Self {}
-
     #[doc(hidden)]
     /// Create a new rust worker from a core worker
     pub fn new_from_core(worker: Arc<dyn CoreWorker>, task_queue: impl Into<String>) -> Self {
@@ -150,7 +158,6 @@ impl Worker {
             workflow_half: WorkflowHalf {
                 workflows: Default::default(),
                 workflow_fns: Default::default(),
-                join_handles: FuturesUnordered::new(),
             },
             activity_half: ActivityHalf {
                 activity_fns: Default::default(),
@@ -180,6 +187,7 @@ impl Worker {
     ) {
         self.workflow_half
             .workflow_fns
+            .get_mut()
             .insert(workflow_type.into(), wf_function.into());
     }
 
@@ -201,67 +209,103 @@ impl Worker {
     /// Runs the worker. Eventually resolves after the worker has been explicitly shut down,
     /// or may return early with an error in the event of some unresolvable problem.
     pub async fn run(&mut self) -> Result<(), anyhow::Error> {
-        let (shutdown_tx, shutdown_rx) = watch::channel(false);
-        let pollers = async move {
-            let (common, wf_half, act_half) = self.split_apart();
-            let (completions_tx, mut completions_rx) = unbounded_channel();
-            let (wf_poll_res, act_poll_res) = tokio::join!(
-                // Workflow polling loop
-                async {
-                    loop {
-                        info!("Polling");
-                        let activation = match common.worker.poll_workflow_activation().await {
-                            Err(PollWfError::ShutDown) => {
-                                break Result::<_, anyhow::Error>::Ok(());
-                            }
-                            o => o?,
-                        };
-                        wf_half
-                            .workflow_activation_handler(
-                                common,
-                                &shutdown_rx,
-                                &completions_tx,
-                                &mut completions_rx,
-                                activation,
-                            )
-                            .await?;
+        let shutdown_token = CancellationToken::new();
+        let (common, wf_half, act_half) = self.split_apart();
+        let (wf_future_tx, wf_future_rx) = unbounded_channel();
+        let (completions_tx, completions_rx) = unbounded_channel();
+        let wf_future_joiner = async {
+            UnboundedReceiverStream::new(wf_future_rx)
+                .map(Result::<_, anyhow::Error>::Ok)
+                .try_for_each_concurrent(
+                    None,
+                    |WorkflowFutureHandle {
+                         join_handle,
+                         run_id,
+                     }| {
+                        let wf_half = &*wf_half;
+                        async move {
+                            join_handle.await??;
+                            wf_half.workflows.borrow_mut().remove(&run_id);
+                            Ok(())
+                        }
+                    },
+                )
+                .await
+        };
+        let wf_completion_processor = async {
+            let r = UnboundedReceiverStream::new(completions_rx)
+                .map(Ok)
+                .try_for_each_concurrent(None, |completion| async {
+                    if let Some(ref i) = common.worker_interceptor {
+                        i.on_workflow_activation_completion(&completion).await;
                     }
-                },
-                // Only poll on the activity queue if activity functions have been registered. This
-                // makes tests which use mocks dramatically more manageable.
-                async {
-                    let mut shutdown_rx = shutdown_rx.clone();
-                    if !act_half.activity_fns.is_empty() {
-                        loop {
-                            tokio::select! {
-                                activity = common.worker.poll_activity_task() => {
-                                    if matches!(activity, Err(PollActivityError::ShutDown)) {
-                                        break;
-                                    }
-                                    act_half.activity_task_handler(common.worker.clone(),
-                                                                   activity?)?;
-                                },
-                                _ = shutdown_rx.changed() => { break }
-                            }
+                    common.worker.complete_workflow_activation(completion).await
+                })
+                .map_err(Into::into)
+                .await;
+            r
+        };
+        tokio::try_join!(
+            // Workflow polling loop
+            async {
+                loop {
+                    let activation = match common.worker.poll_workflow_activation().await {
+                        Err(PollWfError::ShutDown) => {
+                            break;
                         }
+                        o => o?,
                     };
-                    Result::<_, anyhow::Error>::Ok(())
+                    if let Some(wf_fut) = wf_half.workflow_activation_handler(
+                        common,
+                        shutdown_token.clone(),
+                        activation,
+                        &completions_tx,
+                    )? {
+                        if wf_future_tx.send(wf_fut).is_err() {
+                            panic!(
+                                "Receive half of completion processor channel cannot be dropped"
+                            );
+                        }
+                    }
                 }
-            );
-            wf_poll_res?;
-            // TODO: Activity loop errors don't show up until wf loop exits/errors
-            act_poll_res?;
-            Result::<_, anyhow::Error>::Ok(self)
-        };
+                // Tell still-alive workflows to evict themselves
+                shutdown_token.cancel();
+                // It's important to drop these so the future and completion processors will
+                // terminate.
+                drop(wf_future_tx);
+                drop(completions_tx);
+                Result::<_, anyhow::Error>::Ok(())
+            },
+            // Only poll on the activity queue if activity functions have been registered. This
+            // makes tests which use mocks dramatically more manageable.
+            async {
+                if !act_half.activity_fns.is_empty() {
+                    let shutdown_token = shutdown_token.clone();
+                    loop {
+                        tokio::select! {
+                            activity = common.worker.poll_activity_task() => {
+                                if matches!(activity, Err(PollActivityError::ShutDown)) {
+                                    break;
+                                }
+                                act_half.activity_task_handler(common.worker.clone(),
+                                                               common.task_queue.clone(),
+                                                               activity?)?;
+                            },
+                            _ = shutdown_token.cancelled() => { break }
+                        }
+                    }
+                };
+                Result::<_, anyhow::Error>::Ok(())
+            },
+            wf_future_joiner,
+            wf_completion_processor,
+        )?;
 
-        let myself = pollers.await?;
-        info!("Polling loop exited");
-        let _ = shutdown_tx.send(true);
-        while let Some(h) = myself.workflow_half.join_handles.next().await {
-            h??;
+        info!("Polling loops exited");
+        if let Some(i) = self.common.worker_interceptor.as_ref() {
+            i.on_shutdown(self);
         }
-        myself.common.worker.shutdown().await;
-        myself.workflow_half.workflows.clear();
+        self.common.worker.shutdown().await;
         Ok(())
     }
 
@@ -277,6 +321,12 @@ impl Worker {
         self.common.worker = new_core_worker;
     }
 
+    /// Returns number of currently cached workflows as understood by the SDK. Importantly, this
+    /// is not the same as understood by core, though they *should* always be in sync.
+    pub fn cached_workflows(&self) -> usize {
+        self.workflow_half.workflows.borrow().len()
+    }
+
     fn split_apart(&mut self) -> (&mut CommonWorker, &mut WorkflowHalf, &mut ActivityHalf) {
         (
             &mut self.common,
@@ -287,14 +337,19 @@ impl Worker {
 }
 
 impl WorkflowHalf {
-    async fn workflow_activation_handler(
-        &mut self,
+    fn workflow_activation_handler(
+        &self,
         common: &CommonWorker,
-        shutdown_rx: &Receiver<bool>,
-        completions_tx: &UnboundedSender<WorkflowActivationCompletion>,
-        completions_rx: &mut UnboundedReceiver<WorkflowActivationCompletion>,
+        shutdown_token: CancellationToken,
         activation: WorkflowActivation,
-    ) -> Result<(), anyhow::Error> {
+        completions_tx: &UnboundedSender<WorkflowActivationCompletion>,
+    ) -> Result<
+        Option<WorkflowFutureHandle<impl Future<Output = Result<WorkflowResult<()>, JoinError>>>>,
+        anyhow::Error,
+    > {
+        let mut res = None;
+        let run_id = activation.run_id.clone();
+
         // If the activation is to start a workflow, create a new workflow driver for it,
         // using the function associated with that workflow id
         if let Some(WorkflowActivationJob {
@@ -302,8 +357,8 @@ impl WorkflowHalf {
         }) = activation.jobs.get(0)
         {
             let workflow_type = &sw.workflow_type;
-            let wf_function = self
-                .workflow_fns
+            let wf_fns_borrow = self.workflow_fns.borrow();
+            let wf_function = wf_fns_borrow
                 .get(workflow_type)
                 .ok_or_else(|| anyhow!("Workflow type {workflow_type} not found"))?;
 
@@ -314,60 +369,48 @@ impl WorkflowHalf {
                 sw.arguments.clone(),
                 completions_tx.clone(),
             );
-            let mut shutdown_rx = shutdown_rx.clone();
             let jh = tokio::spawn(async move {
                 tokio::select! {
                     r = wff => r,
-                    _ = shutdown_rx.changed() => Ok(WfExitValue::Evicted)
+                    // TODO: This probably shouldn't abort early, as it could cause an in-progress
+                    //  complete to abort. Send synthetic remove activation
+                    _ = shutdown_token.cancelled() => {
+                        Ok(WfExitValue::Evicted)
+                    }
                 }
             });
-            self.workflows
-                .insert(activation.run_id.clone(), activations);
-            self.join_handles.push(jh.boxed());
+            res = Some(WorkflowFutureHandle {
+                join_handle: jh,
+                run_id: run_id.clone(),
+            });
+            self.workflows.borrow_mut().insert(
+                run_id.clone(),
+                WorkflowData {
+                    activation_chan: activations,
+                },
+            );
         }
 
         // The activation is expected to apply to some workflow we know about. Use it to
         // unblock things and advance the workflow.
-        if let Some(tx) = self.workflows.get_mut(&activation.run_id) {
-            tx.send(activation)
+        if let Some(dat) = self.workflows.borrow_mut().get_mut(&run_id) {
+            dat.activation_chan
+                .send(activation)
                 .expect("Workflow should exist if we're sending it an activation");
         } else {
             bail!("Got activation for unknown workflow");
         };
 
-        let completion = completions_rx.recv().await.expect("No workflows left?");
-        if let Some(ref i) = common.worker_interceptor {
-            i.on_workflow_activation_completion(&completion);
-        }
-        common
-            .worker
-            .complete_workflow_activation(completion)
-            .await?;
-        Ok(())
+        Ok(res)
     }
 }
 
-tokio::task_local! {
-    // This works, but maybe just passing a context object for activities like WFs is better
-    static ACT_CANCEL_TOK: CancellationToken
-}
-
-/// Returns a future the completes if and when the activity this was called inside has been
-/// cancelled
-pub async fn act_cancelled() {
-    ACT_CANCEL_TOK.with(|ct| ct.clone()).cancelled().await
-}
-
-/// Returns true if this activity has already been cancelled
-pub fn act_is_cancelled() -> bool {
-    ACT_CANCEL_TOK.with(|ct| ct.is_cancelled())
-}
-
 impl ActivityHalf {
     /// Spawns off a task to handle the provided activity task
     fn activity_task_handler(
         &mut self,
         worker: Arc<dyn CoreWorker>,
+        task_queue: String,
         activity: ActivityTask,
     ) -> Result<(), anyhow::Error> {
         match activity.variant {
@@ -383,13 +426,14 @@ impl ActivityHalf {
                     })?
                     .clone();
                 let ct = CancellationToken::new();
+                let task_token = activity.task_token;
                 self.task_tokens_to_cancels
-                    .insert(activity.task_token.clone().into(), ct.clone());
+                    .insert(task_token.clone().into(), ct.clone());
 
-                tokio::spawn(ACT_CANCEL_TOK.scope(ct, async move {
-                    let mut inputs = start.input;
-                    let arg = inputs.pop().unwrap_or_default();
-                    let output = (act_fn.act_func)(arg).await;
+                let (ctx, arg) =
+                    ActContext::new(worker.clone(), ct, task_queue, task_token.clone(), start);
+                tokio::spawn(async move {
+                    let output = (act_fn.act_func)(ctx, arg).await;
                     let result = match output {
                         Ok(res) => ActivityExecutionResult::ok(res),
                         Err(err) => match err.downcast::<ActivityCancelledError>() {
@@ -399,12 +443,12 @@ impl ActivityHalf {
                     };
                     worker
                         .complete_activity_task(ActivityTaskCompletion {
-                            task_token: activity.task_token,
+                            task_token,
                             result: Some(result),
                         })
                         .await?;
                     Result::<_, anyhow::Error>::Ok(())
-                }));
+                });
             }
             Some(activity_task::Variant::Cancel(_)) => {
                 if let Some(ct) = self.task_tokens_to_cancels.get(&activity.task_token.into()) {
@@ -622,8 +666,9 @@ impl<T: Debug> WfExitValue<T> {
     }
 }
 
-type BoxActFn =
-    Arc<dyn Fn(Payload) -> BoxFuture<'static, Result<Payload, anyhow::Error>> + Send + Sync>;
+type BoxActFn = Arc<
+    dyn Fn(ActContext, Payload) -> BoxFuture<'static, Result<Payload, anyhow::Error>> + Send + Sync,
+>;
 /// Container for user-defined activity functions
 #[derive(Clone)]
 pub struct ActivityFunction {
@@ -650,16 +695,16 @@ pub trait IntoActivityFunc<Args, Res> {
 
 impl<A, Rf, R, F> IntoActivityFunc<A, Rf> for F
 where
-    F: (Fn(A) -> Rf) + Sync + Send + 'static,
+    F: (Fn(ActContext, A) -> Rf) + Sync + Send + 'static,
     A: FromJsonPayloadExt + Send,
     Rf: Future<Output = Result<R, anyhow::Error>> + Send + 'static,
     R: AsJsonPayloadExt,
 {
     fn into_activity_fn(self) -> BoxActFn {
-        let wrapper = move |input: Payload| {
+        let wrapper = move |ctx: ActContext, input: Payload| {
             // Some minor gymnastics are required to avoid needing to clone the function
             match A::from_json_payload(&input) {
-                Ok(deser) => (self)(deser)
+                Ok(deser) => (self)(ctx, deser)
                     .map(|r| r.map(|r| r.as_json_payload())?)
                     .boxed(),
                 Err(e) => async move { Err(e.into()) }.boxed(),

package/sdk-core/sdk-core-protos/src/history_info.rs

@@ -117,17 +117,13 @@ impl HistoryInfo {
     /// Remove events from the beginning of this history such that it looks like what would've been
     /// delivered on a sticky queue where the previously started task was the one before the last
     /// task in this history.
-    ///
-    /// This is not *fully* accurate in that it will include commands that were part of the last
-    /// WFT completion, which the server would typically not include, but it's good enough for
-    /// testing.
     pub fn make_incremental(&mut self) {
         let last_complete_ix = self
             .events
             .iter()
             .rposition(|he| he.event_type() == EventType::WorkflowTaskCompleted)
             .expect("Must be a WFT completed event in history");
-        self.events.drain(0..=last_complete_ix);
+        self.events.drain(0..last_complete_ix);
     }
 
     pub fn events(&self) -> &[HistoryEvent] {
@@ -223,7 +219,7 @@ mod tests {
     fn incremental_works() {
         let t = single_timer("timer1");
         let hi = t.get_one_wft(2).unwrap();
-        assert_eq!(hi.events().len(), 4);
-        assert_eq!(hi.events()[0].event_id, 5);
+        assert_eq!(hi.events().len(), 5);
+        assert_eq!(hi.events()[0].event_id, 4);
     }
 }

package/sdk-core/test-utils/Cargo.toml

@@ -14,6 +14,7 @@ async-trait = "0.1"
 base64 = "0.13"
 futures = "0.3"
 log = "0.4"
+parking_lot = "0.12"
 prost = "0.9"
 prost-types = "0.9"
 rand = "0.8"