Be cautious and wait for possible error causing callbacks before we treat trailing metadata

1 files changed, 17 insertions, 1 deletions

diff --git a/src/core/ext/filters/message_size/message_size_filter.cc b/src/core/ext/filters/message_size/message_size_filter.cc
index c17df86f3d..01d483f45e 100644
--- a/src/core/ext/filters/message_size/message_size_filter.cc
+++ b/src/core/ext/filters/message_size/message_size_filter.cc
@@ -108,6 +108,8 @@ struct call_data {
   grpc_closure* next_recv_message_ready;
   // Original recv_trailing_metadata callback, invoked after our own.
   grpc_closure* original_recv_trailing_metadata_ready;
+  bool seen_recv_trailing_metadata;
+  grpc_error* recv_trailing_metadata_error;
 };
 
 struct channel_data {
@@ -147,7 +149,14 @@ static void recv_message_ready(void* user_data, grpc_error* error) {
     GRPC_ERROR_REF(error);
   }
   // Invoke the next callback.
-  GRPC_CLOSURE_RUN(calld->next_recv_message_ready, error);
+  grpc_closure* closure = calld->next_recv_message_ready;
+  calld->next_recv_message_ready = nullptr;
+  if (calld->seen_recv_trailing_metadata) {
+    GRPC_CALL_COMBINER_START(
+        calld->call_combiner, &calld->recv_trailing_metadata_ready,
+        calld->recv_trailing_metadata_error, "continue recv trailing metadata");
+  }
+  GRPC_CLOSURE_RUN(closure, error);
 }
 
 // Callback invoked on completion of recv_trailing_metadata
@@ -155,6 +164,12 @@ static void recv_message_ready(void* user_data, grpc_error* error) {
 static void recv_trailing_metadata_ready(void* user_data, grpc_error* error) {
   grpc_call_element* elem = static_cast<grpc_call_element*>(user_data);
   call_data* calld = static_cast<call_data*>(elem->call_data);
+  if (calld->next_recv_message_ready) {
+    calld->seen_recv_trailing_metadata = true;
+    calld->recv_trailing_metadata_error = GRPC_ERROR_REF(error);
+    GRPC_CALL_COMBINER_STOP(calld->call_combiner, "wait for recv message");
+    return;
+  }
   error =
       grpc_error_add_child(GRPC_ERROR_REF(error), GRPC_ERROR_REF(calld->error));
   // Invoke the next callback.
@@ -209,6 +224,7 @@ static grpc_error* init_call_elem(grpc_call_element* elem,
   calld->next_recv_message_ready = nullptr;
   calld->original_recv_trailing_metadata_ready = nullptr;
   calld->error = GRPC_ERROR_NONE;
+  calld->seen_recv_trailing_metadata = false;
   GRPC_CLOSURE_INIT(&calld->recv_message_ready, recv_message_ready, elem,
                     grpc_schedule_on_exec_ctx);
   GRPC_CLOSURE_INIT(&calld->recv_trailing_metadata_ready,