From 09f17db9ca24ea069130f49b3c18c9e236509618 Mon Sep 17 00:00:00 2001
From: Lidi Zheng <lidiz@google.com>
Date: Mon, 5 Apr 2021 16:59:23 -0700
Subject: [PATCH] [Aio] Gracefully handle RPCs ends pre-maturely (#25884)

* Gracefully handle RPCs ends pre-maturely

* Fix a typo
---
 .../grpc/_cython/_cygrpc/aio/call.pyx.pxi     | 27 ++++++++++---------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/src/python/grpcio/grpc/_cython/_cygrpc/aio/call.pyx.pxi b/src/python/grpcio/grpc/_cython/_cygrpc/aio/call.pyx.pxi
index 2c2a3ff3f6e..7bce1850dc8 100644
--- a/src/python/grpcio/grpc/_cython/_cygrpc/aio/call.pyx.pxi
+++ b/src/python/grpcio/grpc/_cython/_cygrpc/aio/call.pyx.pxi
@@ -134,37 +134,38 @@ cdef class _AioCall(GrpcCallWrapper):
     cdef void _set_status(self, AioRpcStatus status) except *:
         cdef list waiters
 
+        # No more waiters should be expected since status has been set.
         self._status = status
 
         if self._initial_metadata is None:
             self._set_initial_metadata(_IMMUTABLE_EMPTY_METADATA)
 
-        # No more waiters should be expected since status
-        # has been set.
-        waiters = self._waiters_status
-        self._waiters_status = None
-
-        for waiter in waiters:
+        for waiter in self._waiters_status:
             if not waiter.done():
                 waiter.set_result(None)
+        self._waiters_status = []
 
         for callback in self._done_callbacks:
             callback()
 
     cdef void _set_initial_metadata(self, tuple initial_metadata) except *:
+        if self._initial_metadata is not None:
+            # Some gRPC calls might end before the initial metadata arrived in
+            # the Call object. That causes this method to be invoked twice: 1.
+            # filled with an empty metadata; 2. updated with the actual user
+            # provided metadata.
+            return
+
         cdef list waiters
 
+        # No more waiters should be expected since initial metadata has been
+        # set.
         self._initial_metadata = initial_metadata
 
-        # No more waiters should be expected since initial metadata
-        # has been set.
-        waiters = self._waiters_initial_metadata
-        self._waiters_initial_metadata = None
-
-        for waiter in waiters:
+        for waiter in self._waiters_initial_metadata:
             if not waiter.done():
                 waiter.set_result(None)
-
+        self._waiters_initial_metadata = []
 
     def add_done_callback(self, callback):
         if self.done():