Repository: encode/starlette

Issue:
Background tasks are cancelled if the client closes connection
### Checklist

- [X] The bug is reproducible against the latest release or `master`.
- [X] There are no similar issues or pull requests to fix it yet.

### Describe the bug

When the HTTP client closes the TCP socket immediately after receiving the HTTP response, background tasks are cancelled.

This bug only happens when running the ASGI under uvicorn, and only if at least one HTTP Middleware is defined in the user middleware chain.

### Steps to reproduce the bug

1. Write the following ASGI Starlette application in `repro.py`:
```python
import traceback

import anyio
from starlette.applications import Starlette
from starlette.background import BackgroundTasks
from starlette.middleware import Middleware
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response
from starlette.routing import Route


async def passthrough(request, call_next):
    return await call_next(request)


async def _sleep(identifier, delay):
    print(identifier, "started")
    try:
        await anyio.sleep(delay)
        print(identifier, "completed")
    except BaseException:
        print(identifier, "error")
        traceback.print_exc()
        raise


async def response_with_sleeps(request):
    background_tasks = BackgroundTasks()
    background_tasks.add_task(_sleep, "background task 1", 2)
    background_tasks.add_task(_sleep, "background task 2", 2)
    return Response(background=background_tasks)


application = Starlette(
    middleware=[
        Middleware(BaseHTTPMiddleware, dispatch=passthrough),
    ],
    routes=[
        Route("/", response_with_sleeps),
    ],
)
```
2. Run that application using `uvicorn` (either uvloop or regular asyncio will reproduce the issue) on localhost:8000
```bash
uvicorn repro:application --port 8000
```
3. Run the following client script
```python
#!/usr/bin/env python
import socket

connection = socket.create_connection(("localhost", 8000))
connection.sendall(b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n")
print(connection.recv(10000).decode("utf8"))
connection.close()
```


### Expected behavior

The client script gets the HTTP response, and both background tasks should complete successfully.

The expected behavior will be detectable by the following content in standard output:

```
background task 1 started
background task 1 completed
background task 2 started
background task 2 completed
```

### Actual behavior

Background task 1 is interrupted at the `await` point and background task 2 is never started.

That results in the following content in the output (when running the `repro.py` application):
```
background task 1 started
background task 1 error
Traceback (most recent call last):
  File "/Users/jean/PycharmProjects/starlette-bg-cancelled/./repro.py", line 19, in _sleep
    await anyio.sleep(delay)
  File "/Users/jean/PycharmProjects/starlette-bg-cancelled/venv/lib/python3.9/site-packages/anyio/_core/_eventloop.py", line 69, in sleep
    return await get_asynclib().sleep(delay)
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/asyncio/tasks.py", line 654, in sleep
    return await future
asyncio.exceptions.CancelledError
```

### Debugging material

_No response_

### Environment

- MacOS 10.14.6 / Python 3.9 / Starlette 0.18.0

### Additional context

 * When I remove the `passthrough` middleware, the bug goes away.
 * When I run the same application in `hypercorn`, the bug goes away.
 * There does not seem to be a difference between using `uvloop` or not.
 * If the client script (e.g. with a `time.sleep(10)`) maintains the TCP connection open, the bug goes away.
Shield send "http.response.start" from cancellation (BaseHTTPMiddleware)
Fixes #1634 
- Discussion #1527 
- Caused by #1157 

`await recv_stream.receive()` will raise `anyio.EndOfStream` if request is disconnected, due to:
- `task_group.cancel_scope.cancel()` in `StreamingResponse.__call__.<locals>.wrap` and
- cancellation check in `await checkpoint()` of `MemoryObjectSendStream.send`,

and then `RuntimeError: No response returned.` will be raised in `BaseHTTPMiddleware`.

Let's shield send "http.response.start" from cancellation, since the message is ready to be sent to the receiver.

