@temporalio/core-bridge 0.16.4 → 0.18.0

package/sdk-core/{src → core/src}/worker/activities.rs

@@ -1,10 +1,17 @@
 mod activity_heartbeat_manager;
+mod local_activities;
+
+pub(crate) use local_activities::{
+    DispatchOrTimeoutLA, ExecutingLAId, LACompleteAction, LocalActRequest,
+    LocalActivityExecutionResult, LocalActivityManager, LocalActivityResolution,
+    LocalInFlightActInfo, NewLocalAct,
+};
 
 use crate::{
     pollers::BoxedActPoller,
-    task_token::TaskToken,
     telemetry::metrics::{activity_type, workflow_type, MetricsContext},
-    ActivityHeartbeatError, CompleteActivityError, PollActivityError, ServerGatewayApis,
+    worker::activities::activity_heartbeat_manager::ActivityHeartbeatError,
+    CompleteActivityError, PollActivityError, ServerGatewayApis, TaskToken,
 };
 use activity_heartbeat_manager::ActivityHeartbeatManager;
 use dashmap::DashMap;
@@ -15,7 +22,7 @@ use std::{
 };
 use temporal_sdk_core_protos::{
     coresdk::{
-        activity_result::{self as ar, activity_result},
+        activity_result::{self as ar, activity_execution_result as aer},
         activity_task::{ActivityCancelReason, ActivityTask},
         ActivityHeartbeat,
     },
@@ -24,7 +31,7 @@ use temporal_sdk_core_protos::{
         workflowservice::v1::PollActivityTaskQueueResponse,
     },
 };
