Graceful termination

.github/workflows/ci.yaml

@@ -15,10 +15,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - name: checkout
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: setup python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: 3.9
 
@@ -33,7 +33,7 @@ jobs:
          TOXENV: static
 
   test:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-20.04
     strategy:
       matrix:
         python-version:
@@ -44,10 +44,10 @@ jobs:
 
     steps:
     - name: checkout
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: setup python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
 
@@ -74,10 +74,10 @@ jobs:
 
     steps:
     - name: checkout
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
 
     - name: setup python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: 3.9
 

.gitignore

@@ -3,4 +3,6 @@ __pycache__
 .tox
 *_pb2.py
 *_pb2_grpc.py
-.coverage
+.coverage
+.DS_Store
+.idea

nameko_grpc/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+from nameko_grpc.h2_patch import patch_h2_transitions
+
+
+patch_h2_transitions()

nameko_grpc/channel.py

@@ -111,14 +111,19 @@ class ServerConnectionPool:
     Just accepts new connections and allows them to run until close.
     """
 
-    def __init__(self, host, port, ssl, spawn_thread, handle_request):
+    def __init__(
+        self, host, port, ssl, spawn_thread, handle_request, max_concurrent_streams=100
+    ):
         self.host = host
         self.port = port
         self.ssl = ssl
         self.spawn_thread = spawn_thread
         self.handle_request = handle_request
+        self.max_concurrent_streams = max_concurrent_streams
 
         self.connections = queue.Queue()
+        self.is_accepting = False
+        self.listening_socket = None
 
     def listen(self):
         sock = eventlet.listen((self.host, self.port))
@@ -139,7 +144,9 @@ def run(self):
             sock, _ = self.listening_socket.accept()
             sock.settimeout(60)  # XXX needed and/or correct value?
 
-            connection = ServerConnectionManager(sock, self.handle_request)
+            connection = ServerConnectionManager(
+                sock, self.handle_request, self.max_concurrent_streams
+            )
             self.connections.put(weakref.ref(connection))
             self.spawn_thread(
                 target=connection.run_forever, name=f"grpc server connection [{sock}]"
@@ -165,9 +172,11 @@ def stop(self):
 class ServerChannel:
     """Simple server channel encapsulating incoming connection management."""
 
-    def __init__(self, host, port, ssl, spawn_thread, handle_request):
+    def __init__(
+        self, host, port, ssl, spawn_thread, handle_request, max_concurrent_streams=100
+    ):
         self.conn_pool = ServerConnectionPool(
-            host, port, ssl, spawn_thread, handle_request
+            host, port, ssl, spawn_thread, handle_request, max_concurrent_streams
         )
 
     def start(self):

nameko_grpc/connection.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 import itertools
+import logging
 import select
 import sys
 from collections import deque
@@ -8,7 +9,7 @@
 from threading import Event
 
 from grpc import StatusCode
-from h2.config import H2Configuration
+from h2.config import DummyLogger, H2Configuration
 from h2.connection import H2Connection
 from h2.errors import ErrorCodes
 from h2.events import (
@@ -37,9 +38,30 @@
 log = getLogger(__name__)
 
 
+class H2Logger(DummyLogger):
+    """
+    Provide logger to H2 matching required interface
+    """
+
+    def __init__(self, logger: logging.Logger):
+        super().__init__()
+        self.logger = logger
+
+    def debug(self, *vargs, **kwargs):
+        self.logger.debug(*vargs, **kwargs)
+
+    def trace(self, *vargs, **kwargs):
+        # log level below debug
+        self.logger.log(5, *vargs, **kwargs)
+
+
 SELECT_TIMEOUT = 0.01
 
 
+class ConnectionTerminatingError(Exception):
+    pass
+
+
 class ConnectionManager:
     """
     Base class for managing a single GRPC HTTP/2 connection.
@@ -49,21 +71,24 @@ class ConnectionManager:
     by subclasses.
     """
 
-    def __init__(self, sock, client_side):
+    def __init__(self, sock, client_side, max_concurrent_streams=100):
         self.sock = sock
 
-        config = H2Configuration(client_side=client_side)
+        h2_logger = H2Logger(log.getChild("h2"))
+        config = H2Configuration(client_side=client_side, logger=h2_logger)
         self.conn = H2Connection(config=config)
+        self.conn.local_settings.max_concurrent_streams = max_concurrent_streams
 
         self.receive_streams = {}
         self.send_streams = {}
 
         self.run = True
         self.stopped = Event()
+        self.terminating = False
 
     @property
     def alive(self):
-        return not self.stopped.is_set()
+        return not self.stopped.is_set() and not self.terminating
 
     @contextmanager
     def cleanup_on_exit(self):
@@ -83,23 +108,25 @@ def cleanup_on_exit(self):
                 if send_stream.closed:
                     continue  # stream.close() is idemponent but this prevents the log
                 log.info(
-                    f"Terminating send stream {send_stream}"
+                    f"Terminating send stream {send_stream.stream_id}"
                     f"{f' with error {error}' if error else ''}."
                 )
                 send_stream.close(error)
             for receive_stream in self.receive_streams.values():
                 if receive_stream.closed:
                     continue  # stream.close() is idemponent but this prevents the log
                 log.info(
-                    f"Terminating receive stream {receive_stream}"
+                    f"Terminating receive stream {receive_stream.stream_id}"
                     f"{f' with error {error}' if error else ''}."
                 )
                 receive_stream.close(error)
             self.sock.close()
             self.stopped.set()
+            log.debug(f"connection terminated {self}")
 
     def run_forever(self):
         """Event loop."""
+        log.debug(f"connection initiated {self}")
         self.conn.initiate_connection()
 
         with self.cleanup_on_exit():
@@ -108,6 +135,9 @@ def run_forever(self):
 
                 self.on_iteration()
 
+                if not self.run:
+                    break
+
                 self.sock.sendall(self.conn.data_to_send())
                 ready = select.select([self.sock], [], [], SELECT_TIMEOUT)
                 if not ready[0]:
@@ -142,8 +172,10 @@ def run_forever(self):
                         self.connection_terminated(event)
 
     def stop(self):
-        self.run = False
-        self.stopped.wait()
+        self.conn.close_connection()
+        self.terminating = True
+        log.debug("waiting for connection to terminate (Timeout 5s)")
+        self.stopped.wait(5)
 
     def on_iteration(self):
         """Called on every iteration of the event loop.
