Fix streaming race condition when streaming to multiple viewers

- Use fresh REQ socket for each send instead of caching (prevents ZMQ state errors)
- Standardize both Napari and Fiji on REQ/REP pattern (was PUB/SUB for Napari)
- Remove hardcoded VIEWER_TYPE == 'fiji' check in _get_publisher()
- Both backends now block until viewer acknowledges receipt

Fixes: Cannot stream to both Napari and Fiji in quick succession

Author: trissim

src/polystore/fiji_stream.py

@@ -12,13 +12,17 @@
 """
 
 import logging
+import time
 from pathlib import Path
 from typing import Any, List, Union
 
+import zmq
+
 from .constants import Backend, TransportMode
 from .streaming_constants import StreamingDataType
 from .streaming import StreamingBackend
 from .roi_converters import FijiROIConverter
+from zmqruntime.transport import get_zmq_transport_url, coerce_transport_mode
 
 logger = logging.getLogger(__name__)
 
@@ -81,7 +85,6 @@ def save_batch(self, data_list: List[Any], file_paths: List[Union[str, Path]], *
         port = kwargs['port']
         transport_mode = kwargs['transport_mode']
         transport_config = kwargs.get('transport_config')
-        publisher = self._get_publisher(host, port, transport_mode, transport_config=transport_config)
         display_config = kwargs['display_config']
         microscope_handler = kwargs['microscope_handler']
         source = kwargs.get('source', 'unknown_source')  # Pre-built source value
@@ -132,16 +135,39 @@ def save_batch(self, data_list: List[Any], file_paths: List[Union[str, Path]], *
             transport_config=transport_config,
         )
 
-        # Send with REQ socket (BLOCKING - worker waits for Fiji to acknowledge)
-        # Worker blocks until Fiji receives, copies data from shared memory, and sends ack
-        # This guarantees no messages are lost and shared memory is only closed after Fiji is done
-        logger.info(f"📤 FIJI BACKEND: Sending batch of {len(batch_images)} images to Fiji on port {port} (REQ/REP - blocking until ack)")
-        publisher.send_json(message)  # Blocking send
+        # Create FRESH REQ socket for each send - REQ sockets cannot be reused
+        # This prevents the "Operation cannot be accomplished in current state" error
+        # when multiple streams happen concurrently
+        transport_config = transport_config or self._transport_config
+        url = get_zmq_transport_url(
+            port,
+            host=host,
+            mode=coerce_transport_mode(transport_mode),
+            config=transport_config,
+        )
+
+        if self._context is None:
+            self._context = zmq.Context()
+
+        socket = self._context.socket(zmq.REQ)
+        socket.connect(url)
+        time.sleep(0.1)  # Brief delay for connection to establish
+
+        try:
+            # Send with REQ socket (BLOCKING - worker waits for Fiji to acknowledge)
+            # Worker blocks until Fiji receives, copies data from shared memory, and sends ack
+            # This guarantees no messages are lost and shared memory is only closed after Fiji is done
+            logger.info(f"📤 FIJI BACKEND: Sending batch of {len(batch_images)} images to Fiji on port {port} (REQ/REP - blocking until ack)")
+            socket.send_json(message)  # Blocking send
+
+            # Wait for acknowledgment from Fiji (REP socket)
+            # Fiji will only reply after it has copied all data from shared memory
+            ack_response = socket.recv_json()
+            logger.info(f"✅ FIJI BACKEND: Received ack from Fiji: {ack_response.get('status', 'unknown')}")
 
-        # Wait for acknowledgment from Fiji (REP socket)
-        # Fiji will only reply after it has copied all data from shared memory
-        ack_response = publisher.recv_json()
-        logger.info(f"✅ FIJI BACKEND: Received ack from Fiji: {ack_response.get('status', 'unknown')}")
+        finally:
+            # Always close the socket - never reuse REQ sockets
+            socket.close()
 
         # Clean up publisher's handles after successful send
         # Receiver will unlink the shared memory after copying the data

src/polystore/napari_stream.py

@@ -9,10 +9,11 @@
 - Sender (Worker): Creates shared memory, sends reference via ZMQ, closes handle (does NOT unlink)
 - Receiver (Napari Server): Attaches to shared memory, copies data, closes handle, unlinks
 - Only receiver calls unlink() to prevent FileNotFoundError
-- PUB/SUB socket pattern is non-blocking; receiver must copy data before sender closes handle
+- REQ/REP socket pattern is blocking; worker waits for acknowledgment before closing shared memory
 """
 
 import logging
+import time
 from pathlib import Path
 from typing import Any, List, Union
 
