From 38a21efaa8a769ab3dba7b990139ef0c0e0518ac Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Tue, 22 Sep 2026 00:31:38 -0400 Subject: [PATCH 1/2] Make streaming teardown destructor-safe --- src/polystore/fiji_stream.py | 8 -- src/polystore/napari_stream.py | 6 -- src/polystore/streaming/_streaming_backend.py | 93 +++++++++++------- tests/test_streaming_cleanup.py | 95 +++++++++++++++++++ 4 files changed, 155 insertions(+), 47 deletions(-) create mode 100644 tests/test_streaming_cleanup.py diff --git a/src/polystore/fiji_stream.py b/src/polystore/fiji_stream.py index 11750af..95da974 100644 --- a/src/polystore/fiji_stream.py +++ b/src/polystore/fiji_stream.py @@ -177,11 +177,3 @@ def _prepare_batch_item( item_payload=item_data, streaming_data_type=output_streaming_data_type, ) - - # cleanup() now inherited from ABC - - def __del__(self): - """Cleanup on deletion.""" - logger.info("🔥 FIJI __del__ called, about to call cleanup()") - self.cleanup() - logger.info("🔥 FIJI __del__ cleanup() returned") diff --git a/src/polystore/napari_stream.py b/src/polystore/napari_stream.py index a28c6e7..551a54a 100644 --- a/src/polystore/napari_stream.py +++ b/src/polystore/napari_stream.py @@ -87,9 +87,3 @@ def _prepare_batch_item( item_payload=item_data, streaming_data_type=request.streaming_data_type, ) - - # cleanup() now inherited from ABC - - def __del__(self): - """Cleanup on deletion.""" - self.cleanup() diff --git a/src/polystore/streaming/_streaming_backend.py b/src/polystore/streaming/_streaming_backend.py index e44c14a..f8189a5 100644 --- a/src/polystore/streaming/_streaming_backend.py +++ b/src/polystore/streaming/_streaming_backend.py @@ -814,44 +814,71 @@ def save(self, data: StreamablePayload | str, file_path: FilePath, **kwargs) -> self.save_batch([data], [file_path], **kwargs) def cleanup(self) -> None: + """Clean up owned shared-memory and ZeroMQ resources.""" + + self._cleanup_resources(logger) + + def _cleanup_resources(self, cleanup_logger: logging.Logger | None) -> None: + """Release resources with optional diagnostics. + + The destructor supplies no logger because module globals may already + have been cleared by interpreter shutdown. Explicit cleanup retains + diagnostics and follows this same resource-ownership path. """ - Clean up shared memory and ZeroMQ resources (common for all streaming backends). - """ - logger.info(f"🔥 CLEANUP: Starting cleanup for {self.VIEWER_TYPE}") - # Clean up shared memory blocks - logger.info( - f"🔥 CLEANUP: About to clean {len(self._shared_memory_blocks)} shared memory blocks" - ) - for shm_name, shm in self._shared_memory_blocks.items(): + def info(message: str) -> None: + if cleanup_logger is not None: + cleanup_logger.info(message) + + def warning(message: str) -> None: + if cleanup_logger is not None: + cleanup_logger.warning(message) + + info(f"🔥 CLEANUP: Starting cleanup for {self.VIEWER_TYPE}") + + # Relinquish collection ownership before closing resources so cleanup is + # idempotent even when invoked again during object finalization. + shared_memory_blocks = self._shared_memory_blocks + self._shared_memory_blocks = {} + info(f"🔥 CLEANUP: About to clean {len(shared_memory_blocks)} shared memory blocks") + for shm_name, shm in shared_memory_blocks.items(): try: shm.close() shm.unlink() - except Exception as e: - logger.warning(f"Failed to cleanup shared memory {shm_name}: {e}") - self._shared_memory_blocks.clear() - logger.info("🔥 CLEANUP: Shared memory cleanup complete") - - # Close publishers - logger.info(f"🔥 CLEANUP: About to close {len(self._publishers)} publishers") - for key, publisher in self._publishers.items(): + except Exception as error: + warning(f"Failed to cleanup shared memory {shm_name}: {error}") + info("🔥 CLEANUP: Shared memory cleanup complete") + + publishers = self._publishers + self._publishers = {} + info(f"🔥 CLEANUP: About to close {len(publishers)} publishers") + for key, publisher in publishers.items(): try: - logger.info(f"🔥 CLEANUP: Closing publisher {key}") + info(f"🔥 CLEANUP: Closing publisher {key}") publisher.close() - logger.info(f"🔥 CLEANUP: Publisher {key} closed") - except Exception as e: - logger.warning(f"Failed to close publisher {key}: {e}") - self._publishers.clear() - logger.info("🔥 CLEANUP: Publishers cleanup complete") - - # Terminate context - if self._context: + info(f"🔥 CLEANUP: Publisher {key} closed") + except Exception as error: + warning(f"Failed to close publisher {key}: {error}") + info("🔥 CLEANUP: Publishers cleanup complete") + + context = self._context + self._context = None + if context: try: - logger.info("🔥 CLEANUP: About to terminate ZMQ context") - self._context.term() - logger.info("🔥 CLEANUP: ZMQ context terminated") - except Exception as e: - logger.warning(f"Failed to terminate ZMQ context: {e}") - self._context = None - - logger.info(f"🔥 CLEANUP: {self.VIEWER_TYPE} streaming backend cleaned up") + info("🔥 CLEANUP: About to terminate ZMQ context") + context.term() + info("🔥 CLEANUP: ZMQ context terminated") + except Exception as error: + warning(f"Failed to terminate ZMQ context: {error}") + + info(f"🔥 CLEANUP: {self.VIEWER_TYPE} streaming backend cleaned up") + + def __del__(self) -> None: + """Release resources without relying on interpreter-shutdown globals.""" + + if not all( + hasattr(self, attribute) + for attribute in ("_shared_memory_blocks", "_publishers", "_context") + ): + return + self._cleanup_resources(None) diff --git a/tests/test_streaming_cleanup.py b/tests/test_streaming_cleanup.py new file mode 100644 index 0000000..6958659 --- /dev/null +++ b/tests/test_streaming_cleanup.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +from collections.abc import Iterator + +import pytest + +from polystore import fiji_stream, napari_stream +from polystore.streaming import StreamingBackend, _streaming_backend + + +class _SharedMemoryProbe: + def __init__(self) -> None: + self.close_calls = 0 + self.unlink_calls = 0 + + def close(self) -> None: + self.close_calls += 1 + + def unlink(self) -> None: + self.unlink_calls += 1 + + +class _PublisherProbe: + def __init__(self) -> None: + self.close_calls = 0 + + def close(self) -> None: + self.close_calls += 1 + + +class _ContextProbe: + def __init__(self) -> None: + self.term_calls = 0 + + def term(self) -> None: + self.term_calls += 1 + + +@pytest.mark.parametrize( + "backend_type", + (fiji_stream.FijiStreamingBackend, napari_stream.NapariStreamingBackend), +) +def test_streaming_destructor_cleanup_is_shutdown_safe_and_idempotent( + monkeypatch: pytest.MonkeyPatch, + backend_type: type[StreamingBackend], +) -> None: + backend = backend_type() + shared_memory = _SharedMemoryProbe() + publisher = _PublisherProbe() + context = _ContextProbe() + backend._shared_memory_blocks = {"owned": shared_memory} + backend._publishers = {"owned": publisher} + backend._context = context + + # Python clears module globals in an unspecified order during interpreter + # shutdown. Cleanup must not require the module logger to remain available. + monkeypatch.setattr(_streaming_backend, "logger", None) + + backend.__del__() + backend.__del__() + + assert shared_memory.close_calls == 1 + assert shared_memory.unlink_calls == 1 + assert publisher.close_calls == 1 + assert context.term_calls == 1 + assert backend._shared_memory_blocks == {} + assert backend._publishers == {} + assert backend._context is None + + +def test_streaming_subclasses_share_one_destructor_authority() -> None: + assert fiji_stream.FijiStreamingBackend.__del__ is StreamingBackend.__del__ + assert napari_stream.NapariStreamingBackend.__del__ is StreamingBackend.__del__ + + +def test_streaming_destructor_accepts_partial_construction() -> None: + backend = napari_stream.NapariStreamingBackend.__new__(napari_stream.NapariStreamingBackend) + + backend.__del__() + + +class _BrokenResourceMapping(dict): + def items(self) -> Iterator[tuple[object, object]]: + raise RuntimeError("broken cleanup ownership") + + +def test_streaming_cleanup_does_not_hide_unexpected_ownership_errors() -> None: + backend = napari_stream.NapariStreamingBackend() + backend._shared_memory_blocks = _BrokenResourceMapping() + + try: + with pytest.raises(RuntimeError, match="broken cleanup ownership"): + backend.cleanup() + finally: + backend._shared_memory_blocks = {} From 5c1a46b8e1e30f35a795b3049e6f47bb6c2b6e3b Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Tue, 22 Sep 2026 09:49:19 -0400 Subject: [PATCH 2/2] Synchronize dependency lock with published version --- uv.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/uv.lock b/uv.lock index b81ffa4..fa021aa 100644 --- a/uv.lock +++ b/uv.lock @@ -2905,7 +2905,7 @@ wheels = [ [[package]] name = "polystore" -version = "0.2.18" +version = "0.2.17" source = { editable = "." } dependencies = [ { name = "arraybridge" },