Skip to content
This repository was archived by the owner on Mar 31, 2026. It is now read-only.

feat(experimental): add write resumption strategy #1663

Merged
Pulkit0110 merged 5 commits into from
Jan 2, 2026
Merged

feat(experimental): add write resumption strategy #1663

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
cc7fb6f
integrate retry logic with the MRD
Pulkit0110 Dec 8, 2025
f1f8d65
feat(experimental): add write resumption strategy
Pulkit0110 Dec 15, 2025
9a086e4
address gemini bot comments
Pulkit0110 Dec 29, 2025
7abfae0
addressing comments:
Pulkit0110 Dec 30, 2025
a3f6254
resolve comments
Pulkit0110 Jan 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
resolve comments
  • Loading branch information
@Pulkit0110
Pulkit0110 committed Jan 2, 2026
commit a3f6254f992ad1392366b80935e90ca851632d2f
35 changes: 15 additions & 20 deletions google/cloud/storage/_experimental/asyncio/retry/writes_resumption_strategy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ class _WriteState:

def __init__(
self,
spec: storage_type.AppendObjectSpec,
spec: Union[storage_type.AppendObjectSpec, storage_type.WriteObjectSpec],
chunk_size: int,
user_buffer: IO[bytes],
):
Expand All @@ -48,7 +48,7 @@ def __init__(
self.bytes_sent: int = 0
self.write_handle: Union[bytes, storage_type.BidiWriteHandle, None] = None
self.routing_token: Optional[str] = None
self.is_complete: bool = False
self.is_finalized: bool = False


class _WriteResumptionStrategy(_BaseResumptionStrategy):
Expand All @@ -61,34 +61,29 @@ def generate_requests(

For Appendable Objects, every stream opening should send an
AppendObjectSpec. If resuming, the `write_handle` is added to that spec.

This method is not applicable for `open` methods.
"""
Comment thread
Pulkit0110 marked this conversation as resolved.
write_state: _WriteState = state["write_state"]

if write_state.write_handle:
write_state.spec.write_handle = write_state.write_handle

if write_state.routing_token:
write_state.spec.routing_token = write_state.routing_token

# Initial request of the stream must provide the specification.
# If we have a write_handle, we request a state lookup to verify persisted offset.
do_state_lookup = write_state.write_handle is not None
initial_request = storage_type.BidiWriteObjectRequest()

# Determine if we need to send WriteObjectSpec or AppendObjectSpec
Comment thread
Pulkit0110 marked this conversation as resolved.
initial_request = storage_type.BidiWriteObjectRequest(
state_lookup=do_state_lookup
)

if isinstance(write_state.spec, storage_type.WriteObjectSpec):
initial_request.write_object_spec = write_state.spec
else:
if write_state.write_handle:
write_state.spec.write_handle = write_state.write_handle

if write_state.routing_token:
write_state.spec.routing_token = write_state.routing_token
initial_request.append_object_spec = write_state.spec

yield initial_request
Comment thread
Pulkit0110 marked this conversation as resolved.

# The buffer should already be seeked to the correct position (persisted_size)
# by the `recover_state_on_failure` method before this is called.
while not write_state.is_complete:
while not write_state.is_finalized:
chunk = write_state.user_buffer.read(write_state.chunk_size)

# End of File detection
Expand All @@ -113,16 +108,16 @@ def update_state_from_response(
"""Processes a server response and updates the write state."""
write_state: _WriteState = state["write_state"]

if response.persisted_size is not None:
if response.persisted_size > write_state.persisted_size:
write_state.persisted_size = response.persisted_size
if response.persisted_size:
write_state.persisted_size = response.persisted_size

if response.write_handle:
write_state.write_handle = response.write_handle

if response.resource:
write_state.is_complete = True
write_state.persisted_size = response.resource.size
if response.resource.finalize_time:
write_state.is_finalized = True

async def recover_state_on_failure(
Comment thread
Pulkit0110 marked this conversation as resolved.
self, error: Exception, state: Dict[str, Any]
Expand Down
23 changes: 3 additions & 20 deletions tests/unit/asyncio/retry/test_writes_resumption_strategy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import io
import unittest
import unittest.mock as mock
from datetime import datetime

import pytest
import google_crc32c
Expand Down Expand Up @@ -137,7 +138,6 @@ def test_generate_requests_resumption(self):
self.assertEqual(
requests[0].append_object_spec.write_handle.handle, b"test-handle"
)
self.assertTrue(requests[0].state_lookup)

# Check data starts from offset 4
self.assertEqual(requests[1].write_offset, 4)
Expand Down Expand Up @@ -206,28 +206,11 @@ def test_update_state_from_response(self):
strategy.update_state_from_response(response2, state)
self.assertEqual(write_state.persisted_size, 1024)

final_resource = storage_type.Object(name="test-object", bucket="b", size=2048)
final_resource = storage_type.Object(name="test-object", bucket="b", size=2048, finalize_time=datetime.now())
response3 = storage_type.BidiWriteObjectResponse(resource=final_resource)
strategy.update_state_from_response(response3, state)
self.assertTrue(write_state.is_complete)
self.assertEqual(write_state.persisted_size, 2048)

def test_update_state_from_response_ignores_smaller_persisted_size(self):
strategy = self._make_one()
state = {
"write_state": _WriteState(
mock.Mock(spec=storage_type.AppendObjectSpec),
0,
mock.Mock(spec=io.BytesIO),
),
}
write_state = state["write_state"]
write_state.persisted_size = 2048

response = storage_type.BidiWriteObjectResponse(persisted_size=1024)
strategy.update_state_from_response(response, state)

self.assertEqual(write_state.persisted_size, 2048)
self.assertTrue(write_state.is_finalized)

@pytest.mark.asyncio
async def test_recover_state_on_failure_handles_redirect(self):
Expand Down