Merge pull request #21819 from Skyscanner/close_ongoing_calls

[Aio] Close ongoing calls when the channel is closed

src/python/grpcio/grpc/_cython/_cygrpc/aio/call.pyx.pxi

@@ -132,14 +132,14 @@ cdef class _AioCall(GrpcCallWrapper):
     cdef void _set_status(self, AioRpcStatus status) except *:
         cdef list waiters
 
+        self._status = status
+
         if self._initial_metadata is None:
             self._set_initial_metadata(_IMMUTABLE_EMPTY_METADATA)
 
-        self._status = status
-        waiters = self._waiters_status
-
         # No more waiters should be expected since status
         # has been set.
+        waiters = self._waiters_status
         self._waiters_status = None
 
         for waiter in waiters:
@@ -154,10 +154,9 @@ cdef class _AioCall(GrpcCallWrapper):
 
         self._initial_metadata = initial_metadata
 
-        waiters = self._waiters_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:

src/python/grpcio/grpc/_cython/_cygrpc/aio/channel.pxd.pxi

@@ -15,6 +15,7 @@
 cdef enum AioChannelStatus:
     AIO_CHANNEL_STATUS_UNKNOWN
     AIO_CHANNEL_STATUS_READY
+    AIO_CHANNEL_STATUS_CLOSING
     AIO_CHANNEL_STATUS_DESTROYED
 
 cdef class AioChannel:

src/python/grpcio/grpc/_cython/_cygrpc/aio/channel.pyx.pxi

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+#
 
 
 class _WatchConnectivityFailed(Exception):
@@ -69,9 +70,10 @@ cdef class AioChannel:
         Keeps mirroring the behavior from Core, so we can easily switch to
         other design of API if necessary.
         """
-        if self._status == AIO_CHANNEL_STATUS_DESTROYED:
+        if self._status in (AIO_CHANNEL_STATUS_DESTROYED, AIO_CHANNEL_STATUS_CLOSING):
             # TODO(lidiz) switch to UsageError
             raise RuntimeError('Channel is closed.')
+
         cdef gpr_timespec c_deadline = _timespec_from_time(deadline)
 
         cdef object future = self.loop.create_future()
@@ -92,10 +94,16 @@ cdef class AioChannel:
         else:
             return True
 
+    def closing(self):
+        self._status = AIO_CHANNEL_STATUS_CLOSING
+
     def close(self):
         self._status = AIO_CHANNEL_STATUS_DESTROYED
         grpc_channel_destroy(self.channel)
 
+    def closed(self):
+        return self._status in (AIO_CHANNEL_STATUS_CLOSING, AIO_CHANNEL_STATUS_DESTROYED)
+
     def call(self,
              bytes method,
              object deadline,
@@ -106,7 +114,7 @@ cdef class AioChannel:
         Returns:
           The _AioCall object.
         """
-        if self._status == AIO_CHANNEL_STATUS_DESTROYED:
+        if self.closed():
             # TODO(lidiz) switch to UsageError
             raise RuntimeError('Channel is closed.')
 

src/python/grpcio/grpc/experimental/aio/_channel.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 """Invocation-side implementation of gRPC Asyncio Python."""
 import asyncio
-from typing import Any, AsyncIterable, Optional, Sequence, Text
+from typing import Any, AsyncIterable, Optional, Sequence, AbstractSet, Text
+from weakref import WeakSet
 
+import logging
 import grpc
 from grpc import _common
 from grpc._cython import cygrpc
@@ -30,6 +32,34 @@ from ._utils import _timeout_to_deadline
 
 _IMMUTABLE_EMPTY_TUPLE = tuple()
 
