@temporalio/core-bridge 0.16.4 → 0.18.0

package/sdk-core/{src → core/src}/core_tests/activity_tasks.rs

@@ -1,8 +1,7 @@
 use crate::{
     job_assert,
-    pollers::{MockManualGateway, MockServerGatewayApis},
     test_help::{
-        build_fake_core, canned_histories, fake_sg_opts, gen_assert_and_reply, mock_core,
+        build_fake_core, canned_histories, gen_assert_and_reply, mock_core,
         mock_core_with_opts_no_workers, mock_manual_poller, mock_poller, mock_poller_from_resps,
         poll_and_reply, MockWorker, MocksHolder, TEST_Q,
     },
@@ -11,31 +10,35 @@ use crate::{
 };
 use futures::FutureExt;
 use std::{
+    cell::RefCell,
     collections::{hash_map::Entry, HashMap, VecDeque},
+    rc::Rc,
     sync::{
         atomic::{AtomicUsize, Ordering},
         Arc,
     },
     time::Duration,
 };
+use temporal_client::mocks::{fake_sg_opts, mock_gateway, mock_manual_gateway};
 use temporal_sdk_core_protos::{
     coresdk::{
-        activity_result::{activity_result, ActivityResult},
+        activity_result::{activity_resolution, ActivityExecutionResult, ActivityResolution},
         activity_task::activity_task,
-        workflow_activation::{wf_activation_job, ResolveActivity, WfActivationJob},
+        workflow_activation::{workflow_activation_job, ResolveActivity, WorkflowActivationJob},
         workflow_commands::{
             ActivityCancellationType, CompleteWorkflowExecution, RequestCancelActivity,
-            ScheduleActivity, StartTimer,
+            ScheduleActivity,
         },
         ActivityTaskCompletion,
     },
     temporal::api::workflowservice::v1::{
         PollActivityTaskQueueResponse, RecordActivityTaskHeartbeatResponse,
-        RespondActivityTaskCompletedResponse,
+        RespondActivityTaskCanceledResponse, RespondActivityTaskCompletedResponse,
+        RespondActivityTaskFailedResponse,
     },
 };
-use test_utils::fanout_tasks;
-use tokio::{sync::Notify, time::sleep};
+use temporal_sdk_core_test_utils::{fanout_tasks, start_timer_cmd};
+use tokio::{join, sync::Notify, time::sleep};
 
 #[tokio::test]
 async fn max_activities_respected() {
@@ -57,7 +60,7 @@ async fn max_activities_respected() {
             ..Default::default()
         },
     ]);
-    let mut mock_gateway = MockServerGatewayApis::new();
+    let mut mock_gateway = mock_gateway();
     mock_gateway
         .expect_poll_activity_task()
         .times(3)
@@ -80,7 +83,6 @@ async fn max_activities_respected() {
             .build()
             .unwrap(),
     )
-    .await
     .unwrap();
 
     // We allow two outstanding activities, therefore first two polls should return right away
@@ -93,7 +95,7 @@ async fn max_activities_respected() {
             core.complete_activity_task(ActivityTaskCompletion {
                 task_token: r1.task_token,
                 task_queue: TEST_Q.to_string(),
-                result: Some(ActivityResult::ok(vec![1].into()))
+                result: Some(ActivityExecutionResult::ok(vec![1].into()))
             }).await.unwrap();
             last_finisher.store(1, Ordering::SeqCst);
         },
