feat: make flush size configurable (#1677)

feat: make flush size configurable

Author: chandra-siri

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

@@ -24,6 +24,7 @@
 fromtypingimportOptional, Union
 
 fromgoogle_crc32cimportChecksum
+fromgoogle.api_coreimportexceptions
 
 from ._utilsimportraise_if_no_fast_crc32c
 fromgoogle.cloudimport_storage_v2
@@ -36,7 +37,7 @@
 
 
 _MAX_CHUNK_SIZE_BYTES=2*1024*1024# 2 MiB
-_MAX_BUFFER_SIZE_BYTES=16*1024*1024# 16 MiB
+_DEFAULT_FLUSH_INTERVAL_BYTES=16*1024*1024# 16 MiB
 
 
 classAsyncAppendableObjectWriter:
@@ -49,6 +50,7 @@ def __init__(
 object_name: str,
 generation=None,
 write_handle=None,
+writer_options: Optional[dict] =None,
  ):
 """
  Class for appending data to a GCS Appendable Object.
@@ -125,6 +127,21 @@ def __init__(
 # Please note: `offset` and `persisted_size` are same when the stream is
 # opened.
 self.persisted_size: Optional[int] =None
+ifwriter_optionsisNone:
+writer_options= {}
+self.flush_interval=writer_options.get(
+"FLUSH_INTERVAL_BYTES", _DEFAULT_FLUSH_INTERVAL_BYTES
+ )
+# TODO: add test case for this.
+ifself.flush_interval<_MAX_CHUNK_SIZE_BYTES:
+raiseexceptions.OutOfRange(
+f"flush_interval must be >= {_MAX_CHUNK_SIZE_BYTES} , but provided {self.flush_interval}"
+ )
+ifself.flush_interval%_MAX_CHUNK_SIZE_BYTES!=0:
+raiseexceptions.OutOfRange(
+f"flush_interval must be a multiple of {_MAX_CHUNK_SIZE_BYTES}, but provided {self.flush_interval}"
+ )
+self.bytes_appended_since_last_flush=0
 
 asyncdefstate_lookup(self) ->int:
 """Returns the persisted_size
@@ -193,7 +210,6 @@ async def append(self, data: bytes) -> None:
 self.offset=self.persisted_size
 
 start_idx=0
-bytes_to_flush=0
 whilestart_idx<total_bytes:
 end_idx=min(start_idx+_MAX_CHUNK_SIZE_BYTES, total_bytes)
 data_chunk=data[start_idx:end_idx]
@@ -208,10 +224,10 @@ async def append(self, data: bytes) -> None:
  )
 chunk_size=end_idx-start_idx
 self.offset+=chunk_size
-bytes_to_flush+=chunk_size
-ifbytes_to_flush>=_MAX_BUFFER_SIZE_BYTES:
+self.bytes_appended_since_last_flush+=chunk_size
+ifself.bytes_appended_since_last_flush>=self.flush_interval:
 awaitself.simple_flush()
-bytes_to_flush=0
+self.bytes_appended_since_last_flush=0
 start_idx=end_idx
 
 asyncdefsimple_flush(self) ->None:

tests/system/test_zonal.py

@@ -13,6 +13,7 @@
 fromgoogle.cloud.storage._experimental.asyncio.async_grpc_clientimportAsyncGrpcClient
 fromgoogle.cloud.storage._experimental.asyncio.async_appendable_object_writerimport (
 AsyncAppendableObjectWriter,
+_DEFAULT_FLUSH_INTERVAL_BYTES,
 )
 fromgoogle.cloud.storage._experimental.asyncio.async_multi_range_downloaderimport (
 AsyncMultiRangeDownloader,
@@ -162,6 +163,59 @@ async def test_basic_wrd_in_slices(storage_client, blobs_to_delete, object_size)
 blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+"flush_interval",
+ [2*1024*1024, 4*1024*1024, 8*1024*1024, _DEFAULT_FLUSH_INTERVAL_BYTES],
+)
+asyncdeftest_wrd_with_non_default_flush_interval(
+storage_client,
+blobs_to_delete,
+flush_interval,
+):
+object_name=f"test_basic_wrd-{str(uuid.uuid4())}"
+object_size=9*1024*1024
+
+# Client instantiation; it cannot be part of fixture because.
+# grpc_client's event loop and event loop of coroutine running it
+# (i.e. this test) must be same.
+# Note:
+# 1. @pytest.mark.asyncio ensures new event loop for each test.
+# 2. we can keep the same event loop for entire module but that may
+# create issues if tests are run in parallel and one test hogs the event
+# loop slowing down other tests.
+object_data=os.urandom(object_size)
+object_checksum=google_crc32c.value(object_data)
+grpc_client=AsyncGrpcClient().grpc_client
+
+writer=AsyncAppendableObjectWriter(
+grpc_client,
+_ZONAL_BUCKET,
+object_name,
+writer_options={"FLUSH_INTERVAL_BYTES": flush_interval},
+ )
+awaitwriter.open()
+mark1, mark2=_get_equal_dist(0, object_size)
+awaitwriter.append(object_data[0:mark1])
+awaitwriter.append(object_data[mark1:mark2])
+awaitwriter.append(object_data[mark2:])
+object_metadata=awaitwriter.close(finalize_on_close=True)
+assertobject_metadata.size==object_size
+assertint(object_metadata.checksums.crc32c) ==object_checksum
+
+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()
+assertbuffer.getvalue() ==object_data
+assertmrd.persisted_size==object_size
+
+# Clean up; use json client (i.e. `storage_client` fixture) to delete.
+blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name))
+
+
 @pytest.mark.asyncio
 asyncdeftest_read_unfinalized_appendable_object(storage_client, blobs_to_delete):
 object_name=f"read_unfinalized_appendable_object-{str(uuid.uuid4())[:4]}"

