Own sender-backed shared memory attachment

Author: trissim

src/polystore/streaming/__init__.py

@@ -14,14 +14,15 @@
     FilePath,
     RoiStreamPayload,
     StreamablePayload,
+    StreamingBackend,
     StreamingBatchItemPreparationAuthority,
     StreamingBatchMessageBuilder,
     StreamingBatchMessageRequest,
     StreamingBuiltBatch,
-    StreamingPreparedBatchItems,
-    StreamingBackend,
     StreamingComponentNamesRequest,
     StreamingItemPreparationRequest,
+    StreamingPreparedBatchItems,
+    StreamingSharedMemoryAuthority,
     ViewerDisplayPayloadExtra,
 )
 
@@ -37,5 +38,6 @@
     "StreamingBackend",
     "StreamingComponentNamesRequest",
     "StreamingItemPreparationRequest",
+    "StreamingSharedMemoryAuthority",
     "ViewerDisplayPayloadExtra",
 ]

src/polystore/streaming/_streaming_backend.py

@@ -12,7 +12,8 @@
 import uuid
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
-from multiprocessing import shared_memory
+from multiprocessing import resource_tracker, shared_memory
+from multiprocessing.shared_memory import _USE_POSIX
 from pathlib import Path
 from types import MappingProxyType
 from typing import TypeAlias
@@ -310,7 +311,29 @@ def to_numpy(data: StreamablePayload) -> np.ndarray:
 
 
 class StreamingSharedMemoryAuthority:
-    """Allocate image payloads for viewer transfer through shared memory."""
+    """Own sender allocation and receiver attachment for shared-memory streams."""
+
+    @staticmethod
+    def _release_receiver_tracking(memory: shared_memory.SharedMemory) -> None:
+        """Keep a non-owning POSIX receiver from unlinking sender-owned memory."""
+        if _USE_POSIX:
+            resource_tracker.unregister(memory._name, "shared_memory")
+
+    @classmethod
+    def copy_sender_owned_array(
+        cls,
+        *,
+        name: str,
+        shape: Sequence[int],
+        dtype: str | np.dtype,
+    ) -> np.ndarray:
+        """Copy an attached array without transferring allocation ownership."""
+        memory = shared_memory.SharedMemory(name=name)
+        try:
+            cls._release_receiver_tracking(memory)
+            return np.ndarray(shape, dtype=dtype, buffer=memory.buf).copy()
+        finally:
+            memory.close()
 
     @classmethod
     def create(

tests/test_streaming_shared_memory.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from polystore.streaming import StreamingSharedMemoryAuthority, _streaming_backend
+
+
+class _SharedMemoryProbe:
+    def __init__(self, source: np.ndarray) -> None:
+        self._name = "/sender-owned"
+        self.buf = bytearray(source.tobytes())
+        self.closed = False
+
+    def close(self) -> None:
+        self.closed = True
+
+
+def test_non_posix_receiver_copies_without_resource_tracker_unregister(
+    monkeypatch,
+) -> None:
+    source = np.arange(12, dtype=np.uint16).reshape(3, 4)
+    memory = _SharedMemoryProbe(source)
+    unregister_calls = []
+    monkeypatch.setattr(_streaming_backend, "_USE_POSIX", False)
+    monkeypatch.setattr(
+        _streaming_backend.shared_memory,
+        "SharedMemory",
+        lambda *, name: memory,
+    )
+    monkeypatch.setattr(
+        _streaming_backend.resource_tracker,
+        "unregister",
+        lambda *args: unregister_calls.append(args),
+    )
+
+    copied = StreamingSharedMemoryAuthority.copy_sender_owned_array(
+        name="sender-owned",
+        shape=source.shape,
+        dtype=str(source.dtype),
+    )
+
+    np.testing.assert_array_equal(copied, source)
+    assert copied.flags.owndata
+    assert unregister_calls == []
+    assert memory.closed
+
+
+def test_posix_receiver_releases_tracking_without_unlinking_sender_memory(
+    monkeypatch,
+) -> None:
+    source = np.arange(6, dtype=np.float32).reshape(2, 3)
+    memory = _SharedMemoryProbe(source)
+    unregister_calls = []
+    monkeypatch.setattr(_streaming_backend, "_USE_POSIX", True)
+    monkeypatch.setattr(
+        _streaming_backend.shared_memory,
+        "SharedMemory",
+        lambda *, name: memory,
+    )
+    monkeypatch.setattr(
+        _streaming_backend.resource_tracker,
+        "unregister",
+        lambda *args: unregister_calls.append(args),
+    )
+
+    copied = StreamingSharedMemoryAuthority.copy_sender_owned_array(
+        name="sender-owned",
+        shape=source.shape,
+        dtype=source.dtype,
+    )
+
+    np.testing.assert_array_equal(copied, source)
+    assert unregister_calls == [(memory._name, "shared_memory")]
+    assert memory.closed
+
+
+def test_receiver_closes_attachment_when_array_projection_fails(monkeypatch) -> None:
+    source = np.arange(2, dtype=np.uint8)
+    memory = _SharedMemoryProbe(source)
+    monkeypatch.setattr(_streaming_backend, "_USE_POSIX", False)
+    monkeypatch.setattr(
+        _streaming_backend.shared_memory,
+        "SharedMemory",
+        lambda *, name: memory,
+    )
+
+    with pytest.raises(TypeError, match="buffer is too small"):
+        StreamingSharedMemoryAuthority.copy_sender_owned_array(
+            name="sender-owned",
+            shape=(3,),
+            dtype=source.dtype,
+        )
+
+    assert memory.closed