Add support for notifying end of remote stream (#727)

Author: TingDaoK

awscrt/aio/http.py

@@ -337,6 +337,7 @@ def __init__(self,
                  request: HttpRequest,
                  request_body_generator: AsyncIterator[bytes] = None,
                  loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
+
         # Initialize the parent class
         http2_manual_write = request_body_generator is not None and connection.version is HttpVersion.Http2
         super()._init_common(connection, request, http2_manual_write=http2_manual_write)
@@ -357,6 +358,7 @@ def __init__(self,
 
         # Create futures for async operations
         self._completion_future = Future()
+        self._remote_completion_future = Future()
         self._response_status_future = Future()
         self._response_headers_future = Future()
         self._status_code = None
@@ -386,13 +388,12 @@ def _on_body(self, chunk: bytes) -> None:
         # Set result outside lock (Future is thread-safe)
         future.set_result(chunk)
 
-    def _on_complete(self, error_code: int) -> None:
-        """Set the completion status of the stream."""
-        if error_code == 0:
-            self._completion_future.set_result(self._status_code)
-        else:
-            self._completion_future.set_exception(awscrt.exceptions.from_code(error_code))
+    def _resolve_pending_chunk_futures(self) -> None:
+        """Helper to resolve all pending chunk futures with empty bytes.
 
+        This indicates end of stream to any waiting get_next_response_chunk() calls.
+        Must be called when either the stream completes or remote peer sends END_STREAM.
+        """
         # Resolve all pending chunk futures with lock protection
         with self._deque_lock:
             pending_futures = list(self._chunk_futures)
@@ -402,6 +403,20 @@ def _on_complete(self, error_code: int) -> None:
         for future in pending_futures:
             future.set_result(b"")
 
+    def _on_complete(self, error_code: int) -> None:
+        """Set the completion status of the stream."""
+        if error_code == 0:
+            self._completion_future.set_result(self._status_code)
+        else:
+            self._completion_future.set_exception(awscrt.exceptions.from_code(error_code))
+
+        self._resolve_pending_chunk_futures()
+
+    def _on_h2_remote_end_stream(self) -> None:
+        """Called when the remote peer has finished sending (HTTP/2 only)."""
+        self._remote_completion_future.set_result(None)
+        self._resolve_pending_chunk_futures()
+
     async def _set_request_body_generator(self, body_iterator: AsyncIterator[bytes]):
         ...
 
@@ -431,7 +446,7 @@ async def get_next_response_chunk(self) -> bytes:
         with self._deque_lock:
             if self._received_chunks:
                 return self._received_chunks.popleft()
-            elif self._completion_future.done():
+            elif self._completion_future.done() or self._remote_completion_future.done():
                 return b""
             else:
                 future = Future()

awscrt/http.py

@@ -569,6 +569,14 @@ def _on_complete(self, error_code: int) -> None:
         else:
             self._completion_future.set_exception(awscrt.exceptions.from_code(error_code))
 
+    def _on_h2_remote_end_stream(self) -> None:
+        """Called when remote peer sends END_STREAM (HTTP/2 only).
+
+        This callback is only invoked for HTTP/2 connections. HTTP/1.x streams
+        will never receive this callback. Base implementation does nothing.
+        """
+        pass
+
     def update_window(self, increment_size: int) -> None:
         """
         Update the stream's flow control window.
@@ -613,14 +621,57 @@ def activate(self) -> None:
 
 
 class Http2ClientStream(HttpClientStreamBase):
+    __slots__ = ('_remote_end_stream_future',)
+
     def __init__(self,
                  connection: HttpClientConnection,
                  request: 'HttpRequest',
                  on_response: Optional[Callable[..., None]] = None,
                  on_body: Optional[Callable[..., None]] = None,
                  manual_write: bool = False) -> None:
+        self._remote_end_stream_future = Future()
         self._init_common(connection, request, on_response, on_body, manual_write)
 
+    @property
+    def remote_end_stream_future(self) -> "concurrent.futures.Future":
+        """
+        concurrent.futures.Future: Future that completes when the remote peer has finished
+        sending (HTTP/2 only). This occurs when the server sends an END_STREAM flag.
+
+        The future will contain a result of None on success, or an exception if the stream
+        encounters an error before END_STREAM is received (e.g., RST_STREAM).
+
+        This is different from `completion_future` which completes when both the
+        client and server have finished (bidirectional stream closure).
+
+        Note: This future only applies to HTTP/2 connections. It will complete when the
+        server sends END_STREAM, which may occur before the client finishes sending.
+        In case of stream completed without END_STREAM received, this future will complete
+        with exception.
+        """
+        return self._remote_end_stream_future
+
+    def _on_h2_remote_end_stream(self) -> None:
+        """Internal callback when remote peer sends END_STREAM (HTTP/2 only)."""
+        if not self._remote_end_stream_future.done():
+            self._remote_end_stream_future.set_result(None)
+
+    def _on_complete(self, error_code: int) -> None:
+        # done with HttpRequest, drop reference
+        self._request = None  # type: ignore
+
+        # Ensure remote_completion_future is always resolved
+        if not self._remote_end_stream_future.done():
+            # Stream completed successfully but END_STREAM was never received,
+            # complete `remote_completion_future` with exception.
+            self._remote_end_stream_future.set_exception(
+                RuntimeError("Stream completed without receiving remote END_STREAM"))
+
+        if error_code == 0:
+            self._completion_future.set_result(self._response_status_code)
+        else:
+            self._completion_future.set_exception(awscrt.exceptions.from_code(error_code))
+
     def activate(self) -> None:
         """Begin sending the request.
 

crt/aws-c-event-stream

@@ -213,6 +213,24 @@ static void s_on_stream_complete(struct aws_http_stream *native_stream, int erro
     /*************** GIL RELEASE ***************/
 }
 
+static void s_on_h2_remote_end_stream(struct aws_http_stream *native_stream, void *user_data) {
+    (void)native_stream;
+    struct http_stream_binding *stream = user_data;
+
+    /*************** GIL ACQUIRE ***************/
+    PyGILState_STATE state;
+    if (aws_py_gilstate_ensure(&state)) {
+        return; /* Python has shut down. Nothing matters anymore, but don't crash */
+    }
+
+    PyObject *result = PyObject_CallMethod(stream->self_proxy, "_on_h2_remote_end_stream", "()");
+    if (result) {
+        Py_DECREF(result);
+    }
+    PyGILState_Release(state);
+    /*************** GIL RELEASE ***************/
+}
+
 static void s_stream_capsule_destructor(PyObject *http_stream_capsule) {
     struct http_stream_binding *stream = PyCapsule_GetPointer(http_stream_capsule, s_capsule_name_http_stream);
 
@@ -283,6 +301,7 @@ PyObject *aws_py_http_client_stream_new(PyObject *self, PyObject *args) {
         .on_response_header_block_done = s_on_incoming_header_block_done,
         .on_response_body = s_on_incoming_body,
         .on_complete = s_on_stream_complete,
+        .on_h2_remote_end_stream = s_on_h2_remote_end_stream,
         .user_data = stream,
         .http2_use_manual_data_writes = http2_manual_write,
     };

crt/aws-c-http

@@ -916,5 +916,80 @@ def test_h2_stream_flow_control_blocks_and_resumes(self):
         asyncio.run(self._test_h2_stream_flow_control_blocks_and_resumes())
 
 
+class AIOHttp2RemoteEndStreamTest(NativeResourceTest):
+    """Test suite for HTTP/2 on_h2_remote_end_stream callback in asyncio"""
+    timeout = 10.0
+
+    async def _new_httpbin_h2_connection(self):
+        """Create HTTP/2 connection to httpbin.org"""
+        event_loop_group = EventLoopGroup()
+        host_resolver = DefaultHostResolver(event_loop_group)
+        bootstrap = ClientBootstrap(event_loop_group, host_resolver)
+
+        tls_ctx_options = TlsContextOptions()
+        tls_ctx = ClientTlsContext(tls_ctx_options)
+        tls_conn_opt = tls_ctx.new_connection_options()
+        tls_conn_opt.set_server_name("httpbin.org")
+        tls_conn_opt.set_alpn_list(["h2"])
+
+        connection = await AIOHttp2ClientConnection.new(
+            host_name="httpbin.org",
+            port=443,
+            bootstrap=bootstrap,
+            tls_connection_options=tls_conn_opt)
+
+        return connection
+
+    async def _test_h2_remote_end_stream_ordering(self):
+        """Test that on_h2_remote_end_stream fires before on_complete when server finishes first"""
+        connection = await self._new_httpbin_h2_connection()
+
+        # Use httpbin.org 404 path - server responds immediately
+        request = HttpRequest('POST', '/this-path-does-not-exist-deliberately-404')
+        request.headers.add('host', 'httpbin.org')
+
+        complete_success = asyncio.Event()
+        remote_finished = asyncio.Event()
+        complete_fired = asyncio.Event()
+
+        async def slow_body_generator():
+            # Send first chunk WITHOUT end_stream
+            yield b'chunk1'
+            # Wait for server to finish
+            await remote_finished.wait()
+            if not complete_fired.is_set():
+                # Verify complete hasn't fired yet
+                complete_success.set()
+            # Now finish sending
+            yield b'chunk2'
+
+        stream = connection.request(request, request_body_generator=slow_body_generator())
+
+        # Read response
+        status_code = await stream.get_response_status_code()
+        self.assertEqual(404, status_code)
+
+        # Read all response body
+        while True:
+            chunk = await stream.get_next_response_chunk()
+            if not chunk:
+                break
+
+        # set remove stream
+        remote_finished.set()
+
+        # Wait for stream to complete successfully
+        await stream.wait_for_completion()
+        complete_fired.set()
+
+        self.assertTrue(complete_success.is_set())
+
+        await connection.close()
+
+    def test_h2_remote_end_stream_ordering(self):
+        """Test callback ordering with early server response"""
+        asyncio.run(self._test_h2_remote_end_stream_ordering())
+
+
 if __name__ == '__main__':
     unittest.main()