mccode-plumber 0.6.0__tar.gz → 0.7.0__tar.gz

{mccode_plumber-0.6.0/src/mccode_plumber.egg-info → mccode_plumber-0.7.0}/PKG-INFO

@@ -1,11 +1,12 @@
 Metadata-Version: 2.4
 Name: mccode-plumber
-Version: 0.6.0
+Version: 0.7.0
 Author-email: Gregory Tucker <gregory.tucker@ess.eu>
 Classifier: License :: OSI Approved :: BSD License
 Description-Content-Type: text/markdown
 Requires-Dist: p4p
-Requires-Dist: file-writer-control>=1.3.0
+Requires-Dist: kafka-python>=2.0
+Requires-Dist: ess-streaming-data-types>=0.14.0
 Requires-Dist: restage>=0.4.0
 Requires-Dist: mccode-to-kafka>=0.2.1
 Requires-Dist: moreniius>=0.2.3

{mccode_plumber-0.6.0 → mccode_plumber-0.7.0}/pyproject.toml

@@ -6,7 +6,8 @@ build-backend = 'setuptools.build_meta'
 name = 'mccode-plumber'
 dependencies = [
     'p4p',
-    'file-writer-control>=1.3.0',
+    'kafka-python>=2.0',
+    'ess-streaming-data-types>=0.14.0',
     'restage>=0.4.0',
     'mccode-to-kafka>=0.2.1',
     'moreniius>=0.2.3',

mccode_plumber-0.7.0/src/mccode_plumber/file_writer_control/CommandChannel.py

@@ -0,0 +1,236 @@
+import atexit
+import threading
+from datetime import datetime
+from queue import Queue
+from typing import Dict, List, Optional, Union
+
+from kafka import KafkaConsumer
+from kafka.errors import NoBrokersAvailable
+
+from file_writer_control.CommandStatus import CommandState, CommandStatus
+from file_writer_control.InThreadStatusTracker import (
+    DEAD_ENTITY_TIME_LIMIT,
+    InThreadStatusTracker,
+)
+from file_writer_control.JobStatus import JobStatus
+from file_writer_control.KafkaTopicUrl import KafkaTopicUrl
+from file_writer_control.WorkerStatus import WorkerStatus
+
+
+def thread_function(
+    host_port: str,
+    topic: str,
+    in_queue: Queue,
+    out_queue: Queue,
+    kafka_config: Dict[str, str] = {},
+):
+    """
+    Background thread for consuming Kafka messages.
+    :param host_port: The host + port of the Kafka broker that we are using.
+    :param topic: The Kafka topic that we are listening to.
+    :param in_queue: A queue for sending "exit" messages to the thread.
+    .. note:: The queue will exit upon the reception of the string "exit" on this queue.
+    :param out_queue: The queue to which status updates are published.
+    """
+    status_tracker = InThreadStatusTracker(out_queue)
+    while True:
+        try:
+            consumer = KafkaConsumer(
+                topic,
+                bootstrap_servers=host_port,
+                fetch_max_bytes=52428800 * 6,
+                max_partition_fetch_bytes=52428800 * 10,
+                consumer_timeout_ms=100,
+                **kafka_config
+            )  # Roughly 300MB
+            break
+        except NoBrokersAvailable:
+            pass  # Do not fail if the broker is not immediately available.
+        if not in_queue.empty():
+            new_msg = in_queue.get()
+            if new_msg == "exit":
+                return
+    while True:
+        for message in consumer:
+            status_tracker.process_message(message.value)
+        status_tracker.check_for_lost_connections()
+        if not in_queue.empty():
+            new_msg = in_queue.get()
+            if new_msg == "exit":
+                break
+    consumer.close(True)
+
+
+class CommandChannel(object):
+    """
+    A class that implements the functionality for receiving and interpreting messages that are published to the
+    Kafka command topic of a pool of file-writers.
+    .. note:: This class implements a thread that will continuously attempt to connect to a Kafka broker.
+    """
+
+    def __init__(self, command_topic_url: str, kafka_config: Dict[str, str] = {}):
+        """
+        Constructor.
+        :param command_topic_url: The url of the Kafka topic to where the file-writer status/command messages are published.
+        """
+        kafka_address = KafkaTopicUrl(command_topic_url)
+        self.status_queue = Queue()
+        self.to_thread_queue = Queue()
+        thread_kwargs = {
+            "host_port": kafka_address.host_port,
+            "topic": kafka_address.topic,
+            "in_queue": self.to_thread_queue,
+            "out_queue": self.status_queue,
+            "kafka_config": kafka_config,
+        }
+        self.map_of_workers: Dict[str, WorkerStatus] = {}
+        self.map_of_jobs: Dict[str, JobStatus] = {}
+        self.map_of_commands: Dict[str, CommandStatus] = {}
+        self.run_thread = True
+        self.thread = threading.Thread(
+            target=thread_function, daemon=True, kwargs=thread_kwargs
+        )
+        self.thread.start()
+
+        def do_exit():
+            self.stop_thread()
+
+        atexit.register(do_exit)
+
+    def add_job_id(self, job_id: str):
+        """
+        Add a job identifier to the list of known jobs before it has been encountered on the command topic.
+        :param job_id: The identifier of the new job.
+        """
+        if job_id not in self.map_of_jobs:
+            self.map_of_jobs[job_id] = JobStatus(job_id)
+
+    def add_command_id(self, job_id: str, command_id: str):
+        """
+        Add a command identifier to the list of known commands before it has been encountered on the command topic.
+        :param job_id: The job identifier of the new command.
+        :param command_id: The identifier of the new command.
+        """
+        if command_id not in self.map_of_commands:
+            self.map_of_commands[command_id] = CommandStatus(job_id, command_id)
+            self.map_of_commands[command_id].state = CommandState.WAITING_RESPONSE
+
+    def stop_thread(self):
+        """
+        Stop the thread that is continuously getting command topic messages in the background. Should only be called if
+        we are about to get rid of the current instance of CommandChannel.
+        """
+        self.to_thread_queue.put("exit")
+        try:
+            self.thread.join()
+        except RuntimeError:
+            pass  # Do not throw an exception if the thread has not yet been started.
+
+    def __del__(self):
+        self.stop_thread()
+
+    def update_workers(self, current_time: Optional[datetime] = None):
+        """
+        Update the list of known workers, jobs and commands. This is a non-blocking call but it might take some time
+        to execute if the queue of updates is long. This member function is called by many of the other member functions
+        in this class.
+        """
+        if current_time is None:
+            current_time = datetime.now()
+
+        def handle_worker_status(status_update):
+            if status_update.service_id not in self.map_of_workers:
+                self.map_of_workers[status_update.service_id] = status_update
+            self.map_of_workers[status_update.service_id].update_status(status_update)
+
+        def handle_job_status(status_update):
+            if status_update.job_id not in self.map_of_jobs:
+                self.map_of_jobs[status_update.job_id] = status_update
+            self.map_of_jobs[status_update.job_id].update_status(status_update)
+
+        def handle_command_status(status_update):
+            if status_update.command_id not in self.map_of_commands:
+                self.map_of_commands[status_update.command_id] = status_update
+            self.map_of_commands[status_update.command_id].update_status(status_update)
+
+        status_updater_map = {
+            WorkerStatus: handle_worker_status,
+            CommandStatus: handle_command_status,
+            JobStatus: handle_job_status,
+        }
+        while not self.status_queue.empty():
+            status_update = self.status_queue.get()
+            status_updater_map[type(status_update)](status_update)
+
+        for entity in (
+            list(self.map_of_workers.values())
+            + list(self.map_of_commands.values())
+            + list(self.map_of_jobs.values())
+        ):
+            entity.check_if_outdated(current_time)
+
+        def pruner(entities_dictionary):
+            for key in list(entities_dictionary.keys()):
+                if (
+                    entities_dictionary[key].last_update + DEAD_ENTITY_TIME_LIMIT
+                    < current_time
+                ):
+                    del entities_dictionary[key]
+
+        pruner(self.map_of_commands)
+        pruner(self.map_of_workers)
+        pruner(self.map_of_jobs)
+
+    def list_workers(self) -> List[WorkerStatus]:
+        """
+        :return: A list of the (known) workers with state and status information.
+        """
+        self.update_workers()
+        return list(self.map_of_workers.values())
+
+    def list_jobs(self) -> List[JobStatus]:
+        """
+        :return: A list of the (known) jobs with state and status information.
+        """
+        self.update_workers()
+        return list(self.map_of_jobs.values())
+
+    def list_commands(self) -> List[CommandStatus]:
+        """
+        :return: A list of the (known) commands and their outcomes.
+        """
+        self.update_workers()
+        return list(self.map_of_commands.values())
+
+    def get_job(self, job_id: str) -> Union[JobStatus, None]:
+        """
+        Get the status of a single job.
+        :param job_id: The job identifier of the job we are interested in.
+        :return: The job status or None if the job is not known.
+        """
+        self.update_workers()
+        if job_id in self.map_of_jobs:
+            return self.map_of_jobs[job_id]
+        return None
+
+    def get_worker(self, service_id: str) -> Union[WorkerStatus, None]:
+        """
+        Get the status of a single worker.
+        :param service_id: The service identifier of the worker we are interested in.
+        :return: The worker status or None if the service id is not known.
+        """
+        self.update_workers()
+        if service_id in self.map_of_workers:
+            return self.map_of_workers[service_id]
+        return None
+
+    def get_command(self, command_id: str) -> Union[CommandStatus, None]:
+        """
+        Get the status of a single command.
+        :param command_id: The command identifier of the command we are interested in.
+        :return: The command status/outcome or None if the command is not known.
+        """
+        self.update_workers()
+        if command_id in self.map_of_commands:
+            return self.map_of_commands[command_id]
+        return None

mccode_plumber-0.7.0/src/mccode_plumber/file_writer_control/CommandHandler.py

@@ -0,0 +1,58 @@
+from datetime import timedelta
+
+from file_writer_control.CommandChannel import CommandChannel
+from file_writer_control.CommandStatus import CommandState
+
+
+class CommandHandler:
+    """
+    A stand in for (more easily) checking the state of a command sent to a file-writer.
+    """
+
+    def __init__(self, command_channel: CommandChannel, command_id: str):
+        """
+        Constructor.
+        :param command_channel: The instance of a CommandChannel that this class uses for getting the command status from.
+        :param command_id: The (unique) command identifier.
+        """
+        self.command_id = command_id
+        self.command_channel = command_channel
+
+    def get_state(self) -> CommandState:
+        """
+        Get the command state.
+        :return: The state of the command if possible. CommandState.UNKNOWN if the the state can not be determined.
+        """
+        command = self.command_channel.get_command(self.command_id)
+        if command is None:
+            return CommandState.UNKNOWN
+        return command.state
+
+    def is_done(self) -> bool:
+        """
+        :return: True if the command completed successfully. False otherwise.
+        """
+        current_state = self.command_channel.get_command(self.command_id).state
+        if current_state == CommandState.ERROR:
+            raise RuntimeError(
+                f'Command failed with error message "{self.get_message()}".'
+            )
+        if current_state == CommandState.TIMEOUT_RESPONSE:
+            raise RuntimeError("Timed out while trying to send command.")
+        return current_state == CommandState.SUCCESS
+
+    def get_message(self) -> str:
+        """
+        :return: If there was an error executing the command, this member function will return the error string as
+        sent by the file-writer. Will return an empty string otherwise.
+        """
+        command = self.command_channel.get_command(self.command_id)
+        if command is None:
+            return ""
+        return command.message
+
+    def set_timeout(self, new_timeout: timedelta):
+        self.command_channel.get_command(self.command_id).timeout = new_timeout
+
+    def get_timeout(self):
+        return self.command_channel.get_command(self.command_id).timeout

mccode_plumber-0.7.0/src/mccode_plumber/file_writer_control/CommandStatus.py

@@ -0,0 +1,151 @@
+from datetime import datetime, timedelta
+from enum import Enum, auto
+from typing import Optional
+
+COMMAND_STATUS_TIMEOUT = timedelta(seconds=60)
+
+
+class CommandState(Enum):
+    """
+    The state of a command.
+    """
+
+    UNKNOWN = auto()
+    NO_COMMAND = auto()
+    WAITING_RESPONSE = auto()
+    TIMEOUT_RESPONSE = auto()
+    ERROR = auto()
+    SUCCESS = auto()
+
+
+class CommandStatus(object):
+    """
+    The status of a command.
+    """
+
+    def __init__(
+        self,
+        job_id: str,
+        command_id: str,
+        command_timeout: timedelta = COMMAND_STATUS_TIMEOUT,
+    ):
+        self._command_timeout = command_timeout
+        self._job_id = job_id
+        self._command_id = command_id
+        self._last_update = datetime.now()
+        self._state = CommandState.NO_COMMAND
+        self._message = ""
+        self._response_code = None
+
+    def __eq__(self, other_status: "CommandStatus"):
+        if not isinstance(other_status, CommandStatus):
+            raise NotImplementedError
+        return (
+            other_status.command_id == self.command_id
+            and other_status.job_id == self.job_id
+            and other_status.state == self.state
+            and other_status.response_code == self.response_code
+        )
+
+    def update_status(self, new_status: "CommandStatus"):
+        """
+        Updates the status/state of this instance of the CommandStatus class using another instance.
+        .. note:: The command identifier of both this instance and the other one must be identical.
+        :param new_status: The other instance of the CommandStatus class.
+        """
+        if new_status.command_id != self.command_id:
+            raise RuntimeError(
+                f"Command id of status update is not correct ({self.command_id} vs {new_status.command_id})"
+            )
+        self._state = new_status.state
+        self._response_code = new_status.response_code
+        if new_status.message:
+            self._message = new_status.message
+        self._last_update = new_status.last_update
+
+    def check_if_outdated(self, current_time: datetime):
+        """
+        Given the current time, state and the time of the last update: Have we lost the connection?
+        :param current_time: The current time
+        """
+        if (
+            self.state != CommandState.SUCCESS
+            and self.state != CommandState.ERROR
+            and self.state != CommandState.TIMEOUT_RESPONSE
+            and current_time - self.last_update > self._command_timeout
+        ):
+            self._state = CommandState.TIMEOUT_RESPONSE
+            self._last_update = current_time
+
+    @property
+    def response_code(self) -> Optional[int]:
+        """
+        A code that mirrors the result of a command if a response has been received. Is set to None otherwise.
+        """
+        return self._response_code
+
+    @response_code.setter
+    def response_code(self, new_code: int):
+        """
+        Set the current response code.
+        """
+        self._response_code = new_code
+        self._last_update = datetime.now()
+
+    @property
+    def job_id(self) -> str:
+        """
+        The job identifier under which this command is executed.
+        """
+        return self._job_id
+
+    @property
+    def command_id(self) -> str:
+        """
+        The unique command identifier of this command.
+        """
+        return self._command_id
+
+    @property
+    def message(self) -> str:
+        """
+        A status/error message as returned by the file-writer that responded to the command.
+        :return:
+        """
+        return self._message
+
+    @message.setter
+    def message(self, new_message: str):
+        if new_message:
+            self._message = new_message
+            self._last_update = datetime.now()
+
+    @property
+    def state(self) -> CommandState:
+        """
+        The current state of the command.
+        """
+        return self._state
+
+    @state.setter
+    def state(self, new_state: CommandState):
+        self._state = new_state
+        self._last_update = datetime.now()
+
+    @property
+    def last_update(self) -> datetime:
+        """
+        The local time stamp of the last update of the status of the command.
+        """
+        return self._last_update
+
+    @property
+    def timeout(self) -> timedelta:
+        """
+        Timeout for waiting for response to command.
+        """
+        return self._command_timeout
+
+    @timeout.setter
+    def timeout(self, new_timeout: timedelta):
+        self._command_timeout = new_timeout