| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -188,8 +188,9 @@ async def append(self, data: bytes) -> None: | |||
| 188 | 188 | ie. `self.offset` bytes relative to the begining of the object. | |
| 189 | 189 | ||
| 190 | 190 | This method sends the provided `data` to the GCS server in chunks. | |
| 191 | - and persists data in GCS at every `_MAX_BUFFER_SIZE_BYTES` bytes by | ||
| 192 | - calling `self.simple_flush`. | ||
| 191 | + and persists data in GCS at every `_DEFAULT_FLUSH_INTERVAL_BYTES` bytes | ||
| 192 | + or at the last chunk whichever is earlier. Persisting is done by setting | ||
| 193 | + `flush=True` on request. | ||
| 193 | 194 | ||
| 194 | 195 | :type data: bytes | |
| 195 | 196 | :param data: The bytes to append to the object. | |
@@ -214,20 +215,33 @@ async def append(self, data: bytes) -> None: | |||
| 214 | 215 | while start_idx < total_bytes: | |
| 215 | 216 | end_idx = min(start_idx + _MAX_CHUNK_SIZE_BYTES, total_bytes) | |
| 216 | 217 | data_chunk = data[start_idx:end_idx] | |
| 218 | + is_last_chunk = end_idx == total_bytes | ||
| 219 | + chunk_size = end_idx - start_idx | ||
| 217 | 220 | await self.write_obj_stream.send( | |
| 218 | 221 | _storage_v2.BidiWriteObjectRequest( | |
| 219 | 222 | write_offset=self.offset, | |
| 220 | 223 | checksummed_data=_storage_v2.ChecksummedData( | |
| 221 | 224 | content=data_chunk, | |
| 222 | 225 | crc32c=int.from_bytes(Checksum(data_chunk).digest(), "big"), | |
| 223 | 226 | ), | |
| 227 | + state_lookup=is_last_chunk, | ||
| 228 | + flush=is_last_chunk | ||
| 229 | + or ( | ||
| 230 | + self.bytes_appended_since_last_flush + chunk_size | ||
| 231 | + >= self.flush_interval | ||
| 232 | + ), | ||
| 224 | 233 | ) | |
| 225 | 234 | ) | |
| 226 | - chunk_size = end_idx - start_idx | ||
| 227 | 235 | self.offset += chunk_size | |
| 228 | 236 | self.bytes_appended_since_last_flush += chunk_size | |
| 237 | + | ||
| 229 | 238 | if self.bytes_appended_since_last_flush >= self.flush_interval: | |
| 230 | - await self.simple_flush() | ||
| 239 | + self.bytes_appended_since_last_flush = 0 | ||
| 240 | + | ||
| 241 | + if is_last_chunk: | ||
| 242 | + response = await self.write_obj_stream.recv() | ||
| 243 | + self.persisted_size = response.persisted_size | ||
| 244 | + self.offset = self.persisted_size | ||
| 231 | 245 | self.bytes_appended_since_last_flush = 0 | |
| 232 | 246 | start_idx = end_idx | |
| 233 | 247 | ||
@@ -292,20 +306,24 @@ async def close(self, finalize_on_close=False) -> Union[int, _storage_v2.Object] | |||
| 292 | 306 | raise ValueError("Stream is not open. Call open() before close().") | |
| 293 | 307 | ||
| 294 | 308 | if finalize_on_close: | |
| 295 | - await self.finalize() | ||
| 296 | - else: | ||
| 297 | - await self.flush() | ||
| 309 | + return await self.finalize() | ||
| 298 | 310 | ||
| 299 | 311 | await self.write_obj_stream.close() | |
| 300 | 312 | ||
| 301 | 313 | self._is_stream_open = False | |
| 302 | 314 | self.offset = None | |
| 303 | - return self.object_resource if finalize_on_close else self.persisted_size | ||
| 315 | + return self.persisted_size | ||
| 304 | 316 | ||
| 305 | 317 | async def finalize(self) -> _storage_v2.Object: | |
| 306 | 318 | """Finalizes the Appendable Object. | |
| 307 | 319 | ||
| 308 | 320 | Note: Once finalized no more data can be appended. | |
| 321 | + This method is different from `close`. if `.close()` is called data may | ||
| 322 | + still be appended to object at a later point in time by opening with | ||
| 323 | + generation number. | ||
| 324 | + (i.e. `open(..., generation=<object_generation_number>)`. | ||
| 325 | + However if `.finalize()` is called no more data can be appended to the | ||
| 326 | + object. | ||
| 309 | 327 | ||
| 310 | 328 | rtype: google.cloud.storage_v2.types.Object | |
| 311 | 329 | returns: The finalized object resource. | |
@@ -322,6 +340,10 @@ async def finalize(self) -> _storage_v2.Object: | |||
| 322 | 340 | response = await self.write_obj_stream.recv() | |
| 323 | 341 | self.object_resource = response.resource | |
| 324 | 342 | self.persisted_size = self.object_resource.size | |
| 343 | + await self.write_obj_stream.close() | ||
| 344 | + | ||
| 345 | + self._is_stream_open = False | ||
| 346 | + self.offset = None | ||
| 325 | 347 | return self.object_resource | |
| 326 | 348 | ||
| 327 | 349 | # helper methods. | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -192,7 +192,14 @@ def system(session): | |||
| 192 | 192 | # 2021-05-06: defer installing 'google-cloud-*' to after this package, | |
| 193 | 193 | # in order to work around Python 2.7 googolapis-common-protos | |
| 194 | 194 | # issue. | |
| 195 | - session.install("mock", "pytest", "pytest-rerunfailures", "-c", constraints_path) | ||
| 195 | + session.install( | ||
| 196 | + "mock", | ||
| 197 | + "pytest", | ||
| 198 | + "pytest-rerunfailures", | ||
| 199 | + "pytest-asyncio", | ||
| 200 | + "-c", | ||
| 201 | + constraints_path, | ||
| 202 | + ) | ||
| 196 | 203 | session.install("-e", ".", "-c", constraints_path) | |
| 197 | 204 | session.install( | |
| 198 | 205 | "google-cloud-testutils", | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -307,3 +307,56 @@ async def _read_and_verify(expected_content, generation=None): | |||
| 307 | 307 | del mrd | |
| 308 | 308 | del mrd_2 | |
| 309 | 309 | gc.collect() | |
| 310 | + | ||
| 311 | + | ||
| 312 | + @pytest.mark.asyncio | ||
| 313 | + async def test_append_flushes_and_state_lookup(storage_client, blobs_to_delete): | ||
| 314 | + """ | ||
| 315 | + System test for AsyncAppendableObjectWriter, verifying flushing behavior | ||
| 316 | + for both small and large appends. | ||
| 317 | + """ | ||
| 318 | + object_name = f"test-append-flush-varied-size-{uuid.uuid4()}" | ||
| 319 | + grpc_client = AsyncGrpcClient().grpc_client | ||
| 320 | + writer = AsyncAppendableObjectWriter(grpc_client, _ZONAL_BUCKET, object_name) | ||
| 321 | + | ||
| 322 | + # Schedule for cleanup | ||
| 323 | + blobs_to_delete.append(storage_client.bucket(_ZONAL_BUCKET).blob(object_name)) | ||
| 324 | + | ||
| 325 | + # --- Part 1: Test with small data --- | ||
| 326 | + small_data = b"small data" | ||
| 327 | + | ||
| 328 | + await writer.open() | ||
| 329 | + assert writer._is_stream_open | ||
| 330 | + | ||
| 331 | + await writer.append(small_data) | ||
| 332 | + persisted_size = await writer.state_lookup() | ||
| 333 | + assert persisted_size == len(small_data) | ||
| 334 | + | ||
| 335 | + # --- Part 2: Test with large data --- | ||
| 336 | + large_data = os.urandom(38 * 1024 * 1024) | ||
| 337 | + | ||
| 338 | + # Append data larger than the default flush interval (16 MiB). | ||
| 339 | + # This should trigger the interval-based flushing logic. | ||
| 340 | + await writer.append(large_data) | ||
| 341 | + | ||
| 342 | + # Verify the total data has been persisted. | ||
| 343 | + total_size = len(small_data) + len(large_data) | ||
| 344 | + persisted_size = await writer.state_lookup() | ||
| 345 | + assert persisted_size == total_size | ||
| 346 | + | ||
| 347 | + # --- Part 3: Finalize and verify --- | ||
| 348 | + final_object = await writer.close(finalize_on_close=True) | ||
| 349 | + | ||
| 350 | + assert not writer._is_stream_open | ||
| 351 | + assert final_object.size == total_size | ||
| 352 | + | ||
| 353 | + # Verify the full content of the object. | ||
| 354 | + full_data = small_data + large_data | ||
| 355 | + mrd = AsyncMultiRangeDownloader(grpc_client, _ZONAL_BUCKET, object_name) | ||
| 356 | + buffer = BytesIO() | ||
| 357 | + await mrd.open() | ||
| 358 | + # (0, 0) means read the whole object | ||
| 359 | + await mrd.download_ranges([(0, 0, buffer)]) | ||
| 360 | + await mrd.close() | ||
| 361 | + content = buffer.getvalue() | ||
| 362 | + assert content == full_data | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -206,7 +206,9 @@ def test_update_state_from_response(self): | |||
| 206 | 206 | strategy.update_state_from_response(response2, state) | |
| 207 | 207 | self.assertEqual(write_state.persisted_size, 1024) | |
| 208 | 208 | ||
| 209 | - final_resource = storage_type.Object(name="test-object", bucket="b", size=2048, finalize_time=datetime.now()) | ||
| 209 | + final_resource = storage_type.Object( | ||
| 210 | + name="test-object", bucket="b", size=2048, finalize_time=datetime.now() | ||
| 211 | + ) | ||
| 210 | 212 | response3 = storage_type.BidiWriteObjectResponse(resource=final_resource) | |
| 211 | 213 | strategy.update_state_from_response(response3, state) | |
| 212 | 214 | self.assertEqual(write_state.persisted_size, 2048) | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -364,6 +364,7 @@ async def test_close(mock_write_object_stream, mock_client): | |||
| 364 | 364 | writer = AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT) | |
| 365 | 365 | writer._is_stream_open = True | |
| 366 | 366 | writer.offset = 1024 | |
| 367 | + writer.persisted_size = 1024 | ||
| 367 | 368 | mock_stream = mock_write_object_stream.return_value | |
| 368 | 369 | mock_stream.send = mock.AsyncMock() | |
| 369 | 370 | mock_stream.recv = mock.AsyncMock( | |
@@ -435,16 +436,20 @@ async def test_finalize(mock_write_object_stream, mock_client): | |||
| 435 | 436 | mock_stream.recv = mock.AsyncMock( | |
| 436 | 437 | return_value=_storage_v2.BidiWriteObjectResponse(resource=mock_resource) | |
| 437 | 438 | ) | |
| 439 | + mock_stream.close = mock.AsyncMock() | ||
| 438 | 440 | ||
| 439 | 441 | gcs_object = await writer.finalize() | |
| 440 | 442 | ||
| 441 | 443 | mock_stream.send.assert_awaited_once_with( | |
| 442 | 444 | _storage_v2.BidiWriteObjectRequest(finish_write=True) | |
| 443 | 445 | ) | |
| 444 | 446 | mock_stream.recv.assert_awaited_once() | |
| 447 | + mock_stream.close.assert_awaited_once() | ||
| 445 | 448 | assert writer.object_resource == mock_resource | |
| 446 | 449 | assert writer.persisted_size == 123 | |
| 447 | 450 | assert gcs_object == mock_resource | |
| 451 | + assert writer._is_stream_open is False | ||
| 452 | + assert writer.offset is None | ||
| 448 | 453 | ||
| 449 | 454 | ||
| 450 | 455 | @pytest.mark.asyncio | |
@@ -501,30 +506,39 @@ async def test_append_sends_data_in_chunks(mock_write_object_stream, mock_client | |||
| 501 | 506 | writer.persisted_size = 100 | |
| 502 | 507 | mock_stream = mock_write_object_stream.return_value | |
| 503 | 508 | mock_stream.send = mock.AsyncMock() | |
| 504 | - writer.simple_flush = mock.AsyncMock() | ||
| 505 | 509 | ||
| 506 | 510 | data = b"a" * (_MAX_CHUNK_SIZE_BYTES + 1) | |
| 511 | + mock_stream.recv = mock.AsyncMock( | ||
| 512 | + return_value=_storage_v2.BidiWriteObjectResponse( | ||
| 513 | + persisted_size=100 + len(data) | ||
| 514 | + ) | ||
| 515 | + ) | ||
| 516 | + | ||
| 507 | 517 | await writer.append(data) | |
| 508 | 518 | ||
| 509 | 519 | assert mock_stream.send.await_count == 2 | |
| 510 | - first_call = mock_stream.send.await_args_list[0] | ||
| 511 | - second_call = mock_stream.send.await_args_list[1] | ||
| 520 | + first_request = mock_stream.send.await_args_list[0].args[0] | ||
| 521 | + second_request = mock_stream.send.await_args_list[1].args[0] | ||
| 512 | 522 | ||
| 513 | 523 | # First chunk | |
| 514 | - assert first_call[0][0].write_offset == 100 | ||
| 515 | - assert len(first_call[0][0].checksummed_data.content) == _MAX_CHUNK_SIZE_BYTES | ||
| 516 | - assert first_call[0][0].checksummed_data.crc32c == int.from_bytes( | ||
| 524 | + assert first_request.write_offset == 100 | ||
| 525 | + assert len(first_request.checksummed_data.content) == _MAX_CHUNK_SIZE_BYTES | ||
| 526 | + assert first_request.checksummed_data.crc32c == int.from_bytes( | ||
| 517 | 527 | Checksum(data[:_MAX_CHUNK_SIZE_BYTES]).digest(), byteorder="big" | |
| 518 | 528 | ) | |
| 519 | - # Second chunk | ||
| 520 | - assert second_call[0][0].write_offset == 100 + _MAX_CHUNK_SIZE_BYTES | ||
| 521 | - assert len(second_call[0][0].checksummed_data.content) == 1 | ||
| 522 | - assert second_call[0][0].checksummed_data.crc32c == int.from_bytes( | ||
| 529 | + assert not first_request.flush | ||
| 530 | + assert not first_request.state_lookup | ||
| 531 | + | ||
| 532 | + # Second chunk (last chunk) | ||
| 533 | + assert second_request.write_offset == 100 + _MAX_CHUNK_SIZE_BYTES | ||
| 534 | + assert len(second_request.checksummed_data.content) == 1 | ||
| 535 | + assert second_request.checksummed_data.crc32c == int.from_bytes( | ||
| 523 | 536 | Checksum(data[_MAX_CHUNK_SIZE_BYTES:]).digest(), byteorder="big" | |
| 524 | 537 | ) | |
| 538 | + assert second_request.flush | ||
| 539 | + assert second_request.state_lookup | ||
| 525 | 540 | ||
| 526 | 541 | assert writer.offset == 100 + len(data) | |
| 527 | - writer.simple_flush.assert_not_awaited() | ||
| 528 | 542 | ||
| 529 | 543 | ||
| 530 | 544 | @pytest.mark.asyncio | |
@@ -541,12 +555,25 @@ async def test_append_flushes_when_buffer_is_full( | |||
| 541 | 555 | writer.persisted_size = 0 | |
| 542 | 556 | mock_stream = mock_write_object_stream.return_value | |
| 543 | 557 | mock_stream.send = mock.AsyncMock() | |
| 544 | - writer.simple_flush = mock.AsyncMock() | ||
| 558 | + mock_stream.recv = mock.AsyncMock() | ||
| 545 | 559 | ||
| 546 | 560 | data = b"a" * _DEFAULT_FLUSH_INTERVAL_BYTES | |
| 547 | 561 | await writer.append(data) | |
| 548 | 562 | ||
| 549 | - writer.simple_flush.assert_awaited_once() | ||
| 563 | + num_chunks = _DEFAULT_FLUSH_INTERVAL_BYTES // _MAX_CHUNK_SIZE_BYTES | ||
| 564 | + assert mock_stream.send.await_count == num_chunks | ||
| 565 | + | ||
| 566 | + # All but the last request should not have flush or state_lookup set. | ||
| 567 | + for i in range(num_chunks - 1): | ||
| 568 | + request = mock_stream.send.await_args_list[i].args[0] | ||
| 569 | + assert not request.flush | ||
| 570 | + assert not request.state_lookup | ||
| 571 | + | ||
| 572 | + # The last request should have flush and state_lookup set. | ||
| 573 | + last_request = mock_stream.send.await_args_list[-1].args[0] | ||
| 574 | + assert last_request.flush | ||
| 575 | + assert last_request.state_lookup | ||
| 576 | + assert writer.bytes_appended_since_last_flush == 0 | ||
| 550 | 577 | ||
| 551 | 578 | ||
| 552 | 579 | @pytest.mark.asyncio | |
@@ -561,12 +588,18 @@ async def test_append_handles_large_data(mock_write_object_stream, mock_client): | |||
| 561 | 588 | writer.persisted_size = 0 | |
| 562 | 589 | mock_stream = mock_write_object_stream.return_value | |
| 563 | 590 | mock_stream.send = mock.AsyncMock() | |
| 564 | - writer.simple_flush = mock.AsyncMock() | ||
| 591 | + mock_stream.recv = mock.AsyncMock() | ||
| 565 | 592 | ||
| 566 | 593 | data = b"a" * (_DEFAULT_FLUSH_INTERVAL_BYTES * 2 + 1) | |
| 567 | 594 | await writer.append(data) | |
| 568 | 595 | ||
| 569 | - assert writer.simple_flush.await_count == 2 | ||
| 596 | + flushed_requests = [ | ||
| 597 | + call.args[0] for call in mock_stream.send.await_args_list if call.args[0].flush | ||
| 598 | + ] | ||
| 599 | + assert len(flushed_requests) == 3 | ||
| 600 | + | ||
| 601 | + last_request = mock_stream.send.await_args_list[-1].args[0] | ||
| 602 | + assert last_request.state_lookup | ||
| 570 | 603 | ||
| 571 | 604 | ||
| 572 | 605 | @pytest.mark.asyncio | |
@@ -584,17 +617,35 @@ async def test_append_data_two_times(mock_write_object_stream, mock_client): | |||
| 584 | 617 | writer.persisted_size = 0 | |
| 585 | 618 | mock_stream = mock_write_object_stream.return_value | |
| 586 | 619 | mock_stream.send = mock.AsyncMock() | |
| 587 | - writer.simple_flush = mock.AsyncMock() | ||
| 588 | 620 | ||
| 589 | 621 | data1 = b"a" * (_MAX_CHUNK_SIZE_BYTES + 10) | |
| 622 | + mock_stream.recv = mock.AsyncMock( | ||
| 623 | + return_value=_storage_v2.BidiWriteObjectResponse( | ||
| 624 | + persisted_size= len(data1) | ||
| 625 | + ) | ||
| 626 | + ) | ||
| 590 | 627 | await writer.append(data1) | |
| 591 | 628 | ||
| 629 | + assert mock_stream.send.await_count == 2 | ||
| 630 | + last_request_data1 = mock_stream.send.await_args_list[-1].args[0] | ||
| 631 | + assert last_request_data1.flush | ||
| 632 | + assert last_request_data1.state_lookup | ||
| 633 | + | ||
| 592 | 634 | data2 = b"b" * (_MAX_CHUNK_SIZE_BYTES + 20) | |
| 635 | + mock_stream.recv = mock.AsyncMock( | ||
| 636 | + return_value=_storage_v2.BidiWriteObjectResponse( | ||
| 637 | + persisted_size= len(data2) + len(data1) | ||
| 638 | + ) | ||
| 639 | + ) | ||
| 593 | 640 | await writer.append(data2) | |
| 594 | 641 | ||
| 642 | + assert mock_stream.send.await_count == 4 | ||
| 643 | + last_request_data2 = mock_stream.send.await_args_list[-1].args[0] | ||
| 644 | + assert last_request_data2.flush | ||
| 645 | + assert last_request_data2.state_lookup | ||
| 646 | + | ||
| 595 | 647 | total_data_length = len(data1) + len(data2) | |
| 596 | 648 | assert writer.offset == total_data_length | |
| 597 | - assert writer.simple_flush.await_count == 0 | ||
| 598 | 649 | ||
| 599 | 650 | ||
| 600 | 651 | @pytest.mark.asyncio | |
| Back | FazBrowse Home | New Git URL |
0 commit comments