Clean up minor issues in ZMQ broker

- Remove AiiDA import from server.py, default to json encoder/decoder;
  service.py passes YAML encoder/decoder explicitly
- Remove unused threading lock from server
- Remove unused dataclasses from protocol.py (all construction uses
  make_* factory functions)
- Move handler dispatch dict to instance attribute (avoid per-message
  allocation)
- Replace hasattr broker checks with isinstance(broker, ZmqBroker)

Author: agoscinski

src/aiida/brokers/zmq/protocol.py

@@ -47,7 +47,6 @@
 
 import json
 import uuid
-from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any
 
@@ -77,71 +76,6 @@ class MessageType(str, Enum):
     PING = 'ping'
 
 
-@dataclass
-class Message:
-    """Base message structure."""
-
-    type: MessageType
-    id: str = field(default_factory=lambda: uuid.uuid4().hex)
-    sender: str = ''
-
-
-@dataclass
-class TaskMessage(Message):
-    """Task message sent to the broker for processing."""
-
-    type: MessageType = field(default=MessageType.TASK, init=False)
-    body: Any = None
-    no_reply: bool = False
-
-
-@dataclass
-class TaskResponse(Message):
-    """Response to a task message."""
-
-    type: MessageType = field(default=MessageType.TASK_RESPONSE, init=False)
-    task_id: str = ''
-    result: Any = None
-    error: str | None = None
-
-
-@dataclass
-class RpcMessage(Message):
-    """RPC message sent to a specific recipient."""
-
-    type: MessageType = field(default=MessageType.RPC, init=False)
-    recipient: str = ''
-    body: Any = None
-
-
-@dataclass
-class RpcResponse(Message):
-    """Response to an RPC message."""
-
-    type: MessageType = field(default=MessageType.RPC_RESPONSE, init=False)
-    rpc_id: str = ''
-    result: Any = None
-    error: str | None = None
-
-
-@dataclass
-class BroadcastMessage(Message):
-    """Broadcast message sent to all subscribers."""
-
-    type: MessageType = field(default=MessageType.BROADCAST, init=False)
-    body: Any = None
-    subject: str | None = None
-    correlation_id: str | None = None
-
-
-@dataclass
-class SubscribeMessage(Message):
-    """Subscription request message."""
-
-    type: MessageType = MessageType.SUBSCRIBE_TASK
-    identifier: str | None = None
-
-
 def encode_message(msg: dict, encoder=json.dumps) -> bytes:
     """Encode a message dictionary to bytes."""
     return encoder(msg).encode('utf-8')

src/aiida/brokers/zmq/server.py

@@ -1,25 +1,22 @@
-"""ZeroMQ Broker Server - standalone message broker.
+"""ZeroMQ Broker Server.
 
-This module is completely independent of AiiDA and can be used as a standalone
-message broker server. It handles:
+Can be started as a standalone message broker process. It handles:
 - Task queue management with persistence
 - Request/reply routing for RPC
 - Broadcast distribution
 """
 
 from __future__ import annotations
 
+import json
 import logging
-import threading
 import time
 from collections import deque
 from pathlib import Path
 from typing import Any, Callable
 
 import zmq
 
-from aiida.brokers.utils import YAML_DECODER, YAML_ENCODER
-
 from .protocol import MessageType, decode_message, encode_message
 from .queue import PersistentQueue
 
@@ -54,8 +51,8 @@ def __init__(
         :param encoder: Function to encode messages (default: yaml.dump)
         :param decoder: Function to decode messages (default: yaml.load)
         """
-        encoder = encoder if encoder is not None else YAML_ENCODER
-        decoder = decoder if decoder is not None else YAML_DECODER
+        encoder = encoder if encoder is not None else json.dumps
+        decoder = decoder if decoder is not None else json.loads
         self._storage_path = Path(storage_path)
         self._sockets_path = Path(sockets_path)
 
