From 02aa6eb9588f64c0be331ea118598fdc7d13b6d5 Mon Sep 17 00:00:00 2001
From: Weiliang Li <to.be.impressive@gmail.com>
Date: Wed, 22 Jun 2022 13:40:16 +0900
Subject: [PATCH] Shield bg task in Response

---
 starlette/background.py       | 11 ++++-------
 starlette/middleware/base.py  |  9 ++-------
 starlette/responses.py        | 16 ++++++++++------
 tests/middleware/test_base.py |  8 ++++++++
 4 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/starlette/background.py b/starlette/background.py
index db9b38af8..4aaf7ae3c 100644
--- a/starlette/background.py
+++ b/starlette/background.py
@@ -1,8 +1,6 @@
 import sys
 import typing
 
-import anyio
-
 if sys.version_info >= (3, 10):  # pragma: no cover
     from typing import ParamSpec
 else:  # pragma: no cover
@@ -24,11 +22,10 @@ def __init__(
         self.is_async = is_async_callable(func)
 
     async def __call__(self) -> None:
-        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)
+        if self.is_async:
+            await self.func(*self.args, **self.kwargs)
+        else:
+            await run_in_threadpool(self.func, *self.args, **self.kwargs)
 
 
 class BackgroundTasks(BackgroundTask):
diff --git a/starlette/middleware/base.py b/starlette/middleware/base.py
index b77bb41f9..2ea98a4bc 100644
--- a/starlette/middleware/base.py
+++ b/starlette/middleware/base.py
@@ -33,7 +33,7 @@ async def coro() -> None:
 
                 async with send_stream:
                     try:
-                        # may block in send if body_stream not consumed
+                        # `send_stream.send` blocks until `body_stream` is called
                         await self.app(scope, request.receive, send_stream.send)
                     except Exception as exc:
                         app_exc = exc
@@ -72,12 +72,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)
-
-            t = anyio.get_current_task()
-            if t.name == "anyio.from_thread.BlockingPortal._call_func":
-                # cancel stuck task due to discarded response
-                # see: https://github.com/encode/starlette/issues/1022
-                task_group.cancel_scope.cancel()
+            task_group.cancel_scope.cancel()
 
     async def dispatch(
         self, request: Request, call_next: RequestResponseEndpoint
diff --git a/starlette/responses.py b/starlette/responses.py
index 98c8caf1b..c15fd2848 100644
--- a/starlette/responses.py
+++ b/starlette/responses.py
@@ -166,8 +166,9 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         )
         await send({"type": "http.response.body", "body": self.body})
 
-        if self.background is not None:
-            await self.background()
+        with anyio.CancelScope(shield=True):
+            if self.background is not None:
+                await self.background()
 
 
 class HTMLResponse(Response):
@@ -264,8 +265,9 @@ async def wrap(func: typing.Callable[[], typing.Coroutine]) -> None:
             task_group.start_soon(wrap, partial(self.stream_response, send))
             await wrap(partial(self.listen_for_disconnect, receive))
 
-        if self.background is not None:
-            await self.background()
+        with anyio.CancelScope(shield=True):
+            if self.background is not None:
+                await self.background()
 
 
 class FileResponse(Response):
@@ -350,5 +352,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
                             "more_body": more_body,
                         }
                     )
-        if self.background is not None:
-            await self.background()
+
+        with anyio.CancelScope(shield=True):
+            if self.background is not None:
+                await self.background()
diff --git a/tests/middleware/test_base.py b/tests/middleware/test_base.py
index 98ce5e8bf..c7b65e8e4 100644
--- a/tests/middleware/test_base.py
+++ b/tests/middleware/test_base.py
@@ -92,10 +92,18 @@ async def _sleep(identifier, delay):
         await anyio.sleep(delay)
         print(identifier, "completed")
 
+    def _sleep_sync(identifier, delay):
+        import time
+
+        print(identifier, "started")
+        time.sleep(delay)
+        print(identifier, "completed")
+
     async def bg_task(request):
         background_tasks = BackgroundTasks()
         background_tasks.add_task(_sleep, "background task 1", 2)
         background_tasks.add_task(_sleep, "background task 2", 2)
+        background_tasks.add_task(_sleep_sync, "background task sync", 2)
         return Response(background=background_tasks)
 
     app = Starlette(