Merge pull request #110 from codefly-dev/fix/109-concurrency

fix(connect): a call past the stream limit is refused alone; the scaffold advertises the limit

Author: antoinetoussaint-byte

connect_test.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"encoding/binary"
 	"fmt"
 	"io"
 	"net"
@@ -202,6 +203,8 @@ func TestGeneratedMainServesConnect(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, http.StatusOK, response.StatusCode, "%s", body)
 	require.JSONEq(t, `{"message":"over connect"}`, string(body))
+	require.Equal(t, uint32(1000), advertisedMaxStreams(t, app.grpc),
+		"the scaffold's gRPC server advertises the limit the face enforces, so a burst queues in the client")
 	app.stop(t)
 	require.False(t, listens(app.connect), "the face outlived the app")
 
@@ -316,6 +319,35 @@ func (o *lockedOutput) String() string {
 	return o.buf.String()
 }
 
+// advertisedMaxStreams reads SETTINGS_MAX_CONCURRENT_STREAMS from the first
+// SETTINGS frame the gRPC server on port sends; 0 when it advertises none.
+func advertisedMaxStreams(t *testing.T, port uint16) uint32 {
+	t.Helper()
+	conn, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", port), 2*time.Second)
+	require.NoError(t, err)
+	defer conn.Close()
+	require.NoError(t, conn.SetDeadline(time.Now().Add(5*time.Second)))
+	_, err = conn.Write(append([]byte("PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"), 0, 0, 0, 4, 0, 0, 0, 0, 0))
+	require.NoError(t, err)
+	for {
+		header := make([]byte, 9)
+		_, err = io.ReadFull(conn, header)
+		require.NoError(t, err)
+		payload := make([]byte, int(header[0])<<16|int(header[1])<<8|int(header[2]))
+		_, err = io.ReadFull(conn, payload)
+		require.NoError(t, err)
+		if header[3] != 4 || header[4]&1 != 0 { // not a SETTINGS, or an ACK
+			continue
+		}
+		for at := 0; at+6 <= len(payload); at += 6 {
+			if binary.BigEndian.Uint16(payload[at:]) == 3 {
+				return binary.BigEndian.Uint32(payload[at+2:])
+			}
+		}
+		return 0
+	}
+}
+
 func listens(port uint16) bool {
 	conn, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", port), time.Second)
 	if err != nil {

templates/connect/connect_gen.py

@@ -63,13 +63,28 @@
 _PREFACE = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"
 # An HTTP/2 frame header: 24-bit length, type, flags, stream id (RFC 9113 section 4.1).
 _FRAME = struct.Struct(">BHBBI")
-_DATA, _HEADERS, _RST_STREAM, _SETTINGS, _GOAWAY = 0x0, 0x1, 0x3, 0x4, 0x7
-_END_STREAM_FLAG, _ACK_FLAG = 0x1, 0x1
+_DATA, _HEADERS, _RST_STREAM, _SETTINGS, _GOAWAY, _CONTINUATION = 0x0, 0x1, 0x3, 0x4, 0x7, 0x9
+_END_STREAM_FLAG, _ACK_FLAG, _END_HEADERS_FLAG = 0x1, 0x1, 0x4
 _MAX_CONCURRENT_STREAMS = 0x3
+_PUSH_PROMISE = 0x5
+_PROTOCOL_ERROR, _CANCEL, _ENHANCE_YOUR_CALM = 0x1, 0x8, 0xB
 # Calls one relayed connection may hold open when the gRPC server advertises no
-# SETTINGS_MAX_CONCURRENT_STREAMS (grpcio advertises none); past it the face closes the
-# connection, so what it tracks per connection is bounded.
+# SETTINGS_MAX_CONCURRENT_STREAMS (grpcio advertises none unless the service sets
+# grpc.max_concurrent_streams, as the scaffold's server does). Past it the face resets
+# the call alone: ENHANCE_YOUR_CALM to the client (gRPC's RESOURCE_EXHAUSTED, which a
+# client does not retry on its own), and CANCEL to the server. The server still receives
+# the call's header block, which keeps the connection's header compression in step, so it
+# may have started the call before the cancel: the face never promises it did not. What
+# the face tracks per connection is bounded by it. When the server advertises a limit the
+# face enforces nothing: the server refuses what is past it (REFUSED_STREAM), truthfully.
 DEFAULT_MAX_OPEN_CALLS = 1000
+# Bytes of frames for the client (refusals, a GOAWAY) the relay holds while the server's
+# frame in flight completes; past it the relay reads nothing more from the client until
+# they are written, so a client that opens calls faster than it reads them stalls itself.
+_MAX_QUEUED_FOR_CLIENT = 16 * 1024
+# How long a closing relay waits for the server's frame in flight to complete, so the
+# client's GOAWAY lands between two frames.
+_GOAWAY_GRACE_S = 5.0
 # The most of a SETTINGS or GOAWAY payload the relay reads; the rest is skipped.
 _CONTROL_PAYLOAD = 256
 _ENVELOPE = struct.Struct(">BI")
@@ -224,55 +239,107 @@ class _ProtocolError(Exception):
     """A client frame no HTTP/2 client may send: the relay closes the connection."""
 
 
+def _h2_frame(kind: int, flags: int, stream: int, payload: bytes = b"") -> bytes:
+    size = len(payload)
+    return _FRAME.pack(size >> 16, size & 0xFFFF, kind, flags, stream) + payload
+
+
+def _rst_stream(stream: int, code: int) -> bytes:
+    return _h2_frame(_RST_STREAM, 0, stream, struct.pack(">I", code))
+
+
+def _goaway(last_stream: int, code: int, reason: str) -> bytes:
+    return _h2_frame(_GOAWAY, 0, 0, struct.pack(">II", last_stream, code) + reason.encode())
+
+
 class _Frames:
-    """The HTTP/2 frames in one direction of a relayed connection, after the preface:
-    ``on_frame(type, flags, stream, payload)`` for each. The payload is read for SETTINGS
-    and GOAWAY only (at most ``_CONTROL_PAYLOAD`` bytes of it) and skipped otherwise."""
+    """The HTTP/2 frames in one direction of a relayed connection, after the preface.
+
+    ``feed`` calls ``on_frame(type, flags, stream, payload)`` for each frame and returns
+    where frames end in the data, each with the bytes ``on_frame`` returned, to be written
+    right after that frame in this direction. The payload is read for SETTINGS and GOAWAY
+    only (at most ``_CONTROL_PAYLOAD`` bytes of it) and skipped otherwise. A
+    ``_ProtocolError`` from ``on_frame`` stops the feed; it is kept in ``error`` and the
+    ends before the offending frame are returned."""
 
     def __init__(self, on_frame):
         self._on_frame = on_frame
         self._header = bytearray()
         self._frame = None  # (type, flags, stream) while its payload is read
+        self._after = b""  # what on_frame returned for it
         self._payload = bytearray()
         self._keep = 0  # payload bytes still to read
         self._skip = 0  # payload bytes still to skip
-
-    def feed(self, data: bytes) -> None:
-        at = 0
-        while at < len(data):
-            if self._keep:
-                step = min(self._keep, len(data) - at)
-                self._payload += data[at : at + step]
-                self._keep -= step
-                at += step
-                if not self._keep:
-                    self._emit()
-                continue
-            if self._skip:
-                step = min(self._skip, len(data) - at)
-                self._skip -= step
-                at += step
-                continue
-            step = min(_FRAME.size - len(self._header), len(data) - at)
-            self._header += data[at : at + step]
-            at += step
-            if len(self._header) == _FRAME.size:
-                high, low, kind, flags, stream = _FRAME.unpack(self._header)
-                self._header.clear()
-                length = high << 16 | low
-                self._frame = (kind, flags, stream & 0x7FFFFFFF)
-                if kind in (_SETTINGS, _GOAWAY):
-                    self._keep = min(length, _CONTROL_PAYLOAD)
-                    self._skip = length - self._keep
-                else:
-                    self._skip = length
-                if not self._keep:
+        self._in_block = False  # inside a header block: HEADERS ... CONTINUATION
+        self.error: _ProtocolError | None = None
+
+    @property
+    def at_boundary(self) -> bool:
+        """Between two frames, and outside a header block: where another frame may go
+        (RFC 9113 section 6.10: nothing comes between a HEADERS and its CONTINUATION)."""
+        return not self._header and not self._keep and not self._skip and not self._in_block
+
+    def feed(self, data: bytes, every_end: bool = False) -> list[tuple[int, bytes]]:
+        """Where frames end in ``data``: each end where ``on_frame`` returned bytes to
+        write after the frame, and, with ``every_end``, each end where another frame may
+        go. After a ``_ProtocolError``, the last end is the last place another frame
+        may go before the offending frame."""
+        ends: list[tuple[int, bytes]] = []
+        at, size = 0, len(data)
+        clear = 0 if self.at_boundary else -1  # the last such place in ``data``
+        try:
+            while at < size:
+                if self._keep:
+                    step = min(self._keep, size - at)
+                    self._payload += data[at : at + step]
+                    self._keep -= step
+                    at += step
+                    if self._keep:
+                        break
                     self._emit()
+                elif self._skip:
+                    step = min(self._skip, size - at)
+                    self._skip -= step
+                    at += step
+                    if self._skip:
+                        break
+                elif self._header or size - at < _FRAME.size:
+                    step = min(_FRAME.size - len(self._header), size - at)
+                    self._header += data[at : at + step]
+                    at += step
+                    if len(self._header) < _FRAME.size:
+                        break
+                    self._start(*_FRAME.unpack(self._header))
+                    self._header.clear()
+                else:  # a whole frame header in this chunk: read it in place
+                    self._start(*_FRAME.unpack_from(data, at))
+                    at += _FRAME.size
+                if not (self._keep or self._skip or self._in_block):
+                    clear = at
+                    if every_end or self._after:
+                        ends.append((at, self._after))
+                        self._after = b""
+        except _ProtocolError as refusal:
+            self.error = refusal
+            if clear > (ends[-1][0] if ends else 0):
+                ends.append((clear, b""))
+        return ends
+
+    def _start(self, high, low, kind, flags, stream):
+        """A frame header read: what of its payload to keep, and its frame if none."""
+        length = high << 16 | low
+        self._frame = (kind, flags, stream & 0x7FFFFFFF)
+        if kind in (_HEADERS, _PUSH_PROMISE, _CONTINUATION):
+            self._in_block = not flags & _END_HEADERS_FLAG
+        self._keep = min(length, _CONTROL_PAYLOAD) if kind in (_SETTINGS, _GOAWAY) else 0
+        self._skip = length - self._keep
+        if not self._keep:
+            self._emit()
 
     def _emit(self):
         payload = bytes(self._payload)
         self._payload.clear()
-        self._on_frame(*self._frame, payload)
+        self._after = self._on_frame(*self._frame, payload) or b""
 
 
 class _Connection:
@@ -299,60 +366,133 @@ async def _relay(self, reader, writer, head):
 
         The frame headers say which calls are open: a client HEADERS on a new stream (odd,
         above every stream before it) opens one, and the server's END_STREAM, either
-        side's RST_STREAM, or the server's GOAWAY above the stream ends it. A client that
-        opens an even stream, or more calls than the server allows at once
-        (SETTINGS_MAX_CONCURRENT_STREAMS, ``DEFAULT_MAX_OPEN_CALLS`` when it advertises
-        none), is closed: no HTTP/2 client does either."""
+        side's RST_STREAM, or the server's GOAWAY above the stream ends it. A call past
+        the server's own SETTINGS_MAX_CONCURRENT_STREAMS is relayed as it came, for the
+        server to refuse. When the server advertises none, a call past
+        ``DEFAULT_MAX_OPEN_CALLS`` is reset alone (see there). The face's frames for the
+        client are written only between two of the server's frames and outside a header
+        block; while more than ``_MAX_QUEUED_FOR_CLIENT`` bytes of them wait, the client
+        is not read. A client that opens an even stream, which no HTTP/2 client does, gets
+        a GOAWAY naming why, and the connection is closed."""
         upstream_reader, upstream_writer = await asyncio.open_connection(*self.grpc_address)
         upstream_writer.write(head)
         loop = asyncio.get_running_loop()
         last = loop.time()
         calls: set[int] = set()
-        highest = 0
+        highest = accepted = 0
+        refusing = 0  # the refused stream whose header block is still arriving
         max_calls = DEFAULT_MAX_OPEN_CALLS
+        advertised = False  # the server sent SETTINGS_MAX_CONCURRENT_STREAMS
+        to_client = bytearray()  # frames for the client, written at a frame boundary
+        flushed = asyncio.Event()  # set when the server pump wrote them
+        told = False  # the face wrote to the client since the client pump last drained
 
         def from_client(kind, flags, stream, payload):
-            nonlocal highest
-            if kind == _HEADERS:
+            nonlocal highest, accepted, refusing
+            if kind == _HEADERS and stream > highest:
                 if stream % 2 == 0:
                     raise _ProtocolError(f"HEADERS on stream {stream}, which a server opens")
-                if stream <= highest:
-                    # Trailers on a call, or a stream the server refuses as closed:
-                    # never a new call.
-                    return
                 highest = stream
-                if len(calls) >= max_calls:
-                    raise _ProtocolError(f"more than {max_calls} calls open at once")
-                calls.add(stream)
+                if len(calls) < max_calls:
+                    calls.add(stream)
+                    accepted = stream
+                elif not advertised:
+                    refusing = stream
+                    tell_client(_rst_stream(stream, _ENHANCE_YOUR_CALM))
+                # Past an advertised limit the server refuses the call itself.
+            elif kind == _HEADERS and stream % 2 == 0:
+                raise _ProtocolError(f"HEADERS on stream {stream}, which a server opens")
             elif kind == _RST_STREAM:
                 calls.discard(stream)
+            # A lower odd stream id opens nothing: trailers on a call, or a stream the
+            # server refuses as closed.
+            if refusing and stream == refusing and kind in (_HEADERS, _CONTINUATION):
+                if flags & _END_HEADERS_FLAG:
+                    refusing = 0
+                    return _rst_stream(stream, _CANCEL)  # after its header block
+            return None
 
         def from_server(kind, flags, stream, payload):
-            nonlocal max_calls
+            nonlocal max_calls, advertised
             if kind == _RST_STREAM or (kind in (_DATA, _HEADERS) and flags & _END_STREAM_FLAG):
                 calls.discard(stream)
             elif kind == _SETTINGS and not flags & _ACK_FLAG:
                 for at in range(0, len(payload) - 5, 6):
                     key, value = struct.unpack_from(">HI", payload, at)
                     if key == _MAX_CONCURRENT_STREAMS:
-                        max_calls = value
+                        max_calls, advertised = value, True
             elif kind == _GOAWAY and len(payload) >= 4:
                 # Streams above the last one the server processes are never answered.
                 last_stream = struct.unpack_from(">I", payload)[0] & 0x7FFFFFFF
                 calls.difference_update({call for call in calls if call > last_stream})
+            return None
 
-        async def pump(source, sink, frames):
-            nonlocal last
+        client, server = _Frames(from_client), _Frames(from_server)
+
+        def tell_client(frame):
+            nonlocal told
+            # Only between two of the server's frames: queued while one is half relayed,
+            # and written by the server pump at that frame's end. Nothing to a client gone.
+            if writer.is_closing():
+                return
+            if server.at_boundary and not to_client:
+                writer.write(frame)
+                told = True
+            else:
+                to_client.extend(frame)
+
+        async def written():
+            """Wait until the server pump has written what is queued for the client."""
+            while to_client:
+                flushed.clear()
+                await flushed.wait()
+
+        def relay(chunk, sink, frames, ends, queued):
+            # As few writes as the frames to insert allow: one per insertion point.
+            at = 0
+            for end, after in ends:
+                if not after and not queued:
+                    continue
+                sink.write(chunk[at:end] + after)
+                at = end
+                if queued:
+                    sink.write(bytes(queued))
+                    queued.clear()
+                    flushed.set()
+            stop = len(chunk) if frames.error is None else (ends[-1][0] if ends else at)
+            if stop > at:
+                sink.write(chunk[at:stop])
+
+        async def pump(source, sink, frames, queued):
+            nonlocal last, told
             try:
                 while chunk := await source.read(65536):
+                    if writer.is_closing():
+                        return  # the client is gone: nothing left to relay
                     last = loop.time()
-                    frames.feed(chunk)
-                    sink.write(chunk)
+                    relay(chunk, sink, frames, frames.feed(chunk, bool(queued)), queued)
+                    if frames.error is not None:
+                        _log.info("connect face: closing a gRPC connection: %s", frames.error)
+                        tell_client(_goaway(accepted, _PROTOCOL_ERROR, str(frames.error)))
+                        try:
+                            await asyncio.wait_for(written(), _GOAWAY_GRACE_S)
+                        except TimeoutError:
+                            _log.warning(
+                                "connect face: the GOAWAY was not delivered: the server's "
+                                "frame in flight did not complete within %s s",
+                                _GOAWAY_GRACE_S,
+                            )
+                        return
                     await sink.drain()
+                    if sink is not writer:
+                        # The refusals: a client must read them before it is read again.
+                        if len(to_client) >= _MAX_QUEUED_FOR_CLIENT:
+                            await written()
+                        if told:
+                            told = False
+                            await writer.drain()
             except (ConnectionError, OSError):
                 pass
-            except _ProtocolError as refusal:
-                _log.info("connect face: closing a gRPC connection: %s", refusal)
 
         async def idle():
             # The clock runs from the last byte either way, and only ends the connection
@@ -364,8 +504,8 @@ async def idle():
                 await asyncio.sleep(remaining if remaining > 0 else self.idle_timeout)
 
         tasks = (
-            asyncio.ensure_future(pump(reader, upstream_writer, _Frames(from_client))),
-            asyncio.ensure_future(pump(upstream_reader, writer, _Frames(from_server))),
+            asyncio.ensure_future(pump(reader, upstream_writer, client, bytearray())),
+            asyncio.ensure_future(pump(upstream_reader, writer, server, to_client)),
             asyncio.ensure_future(idle()),
         )
         try: