@temporalio/core-bridge 0.14.0 → 0.16.4

package/sdk-core/src/pollers/gateway.rs

@@ -165,7 +165,7 @@ impl Default for RetryConfig {
 }
 
 impl RetryConfig {
-    pub(crate) fn poll_retry_policy() -> Self {
+    pub(crate) const fn poll_retry_policy() -> Self {
         Self {
             initial_interval: Duration::from_millis(200),
             randomization_factor: 0.2,
@@ -179,7 +179,7 @@ impl RetryConfig {
 
 impl From<RetryConfig> for ExponentialBackoff {
     fn from(c: RetryConfig) -> Self {
-        ExponentialBackoff {
+        Self {
             current_interval: c.initial_interval,
             initial_interval: c.initial_interval,
             randomization_factor: c.randomization_factor,
@@ -269,7 +269,7 @@ impl Interceptor for ServiceCallInterceptor {
                 .parse()
                 .unwrap_or_else(|_| MetadataValue::from_static("")),
         );
-        for (k, v) in self.opts.static_headers.iter() {
+        for (k, v) in &self.opts.static_headers {
             if let (Ok(k), Ok(v)) = (MetadataKey::from_str(k), MetadataValue::from_str(v)) {
                 metadata.insert(k, v);
             }
@@ -355,8 +355,11 @@ pub trait ServerGatewayApis {
 
     /// Fetch new workflow tasks from the provided queue. Should block indefinitely if there is no
     /// work.
-    async fn poll_workflow_task(&self, task_queue: String)
-        -> Result<PollWorkflowTaskQueueResponse>;
+    async fn poll_workflow_task(
+        &self,
+        task_queue: String,
+        is_sticky: bool,
+    ) -> Result<PollWorkflowTaskQueueResponse>;
 
     /// Fetch new activity tasks from the provided queue. Should block indefinitely if there is no
     /// work.
@@ -516,12 +519,17 @@ impl ServerGatewayApis for ServerGateway {
     async fn poll_workflow_task(
         &self,
         task_queue: String,
+        is_sticky: bool,
     ) -> Result<PollWorkflowTaskQueueResponse> {
         let mut request = PollWorkflowTaskQueueRequest {
             namespace: self.opts.namespace.clone(),
             task_queue: Some(TaskQueue {
                 name: task_queue,
-                kind: TaskQueueKind::Unspecified as i32,
+                kind: if is_sticky {
+                    TaskQueueKind::Sticky
+                } else {
+                    TaskQueueKind::Normal
+                } as i32,
             }),
             identity: self.opts.identity.clone(),
             binary_checksum: self.opts.worker_binary_id.clone(),
@@ -894,7 +902,7 @@ mockall::mock! {
         ) -> impl Future<Output = Result<StartWorkflowExecutionResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 
-        fn poll_workflow_task<'a, 'b>(&'a self, task_queue: String)
+        fn poll_workflow_task<'a, 'b>(&'a self, task_queue: String, is_sticky: bool)
             -> impl Future<Output = Result<PollWorkflowTaskQueueResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 

package/sdk-core/src/pollers/poll_buffer.rs

@@ -152,7 +152,7 @@ where
 
     async fn shutdown_box(self: Box<Self>) {
         let this = *self;
-        this.shutdown().await
+        this.shutdown().await;
     }
 }
 
@@ -192,7 +192,7 @@ impl Poller<PollWorkflowTaskQueueResponse> for WorkflowTaskPoller {
 
     async fn shutdown_box(self: Box<Self>) {
         let this = *self;
-        this.shutdown().await
+        this.shutdown().await;
     }
 }
 
@@ -200,6 +200,7 @@ pub type PollWorkflowTaskBuffer = LongPollBuffer<PollWorkflowTaskQueueResponse>;
 pub fn new_workflow_task_buffer(
     sg: Arc<impl ServerGatewayApis + Send + Sync + 'static + ?Sized>,
     task_queue: String,
+    is_sticky: bool,
     concurrent_pollers: usize,
     buffer_size: usize,
 ) -> PollWorkflowTaskBuffer {
@@ -207,7 +208,7 @@ pub fn new_workflow_task_buffer(
         move || {
             let sg = sg.clone();
             let task_queue = task_queue.clone();
-            async move { sg.poll_workflow_task(task_queue).await }
+            async move { sg.poll_workflow_task(task_queue, is_sticky).await }
         },
         concurrent_pollers,
         buffer_size,
@@ -246,7 +247,7 @@ mod tests {
         mock_gateway
             .expect_poll_workflow_task()
             .times(2)
-            .returning(move |_| {
+            .returning(move |_, _| {
                 async {
                     tokio::time::sleep(Duration::from_millis(100)).await;
                     Ok(Default::default())
@@ -255,7 +256,7 @@ mod tests {
             });
         let mock_gateway = Arc::new(mock_gateway);
 
-        let pb = new_workflow_task_buffer(mock_gateway, "someq".to_string(), 1, 1);
+        let pb = new_workflow_task_buffer(mock_gateway, "someq".to_string(), false, 1, 1);
 
         // Poll a bunch of times, "interrupting" it each time, we should only actually have polled
         // once since the poll takes a while

package/sdk-core/src/pollers/retry.rs

@@ -27,7 +27,7 @@ pub struct RetryGateway<SG> {
 
 impl<SG> RetryGateway<SG> {
     /// Use the provided retry config with the provided gateway
-    pub fn new(gateway: SG, retry_config: RetryConfig) -> Self {
+    pub const fn new(gateway: SG, retry_config: RetryConfig) -> Self {
         Self {
             gateway,
             retry_config,
@@ -36,23 +36,12 @@ impl<SG> RetryGateway<SG> {
 }
 
 impl<SG: ServerGatewayApis + Send + Sync + 'static> RetryGateway<SG> {
-    async fn call_with_retry<R, F, Fut>(&self, factory: F) -> Result<R>
-    where
-        F: Fn() -> Fut + Unpin,
-        Fut: Future<Output = Result<R>>,
-    {
-        self.call_type_with_retry(factory, CallType::Normal).await
-    }
-
-    async fn long_poll_call_with_retry<R, F, Fut>(&self, factory: F) -> Result<R>
-    where
-        F: Fn() -> Fut + Unpin,
-        Fut: Future<Output = Result<R>>,
-    {
-        self.call_type_with_retry(factory, CallType::LongPoll).await
-    }
-
-    async fn call_type_with_retry<R, F, Fut>(&self, factory: F, ct: CallType) -> Result<R>
+    async fn call_with_retry<R, F, Fut>(
+        &self,
+        factory: F,
+        ct: CallType,
+        call_name: &'static str,
+    ) -> Result<R>
     where
         F: Fn() -> Fut + Unpin,
         Fut: Future<Output = Result<R>>,
@@ -61,10 +50,12 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> RetryGateway<SG> {
             CallType::Normal => self.retry_config.clone(),
             CallType::LongPoll => RetryConfig::poll_retry_policy(),
         };
-        Ok(FutureRetry::new(factory, TonicErrorHandler::new(rtc, ct))
-            .await
-            .map_err(|(e, _attempt)| e)?
-            .0)
+        Ok(
+            FutureRetry::new(factory, TonicErrorHandler::new(rtc, ct, call_name))
+                .await
+                .map_err(|(e, _attempt)| e)?
+                .0,
+        )
     }
 }
 
@@ -73,17 +64,19 @@ struct TonicErrorHandler {
     backoff: ExponentialBackoff,
     max_retries: usize,
     call_type: CallType,
+    call_name: &'static str,
 }
 impl TonicErrorHandler {
-    fn new(cfg: RetryConfig, call_type: CallType) -> Self {
+    fn new(cfg: RetryConfig, call_type: CallType, call_name: &'static str) -> Self {
         Self {
             max_retries: cfg.max_retries,
             backoff: cfg.into(),
             call_type,
+            call_name,
         }
     }
 
-    fn should_log_retry_warning(&self, cur_attempt: usize) -> bool {
+    const fn should_log_retry_warning(&self, cur_attempt: usize) -> bool {
         // Warn on more than 5 retries for unlimited retrying
         if self.max_retries == 0 && cur_attempt > 5 {
             return true;
@@ -111,9 +104,9 @@ impl ErrorHandler<tonic::Status> for TonicErrorHandler {
         }
 
         if current_attempt == 1 {
-            debug!(error=?e, "gRPC call failed on first attempt")
+            debug!(error=?e, "gRPC call {} failed on first attempt", self.call_name);
         } else if self.should_log_retry_warning(current_attempt) {
-            warn!(error=?e, "gRPC call retried {} times", current_attempt)
+            warn!(error=?e, "gRPC call {} retried {} times", self.call_name, current_attempt);
         }
 
         // Long polls are OK with being cancelled or running into the timeout because there's
@@ -140,6 +133,13 @@ impl ErrorHandler<tonic::Status> for TonicErrorHandler {
     }
 }
 
+macro_rules! retry_call {
+    ($myself:ident, $ctype:expr, $call_name:ident, $($args:expr),*) => {{
+        let fact = move || { $myself.gateway.$call_name($($args,)*)};
+        $myself.call_with_retry(fact, $ctype, stringify!($call_name)).await
+    }};
+}
+
 #[async_trait::async_trait]
 impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryGateway<SG> {
     async fn start_workflow(
@@ -150,33 +150,42 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         workflow_type: String,
         task_timeout: Option<Duration>,
     ) -> Result<StartWorkflowExecutionResponse> {
-        let input = &input;
-        let factory = move || {
-            self.gateway.start_workflow(
-                input.to_vec(),
-                task_queue.clone(),
-                workflow_id.clone(),
-                workflow_type.clone(),
-                task_timeout,
-            )
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            start_workflow,
+            input.clone(),
+            task_queue.clone(),
+            workflow_id.clone(),
+            workflow_type.clone(),
+            task_timeout
+        )
     }
 
     async fn poll_workflow_task(
         &self,
         task_queue: String,
+        is_sticky: bool,
     ) -> Result<PollWorkflowTaskQueueResponse> {
-        let factory = move || self.gateway.poll_workflow_task(task_queue.clone());
-        self.long_poll_call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::LongPoll,
+            poll_workflow_task,
+            task_queue.clone(),
+            is_sticky
+        )
     }
 
     async fn poll_activity_task(
         &self,
         task_queue: String,
     ) -> Result<PollActivityTaskQueueResponse> {
-        let factory = move || self.gateway.poll_activity_task(task_queue.clone());
-        self.long_poll_call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::LongPoll,
+            poll_activity_task,
+            task_queue.clone()
+        )
     }
 
     async fn reset_sticky_task_queue(
@@ -184,19 +193,25 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         workflow_id: String,
         run_id: String,
     ) -> Result<ResetStickyTaskQueueResponse> {
-        let factory = move || {
-            self.gateway
-                .reset_sticky_task_queue(workflow_id.clone(), run_id.clone())
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            reset_sticky_task_queue,
+            workflow_id.clone(),
+            run_id.clone()
+        )
     }
 
     async fn complete_workflow_task(
         &self,
         request: WorkflowTaskCompletion,
     ) -> Result<RespondWorkflowTaskCompletedResponse> {
-        let factory = move || self.gateway.complete_workflow_task(request.clone());
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            complete_workflow_task,
+            request.clone()
+        )
     }
 
     async fn complete_activity_task(
@@ -204,11 +219,13 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         task_token: TaskToken,
         result: Option<Payloads>,
     ) -> Result<RespondActivityTaskCompletedResponse> {
-        let factory = move || {
-            self.gateway
-                .complete_activity_task(task_token.clone(), result.clone())
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            complete_activity_task,
+            task_token.clone(),
+            result.clone()
+        )
     }
 
     async fn record_activity_heartbeat(
@@ -216,11 +233,13 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         task_token: TaskToken,
         details: Option<Payloads>,
     ) -> Result<RecordActivityTaskHeartbeatResponse> {
-        let factory = move || {
-            self.gateway
-                .record_activity_heartbeat(task_token.clone(), details.clone())
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            record_activity_heartbeat,
+            task_token.clone(),
+            details.clone()
+        )
     }
 
     async fn cancel_activity_task(
@@ -228,11 +247,13 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         task_token: TaskToken,
         details: Option<Payloads>,
     ) -> Result<RespondActivityTaskCanceledResponse> {
-        let factory = move || {
-            self.gateway
-                .cancel_activity_task(task_token.clone(), details.clone())
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            cancel_activity_task,
+            task_token.clone(),
+            details.clone()
+        )
     }
 
     async fn fail_activity_task(
@@ -240,11 +261,13 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         task_token: TaskToken,
         failure: Option<Failure>,
     ) -> Result<RespondActivityTaskFailedResponse> {
-        let factory = move || {
-            self.gateway
-                .fail_activity_task(task_token.clone(), failure.clone())
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            fail_activity_task,
+            task_token.clone(),
+            failure.clone()
+        )
     }
 
     async fn fail_workflow_task(
@@ -253,11 +276,14 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         cause: WorkflowTaskFailedCause,
         failure: Option<Failure>,
     ) -> Result<RespondWorkflowTaskFailedResponse> {
-        let factory = move || {
-            self.gateway
-                .fail_workflow_task(task_token.clone(), cause, failure.clone())
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            fail_workflow_task,
+            task_token.clone(),
+            cause,
+            failure.clone()
+        )
     }
 
     async fn signal_workflow_execution(
@@ -267,15 +293,15 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         signal_name: String,
         payloads: Option<Payloads>,
     ) -> Result<SignalWorkflowExecutionResponse> {
-        let factory = move || {
-            self.gateway.signal_workflow_execution(
-                workflow_id.clone(),
-                run_id.clone(),
-                signal_name.clone(),
-                payloads.clone(),
-            )
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            signal_workflow_execution,
+            workflow_id.clone(),
+            run_id.clone(),
+            signal_name.clone(),
+            payloads.clone()
+        )
     }
 
     async fn query_workflow_execution(
@@ -284,14 +310,14 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         run_id: String,
         query: WorkflowQuery,
     ) -> Result<QueryWorkflowResponse> {
-        let factory = move || {
-            self.gateway.query_workflow_execution(
-                workflow_id.clone(),
-                run_id.clone(),
-                query.clone(),
-            )
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            query_workflow_execution,
+            workflow_id.clone(),
+            run_id.clone(),
+            query.clone()
+        )
     }
 
     async fn describe_workflow_execution(
@@ -299,11 +325,13 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         workflow_id: String,
         run_id: Option<String>,
     ) -> Result<DescribeWorkflowExecutionResponse> {
-        let factory = move || {
-            self.gateway
-                .describe_workflow_execution(workflow_id.clone(), run_id.clone())
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            describe_workflow_execution,
+            workflow_id.clone(),
+            run_id.clone()
+        )
     }
 
     async fn get_workflow_execution_history(
@@ -312,14 +340,14 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         run_id: Option<String>,
         page_token: Vec<u8>,
     ) -> Result<GetWorkflowExecutionHistoryResponse> {
-        let factory = move || {
-            self.gateway.get_workflow_execution_history(
-                workflow_id.clone(),
-                run_id.clone(),
-                page_token.clone(),
-            )
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            get_workflow_execution_history,
+            workflow_id.clone(),
+            run_id.clone(),
+            page_token.clone()
+        )
     }
 
     async fn respond_legacy_query(
@@ -327,11 +355,13 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         task_token: TaskToken,
         query_result: QueryResult,
     ) -> Result<RespondQueryTaskCompletedResponse> {
-        let factory = move || {
-            self.gateway
-                .respond_legacy_query(task_token.clone(), query_result.clone())
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            respond_legacy_query,
+            task_token.clone(),
+            query_result.clone()
+        )
     }
 
     async fn cancel_workflow_execution(
@@ -339,11 +369,13 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         workflow_id: String,
         run_id: Option<String>,
     ) -> Result<RequestCancelWorkflowExecutionResponse> {
-        let factory = move || {
-            self.gateway
-                .cancel_workflow_execution(workflow_id.clone(), run_id.clone())
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            cancel_workflow_execution,
+            workflow_id.clone(),
+            run_id.clone()
+        )
     }
 
     async fn terminate_workflow_execution(
@@ -351,15 +383,16 @@ impl<SG: ServerGatewayApis + Send + Sync + 'static> ServerGatewayApis for RetryG
         workflow_id: String,
         run_id: Option<String>,
     ) -> Result<TerminateWorkflowExecutionResponse> {
-        let factory = move || {
-            self.gateway
-                .terminate_workflow_execution(workflow_id.clone(), run_id.clone())
-        };
-        self.call_with_retry(factory).await
+        retry_call!(
+            self,
+            CallType::Normal,
+            terminate_workflow_execution,
+            workflow_id.clone(),
+            run_id.clone()
+        )
     }
 
     async fn list_namespaces(&self) -> Result<ListNamespacesResponse> {
-        let factory = move || self.gateway.list_namespaces();
-        self.call_with_retry(factory).await
+        retry_call!(self, CallType::Normal, list_namespaces,)
     }
 }