+_LOGGER = logging.getLogger(__name__)
+
+
+class _OngoingCalls:
+    """Internal class used for have visibility of the ongoing calls."""
+
+    _calls: AbstractSet[_base_call.RpcContext]
+
+    def __init__(self):
+        self._calls = WeakSet()
+
+    def _remove_call(self, call: _base_call.RpcContext):
+        self._calls.remove(call)
+
+    @property
+    def calls(self) -> AbstractSet[_base_call.RpcContext]:
+        """Returns the set of ongoing calls."""
+        return self._calls
+
+    def size(self) -> int:
+        """Returns the number of ongoing calls."""
+        return len(self._calls)
+
+    def trace_call(self, call: _base_call.RpcContext):
+        """Adds and manages a new ongoing call."""
+        self._calls.add(call)
+        call.add_done_callback(self._remove_call)
+
 
 class _BaseMultiCallable:
     """Base class of all multi callable objects.
@@ -38,6 +68,7 @@ class _BaseMultiCallable:
     """
     _loop: asyncio.AbstractEventLoop
     _channel: cygrpc.AioChannel
+    _ongoing_calls: _OngoingCalls
     _method: bytes
     _request_serializer: SerializingFunction
     _response_deserializer: DeserializingFunction
@@ -49,9 +80,11 @@ class _BaseMultiCallable:
     _interceptors: Optional[Sequence[UnaryUnaryClientInterceptor]]
     _loop: asyncio.AbstractEventLoop
 
+    # pylint: disable=too-many-arguments
     def __init__(
             self,
             channel: cygrpc.AioChannel,
+            ongoing_calls: _OngoingCalls,
             method: bytes,
             request_serializer: SerializingFunction,
             response_deserializer: DeserializingFunction,
@@ -60,6 +93,7 @@ class _BaseMultiCallable:
     ) -> None:
         self._loop = loop
         self._channel = channel
+        self._ongoing_calls = ongoing_calls
         self._method = method
         self._request_serializer = request_serializer
         self._response_deserializer = response_deserializer
@@ -108,18 +142,21 @@ class UnaryUnaryMultiCallable(_BaseMultiCallable):
             metadata = _IMMUTABLE_EMPTY_TUPLE
 
         if not self._interceptors:
-            return UnaryUnaryCall(request, _timeout_to_deadline(timeout),
+            call = UnaryUnaryCall(request, _timeout_to_deadline(timeout),
                                   metadata, credentials, wait_for_ready,
                                   self._channel, self._method,
                                   self._request_serializer,
                                   self._response_deserializer, self._loop)
         else:
-            return InterceptedUnaryUnaryCall(
+            call = InterceptedUnaryUnaryCall(
                 self._interceptors, request, timeout, metadata, credentials,
                 wait_for_ready, self._channel, self._method,
                 self._request_serializer, self._response_deserializer,
                 self._loop)
 
+        self._ongoing_calls.trace_call(call)
+        return call
+
 
 class UnaryStreamMultiCallable(_BaseMultiCallable):
     """Affords invoking a unary-stream RPC from client-side in an asynchronous way."""
@@ -158,10 +195,12 @@ class UnaryStreamMultiCallable(_BaseMultiCallable):
         if metadata is None:
             metadata = _IMMUTABLE_EMPTY_TUPLE
 
-        return UnaryStreamCall(request, deadline, metadata, credentials,
+        call = UnaryStreamCall(request, deadline, metadata, credentials,
                                wait_for_ready, self._channel, self._method,
                                self._request_serializer,
                                self._response_deserializer, self._loop)
+        self._ongoing_calls.trace_call(call)
+        return call
 
 
 class StreamUnaryMultiCallable(_BaseMultiCallable):
@@ -205,10 +244,12 @@ class StreamUnaryMultiCallable(_BaseMultiCallable):
         if metadata is None:
             metadata = _IMMUTABLE_EMPTY_TUPLE
 
-        return StreamUnaryCall(request_async_iterator, deadline, metadata,
+        call = StreamUnaryCall(request_async_iterator, deadline, metadata,
                                credentials, wait_for_ready, self._channel,
                                self._method, self._request_serializer,
                                self._response_deserializer, self._loop)
+        self._ongoing_calls.trace_call(call)
+        return call
 
 
 class StreamStreamMultiCallable(_BaseMultiCallable):
@@ -252,10 +293,12 @@ class StreamStreamMultiCallable(_BaseMultiCallable):
         if metadata is None:
             metadata = _IMMUTABLE_EMPTY_TUPLE
 
-        return StreamStreamCall(request_async_iterator, deadline, metadata,
+        call = StreamStreamCall(request_async_iterator, deadline, metadata,
                                 credentials, wait_for_ready, self._channel,
                                 self._method, self._request_serializer,
                                 self._response_deserializer, self._loop)
+        self._ongoing_calls.trace_call(call)
+        return call
 
 
 class Channel:
@@ -266,6 +309,7 @@ class Channel:
     _loop: asyncio.AbstractEventLoop
     _channel: cygrpc.AioChannel
     _unary_unary_interceptors: Optional[Sequence[UnaryUnaryClientInterceptor]]
+    _ongoing_calls: _OngoingCalls
 
     def __init__(self, target: Text, options: Optional[ChannelArgumentType],
                  credentials: Optional[grpc.ChannelCredentials],
@@ -307,6 +351,62 @@ class Channel:
         self._loop = asyncio.get_event_loop()
         self._channel = cygrpc.AioChannel(_common.encode(target), options,
                                           credentials, self._loop)
+        self._ongoing_calls = _OngoingCalls()
+
+    async def __aenter__(self):
+        """Starts an asynchronous context manager.
+
+        Returns:
+          Channel the channel that was instantiated.
+        """
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Finishes the asynchronous context manager by closing the channel.
+
+        Still active RPCs will be cancelled.
+        """
+        await self._close(None)
+
+    async def _close(self, grace):
+        if self._channel.closed():
+            return
+
+        # No new calls will be accepted by the Cython channel.
+        self._channel.closing()
+
+        if grace:
+            # pylint: disable=unused-variable
+            _, pending = await asyncio.wait(self._ongoing_calls.calls,
+                                            timeout=grace,
+                                            loop=self._loop)
+
+            if not pending:
+                return
+
+        # A new set is created acting as a shallow copy because
+        # when cancellation happens the calls are automatically
+        # removed from the originally set.
+        calls = WeakSet(data=self._ongoing_calls.calls)
+        for call in calls:
+            call.cancel()
+
+        self._channel.close()
+
+    async def close(self, grace: Optional[float] = None):
+        """Closes this Channel and releases all resources held by it.
+
+        This method immediately stops the channel from executing new RPCs in
+        all cases.
+
+        If a grace period is specified, this method wait until all active
+        RPCs are finshed, once the grace period is reached the ones that haven't
+        been terminated are cancelled. If a grace period is not specified
+        (by passing None for grace), all existing RPCs are cancelled immediately.
+
+        This method is idempotent.
+        """
+        await self._close(grace)
 
     def get_state(self,
                   try_to_connect: bool = False) -> grpc.ChannelConnectivity:
@@ -372,7 +472,8 @@ class Channel:
         Returns:
           A UnaryUnaryMultiCallable value for the named unary-unary method.
         """
-        return UnaryUnaryMultiCallable(self._channel, _common.encode(method),
+        return UnaryUnaryMultiCallable(self._channel, self._ongoing_calls,
+                                       _common.encode(method),
                                        request_serializer,
                                        response_deserializer,
                                        self._unary_unary_interceptors,
@@ -384,7 +485,8 @@ class Channel:
             request_serializer: Optional[SerializingFunction] = None,
             response_deserializer: Optional[DeserializingFunction] = None
     ) -> UnaryStreamMultiCallable:
-        return UnaryStreamMultiCallable(self._channel, _common.encode(method),
+        return UnaryStreamMultiCallable(self._channel, self._ongoing_calls,
+                                        _common.encode(method),
                                         request_serializer,
                                         response_deserializer, None, self._loop)
 
@@ -394,7 +496,8 @@ class Channel:
             request_serializer: Optional[SerializingFunction] = None,
             response_deserializer: Optional[DeserializingFunction] = None
     ) -> StreamUnaryMultiCallable:
-        return StreamUnaryMultiCallable(self._channel, _common.encode(method),
+        return StreamUnaryMultiCallable(self._channel, self._ongoing_calls,
+                                        _common.encode(method),
                                         request_serializer,
                                         response_deserializer, None, self._loop)
 
@@ -404,33 +507,8 @@ class Channel:
             request_serializer: Optional[SerializingFunction] = None,
             response_deserializer: Optional[DeserializingFunction] = None
     ) -> StreamStreamMultiCallable:
-        return StreamStreamMultiCallable(self._channel, _common.encode(method),
+        return StreamStreamMultiCallable(self._channel, self._ongoing_calls,
+                                         _common.encode(method),
                                          request_serializer,
                                          response_deserializer, None,
                                          self._loop)
-
-    async def _close(self):
-        # TODO: Send cancellation status
-        self._channel.close()
-
-    async def __aenter__(self):
-        """Starts an asynchronous context manager.
-
-        Returns:
-          Channel the channel that was instantiated.
-        """
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        """Finishes the asynchronous context manager by closing gracefully the channel."""
-        await self._close()
-
-    async def close(self):
-        """Closes this Channel and releases all resources held by it.
-
-        Closing the Channel will proactively terminate all RPCs active with the
-        Channel and it is not valid to invoke new RPCs with the Channel.
-
-        This method is idempotent.
-        """
-        await self._close()

src/python/grpcio/grpc/experimental/aio/_interceptor.py

@@ -25,7 +25,7 @@ from . import _base_call
 from ._call import UnaryUnaryCall, AioRpcError
 from ._utils import _timeout_to_deadline
 from ._typing import (RequestType, SerializingFunction, DeserializingFunction,
-                      MetadataType, ResponseType)
+                      MetadataType, ResponseType, DoneCallbackType)
 
 _LOCAL_CANCELLATION_DETAILS = 'Locally cancelled by application!'
 
@@ -103,6 +103,7 @@ class InterceptedUnaryUnaryCall(_base_call.UnaryUnaryCall):
     _intercepted_call: Optional[_base_call.UnaryUnaryCall]
     _intercepted_call_created: asyncio.Event
     _interceptors_task: asyncio.Task
+    _pending_add_done_callbacks: Sequence[DoneCallbackType]
 
     # pylint: disable=too-many-arguments
     def __init__(self, interceptors: Sequence[UnaryUnaryClientInterceptor],
@@ -119,6 +120,9 @@ class InterceptedUnaryUnaryCall(_base_call.UnaryUnaryCall):
             interceptors, method, timeout, metadata, credentials,
             wait_for_ready, request, request_serializer, response_deserializer),
                                                         loop=loop)
+        self._pending_add_done_callbacks = []
+        self._interceptors_task.add_done_callback(
+            self._fire_pending_add_done_callbacks)
 
     def __del__(self):
         self.cancel()
@@ -166,6 +170,17 @@ class InterceptedUnaryUnaryCall(_base_call.UnaryUnaryCall):
         return await _run_interceptor(iter(interceptors), client_call_details,
                                       request)
 
+    def _fire_pending_add_done_callbacks(self,
+                                         unused_task: asyncio.Task) -> None:
+        for callback in self._pending_add_done_callbacks:
+            callback(self)
+
+        self._pending_add_done_callbacks = []
+
+    def _wrap_add_done_callback(self, callback: DoneCallbackType,
+                                unused_task: asyncio.Task) -> None:
+        callback(self)
+
     def cancel(self) -> bool:
         if self._interceptors_task.done():
             return False
@@ -196,8 +211,22 @@ class InterceptedUnaryUnaryCall(_base_call.UnaryUnaryCall):
 
         return call.done()
 
-    def add_done_callback(self, unused_callback) -> None:
-        raise NotImplementedError()
+    def add_done_callback(self, callback: DoneCallbackType) -> None:
+        if not self._interceptors_task.done():
+            self._pending_add_done_callbacks.append(callback)
+            return
+
+        try:
+            call = self._interceptors_task.result()
+        except (AioRpcError, asyncio.CancelledError):
+            callback(self)
+            return
+
+        if call.done():
+            callback(self)
+        else:
+            callback = functools.partial(self._wrap_add_done_callback, callback)
+            call.add_done_callback(self._wrap_add_done_callback)
 
     def time_remaining(self) -> Optional[float]:
         raise NotImplementedError()

src/python/grpcio_tests/tests_aio/tests.json

@@ -10,6 +10,8 @@
   "unit.call_test.TestUnaryUnaryCall",
   "unit.channel_argument_test.TestChannelArgument",
   "unit.channel_test.TestChannel",
+  "unit.close_channel_test.TestCloseChannel",
+  "unit.close_channel_test.TestOngoingCalls",
   "unit.connectivity_test.TestConnectivityState",
   "unit.done_callback_test.TestDoneCallback",
   "unit.init_test.TestInsecureChannel",

src/python/grpcio_tests/tests_aio/unit/channel_test.py

@@ -15,7 +15,6 @@
 
 import logging
 import os
-import threading
 import unittest
 
 import grpc
@@ -227,5 +226,5 @@ class TestChannel(AioTestBase):
 
 
 if __name__ == '__main__':
-    logging.basicConfig(level=logging.DEBUG)
+    logging.basicConfig(level=logging.INFO)
     unittest.main(verbosity=2)

src/python/grpcio_tests/tests_aio/unit/close_channel_test.py

@@ -0,0 +1,186 @@
+# Copyright 2020 The gRPC Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests behavior of closing a grpc.aio.Channel."""
+
+import asyncio
+import logging
+import unittest
+from weakref import WeakSet
+
+import grpc
+from grpc.experimental import aio
+from grpc.experimental.aio import _base_call
+from grpc.experimental.aio._channel import _OngoingCalls
+
+from src.proto.grpc.testing import messages_pb2, test_pb2_grpc
+from tests_aio.unit._constants import UNARY_CALL_WITH_SLEEP_VALUE
+from tests_aio.unit._test_base import AioTestBase
+from tests_aio.unit._test_server import start_test_server
+
+_UNARY_CALL_METHOD_WITH_SLEEP = '/grpc.testing.TestService/UnaryCallWithSleep'
+
+
+class TestOngoingCalls(unittest.TestCase):
+
+    class FakeCall(_base_call.RpcContext):
+
+        def add_done_callback(self, callback):
+            self.callback = callback
+
+        def cancel(self):
+            raise NotImplementedError
+
+        def cancelled(self):
+            raise NotImplementedError
+
+        def done(self):
+            raise NotImplementedError
+
+        def time_remaining(self):
+            raise NotImplementedError
+
+    def test_trace_call(self):
+        ongoing_calls = _OngoingCalls()
+        self.assertEqual(ongoing_calls.size(), 0)
+
+        call = TestOngoingCalls.FakeCall()
+        ongoing_calls.trace_call(call)
+        self.assertEqual(ongoing_calls.size(), 1)
+        self.assertEqual(ongoing_calls.calls, WeakSet([call]))
+
+        call.callback(call)
+        self.assertEqual(ongoing_calls.size(), 0)
+        self.assertEqual(ongoing_calls.calls, WeakSet())
+
+    def test_deleted_call(self):
+        ongoing_calls = _OngoingCalls()
+
+        call = TestOngoingCalls.FakeCall()
+        ongoing_calls.trace_call(call)
+        del (call)
+        self.assertEqual(ongoing_calls.size(), 0)
+
+
+class TestCloseChannel(AioTestBase):
+
+    async def setUp(self):
+        self._server_target, self._server = await start_test_server()
+
+    async def tearDown(self):
+        await self._server.stop(None)
+
+    async def test_graceful_close(self):
+        channel = aio.insecure_channel(self._server_target)
+        UnaryCallWithSleep = channel.unary_unary(
+            _UNARY_CALL_METHOD_WITH_SLEEP,
+            request_serializer=messages_pb2.SimpleRequest.SerializeToString,
+            response_deserializer=messages_pb2.SimpleResponse.FromString,
+        )
+
+        call = UnaryCallWithSleep(messages_pb2.SimpleRequest())
+
+        await channel.close(grace=UNARY_CALL_WITH_SLEEP_VALUE * 4)
+
+        self.assertEqual(grpc.StatusCode.OK, await call.code())
+
+    async def test_none_graceful_close(self):
+        channel = aio.insecure_channel(self._server_target)
+        UnaryCallWithSleep = channel.unary_unary(
+            _UNARY_CALL_METHOD_WITH_SLEEP,
+            request_serializer=messages_pb2.SimpleRequest.SerializeToString,
+            response_deserializer=messages_pb2.SimpleResponse.FromString,
+        )
+
+        call = UnaryCallWithSleep(messages_pb2.SimpleRequest())
+
+        await channel.close(None)
+
+        self.assertEqual(grpc.StatusCode.CANCELLED, await call.code())
+
+    async def test_close_unary_unary(self):
+        channel = aio.insecure_channel(self._server_target)
+        stub = test_pb2_grpc.TestServiceStub(channel)
+
+        calls = [stub.UnaryCall(messages_pb2.SimpleRequest()) for _ in range(2)]
+
+        self.assertEqual(channel._ongoing_calls.size(), 2)
+
+        await channel.close()
+
+        for call in calls:
+            self.assertTrue(call.cancelled())
+
+        self.assertEqual(channel._ongoing_calls.size(), 0)
+
+    async def test_close_unary_stream(self):
+        channel = aio.insecure_channel(self._server_target)
+        stub = test_pb2_grpc.TestServiceStub(channel)
+
+        request = messages_pb2.StreamingOutputCallRequest()
+        calls = [stub.StreamingOutputCall(request) for _ in range(2)]
+
+        self.assertEqual(channel._ongoing_calls.size(), 2)
+
+        await channel.close()
+
+        for call in calls:
+            self.assertTrue(call.cancelled())
+
+        self.assertEqual(channel._ongoing_calls.size(), 0)
+
+    async def test_close_stream_unary(self):
+        channel = aio.insecure_channel(self._server_target)
+        stub = test_pb2_grpc.TestServiceStub(channel)
+
+        calls = [stub.StreamingInputCall() for _ in range(2)]
+
+        await channel.close()
+
+        for call in calls:
+            self.assertTrue(call.cancelled())
+
+        self.assertEqual(channel._ongoing_calls.size(), 0)
+
+    async def test_close_stream_stream(self):
+        channel = aio.insecure_channel(self._server_target)
+        stub = test_pb2_grpc.TestServiceStub(channel)
+
+        calls = [stub.FullDuplexCall() for _ in range(2)]
+
+        self.assertEqual(channel._ongoing_calls.size(), 2)
+
+        await channel.close()
+
+        for call in calls:
+            self.assertTrue(call.cancelled())
+
+        self.assertEqual(channel._ongoing_calls.size(), 0)
+
+    async def test_close_async_context(self):
+        async with aio.insecure_channel(self._server_target) as channel:
+            stub = test_pb2_grpc.TestServiceStub(channel)
+            calls = [
+                stub.UnaryCall(messages_pb2.SimpleRequest()) for _ in range(2)
+            ]
+            self.assertEqual(channel._ongoing_calls.size(), 2)
+
+        for call in calls:
+            self.assertTrue(call.cancelled())
+
+        self.assertEqual(channel._ongoing_calls.size(), 0)
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+    unittest.main(verbosity=2)

src/python/grpcio_tests/tests_aio/unit/interceptor_test.py

@@ -29,6 +29,7 @@ _INITIAL_METADATA_TO_INJECT = (
     (_INITIAL_METADATA_KEY, 'extra info'),
     (_TRAILING_METADATA_KEY, b'\x13\x37'),
 )
+_TIMEOUT_CHECK_IF_CALLBACK_WAS_CALLED = 1.0
 
 
 class TestUnaryUnaryClientInterceptor(AioTestBase):
@@ -577,6 +578,112 @@ class TestInterceptedUnaryUnaryCall(AioTestBase):
 
             self.assertEqual(await call.code(), grpc.StatusCode.OK)
 
+    async def test_add_done_callback_before_finishes(self):
+        called = asyncio.Event()
+        interceptor_can_continue = asyncio.Event()
+
+        def callback(call):
+            called.set()
+
+        class Interceptor(aio.UnaryUnaryClientInterceptor):
+
+            async def intercept_unary_unary(self, continuation,
+                                            client_call_details, request):
+
+                await interceptor_can_continue.wait()
+                call = await continuation(client_call_details, request)
+                return call
+
+        async with aio.insecure_channel(self._server_target,
+                                        interceptors=[Interceptor()
+                                                     ]) as channel:
+
+            multicallable = channel.unary_unary(
+                '/grpc.testing.TestService/UnaryCall',
+                request_serializer=messages_pb2.SimpleRequest.SerializeToString,
+                response_deserializer=messages_pb2.SimpleResponse.FromString)
+            call = multicallable(messages_pb2.SimpleRequest())
+            call.add_done_callback(callback)
+            interceptor_can_continue.set()
+            await call
+
+            try:
+                await asyncio.wait_for(
+                    called.wait(),
+                    timeout=_TIMEOUT_CHECK_IF_CALLBACK_WAS_CALLED)
+            except:
+                self.fail("Callback was not called")
+
+    async def test_add_done_callback_after_finishes(self):
+        called = asyncio.Event()
+
+        def callback(call):
+            called.set()
+
+        class Interceptor(aio.UnaryUnaryClientInterceptor):
+
+            async def intercept_unary_unary(self, continuation,
+                                            client_call_details, request):
+
+                call = await continuation(client_call_details, request)
+                return call
+
+        async with aio.insecure_channel(self._server_target,
+                                        interceptors=[Interceptor()
+                                                     ]) as channel:
+
+            multicallable = channel.unary_unary(
+                '/grpc.testing.TestService/UnaryCall',
+                request_serializer=messages_pb2.SimpleRequest.SerializeToString,
+                response_deserializer=messages_pb2.SimpleResponse.FromString)
+            call = multicallable(messages_pb2.SimpleRequest())
+
+            await call
+
+            call.add_done_callback(callback)
+
+            try:
+                await asyncio.wait_for(
+                    called.wait(),
+                    timeout=_TIMEOUT_CHECK_IF_CALLBACK_WAS_CALLED)
+            except:
+                self.fail("Callback was not called")
+
+    async def test_add_done_callback_after_finishes_before_await(self):
+        called = asyncio.Event()
+
+        def callback(call):
+            called.set()
+
+        class Interceptor(aio.UnaryUnaryClientInterceptor):
+
+            async def intercept_unary_unary(self, continuation,
+                                            client_call_details, request):
+
+                call = await continuation(client_call_details, request)
+                return call
+
+        async with aio.insecure_channel(self._server_target,
+                                        interceptors=[Interceptor()
+                                                     ]) as channel:
+
+            multicallable = channel.unary_unary(
+                '/grpc.testing.TestService/UnaryCall',
+                request_serializer=messages_pb2.SimpleRequest.SerializeToString,
+                response_deserializer=messages_pb2.SimpleResponse.FromString)
+            call = multicallable(messages_pb2.SimpleRequest())
+
+            call.add_done_callback(callback)
+
+            await call
+
+            try:
+                await asyncio.wait_for(
+                    called.wait(),
+                    timeout=_TIMEOUT_CHECK_IF_CALLBACK_WAS_CALLED)
+            except:
+                self.fail("Callback was not called")
+
 
 if __name__ == '__main__':
     logging.basicConfig()