piker/tests/test_shm_cleanup.py

110 lines
3.3 KiB
Python

'''
Pytest shared-memory ownership and leak-cleanup regressions.
'''
from collections.abc import Callable
from uuid import uuid4
import numpy as np
import pytest
from tractor._exceptions import NoRuntime
from tractor.ipc import _shm
from tractor.ipc._mp_bs import disable_mantracker
from piker.data._sharedmem import maybe_open_shm_array
from piker.data._source import def_iohlcv_fields
def test_shm_tracker_catches_pre_actor_failure(
shm_leak_tracker: Callable[[], set[str]],
) -> None:
'''
Clean allocations which fail before actor-lifetime registration.
``open_shm_ndarray()`` creates its data and index segments before it
asks ``current_actor()`` for the lifetime stack. Without a runtime,
that lookup raises after all three OS names exist and the old harness
leaked them. Reproduce that ordering through Piker's real wrapper,
invoke tracked cleanup, and prove the data, first, and last segments
plus their process-local token are all removed.
'''
key = f'test_shm_failure_{uuid4().hex}'
with pytest.raises(NoRuntime):
maybe_open_shm_array(
key=key,
size=4,
dtype=np.dtype(def_iohlcv_fields),
)
token = _shm.NDToken.from_msg(_shm._known_tokens[key])
names = {
token.shm_name,
token.shm_first_index_name,
token.shm_last_index_name,
}
assert shm_leak_tracker() == names
assert key not in _shm._known_tokens
for name in names:
with pytest.raises(FileNotFoundError):
_shm.SharedMemory(
name=name,
create=False,
)
def test_shm_leak_tracker_unlinks_only_owned_segments(
shm_leak_tracker: Callable[[], set[str]],
) -> None:
'''
Test cleanup must unlink creators without harming attachments.
A failed SHM allocation can escape before Tractor registers actor
lifetime callbacks. Broad cache or filesystem cleanup is unsafe
because tests may attach to a live Piker segment they do not own.
Create one external segment through an unpatched factory, attach to
it through Tractor, and create one test-owned segment through the
tracked factory. Invoke cleanup and prove only the owned name was
removed while the external segment remains attachable.
'''
open_external_shm = disable_mantracker()
external = None
attachment = None
owned = None
try:
external = open_external_shm(
create=True,
size=8,
)
attachment = _shm.SharedMemory(
name=external.name,
create=False,
)
owned = _shm.SharedMemory(None, True, 8)
cleaned: set[str] = shm_leak_tracker()
assert cleaned == {owned.name}
with pytest.raises(FileNotFoundError):
_shm.SharedMemory(
name=owned.name,
create=False,
)
survivor = _shm.SharedMemory(
name=external.name,
create=False,
)
survivor.close()
finally:
if owned is not None:
owned.close()
if attachment is not None:
attachment.close()
if external is not None:
try:
external.unlink()
except FileNotFoundError:
pass
external.close()