Compile fixes to lift out test worker -- tests still not passing

Author: Sushisource

core/benches/workflow_replay.rs

@@ -29,10 +29,9 @@ pub fn criterion_benchmark(c: &mut Criterion) {
             tokio_runtime.block_on(async {
                 let func = timers_wf(num_timers);
                 let (worker, _) = init_core_replay_preloaded("replay_bench", &hist);
-                let mut worker = Worker::new(worker, "replay_bench".to_string(), None);
+                let mut worker = Worker::new(worker, "replay_bench".to_string());
                 worker.register_wf(DEFAULT_WORKFLOW_TYPE, func);
-                worker.incr_expected_run_count(1);
-                worker.run_until_done().await.unwrap();
+                worker.run().await.unwrap();
             })
         })
     });
@@ -46,10 +45,9 @@ pub fn criterion_benchmark(c: &mut Criterion) {
             tokio_runtime.block_on(async {
                 let func = big_signals_wf(num_tasks);
                 let (worker, _) = init_core_replay_preloaded("large_hist_bench", &hist);
-                let mut worker = Worker::new(worker, "large_hist_bench".to_string(), None);
+                let mut worker = Worker::new(worker, "large_hist_bench".to_string());
                 worker.register_wf(DEFAULT_WORKFLOW_TYPE, func);
-                worker.incr_expected_run_count(1);
-                worker.run_until_done().await.unwrap();
+                worker.run().await.unwrap();
             })
         })
     });

core/src/core_tests/child_workflows.rs

@@ -7,12 +7,11 @@ use crate::{
 };
 use std::sync::Arc;
 use temporal_client::{mocks::mock_gateway, WorkflowOptions};
-use temporal_sdk::{
-    ChildWorkflowOptions, Signal, WfContext, Worker, WorkflowFunction, WorkflowResult,
-};
+use temporal_sdk::{ChildWorkflowOptions, Signal, WfContext, WorkflowFunction, WorkflowResult};
 use temporal_sdk_core_protos::coresdk::child_workflow::{
     child_workflow_result, ChildWorkflowCancellationType,
 };
+use temporal_sdk_core_test_utils::TestWorker;
 use tokio::join;
 
 const SIGNAME: &str = "SIGNAME";
@@ -29,7 +28,7 @@ async fn signal_child_workflow(#[case] serial: bool) {
     let mh = MockPollCfg::from_resp_batches(wf_id, t, [ResponseType::AllHistory], mock);
     let mock = build_mock_pollers(mh);
     let core = mock_worker(mock);
-    let mut worker = Worker::new(Arc::new(core), TEST_Q.to_string(), None);
+    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
 
     let wf = move |ctx: WfContext| async move {
         let child = ctx.child_workflow(ChildWorkflowOptions {

core/src/core_tests/determinism.rs

@@ -12,8 +12,9 @@ use std::{
     time::Duration,
 };
 use temporal_client::{mocks::mock_gateway, WorkflowOptions};
-use temporal_sdk::{WfContext, Worker, WorkflowResult};
+use temporal_sdk::{WfContext, WorkflowResult};
 use temporal_sdk_core_protos::temporal::api::enums::v1::WorkflowTaskFailedCause;
+use temporal_sdk_core_test_utils::TestWorker;
 
 static DID_FAIL: AtomicBool = AtomicBool::new(false);
 pub async fn timer_wf_fails_once(ctx: WfContext) -> WorkflowResult<()> {
@@ -43,7 +44,7 @@ async fn test_wf_task_rejected_properly() {
         Box::new(|_, cause, _| matches!(cause, WorkflowTaskFailedCause::Unspecified));
     let mock = build_mock_pollers(mh);
     let core = mock_worker(mock);
-    let mut worker = Worker::new(Arc::new(core), TEST_Q.to_string(), None);
+    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
 
     worker.register_wf(wf_type.to_owned(), timer_wf_fails_once);
     worker
@@ -87,7 +88,7 @@ async fn test_wf_task_rejected_properly_due_to_nondeterminism(#[case] use_cache:
         });
     }
     let core = mock_worker(mock);
-    let mut worker = Worker::new(Arc::new(core), TEST_Q.to_string(), None);
+    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
 
     let started_count: &'static _ = Box::leak(Box::new(AtomicUsize::new(0)));
     worker.register_wf(wf_type.to_owned(), move |ctx: WfContext| async move {

core/src/core_tests/local_activities.rs

@@ -12,11 +12,12 @@ use std::{
     time::Duration,
 };
 use temporal_client::{mocks::mock_gateway, WorkflowOptions};
-use temporal_sdk::{LocalActivityOptions, WfContext, Worker, WorkflowResult};
+use temporal_sdk::{LocalActivityOptions, WfContext, WorkflowResult};
 use temporal_sdk_core_protos::{
     coresdk::{common::RetryPolicy, AsJsonPayloadExt},
     temporal::api::{enums::v1::EventType, failure::v1::Failure},
 };
+use temporal_sdk_core_test_utils::TestWorker;
 use tokio::sync::Barrier;
 
 async fn echo(e: String) -> anyhow::Result<String> {
@@ -55,7 +56,7 @@ async fn local_act_two_wfts_before_marker(#[case] replay: bool, #[case] cached:
         mock.worker_cfg(|cfg| cfg.max_cached_workflows = 1);
     }
     let core = mock_worker(mock);
-    let mut worker = Worker::new(Arc::new(core), TEST_Q.to_string(), None);
+    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
 
     worker.register_wf(
         DEFAULT_WORKFLOW_TYPE.to_owned(),
@@ -119,7 +120,7 @@ async fn local_act_many_concurrent() {
     let mh = MockPollCfg::from_resp_batches(wf_id, t, [1, 2, 3], mock);
     let mock = build_mock_pollers(mh);
     let core = mock_worker(mock);
-    let mut worker = Worker::new(Arc::new(core), TEST_Q.to_string(), None);
+    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
 
     worker.register_wf(DEFAULT_WORKFLOW_TYPE.to_owned(), local_act_fanout_wf);
     worker.register_activity("echo", |str: String| async move { Ok(str) });
@@ -168,7 +169,7 @@ async fn local_act_heartbeat(#[case] shutdown_middle: bool) {
     let mut mock = build_mock_pollers(mh);
     mock.worker_cfg(|wc| wc.max_cached_workflows = 1);
     let core = Arc::new(mock_worker(mock));
-    let mut worker = Worker::new(core.clone(), TEST_Q.to_string(), None);
+    let mut worker = TestWorker::new(core.clone(), TEST_Q.to_string());
     let shutdown_barr: &'static Barrier = Box::leak(Box::new(Barrier::new(2)));
 
     worker.register_wf(
@@ -226,7 +227,7 @@ async fn local_act_fail_and_retry(#[case] eventually_pass: bool) {
     let mh = MockPollCfg::from_resp_batches(wf_id, t, [1], mock);
     let mock = build_mock_pollers(mh);
     let core = mock_worker(mock);
-    let mut worker = Worker::new(Arc::new(core), TEST_Q.to_string(), None);
+    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
 
     worker.register_wf(
         DEFAULT_WORKFLOW_TYPE.to_owned(),
@@ -310,7 +311,7 @@ async fn local_act_retry_long_backoff_uses_timer() {
     let mut mock = build_mock_pollers(mh);
     mock.worker_cfg(|w| w.max_cached_workflows = 1);
     let core = mock_worker(mock);
-    let mut worker = Worker::new(Arc::new(core), TEST_Q.to_string(), None);
+    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
 
     worker.register_wf(
         DEFAULT_WORKFLOW_TYPE.to_owned(),

sdk/src/lib.rs

@@ -53,13 +53,9 @@ use std::{
     collections::HashMap,
     fmt::{Debug, Display, Formatter},
     future::Future,
-    sync::{
-        atomic::{AtomicUsize, Ordering},
-        Arc,
-    },
-    time::Duration,
+    sync::Arc,
 };
-use temporal_client::{ServerGatewayApis, ServerGatewayOptionsBuilder, WorkflowOptions};
+use temporal_client::{ServerGatewayApis, ServerGatewayOptionsBuilder};
 use temporal_sdk_core::Url;
 use temporal_sdk_core_api::{
     errors::{PollActivityError, PollWfError},
@@ -112,7 +108,6 @@ pub fn sdk_client_options(url: impl Into<Url>) -> ServerGatewayOptionsBuilder {
 pub struct Worker {
     worker: Arc<dyn CoreWorker>,
     task_queue: String,
-    task_timeout: Option<Duration>,
     workflow_half: WorkflowHalf,
     activity_half: ActivityHalf,
 }
@@ -122,8 +117,6 @@ struct WorkflowHalf {
     workflows: HashMap<String, UnboundedSender<WorkflowActivation>>,
     /// Maps workflow type to the function for executing workflow runs with that ID
     workflow_fns: HashMap<String, WorkflowFunction>,
-    /// Number of live workflows
-    incomplete_workflows: Arc<AtomicUsize>,
     /// 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>>>,
@@ -137,19 +130,13 @@ struct ActivityHalf {
 
 impl Worker {
     /// Create a new rust worker
-    pub fn new(
-        worker: Arc<dyn CoreWorker>,
-        task_queue: impl Into<String>,
-        task_timeout: Option<Duration>,
-    ) -> Self {
+    pub fn new(worker: Arc<dyn CoreWorker>, task_queue: impl Into<String>) -> Self {
         Self {
             worker,
             task_queue: task_queue.into(),
-            task_timeout,
             workflow_half: WorkflowHalf {
                 workflows: Default::default(),
                 workflow_fns: Default::default(),
-                incomplete_workflows: Arc::new(AtomicUsize::new(0)),
                 join_handles: FuturesUnordered::new(),
             },
             activity_half: ActivityHalf {
@@ -169,36 +156,6 @@ impl Worker {
         &self.task_queue
     }
 
-    /// Create a workflow, asking the server to start it with the provided workflow ID and using the
-    /// provided workflow function.
-    ///
-    /// Increments the expected Workflow run count.
-    ///
-    /// Returns the run id of the started workflow
-    pub async fn submit_wf(
-        &self,
-        workflow_id: impl Into<String>,
-        workflow_type: impl Into<String>,
-        input: Vec<Payload>,
-        mut options: WorkflowOptions,
-    ) -> Result<String, tonic::Status> {
-        options.task_timeout = options.task_timeout.or(self.task_timeout);
-        let res = self
-            .worker
-            .server_gateway()
-            .start_workflow(
-                input,
-                self.task_queue.clone(),
-                workflow_id.into(),
-                workflow_type.into(),
-                options,
-            )
-            .await?;
-
-        self.incr_expected_run_count(1);
-        Ok(res.run_id)
-    }
-
     /// Register a Workflow function to invoke when the Worker is asked to run a workflow of
     /// `workflow_type`
     pub fn register_wf<F: Into<WorkflowFunction>>(
@@ -226,25 +183,9 @@ impl Worker {
         );
     }
 
-    // TODO: Should be removed before making this worker prod ready. There can be a test worker
-    //   which wraps this one and implements the workflow counting / run_until_done concepts.
-    //   This worker can expose an interceptor for completions that could be used to assist with
-    //   workflow tracking
-    /// Increment the expected Workflow run count on this Worker. The Worker tracks the run count
-    /// and will resolve `run_until_done` when it goes down to 0.
-    /// You do not have to increment if scheduled a Workflow with `submit_wf`.
-    pub fn incr_expected_run_count(&self, count: usize) {
-        self.workflow_half
-            .incomplete_workflows
-            .fetch_add(count, Ordering::SeqCst);
-    }
-
-    /// See [Self::run_until_done], except calls the provided callback just before performing core
-    /// shutdown.
-    pub async fn run_until_done_shutdown_hook(
-        &mut self,
-        before_shutdown: impl FnOnce(),
-    ) -> Result<(), anyhow::Error> {
+    /// 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 (worker, task_q, wf_half, act_half) = self.split_apart();
@@ -269,13 +210,6 @@ impl Worker {
                                 activation,
                             )
                             .await?;
-                        if wf_half.incomplete_workflows.load(Ordering::SeqCst) == 0 {
-                            info!("All expected workflows complete");
-                            // Die rebel scum - evict all workflows (which are complete now),
-                            // and turn off activity polling.
-                            let _ = shutdown_tx.send(true);
-                            break Result::<_, anyhow::Error>::Ok(());
-                        }
                     }
                 },
                 // Only poll on the activity queue if activity functions have been registered. This
@@ -309,18 +243,11 @@ impl Worker {
         while let Some(h) = myself.workflow_half.join_handles.next().await {
             h??;
         }
-        before_shutdown();
         myself.worker.shutdown().await;
         myself.workflow_half.workflows.clear();
         Ok(())
     }
 
-    /// Drives all workflows & activities until they have all finished, repeatedly polls server to
-    /// fetch work for them.
-    pub async fn run_until_done(&mut self) -> Result<(), anyhow::Error> {
-        self.run_until_done_shutdown_hook(|| {}).await
-    }
-
     /// Turns this rust worker into a new worker with all the same workflows and activities
     /// registered, but with a new underlying core worker. Can be used to swap the worker for
     /// a replay worker, change task queues, etc.
@@ -398,7 +325,6 @@ impl WorkflowHalf {
         let completion = completions_rx.recv().await.expect("No workflows left?");
         if completion.has_execution_ending() {
             debug!("Workflow {} says it's finishing", &completion.run_id);
-            self.incomplete_workflows.fetch_sub(1, Ordering::SeqCst);
         }
         worker.complete_workflow_activation(completion).await?;
         Ok(())

test-utils/Cargo.toml

@@ -24,6 +24,7 @@ temporal-sdk-core = { path = "../core" }
 temporal-sdk-core-api = { path = "../core-api" }
 thiserror = "1.0"
 tokio = "1.1"
+tracing = "0.1"
 url = "2.2"
 uuid = "0.8"