Skip to content

Commit 186d9a8

Browse files
Merge pull request #122 from gael-obin/fix/connect-arrival-cancellation
fix(connect): admit at headers and cancel abandoned calls
2 parents dc3b6e4 + 7e92b3e commit 186d9a8

2 files changed

Lines changed: 434 additions & 24 deletions

File tree

‎templates/connect/connect_gen.py‎

Lines changed: 102 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,9 @@
4545
bounds. The face does not bound how many calls are in flight, as the gRPC server does
4646
not: each holds its answer until its client reads it.
4747
48-
Every read is bounded in time: an HTTP/1.1 read by ``read_timeout`` (``idle_timeout``
49-
between requests), and a relayed gRPC connection is closed once it has had no open call
48+
An HTTP/1.1 header read is bounded by ``read_timeout`` (``idle_timeout`` between
49+
requests); the entire request body has one ``read_timeout`` budget. A relayed
50+
gRPC connection is closed once it has had no open call
5051
for ``idle_timeout``, as Go's HTTP/2 server bounds an idle connection. A call in flight
5152
is never cut by the face, however long its handler takes to answer: its deadline and the
5253
gRPC server's own policy bound it.
@@ -745,12 +746,11 @@ async def _http1(self, reader, writer, head):
745746
return
746747
exchange = _Exchange(conn, writer)
747748
try:
748-
body = await self._body(conn, reader, event)
749-
await self._dispatch(exchange, event, body)
749+
await self._dispatch(exchange, event, conn, reader)
750750
except _Plain as refusal:
751751
await exchange.plain(refusal)
752752
return # the body may be unread: the connection cannot be reused
753-
except (TimeoutError, h11.RemoteProtocolError):
753+
except (TimeoutError, h11.RemoteProtocolError, ConnectionError):
754754
return
755755
if conn.our_state is not h11.DONE or conn.their_state is not h11.DONE:
756756
return
@@ -774,7 +774,17 @@ async def _body(self, conn, reader, request):
774774
else:
775775
raise h11.RemoteProtocolError("connection closed inside a request body")
776776

777-
async def _dispatch(self, exchange, request, body):
777+
async def _disconnected(self, conn, reader):
778+
# Preserve pipelined bytes for h11's next cycle, but never buffer an unbounded
779+
# pipeline while a handler is running. EOF cancels even a silent unary call.
780+
buffered = 0
781+
while data := await reader.read(65536):
782+
buffered += len(data)
783+
if buffered > self.limit + 65536:
784+
raise h11.RemoteProtocolError("pipelined requests exceed the message bound")
785+
conn.receive_data(data)
786+
787+
async def _dispatch(self, exchange, request, conn, reader):
778788
if request.method != b"POST":
779789
raise _Plain(405, "only POST is served", [(b"allow", b"POST")])
780790
path = request.target.split(b"?", 1)[0].decode("latin-1")
@@ -811,7 +821,42 @@ async def _dispatch(self, exchange, request, body):
811821
self.limit,
812822
self.answer_bounds,
813823
)
814-
await protocol.answer(exchange, call, body, headers)
824+
try:
825+
call.refuse_streams(unary=isinstance(protocol, _ConnectUnary))
826+
except CallError:
827+
# Encode the same cardinality refusal without opening a backend RPC.
828+
await protocol.answer(exchange, call, b"", headers)
829+
return
830+
tasks = []
831+
try:
832+
# Start metadata delivery independently of HTTP body arrival. The existing
833+
# server interceptor owns admission and its hold; there is no second check.
834+
call.start()
835+
body_task = asyncio.create_task(self._body(conn, reader, request))
836+
ended = asyncio.create_task(call.ended())
837+
tasks = [body_task, ended]
838+
async with asyncio.timeout(self.read_timeout):
839+
done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
840+
if ended in done:
841+
call.early_error = ended.result()
842+
body = b""
843+
body_task.cancel()
844+
await asyncio.gather(body_task, return_exceptions=True)
845+
else:
846+
body = body_task.result()
847+
answer = asyncio.create_task(protocol.answer(exchange, call, body, headers))
848+
disconnected = asyncio.create_task(self._disconnected(conn, reader))
849+
tasks += [answer, disconnected]
850+
done, _ = await asyncio.wait([answer, disconnected], return_when=asyncio.FIRST_COMPLETED)
851+
if disconnected in done:
852+
disconnected.result()
853+
raise ConnectionError("client disconnected")
854+
await answer
855+
finally:
856+
call.cancel()
857+
for task in tasks:
858+
task.cancel()
859+
await asyncio.gather(*tasks, return_exceptions=True)
815860

816861

