From 6f5fcabe4220a60d5878dbf79037cca51e148a8a Mon Sep 17 00:00:00 2001
From: Neko65536 <kangfx918@163.com>
Date: Mon, 31 Aug 2026 22:32:50 +0800
Subject: [PATCH] fix: serialize normalized stage lock cleanup

---
 src/quant_data_kit/normalized_v3.py  | 44 ++++++++++++++++++++--------
 tests/test_normalized_v3_failures.py | 33 +++++++++++++++++++++
 2 files changed, 64 insertions(+), 13 deletions(-)

diff --git a/src/quant_data_kit/normalized_v3.py b/src/quant_data_kit/normalized_v3.py
index f7c3797..951133e 100644
--- a/src/quant_data_kit/normalized_v3.py
+++ b/src/quant_data_kit/normalized_v3.py
@@ -498,17 +498,34 @@ def _partition_identity(
 def _streaming_stage(root: Path) -> Iterator[Path]:
     staging_root = _mkdir_in_lake(root, root / "normalized" / "staging")
     owners_root = _mkdir_in_lake(root, root / "normalized" / ".stage-owners")
-    for stale in sorted(staging_root.glob("normalized-batch-stream-*")):
-        checked = _validate_lake_path(root, stale, allow_missing=False)
-        if not checked.is_dir():
-            raise ValidationError(f"Normalized staging entry is not a directory: {checked}")
-        owner = owners_root / f"{checked.name}.lock"
-        try:
-            with process_file_lock(owner, timeout_seconds=0.01):
-                if checked.exists():
-                    shutil.rmtree(checked)
-        except TimeoutError:
-            continue
+    gc_lock = owners_root / ".gc.lock"
+    with process_file_lock(gc_lock):
+        for stale in sorted(staging_root.glob("normalized-batch-stream-*")):
+            if not stale.exists():
+                continue
+            checked = _validate_lake_path(root, stale, allow_missing=False)
+            if not checked.is_dir():
+                raise ValidationError(f"Normalized staging entry is not a directory: {checked}")
+            owner = owners_root / f"{checked.name}.lock"
+            try:
+                with process_file_lock(owner, timeout_seconds=0.01):
+                    if checked.exists():
+                        shutil.rmtree(checked)
+                if owner.exists():
+                    owner.unlink()
+            except TimeoutError:
+                continue
+        for owner in sorted(owners_root.glob("normalized-batch-stream-*.lock")):
+            stage = staging_root / owner.name.removesuffix(".lock")
+            if stage.exists():
+                continue
+            try:
+                with process_file_lock(owner, timeout_seconds=0.01):
+                    removable = not stage.exists()
+                if removable and owner.exists():
+                    owner.unlink()
+            except TimeoutError:
+                continue
     operation = uuid.uuid4().hex
     stage = staging_root / f"normalized-batch-stream-{operation}"
     owner = owners_root / f"{stage.name}.lock"
@@ -519,8 +536,9 @@ def _streaming_stage(root: Path) -> Iterator[Path]:
         finally:
             if stage.exists():
                 shutil.rmtree(stage)
-    if owner.exists():
-        owner.unlink()
+    with process_file_lock(gc_lock):
+        if owner.exists():
+            owner.unlink()
 
 
 def _sql_text(value: str | Path) -> str:
diff --git a/tests/test_normalized_v3_failures.py b/tests/test_normalized_v3_failures.py
index b6da245..b89fa5c 100644
--- a/tests/test_normalized_v3_failures.py
+++ b/tests/test_normalized_v3_failures.py
@@ -2,6 +2,7 @@
 
 import json
 from collections.abc import Callable
+from contextlib import contextmanager
 from copy import deepcopy
 from dataclasses import replace
 from datetime import date, datetime, timedelta, timezone
@@ -166,6 +167,38 @@ def test_digest_and_staging_corruption_guards(
         pass
 
 
+def test_streaming_stage_serializes_owner_file_cleanup(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    root = lake_module._resolved_lake_root(tmp_path / "lake", create=True)
+    real_lock = normalized_v3.process_file_lock
+    real_unlink = Path.unlink
+    held_locks: set[Path] = set()
+
+    @contextmanager
+    def tracked_lock(path: Path, *, timeout_seconds: float = 60.0):
+        checked = Path(path)
+        with real_lock(checked, timeout_seconds=timeout_seconds):
+            held_locks.add(checked)
+            try:
+                yield
+            finally:
+                held_locks.remove(checked)
+
+    def guarded_unlink(path: Path, *args: object, **kwargs: object) -> None:
+        if path.parent.name == ".stage-owners" and path.name.startswith("normalized-batch-stream-"):
+            assert any(item.name == ".gc.lock" for item in held_locks)
+        real_unlink(path, *args, **kwargs)
+
+    monkeypatch.setattr(normalized_v3, "process_file_lock", tracked_lock)
+    monkeypatch.setattr(Path, "unlink", guarded_unlink)
+    with normalized_v3._streaming_stage(root) as stage:
+        assert stage.is_dir()
+    owners_root = root / "normalized" / ".stage-owners"
+    assert sorted(item.name for item in owners_root.iterdir()) == [".gc.lock"]
+
+
 def _rewrite_index_manifest(
     manifest_path: Path,
     mutate: Callable[[dict[str, Any]], None],