nv-ingest-api 2025.4.15.dev20250415__py3-none-any.whl → 2025.4.17.dev20250417__py3-none-any.whl

nv_ingest_api/util/message_brokers/simple_message_broker/broker.py

@@ -0,0 +1,465 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import uuid
+import socket
+import socketserver
+import json
+import logging
+import threading
+from typing import Optional
+
+from pydantic import ValidationError
+
+from nv_ingest_api.internal.schemas.message_brokers.request_schema import (
+    PushRequestSchema,
+    PopRequestSchema,
+    SizeRequestSchema,
+)
+from nv_ingest_api.internal.schemas.message_brokers.response_schema import ResponseSchema
+from nv_ingest_api.util.message_brokers.simple_message_broker.ordered_message_queue import OrderedMessageQueue
+
+logger = logging.getLogger(__name__)
+
+
+class SimpleMessageBrokerHandler(socketserver.BaseRequestHandler):
+    """
+    Handles incoming client requests for the SimpleMessageBroker server, processes commands such as
+    PUSH, POP, SIZE, and PING, and manages message queues with thread-safe operations.
+    """
+
+    def handle(self):
+        """
+        Handles incoming client requests, validates the request data, and dispatches to the appropriate
+        command handler.
+
+        Raises
+        ------
+        ValidationError
+            If the incoming request data fails schema validation.
+        Exception
+            If there is an unexpected error while processing the request.
+        """
+
+        client_address = self.client_address
+
+        data_bytes = None
+        try:
+            data_length_bytes = self._recv_exact(8)
+            if not data_length_bytes:
+                logger.debug("No data length received. Closing connection.")
+                return
+            data_length = int.from_bytes(data_length_bytes, "big")
+
+            data_bytes = self._recv_exact(data_length)
+            if not data_bytes:
+                logger.debug("No data received. Closing connection.")
+                return
+
+            data = data_bytes.decode("utf-8").strip()
+            request_data = json.loads(data)
+
+            command = request_data.get("command")
+            if not command:
+                response = ResponseSchema(response_code=1, response_reason="No command specified")
+                self._send_response(response)
+                return
+
+            # Handle the PING command directly
+            if command == "PING":
+                self._handle_ping()
+                return
+
+            # Validate and extract common fields
+            queue_name = request_data.get("queue_name")
+
+            # Initialize the queue and its lock if necessary
+            if queue_name:
+                self.server._initialize_queue(queue_name)
+                queue_lock = self.server.queue_locks[queue_name]
+                queue = self.server.queues[queue_name]
+            else:
+                queue_lock = None
+                queue = None  # For commands that don't require a queue
+
+            # Dispatch to the appropriate handler
+            if command == "PUSH":
+                validated_data = PushRequestSchema(**request_data)
+                self._handle_push(validated_data, transaction_id=str(uuid.uuid4()), queue=queue, queue_lock=queue_lock)
+            elif command == "PUSH_FOR_NV_INGEST":
+                validated_data = PushRequestSchema(**request_data)
+                self._handle_push_for_nv_ingest(validated_data, queue=queue, queue_lock=queue_lock)
+            elif command == "POP":
+                validated_data = PopRequestSchema(**request_data)
+                self._handle_pop(validated_data, queue=queue, queue_lock=queue_lock)
+            elif command == "SIZE":
+                validated_data = SizeRequestSchema(**request_data)
+                response = self._size_of_queue(validated_data, queue, queue_lock)
+                self._send_response(response)
+            else:
+                response = ResponseSchema(response_code=1, response_reason="Unknown command")
+                self._send_response(response)
+
+        except ValidationError as ve:
+            response = ResponseSchema(response_code=1, response_reason=str(ve))
+            self._send_response(response)
+        except Exception as e:
+            logger.error(f"Error processing command from {client_address}: {e}\n{data_bytes}")
+            response = ResponseSchema(response_code=1, response_reason=str(e))
+            try:
+                self._send_response(response)
+            except BrokenPipeError:
+                logger.error("Cannot send error response; client connection closed.")
+
+    def _handle_ping(self):
+        """
+        Responds to a PING command with a PONG response.
+        """
+
+        response = ResponseSchema(response_code=0, response="PONG")
+        self._send_response(response)
+
+    def _handle_push(
+        self, data: PushRequestSchema, transaction_id: str, queue: OrderedMessageQueue, queue_lock: threading.Lock
+    ):
+        """
+        Handles a PUSH command to add a message to the specified queue.
+
+        Parameters
+        ----------
+        data : PushRequestSchema
+            The validated data for the PUSH command.
+        transaction_id : str
+            The unique transaction ID for the operation.
+        queue : OrderedMessageQueue
+            The queue where the message will be pushed.
+        queue_lock : threading.Lock
+            The lock object to ensure thread-safe access to the queue.
+        """
+
+        timeout = data.timeout
+
+        with queue_lock:
+            if queue.full():
+                # Return failure response immediately
+                response = ResponseSchema(response_code=1, response_reason="Queue is full")
+                self._send_response(response)
+                return
+
+        # Proceed with the 3-way handshake
+        initial_response = ResponseSchema(
+            response_code=0, response="Transaction initiated. Waiting for ACK.", transaction_id=transaction_id
+        )
+        self._send_response(initial_response)
+
+        # Wait for ACK
+        if not self._wait_for_ack(transaction_id, timeout):
+            logger.debug(f"Transaction {transaction_id}: ACK not received. Discarding data.")
+            final_response = ResponseSchema(
+                response_code=1, response_reason="ACK not received.", transaction_id=transaction_id
+            )
+        else:
+            # Perform the PUSH operation after ACK
+            with queue_lock:
+                queue.push(data.message)
+            final_response = ResponseSchema(response_code=0, response="Data stored.", transaction_id=transaction_id)
+
+        # Send final response
+        self._send_response(final_response)
+
+    def _handle_push_for_nv_ingest(
+        self, data: PushRequestSchema, queue: OrderedMessageQueue, queue_lock: threading.Lock
+    ):
+        """
+        Handles a PUSH_FOR_NV_INGEST command, which includes generating a unique job ID and
+        updating the message payload accordingly.
+
+        Parameters
+        ----------
+        data : PushRequestSchema
+            The validated data for the PUSH_FOR_NV_INGEST command.
+        queue : OrderedMessageQueue
+            The queue where the message will be pushed.
+        queue_lock : threading.Lock
+            The lock object to ensure thread-safe access to the queue.
+        """
+
+        timeout = data.timeout
+
+        # Deserialize the message
+        try:
+            message_dict = json.loads(data.message)
+        except json.JSONDecodeError:
+            response = ResponseSchema(response_code=1, response_reason="Invalid JSON message")
+            self._send_response(response)
+            return
+
+        # Generate a UUID for 'job_id' and use it as transaction_id
+        transaction_id = str(uuid.uuid4())
+        message_dict["job_id"] = transaction_id
+
+        # Re-serialize the message
+        updated_message = json.dumps(message_dict)
+
+        with queue_lock:
+            if queue.full():
+                # Return failure response immediately
+                response = ResponseSchema(response_code=1, response_reason="Queue is full")
+                self._send_response(response)
+                return
+
+        # Proceed with the 3-way handshake
+        initial_response = ResponseSchema(
+            response_code=0, response="Transaction initiated. Waiting for ACK.", transaction_id=transaction_id
+        )
+        self._send_response(initial_response)
+
+        # Wait for ACK
+        if not self._wait_for_ack(transaction_id, timeout):
+            logger.debug(f"Transaction {transaction_id}: ACK not received. Discarding data.")
+            final_response = ResponseSchema(
+                response_code=1, response_reason="ACK not received.", transaction_id=transaction_id
+            )
+        else:
+            # Perform the PUSH operation after ACK
+            with queue_lock:
+                queue.push(updated_message)
+            final_response = ResponseSchema(response_code=0, response="Data stored.", transaction_id=transaction_id)
+
+        # Send final response
+        self._send_response(final_response)
+
+    def _handle_pop(self, data: PopRequestSchema, queue: OrderedMessageQueue, queue_lock: threading.Lock):
+        """
+        Handles a POP command to retrieve a message from the specified queue.
+
+        Parameters
+        ----------
+        data : PopRequestSchema
+            The validated data for the POP command.
+        queue : OrderedMessageQueue
+            The queue from which the message will be retrieved.
+        queue_lock : threading.Lock
+            The lock object to ensure thread-safe access to the queue.
+        """
+
+        timeout = data.timeout
+        transaction_id = str(uuid.uuid4())
+
+        with queue_lock:
+            if queue.empty():
+                # Return failure response immediately
+                response = ResponseSchema(response_code=1, response_reason="Queue is empty")
+                self._send_response(response)
+                return
+            # Pop the message from the queue
+            message = queue.pop(transaction_id)
+
+        # Proceed with the 3-way handshake
+        initial_response = ResponseSchema(response_code=0, response=message, transaction_id=transaction_id)
+        self._send_response(initial_response)
+
+        # Wait for ACK
+        if not self._wait_for_ack(transaction_id, timeout):
+            logger.debug(f"Transaction {transaction_id}: ACK not received. Returning data to queue.")
+            with queue_lock:
+                queue.return_message(transaction_id)
+            final_response = ResponseSchema(
+                response_code=1, response_reason="ACK not received.", transaction_id=transaction_id
+            )
+        else:
+            with queue_lock:
+                queue.acknowledge(transaction_id)
+            final_response = ResponseSchema(response_code=0, response="Data processed.", transaction_id=transaction_id)
+
+        # Send final response
+        self._send_response(final_response)
+
+    def _size_of_queue(
+        self, data: SizeRequestSchema, queue: OrderedMessageQueue, queue_lock: threading.Lock
+    ) -> ResponseSchema:
+        """
+        Retrieves the size of the specified queue.
+
+        Parameters
+        ----------
+        data : SizeRequestSchema
+            The validated data for the SIZE command.
+        queue : OrderedMessageQueue
+            The queue whose size will be queried.
+        queue_lock : threading.Lock
+            The lock object to ensure thread-safe access to the queue.
+
+        Returns
+        -------
+        ResponseSchema
+            A response containing the size of the queue.
+        """
+
+        with queue_lock:
+            size = queue.qsize()
+        return ResponseSchema(response_code=0, response=str(size))
+
+    def _wait_for_ack(self, transaction_id: str, timeout: Optional[float]) -> bool:
+        """
+        Waits for an acknowledgment (ACK) from the client for a specific transaction.
+
+        Parameters
+        ----------
+        transaction_id : str
+            The unique transaction ID for the operation.
+        timeout : float, optional
+            The timeout period for waiting for the ACK.
+
+        Returns
+        -------
+        bool
+            True if the ACK is received, False otherwise.
+        """
+
+        try:
+            self.request.settimeout(timeout)
+            ack_length_bytes = self._recv_exact(8)
+            if not ack_length_bytes or len(ack_length_bytes) < 8:
+                return False
+            ack_length = int.from_bytes(ack_length_bytes, "big")
+            ack_data_bytes = self._recv_exact(ack_length)
+            if not ack_data_bytes or len(ack_data_bytes) < ack_length:
+                return False
+            ack_data = ack_data_bytes.decode("utf-8")
+            ack_response = json.loads(ack_data)
+            return ack_response.get("transaction_id") == transaction_id and ack_response.get("ack") is True
+        except (socket.timeout, json.JSONDecodeError, ConnectionResetError) as e:
+            logger.error(f"Error waiting for ACK: {e}")
+            return False
+        finally:
+            self.request.settimeout(None)
+
+    def _send_response(self, response: ResponseSchema):
+        """
+        Sends a response back to the client.
+
+        Parameters
+        ----------
+        response : ResponseSchema
+            The response to send to the client.
+
+        Raises
+        ------
+        Exception
+            If there is an error while sending the response.
+        """
+
+        try:
+            response_json = response.model_dump_json().encode("utf-8")
+            total_length = len(response_json)
+            self.request.sendall(total_length.to_bytes(8, "big"))
+            self.request.sendall(response_json)
+        except BrokenPipeError as e:
+            logger.error(f"BrokenPipeError while sending response: {e}")
+            # Handle the broken pipe gracefully
+        except Exception as e:
+            logger.error(f"Unexpected error while sending response: {e}")
+
+    def _recv_exact(self, num_bytes: int) -> Optional[bytes]:
+        """
+        Receives an exact number of bytes from the client connection.
+
+        Parameters
+        ----------
+        num_bytes : int
+            The number of bytes to receive.
+
+        Returns
+        -------
+        Optional[bytes]
+            The received bytes, or None if the connection is closed.
+        """
+
+        data = bytearray()
+        while len(data) < num_bytes:
+            packet = self.request.recv(num_bytes - len(data))
+            if not packet:
+                return None
+            data.extend(packet)
+        return bytes(data)
+
+
+class SimpleMessageBroker(socketserver.ThreadingMixIn, socketserver.TCPServer):
+    """
+    A thread-safe message broker server that manages multiple message queues and supports commands
+    such as PUSH, POP, SIZE, and PING.
+    """
+
+    allow_reuse_address = True
+    _instances = {}
+    _instances_lock = threading.Lock()
+
+    def __new__(cls, host: str, port: int, max_queue_size: int):
+        """
+        Ensures that only one instance of SimpleMessageBroker is created per host and port combination.
+
+        Parameters
+        ----------
+        host : str
+            The hostname or IP address for the server.
+        port : int
+            The port number for the server.
+        max_queue_size : int
+            The maximum size of each message queue.
+
+        Returns
+        -------
+        SimpleMessageBroker
+            The singleton instance of the server.
+        """
+
+        key = (host, port)
+        with cls._instances_lock:
+            if key not in cls._instances:
+                # Create a new instance and store it in the instances dictionary
+                instance = super(SimpleMessageBroker, cls).__new__(cls)
+                cls._instances[key] = instance
+            else:
+                instance = cls._instances[key]
+        return instance
+
+    def __init__(self, host: str, port: int, max_queue_size: int):
+        """
+        Initializes the SimpleMessageBroker server, setting up message queues and locks.
+
+        Parameters
+        ----------
+        host : str
+            The hostname or IP address for the server.
+        port : int
+            The port number for the server.
+        max_queue_size : int
+            The maximum size of each message queue.
+        """
+
+        # Prevent __init__ from running multiple times on the same instance
+        if hasattr(self, "_initialized") and self._initialized:
+            return
+        super().__init__((host, port), SimpleMessageBrokerHandler)
+        self.max_queue_size = max_queue_size
+        self.queues = {}
+        self.queue_locks = {}  # Dictionary to hold locks for each queue
+        self.lock = threading.Lock()  # Global lock to protect access to queues and locks
+        self._initialized = True  # Flag to indicate initialization is complete
+
+    def _initialize_queue(self, queue_name: str):
+        """
+        Initializes a new message queue with the specified name if it doesn't already exist.
+
+        Parameters
+        ----------
+        queue_name : str
+            The name of the queue to initialize.
+        """
+
+        with self.lock:
+            if queue_name not in self.queues:
+                self.queues[queue_name] = OrderedMessageQueue(maxsize=self.max_queue_size)
+                self.queue_locks[queue_name] = threading.Lock()

