diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py index 1a239ed9..b1795434 100644 --- a/tests/test_kv_lock.py +++ b/tests/test_kv_lock.py @@ -14,9 +14,11 @@ # limitations under the License. import asyncio +import gc import importlib import threading import time +import weakref from unittest.mock import MagicMock import pytest @@ -172,6 +174,27 @@ def test_invalid_lock_settings_raise_before_any_actor_starts(fake_lock_actors, m fake_lock_actors.options.assert_not_called() +def test_fork_child_keeps_inherited_lock_handles_referenced(monkeypatch): + class Handle: + pass + + owned, managed, lease = Handle(), Handle(), Handle() + for name in ("_state_lock", "_table", "_registry", "_renewer_running"): + monkeypatch.setattr(kvl, name, getattr(kvl, name)) + monkeypatch.setattr(kvl, "_owned", [owned]) + monkeypatch.setattr(kvl, "_managers", {0: managed}) + monkeypatch.setattr(kvl, "_leases", {"token": lease}) + monkeypatch.setattr(kvl, "_inherited_after_fork", []) + refs = [weakref.ref(obj) for obj in (owned, managed, lease)] + del owned, managed, lease + + kvl._reset_after_fork() + gc.collect() + + assert kvl._owned == [] and kvl._managers == {} and kvl._leases == {} + assert all(ref() is not None for ref in refs) + + def test_lock_settings_are_ignored_while_the_lock_is_off(monkeypatch): monkeypatch.delenv("TQ_ENABLE_KV_GLOBAL_LOCK", raising=False) monkeypatch.setenv("TQ_KV_GLOBAL_LOCK_NUM_SHARDS", "0") diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 4ddff4f1..3586d2af 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -216,6 +216,7 @@ async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: _managers: dict = {} # shard -> TransferQueueLockManager handle, looked up on first use _owned: list = [] # the lock actors this process's tq.init() created, killed by its tq.close() _closes = 0 # this process's tq.close() calls: a lease created before the latest one is given up +_inherited_after_fork: list = [] # Longest one call waits inside a lock actor; waiting longer takes another call that keeps # the queue position. The bound is long because every waiter spends one call per window, and # those calls share the actor's limited throughput with everyone else's, renewals included. @@ -585,7 +586,10 @@ def _close_lock_managers() -> None: def _reset_after_fork() -> None: # The child inherits keys, leases and Ray handles that belong to threads and connections it lacks. + # Releasing an inherited handle runs the parent's core worker in the child: that crashes the child, + # can report the parent's lock actors out of scope, and leaves the parent's GCS calls timing out. global _state_lock, _table, _registry, _leases, _renewer_running, _managers, _owned + _inherited_after_fork.append((_owned, _managers, _leases)) _state_lock, _table = threading.Lock(), {} _registry, _leases, _renewer_running, _managers, _owned = threading.Condition(), {}, False, {}, []