Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Allow background tasks to run with custom BaseHTTPMiddleware's #1441

Closed
wants to merge 9 commits into from
Closed

Allow background tasks to run with custom BaseHTTPMiddleware's #1441

Changes from 3 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
149aa7c
Allow background tasks to run with custom `BaseHTTPMiddleware`'s
Kludex Jan 28, 2022
e161bae
Ignore coverage for now
Kludex Jan 28, 2022
e1af17b
Merge branch 'master' into fix/remove-cancel-scope
Kludex Jan 28, 2022
3011702
Merge branch 'master' into fix/remove-cancel-scope
Kludex Feb 8, 2022
20406e7
Merge branch 'master' into fix/remove-cancel-scope
Kludex Feb 19, 2022
8e6063d
Merge branch 'master' into fix/remove-cancel-scope
Kludex Apr 4, 2022
5843a57
Merge branch 'master' into fix/remove-cancel-scope
Kludex Apr 24, 2022
346f893
Merge branch 'master' into fix/remove-cancel-scope
Kludex May 8, 2022
470b31c
Merge branch 'master' into fix/remove-cancel-scope
adriangb May 12, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
17 changes: 12 additions & 5 deletions starlette/middleware/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,11 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
await self.app(scope, receive, send)
return

call_next_response = None
send_stream, recv_stream = anyio.create_memory_object_stream()

async def call_next(request: Request) -> Response:
app_exc: typing.Optional[Exception] = None
send_stream, recv_stream = anyio.create_memory_object_stream()

async def coro() -> None:
nonlocal app_exc
Expand Down Expand Up @@ -52,17 +54,22 @@ async def body_stream() -> typing.AsyncGenerator[bytes, None]:
assert message["type"] == "http.response.body"
yield message.get("body", b"")

response = StreamingResponse(
nonlocal call_next_response

call_next_response = StreamingResponse(
status_code=message["status"], content=body_stream()
)
response.raw_headers = message["headers"]
return response
call_next_response.raw_headers = message["headers"]
return call_next_response

async with anyio.create_task_group() as task_group:
request = Request(scope, receive=receive)
response = await self.dispatch_func(request, call_next)
if call_next_response and response is not call_next_response:
async with recv_stream:
async for _ in recv_stream:
... # pragma: no cover
Comment on lines +78 to +80
Copy link
Sponsor Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can create a function that performs this logic, if wanted.

await response(scope, receive, send)
Comment on lines +77 to 81
Copy link
Member

@jhominal jhominal Jan 30, 2022 •
edited

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems to me that we will wait for app to make all of its calls to send_stream.send before actually processing the response returned by dispatch_func. If that is the case, and given that all the messages from app are being discarded, I think it would be preferable, either to call await response(scope, receive, send) before draining recv_stream, or to put draining of recv_stream in another task of task_group.

Of course, I do not know the relevant technologies (ASGI/the server implementations/Starlette) as well as you do, so there could very well be a reason that either my analysis is wrong, or that my suggestions are unworkable.

Suggested change
if call_next_response and response is not call_next_response:
async with recv_stream:
async for _ in recv_stream:
... # pragma: no cover
await response(scope, receive, send)
if call_next_response and response is not call_next_response:
async def drain_stream():
async with recv_stream:
async for _ in recv_stream:
... # pragma: no cover
task_group.start_soon(drain_stream)
await response(scope, receive, send)

task_group.cancel_scope.cancel()

async def dispatch(
self, request: Request, call_next: RequestResponseEndpoint
Expand Down