@@ -96,7 +93,21 @@ def __init__(
 
         # Server state
         self._running = False
-        self._lock = threading.Lock()
+
+        # Message type -> handler mapping (built once, not per message)
+        self._handlers: dict[str, Callable] = {
+            MessageType.TASK.value: self._handle_task,
+            MessageType.TASK_RESPONSE.value: self._handle_task_response,
+            MessageType.TASK_ACK.value: self._handle_task_ack,
+            MessageType.TASK_NACK.value: self._handle_task_nack,
+            MessageType.RPC.value: self._handle_rpc,
+            MessageType.RPC_RESPONSE.value: self._handle_rpc_response,
+            MessageType.BROADCAST.value: self._handle_broadcast,
+            MessageType.SUBSCRIBE_TASK.value: self._handle_subscribe_task,
+            MessageType.SUBSCRIBE_RPC.value: self._handle_subscribe_rpc,
+            MessageType.UNSUBSCRIBE_TASK.value: self._handle_unsubscribe_task,
+            MessageType.UNSUBSCRIBE_RPC.value: self._handle_unsubscribe_rpc,
+        }
 
     @property
     def storage_path(self) -> Path:
@@ -257,22 +268,7 @@ def _handle_router_message(self) -> None:
                 _LOGGER.warning('Message missing type field')
                 return
 
-            # Route by message type
-            handlers: dict[str, Any] = {
-                MessageType.TASK.value: self._handle_task,
-                MessageType.TASK_RESPONSE.value: self._handle_task_response,
-                MessageType.TASK_ACK.value: self._handle_task_ack,
-                MessageType.TASK_NACK.value: self._handle_task_nack,
-                MessageType.RPC.value: self._handle_rpc,
-                MessageType.RPC_RESPONSE.value: self._handle_rpc_response,
-                MessageType.BROADCAST.value: self._handle_broadcast,
-                MessageType.SUBSCRIBE_TASK.value: self._handle_subscribe_task,
-                MessageType.SUBSCRIBE_RPC.value: self._handle_subscribe_rpc,
-                MessageType.UNSUBSCRIBE_TASK.value: self._handle_unsubscribe_task,
-                MessageType.UNSUBSCRIBE_RPC.value: self._handle_unsubscribe_rpc,
-            }
-
-            handler = handlers.get(msg_type)
+            handler = self._handlers.get(msg_type)
             if handler:
                 handler(identity, msg)
             else:

src/aiida/brokers/zmq/service.py

@@ -66,9 +66,13 @@ def start(self) -> None:
         self._sockets_path = Path(tempfile.mkdtemp(prefix='aiida_zmq_'))
         self._sockets_file.write_text(str(self._sockets_path))
 
+        from aiida.brokers.utils import YAML_DECODER, YAML_ENCODER
+
         self._server = ZmqBrokerServer(
             storage_path=self._storage_path,
             sockets_path=self._sockets_path,
+            encoder=YAML_ENCODER,
+            decoder=YAML_DECODER,
         )
 
         self._pid_file.write_text(str(os.getpid()))

src/aiida/cmdline/commands/cmd_daemon.py

@@ -151,7 +151,9 @@ def status(ctx, all_profiles, timeout):
         # Build broker status line for managed brokers (e.g., ZMQ)
         broker_line = ''
         broker = get_manager().get_broker()
-        if broker is not None and hasattr(broker, 'get_service_status'):
+        from aiida.brokers.zmq.broker import ZmqBroker
+
+        if isinstance(broker, ZmqBroker):
             if broker.is_running():
                 status_info = broker.get_service_status()
                 if status_info:

src/aiida/cmdline/commands/cmd_status.py

@@ -156,7 +156,9 @@ def verdi_status(print_traceback, no_rmq):
         else:
             daemon_msg = f'Daemon is running with PID {daemon_status["pid"]}'
             # Append broker info for managed brokers (e.g., ZMQ)
-            if hasattr(broker, 'get_service_status'):
+            from aiida.brokers.zmq.broker import ZmqBroker
+
+            if isinstance(broker, ZmqBroker):
                 if broker.is_running():
                     status_info = broker.get_service_status()
                     if status_info: