Avoid spamming grpc message too large task failures (#1050)

Author: Sushisource

crates/sdk-core/src/core_tests/workflow_tasks.rs

@@ -27,6 +27,7 @@ use std::{
     },
     time::Duration,
 };
+use temporalio_client::MESSAGE_TOO_LARGE_KEY;
 use temporalio_common::{
     Worker as WorkerTrait,
     errors::PollError,
@@ -3070,3 +3071,54 @@ async fn both_normal_and_sticky_pollers_poll_concurrently() {
         "At peak, there should be exactly 2 pollers active at the same time"
     );
 }
+
+#[tokio::test]
+async fn grpc_message_too_large_doesnt_spam_task_fails() {
+    let mut t = TestHistoryBuilder::default();
+    t.add_by_type(EventType::WorkflowExecutionStarted);
+    t.add_workflow_task_scheduled_and_started();
+
+    let mut mh = MockPollCfg::from_resp_batches(
+        "fake_wf_id",
+        t,
+        [
+            ResponseType::AllHistory,
+            ResponseType::AllHistory,
+            ResponseType::AllHistory,
+            ResponseType::AllHistory,
+            ResponseType::AllHistory,
+            ResponseType::AllHistory,
+        ],
+        mock_worker_client(),
+    );
+    mh.num_expected_fails = 1;
+    let mut times = 1;
+    mh.completion_mock_fn = Some(Box::new(move |_| {
+        if times <= 5 {
+            let mut err = tonic::Status::new(tonic::Code::ResourceExhausted, "message too large");
+            // This key is what we look for
+            err.metadata_mut().insert(MESSAGE_TOO_LARGE_KEY, 1.into());
+            times += 1;
+            Err(err)
+        } else {
+            Ok(Default::default())
+        }
+    }));
+
+    let mut mock = build_mock_pollers(mh);
+    mock.worker_cfg(|wc| wc.max_cached_workflows = 1);
+    let core = mock_worker(mock);
+
+    // Since the mock makes us fail 5 times, we should succeed on the sixth
+    for _ in 1..=5 {
+        let act = core.poll_workflow_activation().await.unwrap();
+        core.complete_workflow_activation(WorkflowActivationCompletion::empty(&act.run_id))
+            .await
+            .unwrap();
+        core.handle_eviction().await;
+    }
+    let act = core.poll_workflow_activation().await.unwrap();
+    core.complete_execution(&act.run_id).await;
+    core.drain_pollers_and_shutdown().await;
+    // Mock only expects 1 task failure, and would fail here if we spammed
+}

crates/sdk-core/src/worker/workflow/managed_run.rs

@@ -1142,6 +1142,7 @@ impl ManagedRun {
                 }
             }
 
+            let attempt = self.wft.as_ref().map(|t| t.info.attempt).unwrap_or(1);
             ActivationCompleteOutcome::ReportWFTSuccess(ServerCommandsWithWorkflowInfo {
                 task_token: data.task_token,
                 action: ActivationAction::WftComplete {
@@ -1151,6 +1152,7 @@ impl ManagedRun {
                     query_responses,
                     sdk_metadata: machines_wft_response.metadata_for_complete(),
                     versioning_behavior: data.versioning_behavior,
+                    attempt,
                 },
             })
         } else {

crates/sdk-core/src/worker/workflow/mod.rs

@@ -341,13 +341,14 @@ impl Workflows {
                         force_new_wft,
                         sdk_metadata,
                         mut versioning_behavior,
+                        attempt,
                     },
             } => {
                 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");
+                       messages=%messages.display(), force_new_wft,
+                       "Sending responses to server");
                 if let Some(default_vb) = self.default_versioning_behavior.as_ref()
                     && versioning_behavior == VersioningBehavior::Unspecified
                 {
@@ -393,7 +394,11 @@ impl Workflows {
                                 response.activity_tasks,
                             );
                         }
-                        Err(e) if e.metadata().contains_key(MESSAGE_TOO_LARGE_KEY) => {
+                        // Reply with a task failure if we got grpc too large from server, but
+                        // not if this is a nonfirst attempt to avoid spamming.
+                        Err(e)
+                            if e.metadata().contains_key(MESSAGE_TOO_LARGE_KEY) && attempt < 2 =>
+                        {
                             let failure = Failure {
                                 failure: Some(
                                     temporalio_common::protos::temporal::api::failure::v1::Failure {
@@ -417,6 +422,7 @@ impl Workflows {
                             );
                             self.handle_activation_failed(run_id, completion_time, new_outcome)
                                 .await;
+                            return Err(e);
                         }
                         e => {
                             e?;
@@ -992,6 +998,7 @@ pub(crate) enum ActivationAction {
         force_new_wft: bool,
         sdk_metadata: WorkflowTaskCompletedMetadata,
         versioning_behavior: VersioningBehavior,
+        attempt: u32,
     },
     /// We should respond to a legacy query request
     RespondLegacyQuery { result: Box<QueryResult> },