@@ -155,6 +187,16 @@ def on_iteration(self):
             self.send_headers(stream_id)
             self.send_data(stream_id)
 
+        if self.terminating:
+            send_streams_closed = all(
+                stream.exhausted for stream in self.send_streams.values()
+            )
+            receive_streams_closed = all(
+                stream.exhausted for stream in self.receive_streams.values()
+            )
+            if send_streams_closed and receive_streams_closed:
+                self.run = False
+
     def request_received(self, event):
         """Called when a request is received on a stream.
 
@@ -199,6 +241,7 @@ def window_updated(self, event):
         Any data waiting to be sent on the stream may fit in the window now.
         """
         log.debug("window updated, stream %s", event.stream_id)
+        self.send_headers(event.stream_id)
         self.send_data(event.stream_id)
 
     def stream_ended(self, event):
@@ -210,16 +253,22 @@ def stream_ended(self, event):
         receive_stream = self.receive_streams.pop(event.stream_id, None)
         if receive_stream:
             receive_stream.close()
+        # send_stream = self.send_streams.pop(event.stream_id, None)
+        # if send_stream:
+        #     send_stream.close()
 
     def stream_reset(self, event):
         """Called when an incoming stream is reset.
 
-        Close any `ReceiveStream` that was opened for this stream.
+        Close any `ReceiveStream` or `SendStream` that was opened for this stream.
         """
         log.debug("stream reset, stream %s", event.stream_id)
         receive_stream = self.receive_streams.pop(event.stream_id, None)
         if receive_stream:
             receive_stream.close()
+        send_stream = self.send_streams.pop(event.stream_id, None)
+        if send_stream:
+            send_stream.close()
 
     def settings_changed(self, event):
         log.debug("settings changed")
@@ -235,9 +284,22 @@ def trailers_received(self, event):
 
         receive_stream.trailers.set(*event.headers, from_wire=True)
 
-    def connection_terminated(self, event):
-        log.debug("connection terminated")
-        self.run = False
+    def connection_terminated(self, event: ConnectionTerminated):
+        """H2 signals a connection terminated event after receiving a GOAWAY frame
+
+        If no error has occurred, flag termination and initiate a graceful termination
+        allowing existing streams to finish sending/receiving.
+
+        If an error has occurred then close down immediately.
+        """
+        log.debug(f"received GOAWAY with error code {event.error_code}")
+        if event.error_code not in (ErrorCodes.NO_ERROR, ErrorCodes.ENHANCE_YOUR_CALM):
+            log.debug("connection terminating immediately")
+            self.terminating = True
+            self.run = False
+        else:
+            log.debug("connection terminating")
+            self.terminating = True
 
     def send_headers(self, stream_id, immediate=False):
         """Attempt to send any headers on a stream.
@@ -269,6 +331,14 @@ def send_data(self, stream_id):
             # has been completely sent
             return
 
+        # When a stream is closed, a STREAM_END item or ERROR is placed in the queue.
+        # If we never read from the stream again, these are not surfaced, and the
+        # stream is never exhausted.
+        # Because we shortcut sending data if headers haven't been set yet, we need
+        # to manually flush the queue, surfacing the end/error, and ensuring the
+        # queue exhausts (and we can terminate).
+        send_stream.flush_queue_to_buffer()
+
         if not send_stream.headers_sent:
             # don't attempt to send any data until the headers have been sent
             return
@@ -331,7 +401,19 @@ def send_request(self, request_headers):
         over the response.
 
         Invocations are queued and sent on the next iteration of the event loop.
+
+        raises ConnectionTerminatingError if connection is terminating. Check
+         connection .is_alive() before initiating send_request
+
+        Note:
+            We are handling termination and raising TerminatingError here as the
+            underlying library H2 doesn't do this. If H2 ever begins handling graceful
+            shutdowns, this logic will need altering.
         """
+        if self.terminating:
+            raise ConnectionTerminatingError(
+                "Connection is terminating. No new streams can be initiated"
+            )
         stream_id = next(self.counter)
 
         request_stream = SendStream(stream_id)
@@ -427,8 +509,10 @@ class ServerConnectionManager(ConnectionManager):
     Extends the base `ConnectionManager` to handle incoming GRPC requests.
     """
 
-    def __init__(self, sock, handle_request):
-        super().__init__(sock, client_side=False)
+    def __init__(self, sock, handle_request, max_concurrent_streams=100):
+        super().__init__(
+            sock, client_side=False, max_concurrent_streams=max_concurrent_streams
+        )
         self.handle_request = handle_request
 
     def request_received(self, event):