Drop response handler if request is canceled/dropped

[imabdulbasit]

When a client disconnects before a response is sent, the handler future previously ran to completion. This was problematic for handlers that acquire resources such as database. This would mean that we are unnecessarily causing load on database, and in some cases can not serve new requests.

This PR adds client-disconnect detection in the listener and cancels/drops the response handler when the socket closes.

[gemini-code-assist]

Summary of Changes

Hello @imabdulbasit, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the RateLimitListener to gracefully handle client disconnections. Previously, if a client disconnected while a request was being processed, the handler might continue to run to completion, wasting resources. The changes introduce a mechanism to detect such disconnections and immediately terminate the handler, ensuring efficient resource utilization and preventing stale computations.

Highlights

• Client Disconnection Handling: Implemented logic to detect client disconnections during request processing and immediately drop the associated handler future, preventing unnecessary resource consumption.
• Asynchronous Stream Monitoring: Introduced a new asynchronous function wait_for_disconnect that uses TcpStream::peek to non-invasively monitor the client connection for closure or errors.
• Concurrent Future Selection: Utilized futures::select! macro within the request handling loop to concurrently await either the completion of the application's response or the client's disconnection.
• New Test Case: Added test_handler_dropped_on_client_disconnect to verify that handler futures are correctly cancelled and dropped when a client prematurely closes the connection.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• src/listener.rs
  • Imported TcpStream, select, and FutureExt for advanced asynchronous operations.
  • Modified the async_h1::accept callback to use select! for concurrent monitoring of handler completion and client disconnects.
  • Added a new wait_for_disconnect async function to detect client disconnections using TcpStream::peek.
  • Included AtomicBool and Arc imports for the new test.
  • Added test_handler_dropped_on_client_disconnect to validate the new disconnection handling.

Activity

• No specific activity has been recorded for this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to cancel request handlers when the client disconnects. This is achieved by racing the handler future against a new wait_for_disconnect future that polls the TCP stream for closure. The implementation is sound and includes a new test case to verify the behavior. I have a couple of suggestions to improve code quality and maintainability in the new wait_for_disconnect function and its corresponding test.

[sveitser]

LGTM

[jbearer]

IIUC this would be a breaking change because it requires that handler futures are cancel safe, which wouldn't have been required before. Are we confident this isn't gonna break anything downstream?

[jbearer]

Does an error in peek necessarily mean a client disconnect? If it's possible that future reads from the socket would succeed even after this failure (ie a transient failure), then we should not disconnect here

[imabdulbasit]

I am not sure, but I can match the exact error types

[jbearer]

So we're only doing this for RateLimitListener. Do you know if the standard listener (whatever you get if you just pass in a &str like "0.0.0.0:{port}") already has this behavior? If not, it might be better to do this as a mix-in adapter, that we could use like HandleDisconnectListener<RateLimitListener<State>> or HandleDisconnectListener<StandardListener>

[imabdulbasit]

yeah we have the same issue with standard listener, but the confirmation layer uses RateLimitListener, so I thought to make the change to RateLimitListener only.

[imabdulbasit]

IIUC this would be a breaking change because it requires that handler futures are cancel safe, which wouldn't have been required before. Are we confident this isn't gonna break anything downstream?

I tried compiling the confirmation layer and it worked without any issues.

[jbearer]

It's not a compilation issue. Things can go wrong at runtime if futures that aren't cancel-safe are cancelled (e.g. some futures don't/can't clean up all their resources properly in Drop, and need to be driven to completion). Here's an example