Don't retry on grpc message too large errors

Author: tconley1428

client/src/retry.rs

@@ -201,7 +201,11 @@ where
 {
     type OutError = tonic::Status;
 
-    fn handle(&mut self, current_attempt: usize, e: tonic::Status) -> RetryPolicy<tonic::Status> {
+    fn handle(
+        &mut self,
+        current_attempt: usize,
+        mut e: tonic::Status,
+    ) -> RetryPolicy<tonic::Status> {
         // 0 max retries means unlimited retries
         if self.max_retries > 0 && current_attempt >= self.max_retries {
             return RetryPolicy::ForwardError(e);
@@ -213,6 +217,17 @@ where
             }
         }
 
+        // Short circuit if message is too large - this is not retryable
+        if e.code() == Code::ResourceExhausted
+            && e.message()
+                .contains("grpc: received message larger than max")
+        {
+            // Leave a marker so we don't have duplicate detection logic in the workflow
+            e.metadata_mut()
+                .insert("message-too-large", tonic::metadata::MetadataValue::from(0));
+            return RetryPolicy::ForwardError(e);
+        }
+
         // Task polls are OK with being cancelled or running into the timeout because there's
         // nothing to do but retry anyway
         let long_poll_allowed = self.call_type == CallType::TaskLongPoll
@@ -423,6 +438,26 @@ mod tests {
         assert_matches!(result, RetryPolicy::ForwardError(_))
     }
 
+    #[tokio::test]
+    async fn message_too_large_not_retried() {
+        let mut err_handler = TonicErrorHandler::new_with_clock(
+            CallInfo {
+                call_type: CallType::TaskLongPoll,
+                call_name: POLL_WORKFLOW_METH_NAME,
+                retry_cfg: TEST_RETRY_CONFIG,
+                retry_short_circuit: None,
+            },
+            TEST_RETRY_CONFIG,
+            FixedClock(Instant::now()),
+            FixedClock(Instant::now()),
+        );
+        let result = err_handler.handle(
+            1,
+            Status::new(Code::ResourceExhausted, "received message larger than max"),
+        );
+        assert_matches!(result, RetryPolicy::ForwardError(_));
+    }
+
     #[rstest::rstest]
     #[tokio::test]
     async fn task_poll_retries_forever<R>(

core/src/worker/workflow/mod.rs

@@ -14,6 +14,7 @@ mod workflow_stream;
 pub(crate) use driven_workflow::DrivenWorkflow;
 pub(crate) use history_update::HistoryUpdate;
 
+use crate::protosext::ValidPollWFTQResponse;
 use crate::{
     MetricsContext,
     abstractions::{
@@ -317,6 +318,180 @@ impl Workflows {
         }
     }
 
+    async fn handle_activation_completed_success(
+        &self,
+        run_id: &str,
+        wft_from_complete: &mut Option<ValidPollWFTQResponse>,
+        completion_time: Instant,
+        report: ServerCommandsWithWorkflowInfo,
+    ) -> WFTReportStatus {
+        match report {
+            ServerCommandsWithWorkflowInfo {
+                task_token,
+                action:
+                    ActivationAction::WftComplete {
+                        mut commands,
+                        messages,
+                        query_responses,
+                        force_new_wft,
+                        sdk_metadata,
+                        mut versioning_behavior,
+                    },
+            } => {
+                let reserved_act_permits =
+                    self.reserve_activity_slots_for_outgoing_commands(commands.as_mut_slice());
+                debug!(commands=%commands.display(), query_responses=%query_responses.display(),
+                           messages=%messages.display(), force_new_wft,
+                           "Sending responses to server");
+                if let Some(default_vb) = self.default_versioning_behavior.as_ref() {
+                    if versioning_behavior == VersioningBehavior::Unspecified {
+                        versioning_behavior = *default_vb;
+                    }
+                }
+                let mut completion = WorkflowTaskCompletion {
+                    task_token: task_token.clone(),
+                    commands,
+                    messages,
+                    query_responses,
+                    sticky_attributes: None,
+                    return_new_workflow_task: true,
+                    force_create_new_workflow_task: force_new_wft,
+                    sdk_metadata,
+                    metering_metadata: MeteringMetadata {
+                        nonfirst_local_activity_execution_attempts: self
+                            .local_act_mgr
+                            .get_nonfirst_attempt_count(run_id)
+                            as u32,
+                    },
+                    versioning_behavior,
+                };
+                let sticky_attrs = self.sticky_attrs.clone();
+                // Do not return new WFT if we would not cache, because returned new WFTs are
+                // always partial.
+                if sticky_attrs.is_none() {
+                    completion.return_new_workflow_task = false;
+                }
+                completion.sticky_attributes = sticky_attrs;
+
+                let mut reset_last_started_to = None;
+                self.handle_wft_reporting_errs(run_id, || async {
+                    match self.client.complete_workflow_task(completion).await {
+                        Ok(response) => {
+                            if response.reset_history_event_id > 0 {
+                                reset_last_started_to = Some(response.reset_history_event_id);
+                            }
+                            if let Some(wft) = response.workflow_task {
+                                *wft_from_complete = Some(validate_wft(wft)?);
+                            }
+                            self.handle_eager_activities(
+                                reserved_act_permits,
+                                response.activity_tasks,
+                            );
+                        }
+                        Err(e) if e.metadata().contains_key("message-too-large") => {
+                            let failure = Failure {
+                                failure: Some(
+                                    temporal_sdk_core_protos::temporal::api::failure::v1::Failure {
+                                        message: "GRPC Message too large".to_string(),
+                                        ..Default::default()
+                                    },
+                                ),
+                                force_cause: 0,
+                            };
+                            // TODO: tim - Update workflow cause when API is ready
+                            let new_outcome = FailedActivationWFTReport::Report(
+                                task_token,
+                                WorkflowTaskFailedCause::WorkflowWorkerUnhandledFailure,
+                                failure,
+                            );
+                            self.handle_activation_failed(run_id, completion_time, new_outcome)
+                                .await;
+                        }
+                        e => {
+                            e?;
+                        }
+                    };
+
+                    Ok(())
+                })
+                .await;
+                WFTReportStatus::Reported {
+                    reset_last_started_to,
+                    completion_time,
+                }
+            }
+            ServerCommandsWithWorkflowInfo {
+                task_token,
+                action: ActivationAction::RespondLegacyQuery { result },
+            } => {
+                self.respond_legacy_query(task_token, *result).await;
+                WFTReportStatus::Reported {
+                    reset_last_started_to: None,
+                    completion_time,
+                }
+            }
+        }
+    }
+
+    async fn handle_activation_failed(
+        &self,
+        run_id: &str,
+        completion_time: Instant,
+        outcome: FailedActivationWFTReport,
+    ) -> WFTReportStatus {
+        match outcome {
+            FailedActivationWFTReport::Report(tt, cause, failure) => {
+                warn!(run_id=%run_id, failure=?failure, "Failing workflow task");
+                self.handle_wft_reporting_errs(run_id, || async {
+                    self.client
+                        .fail_workflow_task(tt, cause, failure.failure)
+                        .await
+                })
+                .await;
+                WFTReportStatus::Reported {
+                    reset_last_started_to: None,
+                    completion_time,
+                }
+            }
+            FailedActivationWFTReport::ReportLegacyQueryFailure(task_token, failure) => {
+                warn!(run_id=%run_id, failure=?failure, "Failing legacy query request");
+                self.respond_legacy_query(task_token, legacy_query_failure(failure))
+                    .await;
+                WFTReportStatus::Reported {
+                    reset_last_started_to: None,
+                    completion_time,
+                }
+            }
+        }
+    }
+
+    async fn handle_activation_completed_result(
+        &self,
+        run_id: &str,
+        wft_from_complete: &mut Option<ValidPollWFTQResponse>,
+        completion_outcome: ActivationCompleteResult,
+    ) -> WFTReportStatus {
+        let completion_time = Instant::now();
+
+        match completion_outcome.outcome {
+            ActivationCompleteOutcome::ReportWFTSuccess(report) => {
+                self.handle_activation_completed_success(
+                    run_id,
+                    wft_from_complete,
+                    completion_time,
+                    report,
+                )
+                .await
+            }
+            ActivationCompleteOutcome::ReportWFTFail(outcome) => {
+                self.handle_activation_failed(run_id, completion_time, outcome)
+                    .await
+            }
+            ActivationCompleteOutcome::WFTFailedDontReport => WFTReportStatus::DropWft,
+            ActivationCompleteOutcome::DoNothing => WFTReportStatus::NotReported,
+        }
+    }
+
     /// Queue an activation completion for processing, returning a future that will resolve with
     /// the outcome of that completion. See [ActivationCompletedOutcome].
     ///
@@ -357,114 +532,12 @@ impl Workflows {
             );
             return Ok(());
         };
+        let replaying = completion_outcome.replaying;
 
         let mut wft_from_complete = None;
-        let completion_time = Instant::now();
-        let wft_report_status = match completion_outcome.outcome {
-            ActivationCompleteOutcome::ReportWFTSuccess(report) => match report {
-                ServerCommandsWithWorkflowInfo {
-                    task_token,
-                    action:
-                        ActivationAction::WftComplete {
-                            mut commands,
-                            messages,
-                            query_responses,
-                            force_new_wft,
-                            sdk_metadata,
-                            mut versioning_behavior,
-                        },
-                } => {
-                    let reserved_act_permits =
-                        self.reserve_activity_slots_for_outgoing_commands(commands.as_mut_slice());
-                    debug!(commands=%commands.display(), query_responses=%query_responses.display(),
-                           messages=%messages.display(), force_new_wft,
-                           "Sending responses to server");
-                    if let Some(default_vb) = self.default_versioning_behavior.as_ref() {
-                        if versioning_behavior == VersioningBehavior::Unspecified {
-                            versioning_behavior = *default_vb;
-                        }
-                    }
-                    let mut completion = WorkflowTaskCompletion {
-                        task_token,
-                        commands,
-                        messages,
-                        query_responses,
-                        sticky_attributes: None,
-                        return_new_workflow_task: true,
-                        force_create_new_workflow_task: force_new_wft,
-                        sdk_metadata,
-                        metering_metadata: MeteringMetadata {
-                            nonfirst_local_activity_execution_attempts: self
-                                .local_act_mgr
-                                .get_nonfirst_attempt_count(&run_id)
-                                as u32,
-                        },
-                        versioning_behavior,
-                    };
-                    let sticky_attrs = self.sticky_attrs.clone();
-                    // Do not return new WFT if we would not cache, because returned new WFTs are
-                    // always partial.
-                    if sticky_attrs.is_none() {
-                        completion.return_new_workflow_task = false;
-                    }
-                    completion.sticky_attributes = sticky_attrs;
-
-                    let mut reset_last_started_to = None;
-                    self.handle_wft_reporting_errs(&run_id, || async {
-                        let response = self.client.complete_workflow_task(completion).await?;
-                        if response.reset_history_event_id > 0 {
-                            reset_last_started_to = Some(response.reset_history_event_id);
-                        }
-                        if let Some(wft) = response.workflow_task {
-                            wft_from_complete = Some(validate_wft(wft)?);
-                        }
-                        self.handle_eager_activities(reserved_act_permits, response.activity_tasks);
-                        Ok(())
-                    })
-                    .await;
-                    WFTReportStatus::Reported {
-                        reset_last_started_to,
-                        completion_time,
-                    }
-                }
-                ServerCommandsWithWorkflowInfo {
-                    task_token,
-                    action: ActivationAction::RespondLegacyQuery { result },
-                } => {
-                    self.respond_legacy_query(task_token, *result).await;
-                    WFTReportStatus::Reported {
-                        reset_last_started_to: None,
-                        completion_time,
-                    }
-                }
-            },
-            ActivationCompleteOutcome::ReportWFTFail(outcome) => match outcome {
-                FailedActivationWFTReport::Report(tt, cause, failure) => {
-                    warn!(run_id=%run_id, failure=?failure, "Failing workflow task");
-                    self.handle_wft_reporting_errs(&run_id, || async {
-                        self.client
-                            .fail_workflow_task(tt, cause, failure.failure)
-                            .await
-                    })
-                    .await;
-                    WFTReportStatus::Reported {
-                        reset_last_started_to: None,
-                        completion_time,
-                    }
-                }
-                FailedActivationWFTReport::ReportLegacyQueryFailure(task_token, failure) => {
-                    warn!(run_id=%run_id, failure=?failure, "Failing legacy query request");
-                    self.respond_legacy_query(task_token, legacy_query_failure(failure))
-                        .await;
-                    WFTReportStatus::Reported {
-                        reset_last_started_to: None,
-                        completion_time,
-                    }
-                }
-            },
-            ActivationCompleteOutcome::WFTFailedDontReport => WFTReportStatus::DropWft,
-            ActivationCompleteOutcome::DoNothing => WFTReportStatus::NotReported,
-        };
+        let wft_report_status = self
+            .handle_activation_completed_result(&run_id, &mut wft_from_complete, completion_outcome)
+            .await;
 
         let maybe_pwft = if let Some(wft) = wft_from_complete {
             match HistoryPaginator::from_poll(wft, self.client.clone()).await {
@@ -485,7 +558,7 @@ impl Workflows {
         if let Some(h) = post_activate_hook {
             h(PostActivateHookData {
                 run_id: &run_id,
-                replaying: completion_outcome.replaying,
+                replaying,
             });
         }