Repository: encode/starlette

Issue:
Background tasks don't work with middleware that subclasses `BaseHTTPMiddleware`
When using background tasks with middleware, requests are not processed until the background task has finished.

1. Use the example below
2. Make several requests in a row - works as expected
3. Uncomment `app.add_middleware(TransparentMiddleware)` and re-run
4. Make several requests in a row - subsequent ones are not processed until the 10 second sleep has finished (the first request returns before then though).

The same behaviour occurs with `asyncio.sleep` (async) and `time.sleep` (sync, run in threadpool)

```python
import asyncio
import uvicorn
from starlette.applications import Starlette
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.background import BackgroundTask
from starlette.responses import JSONResponse


class TransparentMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        # simple middleware that does absolutely nothing
        response = await call_next(request)
        return response


app = Starlette(debug=True)
# uncomment to break
# app.add_middleware(TransparentMiddleware)


@app.route("/")
async def test(_):
    task = BackgroundTask(asyncio.sleep, 10)
    return JSONResponse("hello", background=task)


if __name__ == '__main__':
    uvicorn.run(app)
```

Discussion:
florimondmanca: @jhominal Hey, I think I would be very interested in seeing what the tests look like. Might help me better grok what practical situations this PR would help address.

About:

>  @florimondmanca expresses his opinion that "AFAIK there’s no way in the ASGI spec to just discard a streaming body even though more_body is True. In fact I believe the spec says servers MUST pull until that becomes False." - in other words, using task cancellation on a downstream app may not be permitted by ASGI.

This came from what I would call an _interpretation_ of the spec, which reads (emphasis mine):

https://asgi.readthedocs.io/en/latest/specs/www.html#request-receive-event

> more_body (bool) – Signifies if there is additional content to come (as part of a Request message). If True, **the consuming application _should_ wait until it gets a chunk with this set to False**. If False, the request is complete and should be processed. Optional; if missing defaults to False.

I interpret "you should consume as long as `more_body` is `True`" as implying that "you should not interrupt (cancel) the app until `more_body` becomes `False`".
florimondmanca: @jhominal Am I correct saying this would be an alternative to #1710?
jhominal: @florimondmanca Thank you for your reply! I am going to work on the test cases - at least a few of them will be similar to existing issues/pull requests.

But yes, it is my belief that this proposal has the potential to be an alternative to #1710, #1700 (which fixes the "background tasks are canceled" issue by pushing them outside of the response processing), #1699, #1441 (which was closed yesterday).

I believe it would fix at least #1438, which is the issue I raised a little time ago.
adriangb: This does look very nice, I do think removing the cancellation is a step in the right direction.

> this proposal has the potential to be an alternative to https://github.com/encode/starlette/pull/1700 (which fixes the "background tasks are canceled" issue by pushing them outside of the response processing)

Did you identify any issues with that proposal so that they are codified for posteriority / we can close that PR?
jhominal: I will just comment on the two issues that I believe are still open and would be directly affected by this PR:

 * #1438 would be fixed by this PR, as shown by the added test (thanks to @adriangb for his own version of that test in #1700);
 * For #919, I have not been able to add a substantial test case, as I have not been able to reproduce the headline issue as described by the issue reporter, either on current starlette versions, or by installing `starlette==0.13.2`+`uvicorn==0.11.4` - I will make a more detailed comment on that issue;
adriangb: > * For [Background tasks don't work with middleware that subclasses `BaseHTTPMiddleware` #919](https://github.com/encode/starlette/issues/919), I have not been able to add a substantial test case, as I have not been able to reproduce the headline issue as described by the issue reporter, either on current starlette versions, or by installing `starlette==0.13.2`+`uvicorn==0.11.4` - I will make a more detailed comment on that issue;

I tried to add a test case for this in https://github.com/encode/starlette/blob/f5cb1ee53d569bba5ecaf8ca449fa0b59ea21e0f/tests/middleware/test_base.py#L265-L293

The original issue says "subsequent ones are not processed until the 10 second sleep has finished (the first request returns before then though)". I was not able to reproduce _that_ specific outcome, but I was able to reproduce the BackgroundTask not running when there was a BaseHTTPMiddleware present. I think it is the same bug, the behavior has just changed slightly over time.
jhominal: @adriangb: I have looked at your test cases while implementing mine. However, after porting and reviewing a copy of the #919 test, I saw that the only substantial difference between my test cases for #1438 and #919 was that one test was expressed in pure ASGI terms, while the other was expressed using the `TestClient` (and relies on the fact that `TestClient` sends `http.disconnect` immediately after receiving the response - but I am of the opinion that `TestClient` should not be relied on for ASGI-level details). Because of that, I took out the test for #919.

As some of the issues reported on #919 were at a time when Starlette was implemented directly on top of `asyncio`, and given the difference that using `anyio` makes to the usable primitives, it is extremely likely that these issues cannot be reproduced anymore.
adriangb: Hmm good point, looking at my test now I think that may be the case. I'll dig deeper tomorrow, thank you for combing through them.
jhominal: > Hmm good point, looking at my test now I think that may be the case. I'll dig deeper tomorrow, thank you for combing through them.

I just reread your test for #1438 and I realize that actually, it has a difference (which is that the background task waits on the `disconnected` event). I will admit that, as I had already implemented the test for #1438 before implementing yours, I did not read that test case that closely - I mostly compared my version of #1438 and my version of #919 (based on yours).
jhominal: I have just added a test to check for the issue reported by @kissgyorgy on https://github.com/encode/starlette/issues/1678#issuecomment-1172916042 - about the way that `BaseHTTPMiddleware` cancellation prevents the context manager defined in the middleware from running its exit function.

Fixing this test did not require any modification to the PR's modifications of `base.py`.
jhominal: @florimondmanca As you previously expressed interest in this PR, I chose to request you as a reviewer. Please feel free to suggest anyone else who you think should also take a look at this.
teamhide: Any update here?
jhominal: I agree that this PR is languishing a bit, so I have decided to request review from all encode maintainers, in the hope that more people can take a look at this.
Kludex: I'll check tomorrow. I've added this to the checklist for the next release.
jhominal: @adriangb @kludex Thank you very much for taking the time to look at it! (I have not been in Gitter much recently so I do not know if you have talked about this PR over there)
Kludex: I'll let you merge @jhominal :+1: 

Thanks for this, and sorry for taking long to review.

Beautiful code. :+1: 
Kludex: I'll merge this, since it's the only PR missing for the release.

Patch:
diff --git a/starlette/middleware/base.py b/starlette/middleware/base.py
index 49a5e3e..586c987 100644
--- a/starlette/middleware/base.py
+++ b/starlette/middleware/base.py
@@ -4,12 +4,13 @@ import anyio
 
 from starlette.requests import Request
 from starlette.responses import Response, StreamingResponse
-from starlette.types import ASGIApp, Receive, Scope, Send
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
 RequestResponseEndpoint = typing.Callable[[Request], typing.Awaitable[Response]]
 DispatchFunction = typing.Callable[
     [Request, RequestResponseEndpoint], typing.Awaitable[Response]
 ]
+T = typing.TypeVar("T")
 
 
 class BaseHTTPMiddleware:
@@ -24,19 +25,52 @@ class BaseHTTPMiddleware:
             await self.app(scope, receive, send)
             return
 
+        response_sent = anyio.Event()
+
         async def call_next(request: Request) -> Response:
             app_exc: typing.Optional[Exception] = None
             send_stream, recv_stream = anyio.create_memory_object_stream()
 
+            async def receive_or_disconnect() -> Message:
+                if response_sent.is_set():
+                    return {"type": "http.disconnect"}
+
+                async with anyio.create_task_group() as task_group:
+
+                    async def wrap(func: typing.Callable[[], typing.Awaitable[T]]) -> T:
+                        result = await func()
+                        task_group.cancel_scope.cancel()
+                        return result
+
+                    task_group.start_soon(wrap, response_sent.wait)
+                    message = await wrap(request.receive)
+
+                if response_sent.is_set():
+                    return {"type": "http.disconnect"}
+
+                return message
+
+            async def close_recv_stream_on_response_sent() -> None:
+                await response_sent.wait()
+                recv_stream.close()
+
+            async def send_no_error(message: Message) -> None:
+                try:
+                    await send_stream.send(message)
+                except anyio.BrokenResourceError:
+                    # recv_stream has been closed, i.e. response_sent has been set.
+                    return
+
             async def coro() -> None:
                 nonlocal app_exc
 
                 async with send_stream:
                     try:
-                        await self.app(scope, request.receive, send_stream.send)
+                        await self.app(scope, receive_or_disconnect, send_no_error)
                     except Exception as exc:
                         app_exc = exc
 
+            task_group.start_soon(close_recv_stream_on_response_sent)
             task_group.start_soon(coro)
 
             try:
@@ -71,7 +105,7 @@ class BaseHTTPMiddleware:
             request = Request(scope, receive=receive)
             response = await self.dispatch_func(request, call_next)
             await response(scope, receive, send)
-            task_group.cancel_scope.cancel()
+            response_sent.set()
 
     async def dispatch(
         self, request: Request, call_next: RequestResponseEndpoint


Test patch:
diff --git a/tests/middleware/test_base.py b/tests/middleware/test_base.py
index 976d77b..ed0734b 100644
--- a/tests/middleware/test_base.py
+++ b/tests/middleware/test_base.py
@@ -1,8 +1,11 @@
 import contextvars
+from contextlib import AsyncExitStack
 
+import anyio
 import pytest
 
 from starlette.applications import Starlette
+from starlette.background import BackgroundTask
 from starlette.middleware import Middleware
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.responses import PlainTextResponse, StreamingResponse
@@ -206,3 +209,207 @@ def test_contextvars(test_client_factory, middleware_cls: type):
     client = test_client_factory(app)
     response = client.get("/")
     assert response.status_code == 200, response.content
+
+
+@pytest.mark.anyio
+async def test_run_background_tasks_even_if_client_disconnects():
+    # test for https://github.com/encode/starlette/issues/1438
+    request_body_sent = False
+    response_complete = anyio.Event()
+    background_task_run = anyio.Event()
+
+    async def sleep_and_set():
+        # small delay to give BaseHTTPMiddleware a chance to cancel us
+        # this is required to make the test fail prior to fixing the issue
+        # so do not be surprised if you remove it and the test still passes
+        await anyio.sleep(0.1)
+        background_task_run.set()
+
+    async def endpoint_with_background_task(_):
+        return PlainTextResponse(background=BackgroundTask(sleep_and_set))
+
+    async def passthrough(request, call_next):
+        return await call_next(request)
+
+    app = Starlette(
+        middleware=[Middleware(BaseHTTPMiddleware, dispatch=passthrough)],
+        routes=[Route("/", endpoint_with_background_task)],
+    )
+
+    scope = {
+        "type": "http",
+        "version": "3",
+        "method": "GET",
+        "path": "/",
+    }
+
+    async def receive():
+        nonlocal request_body_sent
+        if not request_body_sent:
+            request_body_sent = True
+            return {"type": "http.request", "body": b"", "more_body": False}
+        # We simulate a client that disconnects immediately after receiving the response
+        await response_complete.wait()
+        return {"type": "http.disconnect"}
+
+    async def send(message):
+        if message["type"] == "http.response.body":
+            if not message.get("more_body", False):
+                response_complete.set()
+
+    await app(scope, receive, send)
+
+    assert background_task_run.is_set()
+
+
+@pytest.mark.anyio
+async def test_run_context_manager_exit_even_if_client_disconnects():
+    # test for https://github.com/encode/starlette/issues/1678#issuecomment-1172916042
+    request_body_sent = False
+    response_complete = anyio.Event()
+    context_manager_exited = anyio.Event()
+
+    async def sleep_and_set():
+        # small delay to give BaseHTTPMiddleware a chance to cancel us
+        # this is required to make the test fail prior to fixing the issue
+        # so do not be surprised if you remove it and the test still passes
+        await anyio.sleep(0.1)
+        context_manager_exited.set()
+
+    class ContextManagerMiddleware:
+        def __init__(self, app):
+            self.app = app
+
+        async def __call__(self, scope: Scope, receive: Receive, send: Send):
+            async with AsyncExitStack() as stack:
+                stack.push_async_callback(sleep_and_set)
+                await self.app(scope, receive, send)
+
+    async def simple_endpoint(_):
+        return PlainTextResponse(background=BackgroundTask(sleep_and_set))
+
+    async def passthrough(request, call_next):
+        return await call_next(request)
+
+    app = Starlette(
+        middleware=[
+            Middleware(BaseHTTPMiddleware, dispatch=passthrough),
+            Middleware(ContextManagerMiddleware),
+        ],
+        routes=[Route("/", simple_endpoint)],
+    )
+
+    scope = {
+        "type": "http",
+        "version": "3",
+        "method": "GET",
+        "path": "/",
+    }
+
+    async def receive():
+        nonlocal request_body_sent
+        if not request_body_sent:
+            request_body_sent = True
+            return {"type": "http.request", "body": b"", "more_body": False}
+        # We simulate a client that disconnects immediately after receiving the response
+        await response_complete.wait()
+        return {"type": "http.disconnect"}
+
+    async def send(message):
+        if message["type"] == "http.response.body":
+            if not message.get("more_body", False):
+                response_complete.set()
+
+    await app(scope, receive, send)
+
+    assert context_manager_exited.is_set()
+
+
+def test_app_receives_http_disconnect_while_sending_if_discarded(test_client_factory):
+    class DiscardingMiddleware(BaseHTTPMiddleware):
+        async def dispatch(self, request, call_next):
+            await call_next(request)
+            return PlainTextResponse("Custom")
+
+    async def downstream_app(scope, receive, send):
+        await send(
+            {
+                "type": "http.response.start",
+                "status": 200,
+                "headers": [
+                    (b"content-type", b"text/plain"),
+                ],
+            }
+        )
+        async with anyio.create_task_group() as task_group:
+
+            async def cancel_on_disconnect():
+                while True:
+                    message = await receive()
+                    if message["type"] == "http.disconnect":
+                        task_group.cancel_scope.cancel()
+                        break
+
+            task_group.start_soon(cancel_on_disconnect)
+
+            # A timeout is set for 0.1 second in order to ensure that
+            # cancel_on_disconnect is scheduled by the event loop
+            with anyio.move_on_after(0.1):
+                while True:
+                    await send(
+                        {
+                            "type": "http.response.body",
+                            "body": b"chunk ",
+                            "more_body": True,
+                        }
+                    )
+
+            pytest.fail(
+                "http.disconnect should have been received and canceled the scope"
+            )  # pragma: no cover
+
+    app = DiscardingMiddleware(downstream_app)
+
+    client = test_client_factory(app)
+    response = client.get("/does_not_exist")
+    assert response.text == "Custom"
+
+
+def test_app_receives_http_disconnect_after_sending_if_discarded(test_client_factory):
+    class DiscardingMiddleware(BaseHTTPMiddleware):
+        async def dispatch(self, request, call_next):
+            await call_next(request)
+            return PlainTextResponse("Custom")
+
+    async def downstream_app(scope, receive, send):
+        await send(
+            {
+                "type": "http.response.start",
+                "status": 200,
+                "headers": [
+                    (b"content-type", b"text/plain"),
+                ],
+            }
+        )
+        await send(
+            {
+                "type": "http.response.body",
+                "body": b"first chunk, ",
+                "more_body": True,
+            }
+        )
+        await send(
+            {
+                "type": "http.response.body",
+                "body": b"second chunk",
+                "more_body": True,
+            }
+        )
+        message = await receive()
+        assert message["type"] == "http.disconnect"
+
+    app = DiscardingMiddleware(downstream_app)
+
+    client = test_client_factory(app)
+    response = client.get("/does_not_exist")
+    assert response.text == "Custom"