Add graceful period for closing the channel

Separated the tests for testing how a channel is closed into another
test file.

Fixed a bug introduced into the interceptors code in the previous
commit.

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

@@ -101,6 +101,9 @@ cdef class AioChannel:
         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,
@@ -110,7 +113,7 @@ cdef class AioChannel:
         Returns:
           The _AioCall object.
         """
-        if self._status in (AIO_CHANNEL_STATUS_CLOSING, 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

@@ -377,25 +377,34 @@ class Channel:
         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()
+        """Finishes the asynchronous context manager by closing the channel.
 
-    async def _wait_for_close_ongoing_calls(self):
-        sleep_iterations_sec = 0.001
+        Still active RPCs will be cancelled.
+        """
+        await self._close(None)
 
-        while self._ongoing_calls.size() > 0:
-            await asyncio.sleep(sleep_iterations_sec)
+    async def _close(self, grace):
+        if self._channel.closed():
+            return
 
-    async def _close(self):
         # No new calls will be accepted by the Cython channel.
         self._channel.closing()
 
+        if grace:
+            _, pending = await asyncio.wait(self._ongoing_calls.calls,
+                                            timeout=grace,
+                                            loop=self._loop)
+
+            if not pending:
+                return
+
         calls = self._ongoing_calls.calls
         for call in calls:
             call.cancel()
 
         try:
-            await asyncio.wait_for(self._wait_for_close_ongoing_calls(),
+            await asyncio.wait_for(asyncio.gather(*self._ongoing_calls.calls,
+                                                  loop=self._loop),
                                    _TIMEOUT_WAIT_FOR_CLOSE_ONGOING_CALLS_SEC,
                                    loop=self._loop)
         except asyncio.TimeoutError:
@@ -404,15 +413,20 @@ class Channel:
 
         self._channel.close()
 
-    async def close(self):
+    async def close(self, grace: Optional[float] = None):
         """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 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()
+        await self._close(grace)
 
     def get_state(self,
                   try_to_connect: bool = False) -> grpc.ChannelConnectivity:

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

@@ -201,7 +201,11 @@ class InterceptedUnaryUnaryCall(_base_call.UnaryUnaryCall):
         if not self._interceptors_task.done():
             return False
 
-        call = self._interceptors_task.result()
+        try:
+            call = self._interceptors_task.result()
+        except (AioRpcError, asyncio.CancelledError):
+            return True
+
         return call.done()
 
     def add_done_callback(self, callback: DoneCallbackType) -> None:
@@ -209,7 +213,11 @@ class InterceptedUnaryUnaryCall(_base_call.UnaryUnaryCall):
             self._pending_add_done_callbacks.append(callback)
             return
 
-        call = self._interceptors_task.result()
+        try:
+            call = self._interceptors_task.result()
+        except (AioRpcError, asyncio.CancelledError):
+            callback(self)
+            return
 
         if call.done():
             callback(self)

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

@@ -44,43 +44,6 @@ _REQUEST_PAYLOAD_SIZE = 7
 _RESPONSE_PAYLOAD_SIZE = 42
 
 
-class Test_OngoingCalls(unittest.TestCase):
-
-    def test_trace_call(self):
-
-        class FakeCall(_base_call.RpcContext):
-
-            def __init__(self):
-                self.callback = None
-
-            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
-
-        ongoing_calls = _OngoingCalls()
-        self.assertEqual(ongoing_calls.size(), 0)
-
-        call = FakeCall()
-        ongoing_calls.trace_call(call)
-        self.assertEqual(ongoing_calls.size(), 1)
-        self.assertEqual(ongoing_calls.calls, [call])
-
-        call.callback(call)
-        self.assertEqual(ongoing_calls.size(), 0)
-        self.assertEqual(ongoing_calls.calls, [])
-
-
 class TestChannel(AioTestBase):
 
     async def setUp(self):
@@ -264,65 +227,6 @@ class TestChannel(AioTestBase):
         self.assertEqual(grpc.StatusCode.OK, await call.code())
         await channel.close()
 
-    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_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)

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

@@ -0,0 +1,185 @@
+# Copyright 2019 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 os
+import threading
+import unittest
+
+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.unit.framework.common import test_constants
+from tests_aio.unit._constants import (UNARY_CALL_WITH_SLEEP_VALUE,
+                                       UNREACHABLE_TARGET)
+from tests_aio.unit._test_base import AioTestBase
+from tests_aio.unit._test_server import start_test_server
+
+_UNARY_CALL_METHOD = '/grpc.testing.TestService/UnaryCall'
+_UNARY_CALL_METHOD_WITH_SLEEP = '/grpc.testing.TestService/UnaryCallWithSleep'
+_STREAMING_OUTPUT_CALL_METHOD = '/grpc.testing.TestService/StreamingOutputCall'
+
+_INVOCATION_METADATA = (
+    ('initial-md-key', 'initial-md-value'),
+    ('trailing-md-key-bin', b'\x00\x02'),
+)
+
+_NUM_STREAM_RESPONSES = 5
+_REQUEST_PAYLOAD_SIZE = 7
+_RESPONSE_PAYLOAD_SIZE = 42
+
+
+class TestOngoingCalls(unittest.TestCase):
+
+    def test_trace_call(self):
+
+        class FakeCall(_base_call.RpcContext):
+
+            def __init__(self):
+                self.callback = None
+
+            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
+
+        ongoing_calls = _OngoingCalls()
+        self.assertEqual(ongoing_calls.size(), 0)
+
+        call = FakeCall()
+        ongoing_calls.trace_call(call)
+        self.assertEqual(ongoing_calls.size(), 1)
+        self.assertEqual(ongoing_calls.calls, [call])
+
+        call.callback(call)
+        self.assertEqual(ongoing_calls.size(), 0)
+        self.assertEqual(ongoing_calls.calls, [])
+
+
+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())
+        task = asyncio.ensure_future(call)
+
+        await channel.close(grace=UNARY_CALL_WITH_SLEEP_VALUE * 2)
+
+        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())
+        task = asyncio.ensure_future(call)
+
+        await channel.close(grace=UNARY_CALL_WITH_SLEEP_VALUE / 2)
+
+        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_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)