feat(experimental): add write resumption strategy

[Pulkit0110]

Adding writes resumption strategy which will be used for error handling of bidi writes operation.

[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 significantly enhances the reliability of asynchronous bidirectional streaming operations by introducing a comprehensive write resumption strategy and a generic retry manager. This refactoring centralizes error handling and retry logic, making both read and write operations more resilient to transient failures and network interruptions. The changes ensure that data transfers can seamlessly resume from the last known good state, improving the overall robustness of the system.

Highlights

• Write Resumption Strategy: Introduced a new _WriteResumptionStrategy to provide robust error handling and resumption capabilities for bidirectional streaming write operations.
• Generic Bidi Stream Retry Manager: Implemented a generic _BidiStreamRetryManager to centralize retry logic for all bidirectional streaming operations, making it reusable for both reads and writes.
• Refactored Read Resumption: The AsyncMultiRangeDownloader was refactored to utilize the new _BidiStreamRetryManager and an updated _ReadResumptionStrategy, moving checksum and offset validation into the strategy.
• Retry Policy Integration: The open and download_ranges methods in AsyncMultiRangeDownloader now accept an AsyncRetry policy, allowing callers to customize retry behavior.
• Improved Data Integrity Checks: Checksum and offset validation logic for read operations has been enhanced and moved into the _ReadResumptionStrategy, ensuring data integrity during stream resumption.

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 generic retry manager for bidi streams and applies it to read operations, which is a great improvement for resiliency. It also adds a resumption strategy for write operations. The overall structure is well-designed and the new tests are comprehensive. I've found one critical issue that will cause a TypeError at runtime due to a method being called with an unexpected argument. I've also included a couple of medium-severity comments to improve documentation clarity and maintainability in the new strategy classes.

[Pulkit0110]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a write resumption strategy for handling bidi write operations, which is a valuable addition for improving error handling in experimental async features. The implementation is well-structured and accompanied by a comprehensive set of unit tests that cover various scenarios, including initial uploads, resumption, and failure recovery. My feedback focuses on enhancing type safety, improving code clarity by addressing a potentially unused flag, and increasing the robustness of the tests by using mocks instead of None.

[gemini-code-assist]

The type hint for write_handle uses Any, which is too broad and reduces type safety. The docstring on line 34 correctly specifies the type as bytes | BidiWriteHandle | None. To improve clarity and type safety, please use the more specific storage_type.BidiWriteHandle instead of Any.

Suggested change

	self.write_handle: Union[bytes, Any, None] = None
	self.write_handle: Union[bytes, storage_type.BidiWriteHandle, None] = None

[gemini-code-assist]

The state['first_request'] flag is set to False here and in recover_state_on_failure, but its value is never read within the _WriteResumptionStrategy class. This makes its purpose unclear and could be confusing for future maintenance. If this flag is no longer used, it should be removed. If it is used by an external component, its purpose should be documented.

[gemini-code-assist]

Initializing _WriteState with None for arguments that expect objects (spec and user_buffer) makes this test brittle. If the implementation of update_state_from_response changes in the future to access attributes of spec or user_buffer, this test will fail with an AttributeError. It's safer to use mock objects to ensure the test is robust against such changes.

Suggested change

	"write_state": _WriteState(None, 0, None),
	"write_state": _WriteState(mock.Mock(spec=storage_type.AppendObjectSpec), 0, mock.Mock(spec=io.BytesIO)),