fix: receive eof while closing reads stream

[Pulkit0110]

When writer.close() is called without setting finalize_on_close flag, we need to get two responses:

1. to get the persisted_size
2. eof response

That's why added a check if the first response is not eof, then again receive the response from the stream.

[gemini-code-assist]

Summary of Changes

Hello @Pulkit0110, 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 addresses a critical issue in the asynchronous write stream where the stream might not correctly process all responses from the server during the closing sequence, particularly when the finalize_on_close flag is not set. The changes ensure that both intermediate persisted_size responses and the final EOF signal are properly received and handled, preventing potential data inconsistencies or premature stream closure by explicitly waiting for the complete response sequence.

Highlights

• Enhanced Stream Closure Logic: The requests_done method in _AsyncWriteObjectStream has been updated to correctly handle scenarios where closing an asynchronous write stream might yield two responses from the server: an intermediate response containing persisted_size followed by an EOF signal. This ensures all necessary data is received.
• Improved EOF Detection: New logic was introduced to differentiate between an immediate EOF response and an intermediate response. If the first response contains persisted_size (and is not EOF), the system now explicitly waits for a second response, which is expected to be the final EOF.
• Comprehensive Unit Tests: Two new asynchronous unit tests (test_close_with_persisted_size_then_eof and test_close_with_eof_response) were added to validate the new stream closure behavior, covering both cases where an intermediate persisted_size is received before EOF and when EOF is received immediately.

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 addresses an issue with handling the end-of-file (EOF) signal when closing an asynchronous write stream. The changes correctly manage scenarios where an intermediate response containing persisted_size is received before the final EOF. The accompanying tests effectively validate this new logic. My review includes a suggestion to refactor a part of the implementation for improved clarity and maintainability, which aligns with general best practices and does not conflict with any specific repository rules.

[chandra-siri]

why do we want to call _utils.update_write_handle_if_exists(self, second_resp) for second_resp. it should be EOF. We may add assert second_resp == grpc.aio.EOF

[chandra-siri]

move this to top.