diff --git a/dataflow/utils/storage.py b/dataflow/utils/storage.py index 09de9141..62659b83 100644 --- a/dataflow/utils/storage.py +++ b/dataflow/utils/storage.py @@ -391,7 +391,8 @@ def flush_step(self, step: int): os.replace(tmp_path, file_path) with self._lock: - self._dirty_steps.discard(step) + if self._buffers.get(step) is df: + self._dirty_steps.discard(step) self.logger.success(f"[flush] Persisted step {step} -> {file_path}") finally: try: diff --git a/test/cpu_only/test_lazy_file_storage_flush_concurrency.py b/test/cpu_only/test_lazy_file_storage_flush_concurrency.py new file mode 100644 index 00000000..2e444fbb --- /dev/null +++ b/test/cpu_only/test_lazy_file_storage_flush_concurrency.py @@ -0,0 +1,68 @@ +"""Concurrency regression tests for LazyFileStorage flush generation handling.""" + +from pathlib import Path +from threading import Event, Thread + +import pandas as pd +import pytest + +from dataflow.utils import storage as storage_module +from dataflow.utils.storage import LazyFileStorage + +pytestmark = pytest.mark.cpu + + +def _make_storage(tmp_path): + source = tmp_path / "source.csv" + source.write_text("text\nsource\n", encoding="utf-8") + return LazyFileStorage( + str(source), + cache_path=str(tmp_path / "cache"), + cache_type="jsonl", + save_on_exit=False, + ) + + +def test_write_during_older_flush_remains_dirty_and_persistable(tmp_path, monkeypatch): + storage = _make_storage(tmp_path) + step = storage.step() + output = Path(step.write([{"text": "old generation"}])) + + replace_entered = Event() + release_replace = Event() + original_replace = storage_module.os.replace + + def blocking_replace(src, dst): + # Synchronize below LazyFileStorage.flush_step(): the real flush has + # already selected and serialized its old dataframe snapshot here. + replace_entered.set() + assert release_replace.wait(timeout=10), "test failed to release blocked persistence" + return original_replace(src, dst) + + monkeypatch.setattr(storage_module.os, "replace", blocking_replace) + + flush_thread = Thread(target=storage.flush_step, args=(1,)) + flush_thread.start() + assert replace_entered.wait(timeout=10), "flush never reached persistence boundary" + + # Real concurrent write to the same output step while the older flush is paused. + step.write([{"text": "new generation"}]) + + release_replace.set() + flush_thread.join(timeout=10) + assert not flush_thread.is_alive() + + # The completed old flush may persist only its selected snapshot. + assert pd.read_json(output, lines=True).to_dict("records") == [ + {"text": "old generation"} + ] + + # Restore real persistence and use the normal public flush path. The newer + # write must still be considered dirty so flush_all() persists it. + monkeypatch.setattr(storage_module.os, "replace", original_replace) + storage.flush_all() + + assert pd.read_json(output, lines=True).to_dict("records") == [ + {"text": "new generation"} + ] + assert 1 not in storage._dirty_steps