@@ -22,6 +23,7 @@
 from .streaming_constants import StreamingDataType
 from .streaming import StreamingBackend
 from .roi_converters import NapariROIConverter
+from zmqruntime.transport import get_zmq_transport_url, coerce_transport_mode
 
 logger = logging.getLogger(__name__)
 
@@ -82,7 +84,6 @@ def save_batch(self, data_list: List[Any], file_paths: List[Union[str, Path]], *
         port = kwargs['port']
         transport_mode = kwargs['transport_mode']
         transport_config = kwargs.get('transport_config')
-        publisher = self._get_publisher(host, port, transport_mode, transport_config=transport_config)
         display_config = kwargs['display_config']
         microscope_handler = kwargs['microscope_handler']
         source = kwargs.get('source', 'unknown_source')  # Pre-built source value
@@ -114,22 +115,43 @@ def save_batch(self, data_list: List[Any], file_paths: List[Union[str, Path]], *
             transport_config=transport_config,
         )
 
-        # Send non-blocking to prevent hanging if Napari is slow to process (matches Fiji pattern)
-        send_succeeded = False
-        try:
-            publisher.send_json(message, flags=zmq.NOBLOCK)
-            send_succeeded = True
+        # Create FRESH REQ socket for each send - REQ sockets cannot be reused
+        # This prevents the "Operation cannot be accomplished in current state" error
+        # when multiple streams happen concurrently
+        transport_config = transport_config or self._transport_config
+        url = get_zmq_transport_url(
+            port,
+            host=host,
+            mode=coerce_transport_mode(transport_mode),
+            config=transport_config,
+        )
 
-        except zmq.Again:
-            logger.warning(f"Napari viewer busy, dropped batch of {len(batch_images)} images (port {port})")
+        if self._context is None:
+            self._context = zmq.Context()
 
-        except Exception as e:
-            logger.error(f"Failed to send batch to Napari on port {port}: {e}", exc_info=True)
-            raise  # Re-raise the exception so the pipeline knows it failed
+        socket = self._context.socket(zmq.REQ)
+        socket.connect(url)
+        time.sleep(0.1)  # Brief delay for connection to establish
+
+        try:
+            # Send with REQ socket (BLOCKING - worker waits for Napari to acknowledge)
+            # Worker blocks until Napari receives, copies data from shared memory, and sends ack
+            # This guarantees no messages are lost and shared memory is only closed after Napari is done
+            logger.info(f"📤 NAPARI BACKEND: Sending batch of {len(batch_images)} images to Napari on port {port} (REQ/REP - blocking until ack)")
+            socket.send_json(message)  # Blocking send
+
+            # Wait for acknowledgment from Napari (REP socket)
+            # Napari will only reply after it has copied all data from shared memory
+            ack_response = socket.recv_json()
+            logger.info(f"✅ NAPARI BACKEND: Received ack from Napari: {ack_response.get('status', 'unknown')}")
 
         finally:
-            # Unified cleanup: close our handle after successful send, close+unlink after failure
-            self._cleanup_shared_memory_blocks(batch_images, unlink=not send_succeeded)
+            # Always close the socket - never reuse REQ sockets
+            socket.close()
+
+        # Clean up publisher's handles after successful send
+        # Receiver will unlink the shared memory after copying the data
+        self._cleanup_shared_memory_blocks(batch_images, unlink=False)
 
     # cleanup() now inherited from ABC
 

src/polystore/streaming/_streaming_backend.py

@@ -135,16 +135,12 @@ def _get_publisher(self, host: str, port: int, transport_mode: TransportMode, tr
                 if self._context is None:
                     self._context = zmq.Context()
 
-                # Use REQ socket for Fiji (synchronous request/reply - worker blocks until Fiji acks)
-                # Use PUB socket for Napari (broadcast pattern)
-                socket_type = zmq.REQ if self.VIEWER_TYPE == 'fiji' else zmq.PUB
-                publisher = self._context.socket(socket_type)
-
-                if socket_type == zmq.PUB:
-                    publisher.setsockopt(zmq.SNDHWM, 100000)  # Only for PUB sockets
+                # Use REQ socket for all viewers (synchronous request/reply)
+                # All viewers must send acknowledgment after processing
+                publisher = self._context.socket(zmq.REQ)
 
                 publisher.connect(url)
-                socket_name = "REQ" if socket_type == zmq.REQ else "PUB"
+                socket_name = "REQ"
                 logger.info(f"{self.VIEWER_TYPE} streaming {socket_name} socket connected to {url}")
                 time.sleep(0.1)
                 self._publishers[key] = publisher