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
Next Next commit
addressing comments:
  • Loading branch information
@Pulkit0110
Pulkit0110 committed Jan 2, 2026
commit 7abfae05208ff1a1040a6c59dd94f788ab3957ee
53 changes: 32 additions & 21 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 @@ -25,15 +25,14 @@
class _WriteState:
"""A helper class to track the state of a single upload operation.

Attributes:
spec (AppendObjectSpec): The specification for the object to write.
chunk_size (int): The size of chunks to read from the buffer.
user_buffer (IO[bytes]): The data source.
persisted_size (int): The amount of data confirmed as persisted by the server.
bytes_sent (int): The amount of data currently sent in the active stream.
write_handle (bytes | BidiWriteHandle | None): The handle for the append session.
routing_token (str | None): Token for routing to the correct backend.
is_complete (bool): Whether the upload has finished.
:type spec: :class:`google.cloud.storage_v2.types.AppendObjectSpec`
:param spec: The specification for the object to write.

:type chunk_size: int
:param chunk_size: The size of chunks to write to the server.

:type user_buffer: IO[bytes]
:param user_buffer: The data source.
"""

def __init__(
Expand Down Expand Up @@ -71,23 +70,29 @@ def generate_requests(
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
Comment thread
Pulkit0110 marked this conversation as resolved.
Outdated
Comment thread
Pulkit0110 marked this conversation as resolved.
Outdated
yield storage_type.BidiWriteObjectRequest(
append_object_spec=write_state.spec, state_lookup=do_state_lookup

# 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:
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:
chunk = write_state.user_buffer.read(write_state.chunk_size)

# End of File detection
if not chunk:
write_state.is_complete = True
yield storage_type.BidiWriteObjectRequest(
write_offset=write_state.bytes_sent,
finish_write=True,
)
return

checksummed_data = storage_type.ChecksummedData(content=chunk)
Expand Down Expand Up @@ -122,19 +127,25 @@ def update_state_from_response(
async def recover_state_on_failure(
Comment thread
Pulkit0110 marked this conversation as resolved.
self, error: Exception, state: Dict[str, Any]
) -> None:
"""Handles errors, specifically BidiWriteObjectRedirectedError, and rewinds state."""
"""
Handles errors, specifically BidiWriteObjectRedirectedError, and rewinds state.

This method rewinds the user buffer and internal byte tracking to the
last confirmed 'persisted_size' from the server.
"""
write_state: _WriteState = state["write_state"]
cause = getattr(error, "cause", error)

# Extract routing token and potentially a new write handle.
# Extract routing token and potentially a new write handle for redirection.
if isinstance(cause, BidiWriteObjectRedirectedError):
if cause.routing_token:
write_state.routing_token = cause.routing_token

if hasattr(cause, "write_handle") and cause.write_handle:
write_state.write_handle = cause.write_handle
redirect_handle = getattr(cause, "write_handle", None)
if redirect_handle:
Comment thread
chandra-siri marked this conversation as resolved.
write_state.write_handle = redirect_handle

# We must assume any data sent beyond 'persisted_size' was lost.
Comment thread
chandra-siri marked this conversation as resolved.
# Reset the user buffer to the last known good byte.
# Reset the user buffer to the last known good byte confirmed by the server.
write_state.user_buffer.seek(write_state.persisted_size)
write_state.bytes_sent = write_state.persisted_size
133 changes: 53 additions & 80 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 @@ -38,72 +38,86 @@ def test_ctor(self):
strategy = self._make_one()
self.assertIsInstance(strategy, self._get_target_class())

def test_generate_requests_initial(self):
def test_generate_requests_initial_new_object(self):
"""
Verify the initial request sequence for a new upload.
- First request is AppendObjectSpec with state_lookup=False.
- Following requests are data chunks.
- Final request has finish_write=True.
Verify the initial request sequence for a new upload (WriteObjectSpec).
"""
strategy = self._make_one()
mock_buffer = io.BytesIO(b"0123456789")
mock_spec = storage_type.AppendObjectSpec(object_="test-object")
# Use WriteObjectSpec for new objects
mock_spec = storage_type.WriteObjectSpec(
resource=storage_type.Object(name="test-object")
)
state = {
"write_state": _WriteState(
mock_spec, chunk_size=4, user_buffer=mock_buffer
),
"first_request": True,
}

requests = list(strategy.generate_requests(state))

self.assertEqual(requests[0].append_object_spec, mock_spec)
# Check first request (Spec)
self.assertEqual(requests[0].write_object_spec, mock_spec)
self.assertFalse(requests[0].state_lookup)
self.assertFalse(requests[0].append_object_spec.write_handle)

# Check data chunks
self.assertEqual(requests[1].write_offset, 0)
self.assertEqual(requests[1].checksummed_data.content, b"0123")
self.assertEqual(requests[2].write_offset, 4)
self.assertEqual(requests[2].checksummed_data.content, b"4567")
self.assertEqual(requests[3].write_offset, 8)
self.assertEqual(requests[3].checksummed_data.content, b"89")

self.assertEqual(requests[4].write_offset, 10)
self.assertTrue(requests[4].finish_write)

self.assertEqual(len(requests), 5)
# Total requests: 1 Spec + 3 Chunks
self.assertEqual(len(requests), 4)

def test_generate_requests_empty_file(self):
def test_generate_requests_initial_existing_object(self):
"""
Verify the request sequence for an empty file upload.
- First request is AppendObjectSpec.
- Second and final request has finish_write=True.
Verify the initial request sequence for appending to an existing object (AppendObjectSpec).
"""
strategy = self._make_one()
mock_buffer = io.BytesIO(b"")
mock_spec = storage_type.AppendObjectSpec(object_="test-object")
mock_buffer = io.BytesIO(b"0123")
# Use AppendObjectSpec for existing objects
mock_spec = storage_type.AppendObjectSpec(
object_="test-object", bucket="test-bucket"
)
state = {
"write_state": _WriteState(
mock_spec, chunk_size=4, user_buffer=mock_buffer
),
"first_request": True,
}

requests = list(strategy.generate_requests(state))

# Check first request (Spec)
self.assertEqual(requests[0].append_object_spec, mock_spec)
self.assertFalse(requests[0].state_lookup)

# Check data chunk
self.assertEqual(requests[1].write_offset, 0)
self.assertTrue(requests[1].finish_write)
self.assertEqual(requests[1].checksummed_data.content, b"0123")

def test_generate_requests_empty_file(self):
"""
Verify request sequence for an empty file. Should just be the Spec.
"""
strategy = self._make_one()
mock_buffer = io.BytesIO(b"")
mock_spec = storage_type.AppendObjectSpec(object_="test-object")
state = {
"write_state": _WriteState(
mock_spec, chunk_size=4, user_buffer=mock_buffer
),
}

requests = list(strategy.generate_requests(state))

self.assertEqual(len(requests), 2)
self.assertEqual(len(requests), 1)
self.assertEqual(requests[0].append_object_spec, mock_spec)

def test_generate_requests_resumption(self):
"""
Verify request sequence when resuming an upload.
- First request is AppendObjectSpec with write_handle and state_lookup=True.
- Data streaming starts from the persisted_size.
"""
strategy = self._make_one()
mock_buffer = io.BytesIO(b"0123456789")
Expand All @@ -115,34 +129,26 @@ def test_generate_requests_resumption(self):
write_state.write_handle = storage_type.BidiWriteHandle(handle=b"test-handle")
mock_buffer.seek(4)

state = {"write_state": write_state, "first_request": True}
state = {"write_state": write_state}

requests = list(strategy.generate_requests(state))

# Check first request has handle and lookup
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)
self.assertEqual(requests[1].checksummed_data.content, b"4567")
self.assertEqual(requests[2].write_offset, 8)
self.assertEqual(requests[2].checksummed_data.content, b"89")

self.assertEqual(requests[3].write_offset, 10)
self.assertTrue(requests[3].finish_write)

self.assertEqual(len(requests), 4)

@pytest.mark.asyncio
async def test_generate_requests_after_failure_and_recovery(self):
"""
Verify a complex scenario:
1. Start upload.
2. Receive a persisted_size update.
3. Encounter an error.
4. Recover state.
5. Generate new requests for resumption.
Verify recovery and resumption flow.
"""
strategy = self._make_one()
mock_buffer = io.BytesIO(b"0123456789abcdef")
Expand All @@ -157,12 +163,11 @@ async def test_generate_requests_after_failure_and_recovery(self):

strategy.update_state_from_response(
storage_type.BidiWriteObjectResponse(
persisted_size=4, write_handle=b"handle-1"
persisted_size=4,
write_handle=storage_type.BidiWriteHandle(handle=b"handle-1"),
),
state,
)
self.assertEqual(write_state.persisted_size, 4)
self.assertEqual(write_state.write_handle, b"handle-1")

await strategy.recover_state_on_failure(Exception("network error"), state)

Expand All @@ -172,23 +177,15 @@ async def test_generate_requests_after_failure_and_recovery(self):
requests = list(strategy.generate_requests(state))

self.assertTrue(requests[0].state_lookup)
self.assertEqual(requests[0].append_object_spec.write_handle, b"handle-1")
self.assertEqual(
requests[0].append_object_spec.write_handle.handle, b"handle-1"
)

self.assertEqual(requests[1].write_offset, 4)
self.assertEqual(requests[1].checksummed_data.content, b"4567")
self.assertEqual(requests[2].write_offset, 8)
self.assertEqual(requests[2].checksummed_data.content, b"89ab")
self.assertEqual(requests[3].write_offset, 12)
self.assertEqual(requests[3].checksummed_data.content, b"cdef")

self.assertEqual(requests[4].write_offset, 16)
self.assertTrue(requests[4].finish_write)
self.assertEqual(len(requests), 5)

def test_update_state_from_response(self):
"""
Verify that the write state is correctly updated based on server responses.
"""
"""Verify state updates from server responses."""
strategy = self._make_one()
mock_buffer = io.BytesIO(b"0123456789")
mock_spec = storage_type.AppendObjectSpec(object_="test-object")
Expand Down Expand Up @@ -232,30 +229,10 @@ def test_update_state_from_response_ignores_smaller_persisted_size(self):

self.assertEqual(write_state.persisted_size, 2048)

@pytest.mark.asyncio
async def test_recover_state_on_failure_rewinds_state(self):
"""
Verify that on failure, the buffer is seeked to persisted_size
and bytes_sent is reset.
"""
strategy = self._make_one()
mock_buffer = mock.MagicMock(spec=io.BytesIO)
mock_spec = storage_type.AppendObjectSpec(object_="test-object")

write_state = _WriteState(mock_spec, chunk_size=4, user_buffer=mock_buffer)
write_state.persisted_size = 100
write_state.bytes_sent = 200
state = {"write_state": write_state}

await strategy.recover_state_on_failure(Exception("any error"), state)

mock_buffer.seek.assert_called_once_with(100)
self.assertEqual(write_state.bytes_sent, 100)

@pytest.mark.asyncio
async def test_recover_state_on_failure_handles_redirect(self):
"""
Verify that on a redirect error, the routing_token is extracted and stored.
Verify redirection error handling.
"""
strategy = self._make_one()
mock_buffer = mock.MagicMock(spec=io.BytesIO)
Expand All @@ -271,11 +248,11 @@ async def test_recover_state_on_failure_handles_redirect(self):
await strategy.recover_state_on_failure(wrapped_error, state)

self.assertEqual(write_state.routing_token, "new-token-123")
mock_buffer.seek.assert_called_once_with(0)
self.assertEqual(write_state.bytes_sent, 0)
mock_buffer.seek.assert_called_with(0)

@pytest.mark.asyncio
async def test_recover_state_on_failure_handles_redirect_with_handle(self):
"""Verify redirection that includes a write handle."""
strategy = self._make_one()
mock_buffer = mock.MagicMock(spec=io.BytesIO)
mock_spec = storage_type.AppendObjectSpec(object_="test-object")
Expand All @@ -294,8 +271,7 @@ async def test_recover_state_on_failure_handles_redirect_with_handle(self):
self.assertEqual(write_state.routing_token, "new-token-456")
self.assertEqual(write_state.write_handle, b"redirect-handle")

mock_buffer.seek.assert_called_once_with(0)
self.assertEqual(write_state.bytes_sent, 0)
mock_buffer.seek.assert_called_with(0)

def test_generate_requests_sends_crc32c_checksum(self):
strategy = self._make_one()
Expand All @@ -305,13 +281,10 @@ def test_generate_requests_sends_crc32c_checksum(self):
"write_state": _WriteState(
mock_spec, chunk_size=4, user_buffer=mock_buffer
),
"first_request": True,
}

requests = list(strategy.generate_requests(state))

self.assertEqual(len(requests), 3)

expected_crc = google_crc32c.Checksum(b"0123")
expected_crc_int = int.from_bytes(expected_crc.digest(), "big")
self.assertEqual(requests[1].checksummed_data.crc32c, expected_crc_int)
Expand All @@ -323,7 +296,7 @@ def test_generate_requests_with_routing_token(self):

write_state = _WriteState(mock_spec, chunk_size=4, user_buffer=mock_buffer)
write_state.routing_token = "redirected-token"
state = {"write_state": write_state, "first_request": True}
state = {"write_state": write_state}

requests = list(strategy.generate_requests(state))

Expand Down