@temporalio/core-bridge 1.13.0 → 1.13.2

package/sdk-core/core/src/test_help/unit_helpers.rs

@@ -0,0 +1,218 @@
+//! Unit test helpers - only available in unit tests (cfg(test))
+
+use futures_util::{StreamExt, stream::FuturesUnordered};
+use std::{collections::HashSet, future::Future};
+use temporal_sdk_core_api::Worker as CoreWorker;
+use temporal_sdk_core_protos::coresdk::{
+    workflow_activation::workflow_activation_job,
+    workflow_completion::{WorkflowActivationCompletion, workflow_activation_completion},
+};
+
+/// Given a desired number of concurrent executions and a provided function that produces a future,
+/// run that many instances of the future concurrently.
+///
+/// Annoyingly, because of a sorta-bug in the way async blocks work, the async block produced by
+/// the closure must be `async move` if it uses the provided iteration number. On the plus side,
+/// since you're usually just accessing core in the closure, if core is a reference everything just
+/// works. See <https://github.com/rust-lang/rust/issues/81653>
+pub async fn fanout_tasks<FutureMaker, Fut>(num: usize, fm: FutureMaker)
+where
+    FutureMaker: Fn(usize) -> Fut,
+    Fut: Future,
+{
+    let mut tasks = FuturesUnordered::new();
+    for i in 0..num {
+        tasks.push(fm(i));
+    }
+
+    while tasks.next().await.is_some() {}
+}
+
+/// Generate asserts for [poll_and_reply] by passing patterns to match against the job list
+#[macro_export]
+macro_rules! job_assert {
+    ($($pat:pat),+) => {
+        |res| {
+            assert_matches!(
+                res.jobs.as_slice(),
+                [$(WorkflowActivationJob {
+                    variant: Some($pat),
+                }),+]
+            );
+        }
+    };
+}
+
+type AsserterWithReply<'a> = (
+    &'a dyn Fn(&temporal_sdk_core_protos::coresdk::workflow_activation::WorkflowActivation),
+    workflow_activation_completion::Status,
+);
+
+/// Determines when workflows are kept in the cache or evicted for [poll_and_reply] type tests
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub(crate) enum WorkflowCachingPolicy {
+    /// Workflows are evicted after each workflow task completion. Note that this is *not* after
+    /// each workflow activation - there are often multiple activations per workflow task.
+    NonSticky,
+
+    /// Not a real mode, but good for imitating crashes. Evict workflows after *every* reply,
+    /// even if there are pending activations
+    AfterEveryReply,
+}
+
+/// This function accepts a list of asserts and replies to workflow activations to run against the
+/// provided instance of fake core.
+///
+/// It handles the business of re-sending the same activation replies over again in the event
+/// of eviction or workflow activation failure. Activation failures specifically are only run once,
+/// since they clearly can't be returned every time we replay the workflow, or it could never
+/// proceed
+pub(crate) async fn poll_and_reply<'a>(
+    worker: &'a crate::Worker,
+    eviction_mode: WorkflowCachingPolicy,
+    expect_and_reply: &'a [AsserterWithReply<'a>],
+) {
+    poll_and_reply_clears_outstanding_evicts(worker, None, eviction_mode, expect_and_reply).await;
+}
+
+use crate::{pollers::MockPoller, test_help::OutstandingWFTMap};
+
+pub(crate) async fn poll_and_reply_clears_outstanding_evicts<'a>(
+    worker: &'a crate::Worker,
+    outstanding_map: Option<OutstandingWFTMap>,
+    eviction_mode: WorkflowCachingPolicy,
+    expect_and_reply: &'a [AsserterWithReply<'a>],
+) {
+    let mut evictions = 0;
+    let expected_evictions = expect_and_reply.len() - 1;
+    let mut executed_failures = HashSet::new();
+    let expected_fail_count = expect_and_reply
+        .iter()
+        .filter(|(_, reply)| !reply.is_success())
+        .count();
+
+    'outer: loop {
+        let expect_iter = expect_and_reply.iter();
+
+        for (i, interaction) in expect_iter.enumerate() {
+            let (asserter, reply) = interaction;
+            let complete_is_failure = !reply.is_success();
+            // Only send activation failures once
+            if executed_failures.contains(&i) {
+                continue;
+            }
+
+            let mut res = worker.poll_workflow_activation().await.unwrap();
+            if res.jobs.iter().any(|j| {
+                matches!(
+                    j.variant,
+                    Some(workflow_activation_job::Variant::RemoveFromCache(_))
+                )
+            }) && res.jobs.len() > 1
+            {
+                panic!("Saw an activation with an eviction & other work! {res:?}");
+            }
+            let is_eviction = res.is_only_eviction();
+
+            let mut do_release = false;
+
+            if is_eviction {
+                // If the job is an eviction, clear it, since in the tests we don't explicitly
+                // specify evict assertions
+                res.jobs.clear();
+                do_release = true;
+            }
+
+            // TODO: Can remove this if?
+            if !res.jobs.is_empty() {
+                asserter(&res);
+            }
+
+            let reply = if res.jobs.is_empty() {
+                // Just an eviction
+                WorkflowActivationCompletion::empty(res.run_id.clone())
+            } else {
+                // Eviction plus some work, we still want to issue the reply
+                WorkflowActivationCompletion {
+                    run_id: res.run_id.clone(),
+                    status: Some(reply.clone()),
+                }
+            };
+
+            let ends_execution = reply.has_execution_ending();
+
+            worker.complete_workflow_activation(reply).await.unwrap();
+
+            if do_release && let Some(omap) = outstanding_map.as_ref() {
+                omap.release_run(&res.run_id);
+            }
+            // Restart assertions from the beginning if it was an eviction (and workflow execution
+            // isn't over)
+            if is_eviction && !ends_execution {
+                continue 'outer;
+            }
+
+            if complete_is_failure {
+                executed_failures.insert(i);
+            }
+
+            match eviction_mode {
+                WorkflowCachingPolicy::NonSticky => (),
+                WorkflowCachingPolicy::AfterEveryReply => {
+                    if evictions < expected_evictions {
+                        worker.request_workflow_eviction(&res.run_id);
+                        evictions += 1;
+                    }
+                }
+            }
+        }
+
+        break;
+    }
+
+    assert_eq!(expected_fail_count, executed_failures.len());
+    assert_eq!(worker.outstanding_workflow_tasks().await, 0);
+}
+
+pub(crate) fn gen_assert_and_reply(
+    asserter: &dyn Fn(&temporal_sdk_core_protos::coresdk::workflow_activation::WorkflowActivation),
+    reply_commands: Vec<
+        temporal_sdk_core_protos::coresdk::workflow_commands::workflow_command::Variant,
+    >,
+) -> AsserterWithReply<'_> {
+    (
+        asserter,
+        temporal_sdk_core_protos::coresdk::workflow_completion::Success::from_variants(
+            reply_commands,
+        )
+        .into(),
+    )
+}
+
+pub(crate) fn gen_assert_and_fail(
+    asserter: &dyn Fn(&temporal_sdk_core_protos::coresdk::workflow_activation::WorkflowActivation),
+) -> AsserterWithReply<'_> {
+    (
+        asserter,
+        temporal_sdk_core_protos::coresdk::workflow_completion::Failure {
+            failure: Some(
+                temporal_sdk_core_protos::temporal::api::failure::v1::Failure {
+                    message: "Intentional test failure".to_string(),
+                    ..Default::default()
+                },
+            ),
+            ..Default::default()
+        }
+        .into(),
+    )
+}
+
+pub(crate) fn mock_poller<T>() -> MockPoller<T>
+where
+    T: Send + Sync + 'static,
+{
+    let mut mock_poller = MockPoller::new();
+    mock_poller.expect_shutdown_box().return_const(());
+    mock_poller.expect_notify_shutdown().return_const(());
+    mock_poller
+}

package/sdk-core/core/src/worker/activities/activity_heartbeat_manager.rs

@@ -42,6 +42,7 @@ enum HeartbeatAction {
     Evict {
         token: TaskToken,
         on_complete: Arc<Notify>,
+        should_flush: bool,
     },
     CompleteReport(TaskToken),
     CompleteThrottle(TaskToken),
@@ -118,7 +119,7 @@ impl ActivityHeartbeatManager {
                             HeartbeatAction::SendHeartbeat(hb) => hb_states.record(hb),
                             HeartbeatAction::CompleteReport(tt) => hb_states.handle_report_completed(tt),
                             HeartbeatAction::CompleteThrottle(tt) => hb_states.handle_throttle_completed(tt),
-                            HeartbeatAction::Evict{ token, on_complete } => hb_states.evict(token, on_complete),
+                            HeartbeatAction::Evict{ token, on_complete, should_flush } => hb_states.evict(token, on_complete, should_flush),
                         },
                         hb_states,
                     ))
@@ -230,13 +231,14 @@ impl ActivityHeartbeatManager {
     }
 
     /// Tell the heartbeat manager we are done forever with a certain task, so it may be forgotten.
-    /// This will also force-flush the most recently provided details.
+    /// If should_flush is true, will also force-flush the most recently provided details.
     /// Record *should* not be called with the same TaskToken after calling this.
