Repository: LMCache/LMCache
File: lmcache/v1/cache_controller/controller_manager.py (Python)

Code before:
        self.pull_socket_active_requests = 0
        self.reply_socket_active_requests = 0

        prometheus_logger = PrometheusLogger.GetInstanceOrNone()
        if prometheus_logger is not None:
            prometheus_logger.pull_socket_message_count.set_function(
                lambda: self.pull_socket_message_count
            )
            prometheus_logger.reply_socket_message_count.set_function(
                lambda: self.reply_socket_message_count
            )

            # Socket queue/backlog metrics
            prometheus_logger.pull_socket_has_pending.set_function(
                lambda: self._check_socket_has_pending(self.controller_pull_socket)
            )
            if self.controller_urls["reply"] is not None:
                prometheus_logger.rep_socket_has_pending.set_function(
                    lambda: self._check_socket_has_pending(self.controller_rep_socket)
                )

            # Active request metrics
            prometheus_logger.pull_socket_active_requests.set_function(
                lambda: self.pull_socket_active_requests
            )
            prometheus_logger.rep_socket_active_requests.set_function(
                lambda: self.reply_socket_active_requests
            )

    def _check_socket_has_pending(self, socket) -> int:
        """Check if socket has pending messages.

        Returns:
            1 if socket has pending messages, 0 otherwise
        """
        try:
            events = socket.get(zmq.EVENTS)  # type: ignore[attr-defined]
            # Check if POLLIN flag is set (indicates readable/pending messages)
            has_pending = 1 if (events & zmq.POLLIN) else 0  # type: ignore[attr-defined]
            return has_pending
        except Exception as e:
            logger.error(f"Error checking socket pending status: {e}")
            return 0

    async def handle_batched_push_request(self, socket) -> Optional[MsgBase]:
        while True:
            parts = await socket.recv_multipart()
            part_count = len(parts)
            with self._SocketMetricsContext(self, SocketType.PULL, part_count):
                for part in parts:
                    # Parse message based on format

Review comment (suggestion):
Ideally we should stick with `rep` or stick with `reply`. Switching between them is slightly confusing.

Code after:
    async def handle_batched_push_request(self, socket) -> Optional[MsgBase]:
        while True:
            parts = await socket.recv_multipart()
            part_count = len(parts)
            with SocketMetricsContext(self, SocketType.PULL, part_count):
                for part in parts:
                    # Parse message based on format
                    if part.startswith(b"{"):
                        # JSON format - typically from external systems
                        # like Mooncake
                        msg_dict = json.loads(part)
                        msg = msgspec.convert(msg_dict, type=Msg)
                    else:
                        # MessagePack format - internal LMCache communication
                        msg = msgspec.msgpack.decode(part, type=Msg)
                    if isinstance(msg, WorkerMsg):
                        await self.handle_worker_message(msg)

                    # FIXME(Jiayi): The abstraction of control messages
                    # might not be necessary.
                    # elif isinstance(msg, ControlMsg):
                    #    await self.issue_control_message(msg)
                    elif isinstance(msg, OrchMsg):
                        await self.handle_orchestration_message(msg)
                    else:
                        logger.error(f"Unknown message type: {type(msg)}")

    async def handle_batched_req_request(self, socket) -> Optional[MsgBase]:
        while True:
            part = await socket.recv()
            with SocketMetricsContext(self, SocketType.REPLY):
                # Parse message based on format
                if part.startswith(b"{"):
                    # JSON format - typically from external systems like Mooncake
                    msg_dict = json.loads(part)
                    msg = msgspec.convert(msg_dict, type=Msg)
                else:
                    # MessagePack format - internal LMCache communication
                    msg = msgspec.msgpack.decode(part, type=Msg)

                if isinstance(msg, WorkerReqMsg):
                    ret_msg = await self.handle_worker_req_message(msg)
                    await socket.send(msgspec.msgpack.encode(ret_msg))
                else:
                    logger.error(f"Unknown message type: {type(msg)}")
                    err_msg = ErrorMsg(error=f"Unknown message type: {type(msg)}")
                    await socket.send(msgspec.msgpack.encode(err_msg))

    async def health_check(self):
        while True:
            time.sleep(self.health_check_interval)