Refactored server close logic to gracefully exit without using GOAWAY frames

Author: vmagamedov

grpclib/client.py

@@ -62,7 +62,10 @@
 
 
 class Handler(AbstractHandler):
-    connection_lost = False
+    closing = False
+
+    def connection_made(self, connection: Any) -> None:
+        pass
 
     def accept(self, stream: Any, headers: Any, release_stream: Any) -> None:
         raise NotImplementedError('Client connection can not accept requests')
@@ -71,7 +74,7 @@ def cancel(self, stream: Any) -> None:
         pass
 
     def close(self) -> None:
-        self.connection_lost = True
+        self.closing = True
 
 
 class Stream(StreamIterator[_RecvType], Generic[_SendType, _RecvType]):
@@ -737,7 +740,7 @@ async def _create_connection(self) -> H2Protocol:
     @property
     def _connected(self) -> bool:
         return (self._protocol is not None
-                and not self._protocol.handler.connection_lost)
+                and not cast(Handler, self._protocol.handler).closing)
 
     async def __connect__(self) -> H2Protocol:
         if not self._connected:

grpclib/protocol.py

@@ -488,6 +488,10 @@ def closable(self) -> bool:
 
 class AbstractHandler(ABC):
 
+    @abstractmethod
+    def connection_made(self, connection: Connection) -> None:
+        pass
+
     @abstractmethod
     def accept(
         self,
@@ -709,6 +713,7 @@ def connection_made(self, transport: BaseTransport) -> None:
         self.connection.flush()
         self.connection.initialize()
 
+        self.handler.connection_made(self.connection)
         self.processor = EventsProcessor(self.handler, self.connection)
 
     def data_received(self, data: bytes) -> None:

grpclib/server.py

@@ -4,6 +4,7 @@
 import logging
 import asyncio
 import warnings
+from functools import partial
 
 from types import TracebackType
 from typing import TYPE_CHECKING, Optional, Collection, Generic, Type, cast
@@ -12,6 +13,7 @@
 
 import h2.config
 import h2.exceptions
+from h2.errors import ErrorCodes
 
 from multidict import MultiDict
 
@@ -24,7 +26,7 @@
 from .metadata import Deadline, encode_grpc_message, _Metadata
 from .metadata import encode_metadata, decode_metadata, _MetadataLike
 from .metadata import _STATUS_DETAILS_KEY, encode_bin_value
-from .protocol import H2Protocol, AbstractHandler
+from .protocol import H2Protocol, AbstractHandler, Connection
 from .exceptions import GRPCError, ProtocolError, StreamTerminatedError
 from .encoding.base import GRPC_CONTENT_TYPE, CodecBase, StatusDetailsCodecBase
 from .encoding.proto import ProtoCodec, ProtoStatusDetailsCodec
@@ -496,6 +498,7 @@ def __gc_step__(self) -> None:
 class Handler(_GC, AbstractHandler):
     __gc_interval__ = 10
 
+    connection: Connection
     closing = False
 
     def __init__(
@@ -511,44 +514,54 @@ def __init__(
         self.dispatch = dispatch
         self.loop = asyncio.get_event_loop()
         self._tasks: Dict['protocol.Stream', 'asyncio.Task[None]'] = {}
-        self._cancelled: Set['asyncio.Task[None]'] = set()
 
     def __gc_collect__(self) -> None:
-        self._tasks = {s: t for s, t in self._tasks.items()
-                       if not t.done()}
-        self._cancelled = {t for t in self._cancelled
-                           if not t.done()}
+        self._tasks = {s: t for s, t in self._tasks.items() if not t.done()}
+
+    def connection_made(self, connection: Connection) -> None:
+        self.connection = connection
+
+    def handler_done(self, stream: 'protocol.Stream', _: Any) -> None:
+        self._tasks.pop(stream, None)
+        if not self._tasks:
+            self.connection.close()
 
     def accept(
         self,
         stream: 'protocol.Stream',
         headers: _Headers,
         release_stream: Callable[[], Any],
     ) -> None:
-        self.__gc_step__()
-        self._tasks[stream] = self.loop.create_task(request_handler(
-            self.mapping, stream, headers, self.codec,
-            self.status_details_codec, self.dispatch, release_stream,
-        ))
+        if self.closing:
+            stream.reset_nowait(ErrorCodes.REFUSED_STREAM)
+            release_stream()
+        else:
+            self.__gc_step__()
+            self._tasks[stream] = self.loop.create_task(request_handler(
+                self.mapping, stream, headers, self.codec,
+                self.status_details_codec, self.dispatch, release_stream,
+            ))
 
     def cancel(self, stream: 'protocol.Stream') -> None:
-        task = self._tasks.pop(stream)
-        task.cancel()
-        self._cancelled.add(task)
+        self._tasks[stream].cancel()
 
     def close(self) -> None:
-        for task in self._tasks.values():
+        self.__gc_collect__()
+        for stream, task in self._tasks.items():
+            task.add_done_callback(partial(self.handler_done, stream))
             task.cancel()
-        self._cancelled.update(self._tasks.values())
         self.closing = True
 
     async def wait_closed(self) -> None:
-        if self._cancelled:
-            await asyncio.wait(self._cancelled)
+        self.__gc_collect__()
+        if self._tasks:
+            await asyncio.wait(self._tasks.values())
+        else:
+            self.connection.close()
 
     def check_closed(self) -> bool:
         self.__gc_collect__()
-        return not self._tasks and not self._cancelled
+        return not self._tasks
 
 
 class Server(_GC):
@@ -737,11 +750,11 @@ async def wait_closed(self) -> None:
         if self._server is None or self._server_closed_fut is None:
             raise RuntimeError('Server is not started')
         await self._server_closed_fut
-        await self._server.wait_closed()
         if self._handlers:
             await asyncio.wait({
                 self._loop.create_task(h.wait_closed()) for h in self._handlers
             })
+        await self._server.wait_closed()
 
     async def __aenter__(self) -> 'Server':
         return self

tests/stubs.py

@@ -47,6 +47,9 @@ class DummyHandler(AbstractHandler):
     headers = None
     release_stream = None
 
+    def connection_made(self, connection):
+        pass
+
     def accept(self, stream, headers, release_stream):
         self.stream = stream
         self.headers = headers