fix: wait for in-flight BigQuery writes
BatchProcessor.flush() checked Queue.empty(), which only reports whether an item is still waiting, so it could return after a writer dequeued an item but before that write finished. This waits on the queue's unfinished-task count with Queue.join() so flush() blocks until in-flight writes complete. Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 951620765
This commit is contained in:
committed by
Copybara-Service
parent
2dc07457b0
commit
6e43800fcb
@@ -1484,10 +1484,9 @@ class BatchProcessor:
|
||||
}
|
||||
|
||||
async def flush(self) -> None:
|
||||
"""Flushes the queue by waiting for it to be empty."""
|
||||
if self._queue.empty():
|
||||
return
|
||||
# Wait for all items in the queue to be processed
|
||||
"""Flushes the queue, blocking until in-flight writes complete."""
|
||||
# empty() turns true as soon as an item is dequeued, before its write
|
||||
# finishes; join() waits for the unfinished-task count to reach zero.
|
||||
await self._queue.join()
|
||||
|
||||
async def start(self) -> None:
|
||||
|
||||
@@ -8547,6 +8547,19 @@ class TestDropStats:
|
||||
fake_batch.serialize.return_value.to_pybytes.return_value = b"batch"
|
||||
bp._prepare_arrow_batch = mock.MagicMock(return_value=fake_batch)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_waits_for_dequeued_write(self, dummy_arrow_schema):
|
||||
bp = self._make_processor(dummy_arrow_schema)
|
||||
await bp.append({"event": 0})
|
||||
await bp._queue.get()
|
||||
|
||||
flush_task = asyncio.create_task(bp.flush())
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert not flush_task.done()
|
||||
bp._queue.task_done()
|
||||
await flush_task
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queue_full_drops_are_counted(self, dummy_arrow_schema):
|
||||
# Writer is not started, so a size-1 queue fills after one append and the
|
||||
|
||||
Reference in New Issue
Block a user