tests/unit/asyncio/test_async_appendable_object_writer.py

@@ -21,6 +21,9 @@
 fromgoogle.cloud.storage._experimental.asyncio.async_appendable_object_writerimport (
 AsyncAppendableObjectWriter,
 )
+fromgoogle.cloud.storage._experimental.asyncio.async_appendable_object_writerimport (
+_MAX_CHUNK_SIZE_BYTES,
+)
 fromgoogle.cloudimport_storage_v2
 
 
@@ -29,6 +32,7 @@
 GENERATION=123
 WRITE_HANDLE=b"test-write-handle"
 PERSISTED_SIZE=456
+EIGHT_MIB=8*1024*1024
 
 
 @pytest.fixture
@@ -52,6 +56,7 @@ def test_init(mock_write_object_stream, mock_client):
 assertnotwriter._is_stream_open
 assertwriter.offsetisNone
 assertwriter.persisted_sizeisNone
+assertwriter.bytes_appended_since_last_flush==0
 
 mock_write_object_stream.assert_called_once_with(
 client=mock_client,
@@ -78,6 +83,7 @@ def test_init_with_optional_args(mock_write_object_stream, mock_client):
 
 assertwriter.generation==GENERATION
 assertwriter.write_handle==WRITE_HANDLE
+assertwriter.bytes_appended_since_last_flush==0
 
 mock_write_object_stream.assert_called_once_with(
 client=mock_client,
@@ -88,6 +94,60 @@ def test_init_with_optional_args(mock_write_object_stream, mock_client):
  )
 
 
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+deftest_init_with_writer_options(mock_write_object_stream, mock_client):
+"""Test the constructor with optional arguments."""
+writer=AsyncAppendableObjectWriter(
+mock_client,
+BUCKET,
+OBJECT,
+writer_options={"FLUSH_INTERVAL_BYTES": EIGHT_MIB},
+ )
+
+assertwriter.flush_interval==EIGHT_MIB
+assertwriter.bytes_appended_since_last_flush==0
+
+mock_write_object_stream.assert_called_once_with(
+client=mock_client,
+bucket_name=BUCKET,
+object_name=OBJECT,
+generation_number=None,
+write_handle=None,
+ )
+
+
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+deftest_init_with_flush_interval_less_than_chunk_size_raises_error(mock_client):
+"""Test that an OutOfRange error is raised if flush_interval is less than the chunk size."""
+
+withpytest.raises(exceptions.OutOfRange):
+AsyncAppendableObjectWriter(
+mock_client,
+BUCKET,
+OBJECT,
+writer_options={"FLUSH_INTERVAL_BYTES": _MAX_CHUNK_SIZE_BYTES-1},
+ )
+
+
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+deftest_init_with_flush_interval_not_multiple_of_chunk_size_raises_error(mock_client):
+"""Test that an OutOfRange error is raised if flush_interval is not a multiple of the chunk size."""
+
+withpytest.raises(exceptions.OutOfRange):
+AsyncAppendableObjectWriter(
+mock_client,
+BUCKET,
+OBJECT,
+writer_options={"FLUSH_INTERVAL_BYTES": _MAX_CHUNK_SIZE_BYTES+1},
+ )
+
+
 @mock.patch("google.cloud.storage._experimental.asyncio._utils.google_crc32c")
 @mock.patch(
 "google.cloud.storage._experimental.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client"
@@ -477,7 +537,7 @@ async def test_append_flushes_when_buffer_is_full(
 ):
 """Test that append flushes the stream when the buffer size is reached."""
 fromgoogle.cloud.storage._experimental.asyncio.async_appendable_object_writerimport (
-_MAX_BUFFER_SIZE_BYTES,
+_DEFAULT_FLUSH_INTERVAL_BYTES,
  )
 
 writer=AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT)
@@ -487,7 +547,7 @@ async def test_append_flushes_when_buffer_is_full(
 mock_stream.send=mock.AsyncMock()
 writer.simple_flush=mock.AsyncMock()
 
-data=b"a"*_MAX_BUFFER_SIZE_BYTES
+data=b"a"*_DEFAULT_FLUSH_INTERVAL_BYTES
 awaitwriter.append(data)
 
 writer.simple_flush.assert_awaited_once()
@@ -500,7 +560,7 @@ async def test_append_flushes_when_buffer_is_full(
 asyncdeftest_append_handles_large_data(mock_write_object_stream, mock_client):
 """Test that append handles data larger than the buffer size."""
 fromgoogle.cloud.storage._experimental.asyncio.async_appendable_object_writerimport (
-_MAX_BUFFER_SIZE_BYTES,
+_DEFAULT_FLUSH_INTERVAL_BYTES,
  )
 
 writer=AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT)
@@ -510,7 +570,7 @@ async def test_append_handles_large_data(mock_write_object_stream, mock_client):
 mock_stream.send=mock.AsyncMock()
 writer.simple_flush=mock.AsyncMock()
 
-data=b"a"* (_MAX_BUFFER_SIZE_BYTES*2+1)
+data=b"a"* (_DEFAULT_FLUSH_INTERVAL_BYTES*2+1)
 awaitwriter.append(data)
 
 assertwriter.simple_flush.await_count==2