817862
def _method(path):
@@ -864,13 +909,52 @@ def __init__(self, channel, path, method, codec, raw_headers, headers, limit, an
864909
self.output = message_factory.GetMessageClass(method.output_type)
865910
self.metadata = _metadata(raw_headers)
866911
self.timeout = _timeout(headers)
912+
self.rpc = None
913+
self.early_error = None
914+
915+
def start(self):
916+
# The wire still carries exactly one request. Streaming writes let grpc.aio
917+
# send headers without requiring that request to have arrived first.
918+
factory = (
919+
self.channel.stream_stream if self.method.server_streaming else self.channel.stream_unary
920+
)
921+
self.rpc = factory(self.path)(metadata=self.metadata, timeout=self.timeout)
922+
923+
def cancel(self):
924+
if self.rpc is not None:
925+
self.rpc.cancel()
926+
927+
async def ended(self):
928+
code = await self.rpc.code()
929+
return CallError(
930+
code if code is not grpc.StatusCode.OK else grpc.StatusCode.INTERNAL,
931+
await self.rpc.details() or "call ended before its request",
932+
await self.rpc.initial_metadata() or (),
933+
await self.rpc.trailing_metadata() or (),
934+
)
867935

868-
def refuse_streams(self):
936+
async def _write(self, message):
937+
request = self.codec.request(self.input, message)
938+
try:
939+
await self.rpc.write(request)
940+
await self.rpc.done_writing()
941+
except asyncio.InvalidStateError:
942+
# A refusal/deadline may arrive while the HTTP body is being decoded.
943+
raise await self.ended() from None
944+
945+
def refuse_streams(self, *, unary=False):
946+
if self.early_error is not None:
947+
raise self.early_error
869948
if self.method.client_streaming:
870949
raise CallError(
871950
grpc.StatusCode.UNIMPLEMENTED,
872951
f"{self.path} streams its requests: it is served on the gRPC protocol only",
873952
)
953+
if unary and self.method.server_streaming:
954+
raise CallError(
955+
grpc.StatusCode.UNIMPLEMENTED,
956+
f"{self.path} streams its answers: call it with application/connect+",
957+
)
874958

875959
def decompress(self, data: bytes, encoding: str | None) -> bytes:
876960
if encoding in (None, "", "identity"):
@@ -920,11 +1004,9 @@ def _too_large(self, size: int) -> CallError:
9201004

9211005
async def unary(self, message: bytes):
9221006
"""The answer and the header and trailer metadata of a unary call."""
923-
request = self.codec.request(self.input, message)
924-
call = self.channel.unary_unary(self.path)(
925-
request, metadata=self.metadata, timeout=self.timeout
926-
)
1007+
call = self.rpc
9271008
try:
1009+
await self._write(message)
9281010
answer = await call
9291011
return (
9301012
self.answer(answer),
@@ -934,11 +1016,12 @@ async def unary(self, message: bytes):
9341016
except grpc.aio.AioRpcError as exc:
9351017
raise self.failed(exc) from None
9361018

937-
def stream(self, message: bytes):
938-
request = self.codec.request(self.input, message)
939-
return self.channel.unary_stream(self.path)(
940-
request, metadata=self.metadata, timeout=self.timeout
941-
)
1019+
async def stream(self, message: bytes):
1020+
try:
1021+
await self._write(message)
1022+
return self.rpc
1023+
except grpc.aio.AioRpcError as exc:
1024+
raise self.failed(exc) from None
9421025

9431026

9441027
class _Exchange:
@@ -977,12 +1060,7 @@ def __init__(self, media):
9771060
async def answer(self, exchange, call, body, headers):
9781061
response_type = [(b"content-type", self.media.encode())]
9791062
try:
980-
call.refuse_streams()
981-
if call.method.server_streaming:
982-
raise CallError(
983-
grpc.StatusCode.UNIMPLEMENTED,
984-
f"{call.path} streams its answers: call it with application/connect+",
985-
)
1063+
call.refuse_streams(unary=True)
9861064
message = call.decompress(body, headers.get("content-encoding"))
9871065
answer, header, trailer = await call.unary(message)
9881066
except CallError as error:
@@ -1073,7 +1151,7 @@ async def _open(call, body, encoding):
10731151
payload, compressed = _one_message(body)
10741152
message = call.decompress(payload, encoding if compressed else None)
10751153
if call.method.server_streaming:
1076-
stream = call.stream(message)
1154+
stream = await call.stream(message)
10771155
return stream, tuple(await stream.initial_metadata() or ())
10781156
answer, header, trailer = await call.unary(message)
10791157
return _Answered(answer, trailer), header

0 commit comments

Comments
 (0)