This is an alternative implementation of #1706 in `BaseHTTPMiddleware` instead of `StreamingResponse`.
We should not force the shielding in `StreamingResponse`, since the cancellation check is an intended feature of `MemoryObjectSendStream`. `BaseHTTPMiddleware`, which uses both, should be responsible for the compatibility.


Discussion:
Thanks for the precise report @jhominal :)

I've implemented a solution on #1441. Let's see how the discussion around that goes.
I looked the code that `BaseHTTPMiddleware` and `StreamingResponse` use the same `task_group`, when one scope receive `http.disconnect`,  `task_group` cancel all taskes. But I don't know why it can cancel `Response.backgroundtasks`. Do I missing other thing?
The PR that solves the issue has all the details. 

EDIT: My PR had some regressions. I've closed it.
The background tasks need to be shielded.

```diff
diff --git a/starlette/background.py b/starlette/background.py
index 4aaf7ae..db9b38a 100644
--- a/starlette/background.py
+++ b/starlette/background.py
@@ -1,6 +1,8 @@
 import sys
 import typing
 
+import anyio
+
 if sys.version_info >= (3, 10):  # pragma: no cover
     from typing import ParamSpec
 else:  # pragma: no cover
@@ -22,10 +24,11 @@ class BackgroundTask:
         self.is_async = is_async_callable(func)
 
     async def __call__(self) -> None:
-        if self.is_async:
-            await self.func(*self.args, **self.kwargs)
-        else:
-            await run_in_threadpool(self.func, *self.args, **self.kwargs)
+        with anyio.CancelScope(shield=True):
+            if self.is_async:
+                await self.func(*self.args, **self.kwargs)
+            else:
+                await run_in_threadpool(self.func, *self.args, **self.kwargs)
```
I tried my patch to `base.py` from https://github.com/encode/starlette/pull/1706#issuecomment-1166342235 on this branch and the current test (in this branch, not the patch) also passes (I can confirm the test fails on `master` so it's not just that it passes on any branch). I do think your solution is nicer in some sense, but does it fix any edge cases that the raise an exception solution doesn't?
I don't think this solution fixes any additional edge cases.

Both solutions pretty much target this specific edge case, where `StreamingResponse` does `listen_for_disconnect` and handles the disconnection in a way that is not compatible with how `MemorySendStream` is used in `BaseHTTPMiddleware`.
If there are no pros to doing cancellation, I would prefer to just raise an exception, but that's just my opinion and I'm biased, so I'd like others to weigh in.

Otherwise, this PR looks good and I'd be happy to move forward with it, thank you for the contribution and your persistence @acjh !
> If there are no pros to doing cancellation

I believe you mean *not* doing cancellation. More accurately, it's temporarily shielding from cancellation.

> I would prefer to just raise an exception

I'm totally fine with closing this PR in favour of that, btw. 😄 

This PR simply handles the compatibility between `StreamingResponse` and usage of anyio memory streams.
Raising an exception takes control of the flow by forcing exception handling, ignoring the streams' behaviour.
> > If there are no pros to doing cancellation
> 
> I believe you mean _not_ doing cancellation. More accurately, it's temporarily shielding from cancellation.

Yep sorry for being unclear

> > I would prefer to just raise an exception
> 
> I'm totally fine with closing this PR in favour of that, btw. 😄
> 
> This PR simply handles the compatibility between `StreamingResponse` and usage of anyio memory streams. Raising an exception takes control of the flow by forcing exception handling, ignoring the streams' behaviour.

Let's just wait to see what others think, I'm also fine with either 😄 
Sure. Thank you for your patience and prompt responses!

Resolving patch:
diff --git a/starlette/middleware/base.py b/starlette/middleware/base.py
--- a/starlette/middleware/base.py
+++ b/starlette/middleware/base.py
@@ -4,12 +4,13 @@
 
 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 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
             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 @@ async def body_stream() -> typing.AsyncGenerator[bytes, None]:
             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
--- 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 @@ async def homepage(request):
     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"