feat(experimental): flush the last chunk in append method (#1699)

Earlier the last chunk was being flushed while calling the close()
method. Now it will be done inside the append method itself.

Author: Pulkit0110

google/cloud/storage/_experimental/asyncio/async_appendable_object_writer.py

@@ -188,8 +188,9 @@ async def append(self, data: bytes) -> None:
  ie. `self.offset` bytes relative to the begining of the object.
 
  This method sends the provided `data` to the GCS server in chunks.
- and persists data in GCS at every `_MAX_BUFFER_SIZE_BYTES` bytes by
- calling `self.simple_flush`.
+ and persists data in GCS at every `_DEFAULT_FLUSH_INTERVAL_BYTES` bytes
+ or at the last chunk whichever is earlier. Persisting is done by setting
+ `flush=True` on request.
 
  :type data: bytes
  :param data: The bytes to append to the object.
@@ -214,20 +215,33 @@ async def append(self, data: bytes) -> None:
 whilestart_idx<total_bytes:
 end_idx=min(start_idx+_MAX_CHUNK_SIZE_BYTES, total_bytes)
 data_chunk=data[start_idx:end_idx]
+is_last_chunk=end_idx==total_bytes
+chunk_size=end_idx-start_idx
 awaitself.write_obj_stream.send(
 _storage_v2.BidiWriteObjectRequest(
 write_offset=self.offset,
 checksummed_data=_storage_v2.ChecksummedData(
 content=data_chunk,
 crc32c=int.from_bytes(Checksum(data_chunk).digest(), "big"),
  ),
+state_lookup=is_last_chunk,
+flush=is_last_chunk
+or (
+self.bytes_appended_since_last_flush+chunk_size
+>=self.flush_interval
+ ),
  )
  )
-chunk_size=end_idx-start_idx
 self.offset+=chunk_size
 self.bytes_appended_since_last_flush+=chunk_size
+
 ifself.bytes_appended_since_last_flush>=self.flush_interval:
-awaitself.simple_flush()
+self.bytes_appended_since_last_flush=0
+
+ifis_last_chunk:
+response=awaitself.write_obj_stream.recv()
+self.persisted_size=response.persisted_size
+self.offset=self.persisted_size
 self.bytes_appended_since_last_flush=0
 start_idx=end_idx
 
@@ -292,20 +306,24 @@ async def close(self, finalize_on_close=False) -> Union[int, _storage_v2.Object]
 raiseValueError("Stream is not open. Call open() before close().")
 
 iffinalize_on_close:
-awaitself.finalize()
-else:
-awaitself.flush()
+returnawaitself.finalize()
 
 awaitself.write_obj_stream.close()
 
 self._is_stream_open=False
 self.offset=None
-returnself.object_resourceiffinalize_on_closeelseself.persisted_size
+returnself.persisted_size
 
 asyncdeffinalize(self) ->_storage_v2.Object:
 """Finalizes the Appendable Object.
 
  Note: Once finalized no more data can be appended.
+ This method is different from `close`. if `.close()` is called data may
+ still be appended to object at a later point in time by opening with
+ generation number.
+ (i.e. `open(..., generation=<object_generation_number>)`.
+ However if `.finalize()` is called no more data can be appended to the
+ object.
 
  rtype: google.cloud.storage_v2.types.Object
  returns: The finalized object resource.
@@ -322,6 +340,10 @@ async def finalize(self) -> _storage_v2.Object:
 response=awaitself.write_obj_stream.recv()
 self.object_resource=response.resource
 self.persisted_size=self.object_resource.size
+awaitself.write_obj_stream.close()
+
+self._is_stream_open=False
+self.offset=None
 returnself.object_resource
 
 # helper methods.

noxfile.py

@@ -192,7 +192,14 @@ def system(session):
 # 2021-05-06: defer installing 'google-cloud-*' to after this package,
 # in order to work around Python 2.7 googolapis-common-protos
 # issue.
-session.install("mock", "pytest", "pytest-rerunfailures", "-c", constraints_path)
+session.install(
+"mock",
+"pytest",
+"pytest-rerunfailures",
+"pytest-asyncio",
+"-c",
+constraints_path,
+ )
 session.install("-e", ".", "-c", constraints_path)
 session.install(
 "google-cloud-testutils",

tests/system/test_zonal.py

@@ -307,3 +307,56 @@ async def _read_and_verify(expected_content, generation=None):
 delmrd
 delmrd_2
 gc.collect()
+
+
+@pytest.mark.asyncio
+asyncdeftest_append_flushes_and_state_lookup(storage_client, blobs_to_delete):
+"""
+ System test for AsyncAppendableObjectWriter, verifying flushing behavior
+ for both small and large appends.
+ """
+object_name=f"test-append-flush-varied-size-{uuid.uuid4()}"
+grpc_client=AsyncGrpcClient().grpc_client
+writer=AsyncAppendableObjectWriter(grpc_client, _ZONAL_BUCKET, object_name)
+
+# Schedule for cleanup
+blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))
+
+# --- Part 1: Test with small data ---
+small_data=b"small data"
+
+awaitwriter.open()
+assertwriter._is_stream_open
+
+awaitwriter.append(small_data)
+persisted_size=awaitwriter.state_lookup()
+assertpersisted_size==len(small_data)
+
+# --- Part 2: Test with large data ---
+large_data=os.urandom(38*1024*1024)
+
+# Append data larger than the default flush interval (16 MiB).
+# This should trigger the interval-based flushing logic.
+awaitwriter.append(large_data)
+
+# Verify the total data has been persisted.
+total_size=len(small_data) +len(large_data)
+persisted_size=awaitwriter.state_lookup()
+assertpersisted_size==total_size
+
+# --- Part 3: Finalize and verify ---
+final_object=awaitwriter.close(finalize_on_close=True)
+
+assertnotwriter._is_stream_open
+assertfinal_object.size==total_size
+
+# Verify the full content of the object.
+full_data=small_data+large_data
+mrd=AsyncMultiRangeDownloader(grpc_client, _ZONAL_BUCKET, object_name)
+buffer=BytesIO()
+awaitmrd.open()
+# (0, 0) means read the whole object
+awaitmrd.download_ranges([(0, 0, buffer)])
+awaitmrd.close()
+content=buffer.getvalue()
+assertcontent==full_data

tests/unit/asyncio/retry/test_writes_resumption_strategy.py

@@ -206,7 +206,9 @@ 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, finalize_time=datetime.now())
+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.assertEqual(write_state.persisted_size, 2048)

tests/unit/asyncio/test_async_appendable_object_writer.py

@@ -364,6 +364,7 @@ async def test_close(mock_write_object_stream, mock_client):
 writer=AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT)
 writer._is_stream_open=True
 writer.offset=1024
+writer.persisted_size=1024
 mock_stream=mock_write_object_stream.return_value
 mock_stream.send=mock.AsyncMock()
 mock_stream.recv=mock.AsyncMock(
@@ -435,16 +436,20 @@ async def test_finalize(mock_write_object_stream, mock_client):
 mock_stream.recv=mock.AsyncMock(
 return_value=_storage_v2.BidiWriteObjectResponse(resource=mock_resource)
  )
+mock_stream.close=mock.AsyncMock()
 
 gcs_object=awaitwriter.finalize()
 
 mock_stream.send.assert_awaited_once_with(
 _storage_v2.BidiWriteObjectRequest(finish_write=True)
  )
 mock_stream.recv.assert_awaited_once()
+mock_stream.close.assert_awaited_once()
 assertwriter.object_resource==mock_resource
 assertwriter.persisted_size==123
 assertgcs_object==mock_resource
+assertwriter._is_stream_openisFalse
+assertwriter.offsetisNone
 
 
 @pytest.mark.asyncio
@@ -501,30 +506,39 @@ async def test_append_sends_data_in_chunks(mock_write_object_stream, mock_client
 writer.persisted_size=100
 mock_stream=mock_write_object_stream.return_value
 mock_stream.send=mock.AsyncMock()
-writer.simple_flush=mock.AsyncMock()
 
 data=b"a"* (_MAX_CHUNK_SIZE_BYTES+1)
+mock_stream.recv=mock.AsyncMock(
+return_value=_storage_v2.BidiWriteObjectResponse(
+persisted_size=100+len(data)
+ )
+ )
+
 awaitwriter.append(data)
 
 assertmock_stream.send.await_count==2
-first_call=mock_stream.send.await_args_list[0]
-second_call=mock_stream.send.await_args_list[1]
+first_request=mock_stream.send.await_args_list[0].args[0]
+second_request=mock_stream.send.await_args_list[1].args[0]
 
 # First chunk
-assertfirst_call[0][0].write_offset==100
-assertlen(first_call[0][0].checksummed_data.content) ==_MAX_CHUNK_SIZE_BYTES
-assertfirst_call[0][0].checksummed_data.crc32c==int.from_bytes(
+assertfirst_request.write_offset==100
+assertlen(first_request.checksummed_data.content) ==_MAX_CHUNK_SIZE_BYTES
+assertfirst_request.checksummed_data.crc32c==int.from_bytes(
 Checksum(data[:_MAX_CHUNK_SIZE_BYTES]).digest(), byteorder="big"
  )
-# Second chunk
-assertsecond_call[0][0].write_offset==100+_MAX_CHUNK_SIZE_BYTES
-assertlen(second_call[0][0].checksummed_data.content) ==1
-assertsecond_call[0][0].checksummed_data.crc32c==int.from_bytes(
+assertnotfirst_request.flush
+assertnotfirst_request.state_lookup
+
+# Second chunk (last chunk)
+assertsecond_request.write_offset==100+_MAX_CHUNK_SIZE_BYTES
+assertlen(second_request.checksummed_data.content) ==1
+assertsecond_request.checksummed_data.crc32c==int.from_bytes(
 Checksum(data[_MAX_CHUNK_SIZE_BYTES:]).digest(), byteorder="big"
  )
+assertsecond_request.flush
+assertsecond_request.state_lookup
 
 assertwriter.offset==100+len(data)
-writer.simple_flush.assert_not_awaited()
 
 
 @pytest.mark.asyncio
@@ -541,12 +555,25 @@ async def test_append_flushes_when_buffer_is_full(
 writer.persisted_size=0
 mock_stream=mock_write_object_stream.return_value
 mock_stream.send=mock.AsyncMock()
-writer.simple_flush=mock.AsyncMock()
+mock_stream.recv=mock.AsyncMock()
 
 data=b"a"*_DEFAULT_FLUSH_INTERVAL_BYTES
 awaitwriter.append(data)
 
-writer.simple_flush.assert_awaited_once()
+num_chunks=_DEFAULT_FLUSH_INTERVAL_BYTES//_MAX_CHUNK_SIZE_BYTES
+assertmock_stream.send.await_count==num_chunks
+
+# All but the last request should not have flush or state_lookup set.
+foriinrange(num_chunks-1):
+request=mock_stream.send.await_args_list[i].args[0]
+assertnotrequest.flush
+assertnotrequest.state_lookup
+
+# The last request should have flush and state_lookup set.
+last_request=mock_stream.send.await_args_list[-1].args[0]
+assertlast_request.flush
+assertlast_request.state_lookup
+assertwriter.bytes_appended_since_last_flush==0
 
 
 @pytest.mark.asyncio
@@ -561,12 +588,18 @@ async def test_append_handles_large_data(mock_write_object_stream, mock_client):
 writer.persisted_size=0
 mock_stream=mock_write_object_stream.return_value
 mock_stream.send=mock.AsyncMock()
-writer.simple_flush=mock.AsyncMock()
+mock_stream.recv=mock.AsyncMock()
 
 data=b"a"* (_DEFAULT_FLUSH_INTERVAL_BYTES*2+1)
 awaitwriter.append(data)
 
-assertwriter.simple_flush.await_count==2
+flushed_requests= [
+call.args[0] forcallinmock_stream.send.await_args_listifcall.args[0].flush
+ ]
+assertlen(flushed_requests) ==3
+
+last_request=mock_stream.send.await_args_list[-1].args[0]
+assertlast_request.state_lookup
 
 
 @pytest.mark.asyncio
@@ -584,17 +617,35 @@ async def test_append_data_two_times(mock_write_object_stream, mock_client):
 writer.persisted_size=0
 mock_stream=mock_write_object_stream.return_value
 mock_stream.send=mock.AsyncMock()
-writer.simple_flush=mock.AsyncMock()
 
 data1=b"a"* (_MAX_CHUNK_SIZE_BYTES+10)
+mock_stream.recv=mock.AsyncMock(
+return_value=_storage_v2.BidiWriteObjectResponse(
+persisted_size=len(data1)
+ )
+ )
 awaitwriter.append(data1)
 
+assertmock_stream.send.await_count==2
+last_request_data1=mock_stream.send.await_args_list[-1].args[0]
+assertlast_request_data1.flush
+assertlast_request_data1.state_lookup
+
 data2=b"b"* (_MAX_CHUNK_SIZE_BYTES+20)
+mock_stream.recv=mock.AsyncMock(
+return_value=_storage_v2.BidiWriteObjectResponse(
+persisted_size=len(data2) +len(data1)
+ )
+ )
 awaitwriter.append(data2)
 
+assertmock_stream.send.await_count==4
+last_request_data2=mock_stream.send.await_args_list[-1].args[0]
+assertlast_request_data2.flush
+assertlast_request_data2.state_lookup
+
 total_data_length=len(data1) +len(data2)
 assertwriter.offset==total_data_length
-assertwriter.simple_flush.await_count==0
 
 
 @pytest.mark.asyncio