-use tokio::sync::Semaphore;
+use tokio::sync::{Notify, Semaphore};
 
 #[derive(Debug, derive_more::Constructor)]
 struct PendingActivityCancel {
@@ -32,12 +39,18 @@ struct PendingActivityCancel {
     reason: ActivityCancelReason,
 }
 
-/// Contains minimal set of details that core needs to store, while activity is running.
+/// Contains minimal set of details that core needs to store while an activity is running.
 #[derive(Debug)]
-struct InflightActivityDetails {
-    pub activity_id: String,
+struct InFlightActInfo {
     pub activity_type: String,
     pub workflow_type: String,
+    start_time: Instant,
+}
+
+/// Augments [InFlightActInfo] with details specific to remote activities
+#[derive(Debug)]
+struct RemoteInFlightActInfo {
+    pub base: InFlightActInfo,
     /// Used to calculate aggregation delay between activity heartbeats.
     pub heartbeat_timeout: Option<prost_types::Duration>,
     /// Set to true if we have already issued a cancellation activation to lang for this activity
@@ -46,23 +59,22 @@ struct InflightActivityDetails {
     /// we have learned from heartbeating and issued a cancel task, in which case we may simply
     /// discard the reply.
     pub known_not_found: bool,
-    start_time: Instant,
 }
-impl InflightActivityDetails {
+impl RemoteInFlightActInfo {
     fn new(
-        activity_id: String,
         activity_type: String,
         workflow_type: String,
         heartbeat_timeout: Option<prost_types::Duration>,
     ) -> Self {
         Self {
-            activity_id,
-            activity_type,
-            workflow_type,
+            base: InFlightActInfo {
+                activity_type,
+                workflow_type,
+                start_time: Instant::now(),
+            },
             heartbeat_timeout,
             issued_cancel_to_lang: false,
             known_not_found: false,
-            start_time: Instant::now(),
         }
     }
 }
@@ -71,12 +83,14 @@ pub(crate) struct WorkerActivityTasks {
     /// Centralizes management of heartbeat issuing / throttling
     heartbeat_manager: ActivityHeartbeatManager,
     /// Activities that have been issued to lang but not yet completed
-    outstanding_activity_tasks: DashMap<TaskToken, InflightActivityDetails>,
+    outstanding_activity_tasks: DashMap<TaskToken, RemoteInFlightActInfo>,
     /// Buffers activity task polling in the event we need to return a cancellation while a poll is
     /// ongoing.
     poller: BoxedActPoller,
     /// Ensures we stay at or below this worker's maximum concurrent activity limit
     activities_semaphore: Semaphore,
+    /// Wakes every time an activity is removed from the outstanding map
+    complete_notify: Notify,
 
     metrics: MetricsContext,
 
@@ -98,6 +112,7 @@ impl WorkerActivityTasks {
             outstanding_activity_tasks: Default::default(),
             poller,
             activities_semaphore: Semaphore::new(max_activity_tasks),
+            complete_notify: Notify::new(),
             metrics,
             max_heartbeat_throttle_interval,
             default_heartbeat_throttle_interval,
@@ -108,13 +123,21 @@ impl WorkerActivityTasks {
         self.poller.notify_shutdown();
     }
 
+    /// Wait for all outstanding activity tasks to finish
+    pub(crate) async fn wait_all_finished(&self) {
+        while !self.outstanding_activity_tasks.is_empty() {
+            self.complete_notify.notified().await
+        }
+    }
+
     pub(crate) async fn shutdown(self) {
         self.poller.shutdown_box().await;
         self.heartbeat_manager.shutdown().await;
     }
 
-    /// Poll for an activity task. Returns `Ok(None)` if no activity is ready and the overall
-    /// polling loop should be retried.
+    /// Wait until not at the outstanding activity limit, and then poll for an activity task.
+    ///
+    /// Returns `Ok(None)` if no activity is ready and the overall polling loop should be retried.
     pub(crate) async fn poll(&self) -> Result<Option<ActivityTask>, PollActivityError> {
         let poll_with_semaphore = async {
             // Acquire and subsequently forget a permit for an outstanding activity. When they are
@@ -150,8 +173,7 @@ impl WorkerActivityTasks {
 
                         self.outstanding_activity_tasks.insert(
                             work.task_token.clone().into(),
-                            InflightActivityDetails::new(
-                                work.activity_id.clone(),
+                            RemoteInFlightActInfo::new(
                                 work.activity_type.clone().unwrap_or_default().name,
                                 work.workflow_type.clone().unwrap_or_default().name,
                                 work.heartbeat_timeout.clone()
@@ -173,36 +195,37 @@ impl WorkerActivityTasks {
     pub(crate) async fn complete(
         &self,
         task_token: TaskToken,
-        status: activity_result::Status,
+        status: aer::Status,
         gateway: &(dyn ServerGatewayApis + Send + Sync),
     ) -> Result<(), CompleteActivityError> {
         if let Some((_, act_info)) = self.outstanding_activity_tasks.remove(&task_token) {
             let act_metrics = self.metrics.with_new_attrs([
-                activity_type(act_info.activity_type.clone()),
-                workflow_type(act_info.workflow_type.clone()),
+                activity_type(act_info.base.activity_type.clone()),
+                workflow_type(act_info.base.workflow_type.clone()),
             ]);
-            act_metrics.act_execution_latency(act_info.start_time.elapsed());
+            act_metrics.act_execution_latency(act_info.base.start_time.elapsed());
             self.activities_semaphore.add_permits(1);
-            self.heartbeat_manager.evict(task_token.clone());
+            self.heartbeat_manager.evict(task_token.clone()).await;
             let known_not_found = act_info.known_not_found;
             drop(act_info); // TODO: Get rid of dashmap. If we hold ref across await, bad stuff.
+            self.complete_notify.notify_waiters();
 
             // No need to report activities which we already know the server doesn't care about
             if !known_not_found {
                 let maybe_net_err = match status {
-                    activity_result::Status::WillCompleteAsync(_) => None,
-                    activity_result::Status::Completed(ar::Success { result }) => gateway
+                    aer::Status::WillCompleteAsync(_) => None,
+                    aer::Status::Completed(ar::Success { result }) => gateway
                         .complete_activity_task(task_token.clone(), result.map(Into::into))
                         .await
                         .err(),
-                    activity_result::Status::Failed(ar::Failure { failure }) => {
+                    aer::Status::Failed(ar::Failure { failure }) => {
                         act_metrics.act_execution_failed();
                         gateway
                             .fail_activity_task(task_token.clone(), failure.map(Into::into))
                             .await
                             .err()
                     }
-                    activity_result::Status::Cancelled(ar::Cancellation { failure }) => {
+                    aer::Status::Cancelled(ar::Cancellation { failure }) => {
                         let details = if let Some(Failure {
                             failure_info:
                                 Some(FailureInfo::CanceledFailureInfo(CanceledFailureInfo { details })),
@@ -291,11 +314,7 @@ impl WorkerActivityTasks {
                 if reason == ActivityCancelReason::NotFound {
                     details.known_not_found = true;
                 }
-                Ok(Some(ActivityTask::cancel_from_ids(
-                    task_token.0,
-                    details.activity_id.clone(),
-                    reason,
-                )))
+                Ok(Some(ActivityTask::cancel_from_ids(task_token.0, reason)))
             } else {
                 warn!(task_token = ?task_token,
                               "Unknown activity task when issuing cancel");

package/sdk-core/{src → core/src}/worker/dispatcher.rs

@@ -1,5 +1,4 @@
 use crate::{
-    pollers::GatewayRef,
     telemetry::metrics::{task_queue, MetricsContext},
     worker::Worker,
     WorkerConfig, WorkerLookupErr, WorkerRegistrationError,
@@ -7,6 +6,7 @@ use crate::{
 use arc_swap::ArcSwap;
 use futures::future::join_all;
 use std::{collections::HashMap, ops::Deref, option::Option, sync::Arc};
+use temporal_client::ServerGatewayApis;
 use tokio::sync::Notify;
 
 /// Allows access to workers by task queue name
@@ -18,11 +18,11 @@ pub(crate) struct WorkerDispatcher {
 }
 
 impl WorkerDispatcher {
-    pub async fn new_worker(
+    pub fn new_worker(
         &self,
         config: WorkerConfig,
         sticky_queue: Option<String>,
-        gateway: Arc<GatewayRef>,
+        gateway: Arc<dyn ServerGatewayApis + Send + Sync>,
         parent_metrics: MetricsContext,
     ) -> Result<(), WorkerRegistrationError> {
         let tq = config.task_queue.clone();
@@ -96,13 +96,15 @@ impl WorkerDispatcher {
         }))
         .await;
 
-        let mut all_workers = HashMap::new();
+        let mut all_workers = vec![];
         self.workers.rcu(|map| {
             let mut map = HashMap::clone(map);
-            all_workers.extend(map.drain());
+            for worker in map.values_mut() {
+                all_workers.push(worker.take());
+            }
             map
         });
-        join_all(all_workers.into_values().map(|w| async move {
+        join_all(all_workers.into_iter().map(|w| async move {
             if let Some(w) = w {
                 w.destroy().await;
             }