-    pub(super) async fn evict(&self, task_token: TaskToken) {
+    pub(super) async fn evict(&self, task_token: TaskToken, should_flush: bool) {
         let completed = Arc::new(Notify::new());
         let _ = self.heartbeat_tx.send(HeartbeatAction::Evict {
             token: task_token,
             on_complete: completed.clone(),
+            should_flush,
         });
         completed.notified().await;
     }
@@ -397,12 +399,15 @@ impl HeartbeatStreamState {
         &mut self,
         tt: TaskToken,
         on_complete: Arc<Notify>,
+        should_flush: bool,
     ) -> Option<HeartbeatExecutorAction> {
         if let Some(state) = self.tt_to_state.remove(&tt) {
             if let Some(cancel_tok) = state.throttled_cancellation_token {
                 cancel_tok.cancel();
             }
-            if let Some(last_deets) = state.last_recorded_details {
+            if let Some(last_deets) = state.last_recorded_details
+                && should_flush
+            {
                 self.tt_needs_flush.insert(tt.clone(), on_complete);
                 return Some(HeartbeatExecutorAction::Report {
                     task_token: tt,
@@ -524,7 +529,7 @@ mod test {
         record_heartbeat(&hm, fake_task_token.clone(), 0, Duration::from_millis(100));
         // Let it propagate
         sleep(Duration::from_millis(10)).await;
-        hm.evict(fake_task_token.clone().into()).await;
+        hm.evict(fake_task_token.clone().into(), true).await;
         record_heartbeat(&hm, fake_task_token, 0, Duration::from_millis(100));
         // Let it propagate
         sleep(Duration::from_millis(10)).await;
@@ -543,7 +548,38 @@ mod test {
         let hm = ActivityHeartbeatManager::new(Arc::new(mock_client), cancel_tx);
         let fake_task_token = vec![1, 2, 3];
         record_heartbeat(&hm, fake_task_token.clone(), 0, Duration::from_millis(100));
-        hm.evict(fake_task_token.clone().into()).await;
+        hm.evict(fake_task_token.clone().into(), true).await;
+        hm.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn no_flush_on_successful_completion() {
+        let mut mock_client = mock_worker_client();
+        // Should only expect 1 heartbeat call, not 2 (the second would be from evict flushing)
+        mock_client
+            .expect_record_activity_heartbeat()
+            .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
+            .times(1);
+        let (cancel_tx, _cancel_rx) = unbounded_channel();
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client), cancel_tx);
+        let fake_task_token = vec![1, 2, 3];
+
+        // Record initial heartbeat - this should be sent immediately
+        record_heartbeat(&hm, fake_task_token.clone(), 0, Duration::from_millis(100));
+
+        // Wait a bit for initial heartbeat to process and enter throttling phase
+        sleep(Duration::from_millis(50)).await;
+
+        // Record another heartbeat while throttled - this should be stored in last_recorded_details
+        record_heartbeat(&hm, fake_task_token.clone(), 1, Duration::from_millis(100));
+
+        // Wait a bit to ensure the second heartbeat is recorded but not sent
+        sleep(Duration::from_millis(10)).await;
+
+        // Evict the activity with should_flush false
+        // This should NOT send the stored heartbeat details since the activity completed successfully
+        hm.evict(fake_task_token.into(), false).await;
+
         hm.shutdown().await;
     }
 

package/sdk-core/core/src/worker/activities/local_activities.rs

@@ -2,7 +2,6 @@ use crate::{
     MetricsContext, TaskToken,
     abstractions::{MeteredPermitDealer, OwnedMeteredSemPermit, UsedMeteredSemPermit, dbg_panic},
     protosext::ValidScheduleLA,
-    retry_logic::RetryPolicyExt,
     telemetry::metrics::{activity_type, should_record_failure_metric, workflow_type},
     worker::workflow::HeartbeatTimeoutMsg,
 };
@@ -13,6 +12,7 @@ use parking_lot::{Mutex, MutexGuard};
 use std::{
     collections::{HashMap, hash_map::Entry},
     fmt::{Debug, Formatter},
+    num::NonZero,
     pin::Pin,
     task::{Context, Poll},
     time::{Duration, Instant, SystemTime},
@@ -490,7 +490,7 @@ impl LocalActivityManager {
                 dispatch_time: Instant::now(),
                 attempt,
                 _permit: permit.into_used(LocalActivitySlotInfo {
-                    activity_type: new_la.workflow_type.clone(),
+                    activity_type: sa.activity_type.clone(),
                 }),
             },
         );
@@ -525,7 +525,7 @@ impl LocalActivityManager {
                     .or(schedule_to_close)
                     .and_then(|t| t.try_into().ok()),
                 heartbeat_timeout: None,
-                retry_policy: Some(sa.retry_policy),
+                retry_policy: Some(sa.retry_policy.into()),
                 priority: Some(Default::default()),
                 is_local: true,
             })),
@@ -570,7 +570,7 @@ impl LocalActivityManager {
             macro_rules! calc_backoff {
                 ($fail: ident) => {
                     info.la_info.schedule_cmd.retry_policy.should_retry(
-                        info.attempt as usize,
+                        info.attempt.try_into().unwrap_or(NonZero::<u32>::MIN),
                         $fail
                             .failure
                             .as_ref()
@@ -976,7 +976,7 @@ impl Drop for TimeoutBag {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{prost_dur, protosext::LACloseTimeouts};
+    use crate::{prost_dur, protosext::LACloseTimeouts, retry_logic::ValidatedRetryPolicy};
     use futures_util::FutureExt;
     use temporal_sdk_core_protos::temporal::api::{
         common::v1::RetryPolicy,
@@ -1111,13 +1111,13 @@ mod tests {
                 seq: 1,
                 activity_id: 1.to_string(),
                 attempt: 5,
-                retry_policy: RetryPolicy {
+                retry_policy: ValidatedRetryPolicy::from_proto_with_defaults(RetryPolicy {
                     initial_interval: Some(prost_dur!(from_secs(1))),
                     backoff_coefficient: 10.0,
                     maximum_interval: Some(prost_dur!(from_secs(10))),
                     maximum_attempts: 10,
                     non_retryable_error_types: vec![],
-                },
+                }),
                 local_retry_threshold: Duration::from_secs(5),
                 ..Default::default()
             },
@@ -1146,13 +1146,13 @@ mod tests {
                 seq: 1,
                 activity_id: "1".to_string(),
                 attempt: 1,
-                retry_policy: RetryPolicy {
+                retry_policy: ValidatedRetryPolicy::from_proto_with_defaults(RetryPolicy {
                     initial_interval: Some(prost_dur!(from_secs(1))),
                     backoff_coefficient: 10.0,
                     maximum_interval: Some(prost_dur!(from_secs(10))),
                     maximum_attempts: 10,
                     non_retryable_error_types: vec!["TestError".to_string()],
-                },
+                }),
                 local_retry_threshold: Duration::from_secs(5),
                 ..Default::default()
             },
@@ -1190,13 +1190,13 @@ mod tests {
                 seq: 1,
                 activity_id: 1.to_string(),
                 attempt: 5,
-                retry_policy: RetryPolicy {
+                retry_policy: ValidatedRetryPolicy::from_proto_with_defaults(RetryPolicy {
                     initial_interval: Some(prost_dur!(from_secs(10))),
                     backoff_coefficient: 1.0,
                     maximum_interval: Some(prost_dur!(from_secs(10))),
                     maximum_attempts: 10,
                     non_retryable_error_types: vec![],
-                },
+                }),
                 local_retry_threshold: Duration::from_secs(500),
                 ..Default::default()
             },
@@ -1239,11 +1239,11 @@ mod tests {
                 seq: 1,
                 activity_id: 1.to_string(),
                 attempt: 5,
-                retry_policy: RetryPolicy {
+                retry_policy: ValidatedRetryPolicy::from_proto_with_defaults(RetryPolicy {
                     initial_interval: Some(prost_dur!(from_millis(10))),
                     backoff_coefficient: 1.0,
                     ..Default::default()
-                },
+                }),
                 local_retry_threshold: Duration::from_secs(500),
                 ..Default::default()
             },
@@ -1276,11 +1276,11 @@ mod tests {
                 seq: 1,
                 activity_id: 1.to_string(),
                 attempt: 5,
-                retry_policy: RetryPolicy {
+                retry_policy: ValidatedRetryPolicy::from_proto_with_defaults(RetryPolicy {
                     initial_interval: Some(prost_dur!(from_millis(10))),
                     backoff_coefficient: 1.0,
                     ..Default::default()
-                },
+                }),
                 local_retry_threshold: Duration::from_secs(500),
                 schedule_to_start_timeout: Some(timeout),
                 ..Default::default()
@@ -1319,12 +1319,12 @@ mod tests {
                 seq: 1,
                 activity_id: 1.to_string(),
                 attempt: 5,
-                retry_policy: RetryPolicy {
+                retry_policy: ValidatedRetryPolicy::from_proto_with_defaults(RetryPolicy {
                     initial_interval: Some(prost_dur!(from_millis(10))),
                     backoff_coefficient: 1.0,
                     maximum_attempts: 1,
                     ..Default::default()
-                },
+                }),
                 local_retry_threshold: Duration::from_secs(500),
                 close_timeouts,
                 ..Default::default()
@@ -1400,11 +1400,11 @@ mod tests {
             schedule_cmd: ValidScheduleLA {
                 seq: 1,
                 activity_id: 1.to_string(),
-                retry_policy: RetryPolicy {
+                retry_policy: ValidatedRetryPolicy::from_proto_with_defaults(RetryPolicy {
                     initial_interval: Some(prost_dur!(from_millis(1))),
                     backoff_coefficient: 1.0,
                     ..Default::default()
-                },
+                }),
                 local_retry_threshold: Duration::from_secs(500),
                 ..Default::default()
             },

package/sdk-core/core/src/worker/activities.rs

@@ -333,7 +333,14 @@ impl WorkerActivityTasks {
             if let Some(jh) = act_info.local_timeouts_task {
                 jh.abort()
             };
-            self.heartbeat_manager.evict(task_token.clone()).await;
+            let should_flush = !known_not_found
+                && !matches!(
+                    &status,
+                    aer::Status::Completed(_) | aer::Status::WillCompleteAsync(_)
+                );
+            self.heartbeat_manager
+                .evict(task_token.clone(), should_flush)
+                .await;
 
             // No need to report activities which we already know the server doesn't care about
             if !known_not_found {
@@ -400,8 +407,6 @@ impl WorkerActivityTasks {
                     }
                 };
 
-                self.complete_notify.notify_waiters();
-
                 if let Some(e) = maybe_net_err {
                     if e.code() == tonic::Code::NotFound {
                         warn!(task_token=?task_token, details=?e, "Activity not found on \
@@ -418,6 +423,8 @@ impl WorkerActivityTasks {
                 &task_token
             );
         }
+
+        self.complete_notify.notify_waiters();
     }
 
     /// Attempt to record an activity heartbeat

package/sdk-core/core/src/worker/client/mocks.rs

@@ -20,9 +20,9 @@ pub(crate) static DEFAULT_TEST_CAPABILITIES: &Capabilities = &Capabilities {
     nexus: false,
 };
 
-#[cfg(test)]
+#[cfg(any(feature = "test-utilities", test))]
 /// Create a mock client primed with basic necessary expectations
-pub(crate) fn mock_worker_client() -> MockWorkerClient {
+pub fn mock_worker_client() -> MockWorkerClient {
     let mut r = MockWorkerClient::new();
     r.expect_capabilities()
         .returning(|| Some(*DEFAULT_TEST_CAPABILITIES));
@@ -152,7 +152,7 @@ mockall::mock! {
 
         fn record_worker_heartbeat<'a, 'b>(&self, heartbeat: WorkerHeartbeat) -> impl Future<Output = Result<RecordWorkerHeartbeatResponse>> + Send + 'b where 'a: 'b, Self: 'b;
 
-        fn replace_client(&self, new_client: RetryClient<Client>);
+        fn replace_client(&self, new_client: Client);
         fn capabilities(&self) -> Option<Capabilities>;
         fn workers(&self) -> Arc<SlotManager>;
         fn is_mock(&self) -> bool;