@temporalio/core-bridge 0.19.2 → 0.20.2

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

@@ -1,6 +1,9 @@
 //! This crate contains testing functionality that can be useful when building SDKs against Core,
 //! or even when testing workflows written in SDKs that use Core.
 
+#[macro_use]
+extern crate tracing;
+
 pub mod canned_histories;
 
 use futures::{stream::FuturesUnordered, StreamExt};
@@ -8,18 +11,28 @@ use log::LevelFilter;
 use prost::Message;
 use rand::{distributions::Standard, Rng};
 use std::{
-    convert::TryFrom, env, future::Future, net::SocketAddr, path::PathBuf, sync::Arc,
+    cell::RefCell,
+    convert::TryFrom,
+    env,
+    future::Future,
+    net::SocketAddr,
+    path::PathBuf,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
     time::Duration,
 };
-use temporal_sdk::TestRustWorker;
+use temporal_client::{Client, RetryClient, WorkflowClientTrait, WorkflowOptions};
+use temporal_sdk::{interceptors::WorkerInterceptor, IntoActivityFunc, Worker, WorkflowFunction};
 use temporal_sdk_core::{
-    replay::{init_core_replay, ReplayCore},
-    CoreInitOptions, CoreInitOptionsBuilder, ServerGatewayOptions, ServerGatewayOptionsBuilder,
+    init_replay_worker, init_worker, telemetry_init, ClientOptions, ClientOptionsBuilder,
     TelemetryOptions, TelemetryOptionsBuilder, WorkerConfig, WorkerConfigBuilder,
 };
-use temporal_sdk_core_api::Core;
+use temporal_sdk_core_api::Worker as CoreWorker;
 use temporal_sdk_core_protos::{
     coresdk::{
+        common::Payload,
         workflow_commands::{
             workflow_command, ActivityCancellationType, CompleteWorkflowExecution,
             ScheduleActivity, StartTimer,
@@ -28,6 +41,7 @@ use temporal_sdk_core_protos::{
     },
     temporal::api::history::v1::History,
 };
+use tokio::sync::OnceCell;
 use url::Url;
 
 pub const NAMESPACE: &str = "default";
@@ -37,27 +51,28 @@ const OTEL_URL_ENV_VAR: &str = "TEMPORAL_INTEG_OTEL_URL";
 /// If set, enable direct scraping of prom metrics on the specified port
 const PROM_ENABLE_ENV_VAR: &str = "TEMPORAL_INTEG_PROM_PORT";
 
-/// Create a core instance which will use the provided test name to base the task queue and wf id
+/// Create a worker instance which will use the provided test name to base the task queue and wf id
 /// upon. Returns the instance and the task queue name (which is also the workflow id).
-pub async fn init_core_and_create_wf(test_name: &str) -> (Arc<dyn Core>, String) {
+pub async fn init_core_and_create_wf(test_name: &str) -> CoreWfStarter {
     let mut starter = CoreWfStarter::new(test_name);
-    let core = starter.get_core().await;
+    let _ = starter.get_worker().await;
     starter.start_wf().await;
-    (core, starter.get_task_queue().to_string())
+    starter
 }
 
-/// Create a core replay instance preloaded with just one provided history. Returns the core impl
+/// Create a worker replay instance preloaded with a provided history. Returns the worker impl
 /// and the task queue name as in [init_core_and_create_wf].
-pub fn init_core_replay_preloaded(test_name: &str, history: &History) -> (Arc<dyn Core>, String) {
-    let replay_core = init_core_replay(get_integ_telem_options());
+pub fn init_core_replay_preloaded(
+    test_name: &str,
+    history: &History,
+) -> (Arc<dyn CoreWorker>, String) {
     let worker_cfg = WorkerConfigBuilder::default()
+        .namespace(NAMESPACE)
         .task_queue(test_name)
         .build()
         .expect("Configuration options construct properly");
-    replay_core
-        .make_replay_worker(worker_cfg, history)
-        .expect("Worker registration works");
-    (Arc::new(replay_core), test_name.to_string())
+    let worker = init_replay_worker(worker_cfg, history).expect("Replay worker must init properly");
+    (Arc::new(worker), test_name.to_string())
 }
 
 /// Load history from a file containing the protobuf serialization of it
@@ -73,10 +88,14 @@ pub async fn history_from_proto_binary(path_from_root: &str) -> Result<History,
 pub struct CoreWfStarter {
     /// Used for both the task queue and workflow id
     task_queue_name: String,
-    core_options: CoreInitOptions,
+    telemetry_options: TelemetryOptions,
     worker_config: WorkerConfig,
     wft_timeout: Option<Duration>,
-    initted_core: Option<Arc<dyn Core>>,
+    initted_worker: OnceCell<InitializedWorker>,
+}
+struct InitializedWorker {
+    worker: Arc<dyn CoreWorker>,
+    client: Arc<RetryClient<Client>>,
 }
 
 impl CoreWfStarter {
@@ -90,65 +109,61 @@ impl CoreWfStarter {
     pub fn new_tq_name(task_queue: &str) -> Self {
         Self {
             task_queue_name: task_queue.to_owned(),
-            core_options: CoreInitOptionsBuilder::default()
-                .gateway_opts(get_integ_server_options())
-                .telemetry_opts(get_integ_telem_options())
-                .build()
-                .unwrap(),
+            telemetry_options: get_integ_telem_options(),
             worker_config: WorkerConfigBuilder::default()
+                .namespace(NAMESPACE)
                 .task_queue(task_queue)
                 .max_cached_workflows(1000_usize)
                 .build()
                 .unwrap(),
             wft_timeout: None,
-            initted_core: None,
+            initted_worker: OnceCell::new(),
         }
     }
 
-    pub async fn worker(&mut self) -> TestRustWorker {
-        TestRustWorker::new(
-            self.get_core().await,
+    pub async fn worker(&mut self) -> TestWorker {
+        let mut w = TestWorker::new(
+            self.get_worker().await,
             self.worker_config.task_queue.clone(),
-            self.wft_timeout,
-        )
+        );
+        w.client = Some(self.get_client().await);
+
+        w
     }
 
     pub async fn shutdown(&mut self) {
-        self.get_core().await.shutdown().await;
+        self.get_worker().await.shutdown().await;
     }
 
-    pub async fn get_core(&mut self) -> Arc<dyn Core> {
-        if self.initted_core.is_none() {
-            let core = temporal_sdk_core::init(self.core_options.clone())
-                .await
-                .unwrap();
-            // Register a worker for the task queue
-            core.register_worker(self.worker_config.clone()).unwrap();
-            self.initted_core = Some(Arc::new(core));
-        }
-        self.initted_core.as_ref().unwrap().clone()
+    pub async fn get_worker(&mut self) -> Arc<dyn CoreWorker> {
+        self.get_or_init().await.worker.clone()
+    }
+
+    pub async fn get_client(&mut self) -> Arc<RetryClient<Client>> {
+        self.get_or_init().await.client.clone()
     }
 
     /// Start the workflow defined by the builder and return run id
     pub async fn start_wf(&self) -> String {
-        self.start_wf_with_id(self.task_queue_name.clone()).await
+        self.start_wf_with_id(self.task_queue_name.clone(), WorkflowOptions::default())
+            .await
     }
 
-    pub async fn start_wf_with_id(&self, workflow_id: String) -> String {
-        self.initted_core
-            .as_ref()
+    pub async fn start_wf_with_id(&self, workflow_id: String, mut opts: WorkflowOptions) -> String {
+        opts.task_timeout = opts.task_timeout.or(self.wft_timeout);
+        self.initted_worker
+            .get()
             .expect(
-                "Core must be initted before starting a workflow.\
-                             Tests must call `get_core` first.",
+                "Worker must be initted before starting a workflow.\
+                             Tests must call `get_worker` first.",
             )
-            .as_ref()
-            .server_gateway()
+            .client
             .start_workflow(
                 vec![],
                 self.worker_config.task_queue.clone(),
                 workflow_id,
                 self.task_queue_name.clone(),
-                self.wft_timeout,
+                opts,
             )
             .await
             .unwrap()
@@ -161,21 +176,20 @@ impl CoreWfStarter {
         &mut self,
         wf_id: impl Into<String>,
         run_id: impl Into<String>,
-        worker: &mut TestRustWorker,
+        // TODO: Need not be passed in
+        worker: &mut Worker,
     ) -> Result<(), anyhow::Error> {
         // Fetch history and replay it
         let history = self
-            .get_core()
+            .get_client()
             .await
-            .server_gateway()
             .get_workflow_execution_history(wf_id.into(), Some(run_id.into()), vec![])
             .await?
             .history
             .expect("history field must be populated");
-        let (replay_core, _) = init_core_replay_preloaded(worker.task_queue(), &history);
-        let mut replay_worker = worker.swap_core(replay_core);
-        replay_worker.incr_expected_run_count(1);
-        replay_worker.run_until_done().await.unwrap();
+        let (replay_worker, _) = init_core_replay_preloaded(worker.task_queue(), &history);
+        worker.with_new_core_worker(replay_worker);
+        worker.run().await.unwrap();
         Ok(())
     }
 
@@ -216,16 +230,159 @@ impl CoreWfStarter {
         self.wft_timeout = Some(timeout);
         self
     }
+
+    async fn get_or_init(&mut self) -> &InitializedWorker {
+        self.initted_worker
+            .get_or_init(|| async {
+                telemetry_init(&self.telemetry_options).expect("Telemetry inits cleanly");
+                let client = Arc::new(
+                    get_integ_server_options()
+                        .connect(self.worker_config.namespace.clone(), None)
+                        .await
+                        .expect("Must connect"),
+                );
+                let worker = init_worker(self.worker_config.clone(), client.clone());
+                InitializedWorker {
+                    worker: Arc::new(worker),
+                    client,
+                }
+            })
+            .await
+    }
+}
+
+/// Provides conveniences for running integ tests with the SDK
+pub struct TestWorker {
+    inner: Worker,
+    client: Option<Arc<dyn WorkflowClientTrait>>,
+    incomplete_workflows: Arc<AtomicUsize>,
+}
+impl TestWorker {
+    /// Create a new test worker
+    pub fn new(core_worker: Arc<dyn CoreWorker>, task_queue: impl Into<String>) -> Self {
+        let ct = Arc::new(AtomicUsize::new(0));
+        let mut inner = Worker::new_from_core(core_worker, task_queue);
+        let iceptor = WorkflowCompletionCountingInterceptor {
+            incomplete_workflows: ct.clone(),
+            shutdown_handle: Box::new(inner.shutdown_handle()),
+        };
+        inner.set_worker_interceptor(Box::new(iceptor));
+        Self {
+            inner,
+            client: None,
+            incomplete_workflows: ct,
+        }
+    }
+
+    pub fn inner_mut(&mut self) -> &mut Worker {
+        &mut self.inner
+    }
+
+    pub fn incr_expected_run_count(&self, amount: usize) {
+        self.incomplete_workflows
+            .fetch_add(amount, Ordering::AcqRel);
+    }
+
+    // TODO: Maybe trait-ify?
+    pub fn register_wf<F: Into<WorkflowFunction>>(
+        &mut self,
+        workflow_type: impl Into<String>,
+        wf_function: F,
+    ) {
+        self.inner.register_wf(workflow_type, wf_function)
+    }
+
+    pub fn register_activity<A, R>(
+        &mut self,
+        activity_type: impl Into<String>,
+        act_function: impl IntoActivityFunc<A, R>,
+    ) {
+        self.inner.register_activity(activity_type, act_function)
+    }
+
+    /// 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>,
+        options: WorkflowOptions,
+    ) -> Result<String, anyhow::Error> {
+        self.incr_expected_run_count(1);
+        if let Some(c) = self.client.as_ref() {
+            let wfid = workflow_id.into();
+            let res = c
+                .start_workflow(
+                    input,
+                    self.inner.task_queue().to_string(),
+                    wfid.clone(),
+                    workflow_type.into(),
+                    options,
+                )
+                .await?;
+            Ok(res.run_id)
+        } else {
+            Ok("fake_run_id".to_string())
+        }
+    }
+
+    /// Runs until all expected workflows have completed
+    pub async fn run_until_done(&mut self) -> Result<(), anyhow::Error> {
+        self.inner.run().await
+    }
+
+    /// 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() + 'static,
+    ) -> Result<(), anyhow::Error> {
+        // Replace shutdown interceptor with one that calls the before hook first
+        let b4shut = RefCell::new(Some(before_shutdown));
+        let shutdown_handle = self.inner.shutdown_handle();
+        let iceptor = WorkflowCompletionCountingInterceptor {
+            incomplete_workflows: self.incomplete_workflows.clone(),
+            shutdown_handle: Box::new(move || {
+                if let Some(s) = b4shut.borrow_mut().take() {
+                    s();
+                }
+                shutdown_handle();
+            }),
+        };
+        self.inner.set_worker_interceptor(Box::new(iceptor));
+        self.inner.run().await
+    }
+}
+
+struct WorkflowCompletionCountingInterceptor {
+    incomplete_workflows: Arc<AtomicUsize>,
+    shutdown_handle: Box<dyn Fn()>,
+}
+impl WorkerInterceptor for WorkflowCompletionCountingInterceptor {
+    fn on_workflow_activation_completion(&self, completion: &WorkflowActivationCompletion) {
+        if completion.has_execution_ending() {
+            info!("Workflow {} says it's finishing", &completion.run_id);
+            let prev = self.incomplete_workflows.fetch_sub(1, Ordering::SeqCst);
+            if prev <= 1 {
+                // There are now zero, we just subtracted one
+                (self.shutdown_handle)()
+            }
+        }
+    }
 }
 
-pub fn get_integ_server_options() -> ServerGatewayOptions {
+pub fn get_integ_server_options() -> ClientOptions {
     let temporal_server_address = match env::var("TEMPORAL_SERVICE_ADDRESS") {
         Ok(addr) => addr,
         Err(_) => "http://localhost:7233".to_owned(),
     };
     let url = Url::try_from(&*temporal_server_address).unwrap();
-    ServerGatewayOptionsBuilder::default()
-        .namespace(NAMESPACE.to_string())
+    ClientOptionsBuilder::default()
         .identity("integ_tester".to_string())
         .worker_binary_id("fakebinaryid".to_string())
         .target_url(url)
@@ -308,19 +465,18 @@ where
 }
 
 #[async_trait::async_trait]
-pub trait CoreTestHelpers {
-    async fn complete_execution(&self, task_q: &str, run_id: &str);
-    async fn complete_timer(&self, task_q: &str, run_id: &str, seq: u32, duration: Duration);
+pub trait WorkerTestHelpers {
+    async fn complete_execution(&self, run_id: &str);
+    async fn complete_timer(&self, run_id: &str, seq: u32, duration: Duration);
 }
 
 #[async_trait::async_trait]
-impl<T> CoreTestHelpers for T
+impl<T> WorkerTestHelpers for T
 where
-    T: Core + ?Sized,
+    T: CoreWorker + ?Sized,
 {
-    async fn complete_execution(&self, task_q: &str, run_id: &str) {
+    async fn complete_execution(&self, run_id: &str) {
         self.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
-            task_q.to_string(),
             run_id.to_string(),
             vec![CompleteWorkflowExecution { result: None }.into()],
         ))
@@ -328,9 +484,8 @@ where
         .unwrap();
     }
 
-    async fn complete_timer(&self, task_q: &str, run_id: &str, seq: u32, duration: Duration) {
+    async fn complete_timer(&self, run_id: &str, seq: u32, duration: Duration) {
         self.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
-            task_q.to_string(),
             run_id.to_string(),
             vec![StartTimer {
                 seq,

package/sdk-core/tests/integ_tests/client_tests.rs

@@ -1,13 +1,36 @@
 use std::time::Duration;
-use temporal_sdk_core_test_utils::CoreWfStarter;
+use temporal_client::{RetryClient, WorkflowClientTrait, WorkflowService};
+use temporal_sdk_core_protos::temporal::api::workflowservice::v1::DescribeNamespaceRequest;
+use temporal_sdk_core_test_utils::{get_integ_server_options, CoreWfStarter, NAMESPACE};
 
 #[tokio::test]
-async fn can_use_retry_gateway() {
+async fn can_use_retry_client() {
     // Not terribly interesting by itself but can be useful for manually inspecting metrics etc
-    let mut core = CoreWfStarter::new("retry_gateway");
-    let retry_client = core.get_core().await.server_gateway();
+    let mut core = CoreWfStarter::new("retry_client");
+    let retry_client = core.get_client().await;
     for _ in 0..10 {
         retry_client.list_namespaces().await.unwrap();
         tokio::time::sleep(Duration::from_millis(10)).await;
     }
 }
+
+#[tokio::test]
+async fn can_use_retry_raw_client() {
+    let opts = get_integ_server_options();
+    let raw_client = opts.connect_no_namespace(None).await.unwrap();
+    let mut retry_client = RetryClient::new(raw_client, opts.retry_config);
+    retry_client
+        .describe_namespace(DescribeNamespaceRequest {
+            namespace: NAMESPACE.to_string(),
+            ..Default::default()
+        })
+        .await
+        .unwrap();
+}
+
+#[tokio::test]
+async fn calls_get_system_info() {
+    let opts = get_integ_server_options();
+    let raw_client = opts.connect_no_namespace(None).await.unwrap();
+    assert!(raw_client.get_client().capabilities().is_some());
+}

package/sdk-core/tests/integ_tests/heartbeat_tests.rs

@@ -4,41 +4,44 @@ use temporal_sdk_core_protos::coresdk::{
     activity_result::{
         self, activity_resolution as act_res, ActivityExecutionResult, ActivityResolution,
     },
-    activity_task::activity_task as act_task,
-    common::Payload,
+    activity_task::activity_task,
+    common::{Payload, RetryPolicy},
     workflow_activation::{workflow_activation_job, ResolveActivity, WorkflowActivationJob},
-    workflow_commands::ActivityCancellationType,
+    workflow_commands::{ActivityCancellationType, ScheduleActivity},
+    workflow_completion::WorkflowActivationCompletion,
     ActivityHeartbeat, ActivityTaskCompletion, IntoCompletion,
 };
 use temporal_sdk_core_test_utils::{
-    init_core_and_create_wf, schedule_activity_cmd, CoreTestHelpers,
+    init_core_and_create_wf, schedule_activity_cmd, WorkerTestHelpers,
 };
 use tokio::time::sleep;
 
 #[tokio::test]
 async fn activity_heartbeat() {
-    let (core, task_q) = init_core_and_create_wf("activity_heartbeat").await;
+    let mut starter = init_core_and_create_wf("activity_heartbeat").await;
+    let core = starter.get_worker().await;
+    let task_q = starter.get_task_queue();
     let activity_id = "act-1";
-    let task = core.poll_workflow_activation(&task_q).await.unwrap();
+    let task = core.poll_workflow_activation().await.unwrap();
     // Complete workflow task and schedule activity
     core.complete_workflow_activation(
         schedule_activity_cmd(
             0,
-            &task_q,
+            task_q,
             activity_id,
             ActivityCancellationType::TryCancel,
             Duration::from_secs(60),
             Duration::from_secs(1),
         )
-        .into_completion(task_q.to_string(), task.run_id),
+        .into_completion(task.run_id),
     )
     .await
     .unwrap();
     // Poll activity and verify that it's been scheduled with correct parameters
-    let task = core.poll_activity_task(&task_q).await.unwrap();
+    let task = core.poll_activity_task().await.unwrap();
     assert_matches!(
         task.variant,
-        Some(act_task::Variant::Start(start_activity)) => {
+        Some(activity_task::Variant::Start(start_activity)) => {
             assert_eq!(start_activity.activity_type, "test_activity".to_string())
         }
     );
@@ -49,7 +52,6 @@ async fn activity_heartbeat() {
         sleep(Duration::from_millis(100)).await;
         core.record_activity_heartbeat(ActivityHeartbeat {
             task_token: task.task_token.clone(),
-            task_queue: task_q.to_string(),
             details: vec![],
         });
     }
@@ -61,13 +63,12 @@ async fn activity_heartbeat() {
     // Complete activity successfully.
     core.complete_activity_task(ActivityTaskCompletion {
         task_token: task.task_token,
-        task_queue: task_q.to_string(),
         result: Some(ActivityExecutionResult::ok(response_payload.clone())),
     })
     .await
     .unwrap();
     // Poll workflow task and verify that activity has succeeded.
-    let task = core.poll_workflow_activation(&task_q).await.unwrap();
+    let task = core.poll_workflow_activation().await.unwrap();
     assert_matches!(
         task.jobs.as_slice(),
         [
@@ -83,5 +84,83 @@ async fn activity_heartbeat() {
             assert_eq!(r, &response_payload);
         }
     );
-    core.complete_execution(&task_q, &task.run_id).await;
+    core.complete_execution(&task.run_id).await;
+}
+
+#[tokio::test]
+async fn many_act_fails_with_heartbeats() {
+    let mut starter = init_core_and_create_wf("many_act_fails_with_heartbeats").await;
+    let core = starter.get_worker().await;
+    let activity_id = "act-1";
+    let task = core.poll_workflow_activation().await.unwrap();
+    // Complete workflow task and schedule activity
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
+        task.run_id,
+        ScheduleActivity {
+            seq: 0,
+            activity_id: activity_id.to_string(),
+            activity_type: "test_act".to_string(),
+            task_queue: starter.get_task_queue().to_string(),
+            start_to_close_timeout: Some(Duration::from_secs(30).into()),
+            retry_policy: Some(RetryPolicy {
+                initial_interval: Some(Duration::from_millis(10).into()),
+                backoff_coefficient: 1.0,
+                maximum_attempts: 4,
+                ..Default::default()
+            }),
+            heartbeat_timeout: Some(Duration::from_secs(1).into()),
+            ..Default::default()
+        }
+        .into(),
+    ))
+    .await
+    .unwrap();
+
+    // Multiple times, poll for the activity, heartbeat, and then immediately fail
+    // Poll activity and verify that it's been scheduled with correct parameters
+    for i in 0u8..=3 {
+        let task = core.poll_activity_task().await.unwrap();
+        let start_t = assert_matches!(task.variant, Some(activity_task::Variant::Start(s)) => s);
+
+        core.record_activity_heartbeat(ActivityHeartbeat {
+            task_token: task.task_token.clone(),
+            details: vec![[i].into()],
+        });
+
+        let compl = if i == 3 {
+            // Verify last hb was recorded
+            assert_eq!(start_t.heartbeat_details, [[2].into()]);
+            ActivityTaskCompletion {
+                task_token: task.task_token,
+                result: Some(ActivityExecutionResult::ok("passed".into())),
+            }
+        } else {
+            if i != 0 {
+                assert_eq!(start_t.heartbeat_details, [[i - 1].into()]);
+            }
+            ActivityTaskCompletion {
+                task_token: task.task_token,
+                result: Some(ActivityExecutionResult::fail(format!("Die on {i}").into())),
+            }
+        };
+        core.complete_activity_task(compl).await.unwrap();
+    }
+    let task = core.poll_workflow_activation().await.unwrap();
+
+    assert_matches!(
+        task.jobs.as_slice(),
+        [WorkflowActivationJob {
+            variant: Some(workflow_activation_job::Variant::ResolveActivity(
+                ResolveActivity {
+                    result: Some(ActivityResolution {
+                        status: Some(act_res::Status::Completed(activity_result::Success { .. })),
+                        ..
+                    }),
+                    ..
+                }
+            )),
+        },]
+    );
+    core.complete_execution(&task.run_id).await;
+    core.shutdown().await;
 }