nv_ingest_api/util/message_brokers/simple_message_broker/ordered_message_queue.py

@@ -0,0 +1,71 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import threading
+import heapq
+
+
+class OrderedMessageQueue:
+    def __init__(self, maxsize=0):
+        self.queue = []  # List of (index, message) tuples
+        self.maxsize = maxsize
+        self.next_index = 0  # Monotonically increasing message index
+        self.in_flight = {}  # Mapping of transaction_id to (index, message)
+        self.lock = threading.Lock()
+        self.not_empty = threading.Condition(self.lock)
+        self.not_full = threading.Condition(self.lock)
+
+    def can_push(self):
+        """Check if the queue can accept more messages."""
+        with self.lock:
+            return self.maxsize == 0 or (len(self.queue) + len(self.in_flight)) < self.maxsize
+
+    def push(self, message):
+        """Add a message to the queue after it has been acknowledged."""
+        with self.lock:
+            index = self.next_index
+            self.next_index += 1
+            heapq.heappush(self.queue, (index, message))
+            self.not_empty.notify()
+            return index
+
+    def pop(self, transaction_id):
+        """Pop a message from the queue and mark it as in-flight."""
+        with self.lock:
+            while not self.queue:
+                self.not_empty.wait()
+            index, message = heapq.heappop(self.queue)
+            self.in_flight[transaction_id] = (index, message)
+            self.not_full.notify()
+            return message
+
+    def acknowledge(self, transaction_id):
+        """Acknowledge that a message has been processed."""
+        with self.lock:
+            if transaction_id in self.in_flight:
+                del self.in_flight[transaction_id]
+
+    def return_message(self, transaction_id):
+        """Return an unacknowledged message back to the queue."""
+        with self.lock:
+            if transaction_id in self.in_flight:
+                index, message = self.in_flight.pop(transaction_id)
+                heapq.heappush(self.queue, (index, message))
+                self.not_empty.notify()
+
+    def qsize(self):
+        """Get the number of messages currently in the queue."""
+        with self.lock:
+            return len(self.queue)
+
+    def empty(self):
+        """Check if the queue is empty."""
+        with self.lock:
+            return not self.queue
+
+    def full(self):
+        """Check if the queue is full."""
+        with self.lock:
+            return self.maxsize > 0 and (len(self.queue) + len(self.in_flight)) >= self.maxsize