@@ -108,7 +110,7 @@ async fn max_activities_respected() {
 
 #[tokio::test]
 async fn activity_not_found_returns_ok() {
-    let mut mock_gateway = MockServerGatewayApis::new();
+    let mut mock_gateway = mock_gateway();
     // Mock won't even be called, since we weren't tracking activity
     mock_gateway.expect_complete_activity_task().times(0);
 
@@ -121,7 +123,7 @@ async fn activity_not_found_returns_ok() {
     core.complete_activity_task(ActivityTaskCompletion {
         task_token: vec![1],
         task_queue: TEST_Q.to_string(),
-        result: Some(ActivityResult::ok(vec![1].into())),
+        result: Some(ActivityExecutionResult::ok(vec![1].into())),
     })
     .await
     .unwrap();
@@ -130,7 +132,7 @@ async fn activity_not_found_returns_ok() {
 
 #[tokio::test]
 async fn heartbeats_report_cancels_only_once() {
-    let mut mock_gateway = MockServerGatewayApis::new();
+    let mut mock_gateway = mock_gateway();
     mock_gateway
         .expect_record_activity_heartbeat()
         .times(2)
@@ -139,6 +141,14 @@ async fn heartbeats_report_cancels_only_once() {
                 cancel_requested: true,
             })
         });
+    mock_gateway
+        .expect_complete_activity_task()
+        .times(1)
+        .returning(|_, _| Ok(RespondActivityTaskCompletedResponse::default()));
+    mock_gateway
+        .expect_cancel_activity_task()
+        .times(1)
+        .returning(|_, _| Ok(RespondActivityTaskCanceledResponse::default()));
 
     let core = mock_core(MocksHolder::from_gateway_with_responses(
         mock_gateway,
@@ -181,22 +191,39 @@ async fn heartbeats_report_cancels_only_once() {
     // Allow heartbeat delay to elapse
     sleep(Duration::from_millis(10)).await;
     core.record_activity_heartbeat(ActivityHeartbeat {
-        task_token: act.task_token,
+        task_token: act.task_token.clone(),
         task_queue: TEST_Q.to_string(),
         details: vec![vec![1_u8, 2, 3].into()],
     });
+    // Wait delay again to flush heartbeat
     sleep(Duration::from_millis(10)).await;
+    // Now complete it as cancelled
+    core.complete_activity_task(ActivityTaskCompletion {
+        task_token: act.task_token,
+        task_queue: TEST_Q.to_string(),
+        result: Some(ActivityExecutionResult::cancel_from_details(None)),
+    })
+    .await
+    .unwrap();
     // Since cancels always come before new tasks, if we get a new non-cancel task, we did not
     // double-issue cancels.
     let act = core.poll_activity_task(TEST_Q).await.unwrap();
     assert_matches!(
-        act,
+        &act,
         ActivityTask {
             task_token,
             variant: Some(activity_task::Variant::Start(_)),
             ..
-        } => { task_token == vec![2] }
+        } => { task_token == &[2] }
     );
+    // Complete it so shutdown goes through
+    core.complete_activity_task(ActivityTaskCompletion {
+        task_token: act.task_token,
+        task_queue: TEST_Q.to_string(),
+        result: Some(ActivityExecutionResult::ok(vec![1].into())),
+    })
+    .await
+    .unwrap();
     core.shutdown().await;
 }
 
@@ -223,7 +250,7 @@ async fn activity_cancel_interrupts_poll() {
         .times(2)
         .returning(move || poll_resps.pop_front().unwrap());
 
-    let mut mock_gateway = MockManualGateway::new();
+    let mut mock_gateway = mock_manual_gateway();
     mock_gateway
         .expect_record_activity_heartbeat()
         .times(1)
@@ -235,6 +262,10 @@ async fn activity_cancel_interrupts_poll() {
             }
             .boxed()
         });
+    mock_gateway
+        .expect_complete_activity_task()
+        .times(1)
+        .returning(|_, _| async { Ok(RespondActivityTaskCompletedResponse::default()) }.boxed());
 
     let mock_worker = MockWorker {
         act_poller: Some(Box::from(mock_poller)),
@@ -256,7 +287,15 @@ async fn activity_cancel_interrupts_poll() {
             last_finisher.store(1, Ordering::SeqCst);
         },
         async {
-            core.poll_activity_task(TEST_Q).await.unwrap();
+            let act = core.poll_activity_task(TEST_Q).await.unwrap();
+            // Must complete this activity for shutdown to finish
+            core.complete_activity_task(
+                ActivityTaskCompletion {
+                    task_token: act.task_token,
+                    task_queue: TEST_Q.to_string(),
+                    result: Some(ActivityExecutionResult::ok(vec![1].into())),
+                }
+            ).await.unwrap();
             last_finisher.store(2, Ordering::SeqCst);
         }
     };
@@ -267,7 +306,7 @@ async fn activity_cancel_interrupts_poll() {
 
 #[tokio::test]
 async fn activity_poll_timeout_retries() {
-    let mock_gateway = MockServerGatewayApis::new();
+    let mock_gateway = mock_gateway();
     let mut calls = 0;
     let mut mock_act_poller = mock_poller();
     mock_act_poller.expect_poll().times(3).returning(move || {
@@ -296,7 +335,7 @@ async fn many_concurrent_heartbeat_cancels() {
     // them after a few successful heartbeats
     const CONCURRENCY_NUM: usize = 5;
 
-    let mut mock_gateway = MockManualGateway::new();
+    let mut mock_gateway = mock_manual_gateway();
     let mut poll_resps = VecDeque::from(
         (0..CONCURRENCY_NUM)
             .map(|i| {
@@ -362,7 +401,6 @@ async fn many_concurrent_heartbeat_cancels() {
             .build()
             .unwrap(),
     )
-    .await
     .unwrap();
 
     // Poll all activities first so they are registered
@@ -397,7 +435,7 @@ async fn many_concurrent_heartbeat_cancels() {
         core.complete_activity_task(ActivityTaskCompletion {
             task_token: r.task_token.clone(),
             task_queue: TEST_Q.to_string(),
-            result: Some(ActivityResult::cancel_from_details(None)),
+            result: Some(ActivityExecutionResult::cancel_from_details(None)),
         })
         .await
         .unwrap();
@@ -418,7 +456,7 @@ async fn activity_timeout_no_double_resolve() {
         NonSticky,
         &[
             gen_assert_and_reply(
-                &job_assert!(wf_activation_job::Variant::StartWorkflow(_)),
+                &job_assert!(workflow_activation_job::Variant::StartWorkflow(_)),
                 vec![ScheduleActivity {
                     seq: activity_id,
                     activity_id: activity_id.to_string(),
@@ -428,21 +466,17 @@ async fn activity_timeout_no_double_resolve() {
                 .into()],
             ),
             gen_assert_and_reply(
-                &job_assert!(wf_activation_job::Variant::SignalWorkflow(_)),
+                &job_assert!(workflow_activation_job::Variant::SignalWorkflow(_)),
                 vec![
                     RequestCancelActivity { seq: activity_id }.into(),
-                    StartTimer {
-                        seq: 2,
-                        ..Default::default()
-                    }
-                    .into(),
+                    start_timer_cmd(2, Duration::from_secs(1)),
                 ],
             ),
             gen_assert_and_reply(
-                &job_assert!(wf_activation_job::Variant::ResolveActivity(
+                &job_assert!(workflow_activation_job::Variant::ResolveActivity(
                     ResolveActivity {
-                        result: Some(ActivityResult {
-                            status: Some(activity_result::Status::Cancelled(..)),
+                        result: Some(ActivityResolution {
+                            status: Some(activity_resolution::Status::Cancelled(..)),
                         }),
                         ..
                     }
@@ -451,8 +485,8 @@ async fn activity_timeout_no_double_resolve() {
             ),
             gen_assert_and_reply(
                 &job_assert!(
-                    wf_activation_job::Variant::SignalWorkflow(_),
-                    wf_activation_job::Variant::FireTimer(_)
+                    workflow_activation_job::Variant::SignalWorkflow(_),
+                    workflow_activation_job::Variant::FireTimer(_)
                 ),
                 vec![CompleteWorkflowExecution { result: None }.into()],
             ),
@@ -465,7 +499,7 @@ async fn activity_timeout_no_double_resolve() {
 
 #[tokio::test]
 async fn only_returns_cancels_for_desired_queue() {
-    let mut mock_gateway = MockServerGatewayApis::new();
+    let mut mock_gateway = mock_gateway();
     let seen_cancel = Arc::new(Notify::new());
     let sc = seen_cancel.clone();
     mock_gateway
@@ -478,6 +512,10 @@ async fn only_returns_cancels_for_desired_queue() {
                 cancel_requested: true,
             })
         });
+    mock_gateway
+        .expect_cancel_activity_task()
+        .times(1)
+        .returning(|_, _| Ok(RespondActivityTaskCanceledResponse::default()));
 
     let mut w1 = MockWorker::for_queue("q1");
     w1.act_poller = Some(mock_poller_from_resps([PollActivityTaskQueueResponse {
@@ -515,12 +553,130 @@ async fn only_returns_cancels_for_desired_queue() {
         }
     };
     assert_matches!(
-        q1_res,
+        &q1_res,
         ActivityTask {
             variant: Some(activity_task::Variant::Cancel(_)),
             ..
         }
     );
+    // act needs to complete to finalize shutdown
+    core.complete_activity_task(ActivityTaskCompletion {
+        task_token: q1_res.task_token,
+        task_queue: "q1".to_string(),
+        result: Some(ActivityExecutionResult::cancel_from_details(None)),
+    })
+    .await
+    .unwrap();
 
     core.shutdown().await;
 }
+
+#[tokio::test]
+async fn can_heartbeat_acts_during_shutdown() {
+    let mut mock_gateway = mock_gateway();
+    mock_gateway
+        .expect_record_activity_heartbeat()
+        .times(1)
+        .returning(|_, _| {
+            Ok(RecordActivityTaskHeartbeatResponse {
+                cancel_requested: false,
+            })
+        });
+    mock_gateway
+        .expect_complete_activity_task()
+        .times(1)
+        .returning(|_, _| Ok(RespondActivityTaskCompletedResponse::default()));
+
+    let core = mock_core(MocksHolder::from_gateway_with_responses(
+        mock_gateway,
+        [],
+        [PollActivityTaskQueueResponse {
+            task_token: vec![1],
+            activity_id: "act1".to_string(),
+            heartbeat_timeout: Some(Duration::from_millis(1).into()),
+            ..Default::default()
+        }],
+    ));
+
+    let act = core.poll_activity_task(TEST_Q).await.unwrap();
+    let complete_order = RefCell::new(vec![]);
+    // Start shutdown before completing the activity
+    let shutdown_fut = async {
+        core.shutdown().await;
+        complete_order.borrow_mut().push(1);
+    };
+    let complete_fut = async {
+        core.record_activity_heartbeat(ActivityHeartbeat {
+            task_token: act.task_token.clone(),
+            task_queue: TEST_Q.to_string(),
+            details: vec![vec![1_u8, 2, 3].into()],
+        });
+        core.complete_activity_task(ActivityTaskCompletion {
+            task_token: act.task_token,
+            task_queue: TEST_Q.to_string(),
+            result: Some(ActivityExecutionResult::ok(vec![1].into())),
+        })
+        .await
+        .unwrap();
+        complete_order.borrow_mut().push(2);
+    };
+    join!(shutdown_fut, complete_fut);
+    assert_eq!(&complete_order.into_inner(), &[2, 1])
+}
+
+/// Verifies that if a user has tried to record a heartbeat and then immediately after failed the
+/// activity, that we flush those details before reporting the failure completion.
+#[tokio::test]
+async fn complete_act_with_fail_flushes_heartbeat() {
+    let last_hb = 50;
+    let mut mock_gateway = mock_gateway();
+    let last_seen_payload = Rc::new(RefCell::new(None));
+    let lsp = last_seen_payload.clone();
+    mock_gateway
+        .expect_record_activity_heartbeat()
+        // Two times b/c we always record the first heartbeat, and we'll flush the last
+        .times(2)
+        .returning_st(move |_, payload| {
+            *lsp.borrow_mut() = payload;
+            Ok(RecordActivityTaskHeartbeatResponse {
+                cancel_requested: false,
+            })
+        });
+    mock_gateway
+        .expect_fail_activity_task()
+        .times(1)
+        .returning(|_, _| Ok(RespondActivityTaskFailedResponse::default()));
+
+    let core = mock_core(MocksHolder::from_gateway_with_responses(
+        mock_gateway,
+        [],
+        [PollActivityTaskQueueResponse {
+            task_token: vec![1],
+            activity_id: "act1".to_string(),
+            heartbeat_timeout: Some(Duration::from_secs(10).into()),
+            ..Default::default()
+        }],
+    ));
+
+    let act = core.poll_activity_task(TEST_Q).await.unwrap();
+    // Record a bunch of heartbeats
+    for i in 1..=last_hb {
+        core.record_activity_heartbeat(ActivityHeartbeat {
+            task_token: act.task_token.clone(),
+            task_queue: TEST_Q.to_string(),
+            details: vec![vec![i].into()],
+        });
+    }
+    core.complete_activity_task(ActivityTaskCompletion {
+        task_token: act.task_token.clone(),
+        task_queue: TEST_Q.to_string(),
+        result: Some(ActivityExecutionResult::fail("Ahh".into())),
+    })
+    .await
+    .unwrap();
+    core.shutdown().await;
+
+    // Verify the last seen call to record a heartbeat had the last detail payload
+    let last_seen_payload = &last_seen_payload.take().unwrap().payloads[0];
+    assert_eq!(last_seen_payload.data, &[last_hb]);
+}

package/sdk-core/{src → core/src}/core_tests/child_workflows.rs

@@ -1,15 +1,15 @@
 use crate::{
-    pollers::MockServerGatewayApis,
-    prototype_rust_sdk::{
-        ChildWorkflowOptions, TestRustWorker, WfContext, WorkflowFunction, WorkflowResult,
-    },
+    replay::DEFAULT_WORKFLOW_TYPE,
     test_help::{
-        build_mock_pollers, canned_histories, mock_core, MockPollCfg, ResponseType,
-        DEFAULT_WORKFLOW_TYPE, TEST_Q,
+        build_mock_pollers, canned_histories, mock_core, MockPollCfg, ResponseType, TEST_Q,
     },
     workflow::managed_wf::ManagedWFFunc,
 };
 use std::sync::Arc;
+use temporal_client::mocks::mock_gateway;
+use temporal_sdk::{
+    ChildWorkflowOptions, TestRustWorker, WfContext, WorkflowFunction, WorkflowResult,
+};
 use temporal_sdk_core_protos::coresdk::child_workflow::{
     child_workflow_result, ChildWorkflowCancellationType,
 };
@@ -25,13 +25,13 @@ async fn signal_child_workflow(#[case] serial: bool) {
     let wf_id = "fakeid";
     let wf_type = DEFAULT_WORKFLOW_TYPE;
     let t = canned_histories::single_child_workflow_signaled("child-id-1", SIGNAME);
-    let mock = MockServerGatewayApis::new();
+    let mock = mock_gateway();
     let mh = MockPollCfg::from_resp_batches(wf_id, t, [ResponseType::AllHistory], mock);
     let mock = build_mock_pollers(mh);
     let core = mock_core(mock);
     let mut worker = TestRustWorker::new(Arc::new(core), TEST_Q.to_string(), None);
 
-    let wf = move |mut ctx: WfContext| async move {
+    let wf = move |ctx: WfContext| async move {
         let child = ctx.child_workflow(ChildWorkflowOptions {
             workflow_id: "child-id-1".to_string(),
             workflow_type: "child".to_string(),
@@ -39,16 +39,16 @@ async fn signal_child_workflow(#[case] serial: bool) {
         });
 
         let start_res = child
-            .start(&mut ctx)
+            .start(&ctx)
             .await
             .into_started()
             .expect("Child should get started");
         let (sigres, res) = if serial {
-            let sigres = start_res.signal(&mut ctx, SIGNAME, b"Hi!").await;
+            let sigres = start_res.signal(&ctx, SIGNAME, b"Hi!").await;
             let res = start_res.result().await;
             (sigres, res)
         } else {
-            let sigfut = start_res.signal(&mut ctx, SIGNAME, b"Hi!");
+            let sigfut = start_res.signal(&ctx, SIGNAME, b"Hi!");
             let resfut = start_res.result();
             join!(sigfut, resfut)
         };
@@ -65,7 +65,7 @@ async fn signal_child_workflow(#[case] serial: bool) {
     worker.run_until_done().await.unwrap();
 }
 
-async fn parent_cancels_child_wf(mut ctx: WfContext) -> WorkflowResult<()> {
+async fn parent_cancels_child_wf(ctx: WfContext) -> WorkflowResult<()> {
     let child = ctx.child_workflow(ChildWorkflowOptions {
         workflow_id: "child-id-1".to_string(),
         workflow_type: "child".to_string(),
@@ -74,11 +74,11 @@ async fn parent_cancels_child_wf(mut ctx: WfContext) -> WorkflowResult<()> {
     });
 
     let start_res = child
-        .start(&mut ctx)
+        .start(&ctx)
         .await
         .into_started()
         .expect("Child should get started");
-    let cancel_fut = start_res.cancel(&mut ctx);
+    let cancel_fut = start_res.cancel(&ctx);
     let resfut = start_res.result();
     let (cancel_res, res) = join!(cancel_fut, resfut);
     cancel_res.expect("cancel result is ok");

package/sdk-core/{src → core/src}/core_tests/determinism.rs

@@ -1,9 +1,7 @@
 use crate::{
-    pollers::MockServerGatewayApis,
-    prototype_rust_sdk::{TestRustWorker, WfContext, WorkflowResult},
+    replay::DEFAULT_WORKFLOW_TYPE,
     test_help::{
-        build_mock_pollers, canned_histories, mock_core, MockPollCfg, ResponseType,
-        DEFAULT_WORKFLOW_TYPE, TEST_Q,
+        build_mock_pollers, canned_histories, mock_core, MockPollCfg, ResponseType, TEST_Q,
     },
 };
 use std::{
@@ -13,10 +11,12 @@ use std::{
     },
     time::Duration,
 };
+use temporal_client::mocks::mock_gateway;
+use temporal_sdk::{TestRustWorker, WfContext, WorkflowResult};
 use temporal_sdk_core_protos::temporal::api::enums::v1::WorkflowTaskFailedCause;
 
 static DID_FAIL: AtomicBool = AtomicBool::new(false);
-pub async fn timer_wf_fails_once(mut ctx: WfContext) -> WorkflowResult<()> {
+pub async fn timer_wf_fails_once(ctx: WfContext) -> WorkflowResult<()> {
     ctx.timer(Duration::from_secs(1)).await;
     if DID_FAIL
         .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
@@ -34,7 +34,7 @@ async fn test_wf_task_rejected_properly() {
     let wf_id = "fakeid";
     let wf_type = DEFAULT_WORKFLOW_TYPE;
     let t = canned_histories::workflow_fails_with_failure_after_timer("1");
-    let mock = MockServerGatewayApis::new();
+    let mock = mock_gateway();
     let mut mh = MockPollCfg::from_resp_batches(wf_id, t, [1, 2, 2], mock);
     // We should see one wft failure which has unspecified cause, since panics don't have a defined
     // type.
@@ -63,7 +63,7 @@ async fn test_wf_task_rejected_properly_due_to_nondeterminism(#[case] use_cache:
     let wf_id = "fakeid";
     let wf_type = DEFAULT_WORKFLOW_TYPE;
     let t = canned_histories::single_timer_wf_completes("1");
-    let mock = MockServerGatewayApis::new();
+    let mock = mock_gateway();
     let mut mh = MockPollCfg::from_resp_batches(
         wf_id,
         t,
@@ -85,7 +85,7 @@ async fn test_wf_task_rejected_properly_due_to_nondeterminism(#[case] use_cache:
     let mut worker = TestRustWorker::new(Arc::new(core), TEST_Q.to_string(), None);
 
     let started_count: &'static _ = Box::leak(Box::new(AtomicUsize::new(0)));
-    worker.register_wf(wf_type.to_owned(), move |mut ctx: WfContext| async move {
+    worker.register_wf(wf_type.to_owned(), move |ctx: WfContext| async move {
         // The workflow is replaying all of history, so the when it schedules an extra timer it
         // should not have, it causes a nondeterminism error.
         if started_count.fetch_add(1, Ordering::Relaxed) == 0 {