mccode-plumber 0.6.0__py3-none-any.whl → 0.7.1__py3-none-any.whl

mccode_plumber/file_writer_control/InThreadStatusTracker.py

@@ -0,0 +1,228 @@
+import json
+from datetime import datetime, timedelta
+from queue import Queue
+from typing import Dict
+
+from streaming_data_types import deserialise_6s4t as deserialise_stop_time
+from streaming_data_types import deserialise_answ as deserialise_answer
+from streaming_data_types import deserialise_pl72 as deserialise_start
+from streaming_data_types import deserialise_wrdn as deserialise_stopped
+from streaming_data_types import deserialise_x5f2 as deserialise_status
+from streaming_data_types.action_response_answ import FILE_IDENTIFIER as ANSW_IDENTIFIER
+from streaming_data_types.action_response_answ import Response
+from streaming_data_types.finished_writing_wrdn import (
+    FILE_IDENTIFIER as STOPPED_IDENTIFIER,
+)
+from streaming_data_types.finished_writing_wrdn import WritingFinished
+from streaming_data_types.run_start_pl72 import FILE_IDENTIFIER as START_IDENTIFIER
+from streaming_data_types.run_start_pl72 import RunStartInfo
+from streaming_data_types.run_stop_6s4t import FILE_IDENTIFIER as STOP_TIME_IDENTIFIER
+from streaming_data_types.run_stop_6s4t import RunStopInfo
+from streaming_data_types.status_x5f2 import FILE_IDENTIFIER as STAT_IDENTIFIER
+from streaming_data_types.status_x5f2 import StatusMessage
+from streaming_data_types.utils import get_schema
+
+from file_writer_control.CommandStatus import CommandState, CommandStatus
+from file_writer_control.JobStatus import JobState, JobStatus
+from file_writer_control.StateExtractor import (
+    extract_job_state_from_answer,
+    extract_state_from_command_answer,
+    extract_worker_state_from_status,
+)
+from file_writer_control.WorkerStatus import WorkerState, WorkerStatus
+
+DEAD_ENTITY_TIME_LIMIT = timedelta(hours=1)
+
+
+class InThreadStatusTracker:
+    """
+    Implements de-coding of flatbuffer messages and sends updates of worker, job and command state/status back to the
+    "main"-thread if there has been changes.
+    """
+
+    def __init__(self, status_queue: Queue):
+        """
+        Constructor.
+        :param status_queue: The output queue to which state/status updates are pushed.
+        """
+        self.queue = status_queue
+        self.known_workers: Dict[str, WorkerStatus] = {}
+        self.known_jobs: Dict[str, JobStatus] = {}
+        self.known_commands: Dict[str, CommandStatus] = {}
+
+    def process_message(self, message: bytes):
+        """
+        Process a binary message.
+        :param message: The binary message to be processed.
+        """
+        current_schema = get_schema(message).encode("utf-8")
+        update_time = datetime.now()
+        msg_process_map = {
+            ANSW_IDENTIFIER: lambda msg: self.process_answer(deserialise_answer(msg)),
+            STAT_IDENTIFIER: lambda msg: self.process_status(deserialise_status(msg)),
+            STOP_TIME_IDENTIFIER: lambda msg: self.process_set_stop_time(
+                deserialise_stop_time(msg)
+            ),
+            START_IDENTIFIER: lambda msg: self.process_start(deserialise_start(msg)),
+            STOPPED_IDENTIFIER: lambda msg: self.process_stopped(
+                deserialise_stopped(msg)
+            ),
+        }
+        if current_schema in msg_process_map:
+            msg_process_map[current_schema](message)
+
+        self.send_status_if_updated(update_time)
+
+    def send_status_if_updated(self, limit_time: datetime):
+        """
+        Sends status updates of workers, jobs and commands (to the status queue) if there has been any updates on or
+        after the limit_time.
+        :param limit_time: The cut-off time for deciding which updates should be sent to the status queue.
+        """
+        for entity in (
+            list(self.known_workers.values())
+            + list(self.known_jobs.values())
+            + list(self.known_commands.values())
+        ):
+            if entity.last_update >= limit_time:
+                self.queue.put(entity)
+
+    def check_for_worker_presence(self, service_id: str):
+        """
+        Check if a service_id is known and add it to a list of known ones if it is not.
+        :param service_id: The service identifier to look for.
+        """
+        if service_id not in self.known_workers:
+            self.known_workers[service_id] = WorkerStatus(service_id)
+
+    def check_for_job_presence(self, job_id: str):
+        """
+        Check if a job identifier is known and add it to a list of known ones if it is not.
+        :param job_id: The job identifier to look for.
+        """
+        if job_id not in self.known_jobs:
+            new_job = JobStatus(job_id)
+            self.known_jobs[job_id] = new_job
+
+    def check_for_command_presence(self, job_id: str, command_id: str):
+        """
+        Check if a command identifier is known and add it to a list of known ones if it is not.
+        :param job_id: The job identifier of the command that we are looking for. (Only used if we need to add the
+        command.)
+        :param command_id: The command identifier to look for.
+        """
+        if command_id not in self.known_commands:
+            new_command = CommandStatus(job_id, command_id)
+            self.known_commands[command_id] = new_command
+
+    def check_for_lost_connections(self):
+        """
+        Check workers, commands and jobs for the last update time and change the state of these if a timeout has been
+        reached.
+        """
+        now = datetime.now()
+        for entity in (
+            list(self.known_workers.values())
+            + list(self.known_jobs.values())
+            + list(self.known_commands.values())
+        ):
+            entity.check_if_outdated(now)
+        self.send_status_if_updated(now)
+
+    def prune_dead_entities(self, current_time: datetime):
+        """
+        Will remove old jobs, workers and commands that have not been updated recently.
+        :return:
+        """
+
+        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.known_workers)
+        pruner(self.known_commands)
+        pruner(self.known_jobs)
+
+    def process_answer(self, answer: Response):
+        """
+        Update workers, jobs and commands based on information in a response message.
+        :param answer: The response/answer message to use for status updates.
+        """
+        self.check_for_worker_presence(answer.service_id)
+        self.check_for_job_presence(answer.job_id)
+        self.check_for_command_presence(answer.job_id, answer.command_id)
+        new_job_state = extract_job_state_from_answer(answer)
+        if new_job_state is not None:
+            self.known_jobs[answer.job_id].state = new_job_state
+        current_command = self.known_commands[answer.command_id]
+        current_command.state = extract_state_from_command_answer(answer)
+        current_command.message = answer.message
+        current_command.response_code = Response.status_code
+        self.known_jobs[answer.job_id].message = answer.message
+        try:
+            self.known_jobs[answer.job_id].service_id = answer.service_id
+        except RuntimeError:
+            pass
+
+    def process_status(self, status_update: StatusMessage):
+        """
+        Update workers and jobs based on information in a status message.
+        :param status_update: The status message to use for updates.
+        """
+        self.check_for_worker_presence(status_update.service_id)
+        current_state = extract_worker_state_from_status(status_update)
+        self.known_workers[status_update.service_id].state = current_state
+        if current_state == WorkerState.WRITING:
+            try:
+                json_data = json.loads(status_update.status_json)
+                job_id = json_data["job_id"]
+            except (KeyError, json.JSONDecodeError):
+                raise RuntimeError("Unable to extract JSON data from status message.")
+            file_name = json_data["file_being_written"]
+            self.check_for_job_presence(job_id)
+            self.known_jobs[job_id].state = JobState.WRITING
+            self.known_jobs[job_id].file_name = file_name
+            self.known_jobs[job_id].metadata = json_data
+            # For some jobs, we will only know the service-id when a worker starts working on a job.
+            # Thus we need the following statement to update the (known) service-id of a job.
+            try:
+                self.known_jobs[job_id].service_id = status_update.service_id
+            except RuntimeError:
+                pass  # Expected error (i.e. the job is not known), do nothing
+
+    def process_set_stop_time(self, stop_time: RunStopInfo):
+        """
+        Update commands and jobs based on information in a "set stop time" message.
+        :param stop_time: The "stop" message to use for updates.
+        """
+        self.check_for_command_presence(stop_time.job_id, stop_time.command_id)
+        self.known_commands[stop_time.command_id].state = CommandState.WAITING_RESPONSE
+
+    def process_start(self, start: RunStartInfo):
+        """
+        Update commands and jobs based on information in a "start" message.
+        :param start: The "start" message to use for updates.
+        """
+        self.check_for_job_presence(start.job_id)
+        self.check_for_command_presence(start.job_id, start.job_id)
+        self.known_commands[start.job_id].state = CommandState.WAITING_RESPONSE
+
+    def process_stopped(self, stopped: WritingFinished):
+        """
+        Update workers and jobs based on information in a "has stopped" message.
+        :param stopped: The "stopped" message to use for updates.
+        """
+        self.check_for_job_presence(stopped.job_id)
+        self.check_for_worker_presence(stopped.service_id)
+        current_job = self.known_jobs[stopped.job_id]
+        if stopped.error_encountered:
+            current_job.state = JobState.ERROR
+        else:
+            current_job.state = JobState.DONE
+            current_job.metadata = json.loads(stopped.metadata)
+        current_job.message = stopped.message
+        self.known_workers[stopped.service_id].state = WorkerState.IDLE

mccode_plumber/file_writer_control/JobHandler.py

@@ -0,0 +1,102 @@
+from datetime import datetime
+
+from file_writer_control.CommandHandler import CommandHandler
+from file_writer_control.JobStatus import JobState
+from file_writer_control.WorkerFinder import WorkerFinder
+from file_writer_control.WriteJob import WriteJob
+
+
+class JobHandler:
+    """
+    A stand in for controlling and checking the state of a job running on a file-writer instance.
+    """
+
+    def __init__(self, worker_finder: WorkerFinder, job_id=""):
+        """
+        Constructor.
+        :param worker_finder: An instance of a class that inherits from WorkerFinder and implements the member function
+        try_start_job of that class.
+        :param job_id: (Optional) The job identifier of an existing job.
+        """
+        self.worker_finder = worker_finder
+        self._job_id = job_id
+
+    def start_job(self, job: WriteJob) -> CommandHandler:
+        """
+        Start a write job. This call is not blocking. It does not guarantee that the write job will actually be started.
+        :param job: The write to be started.
+        .. note:: Starting a new job will cause the current instance of this class to no longer being able to track or
+        control previous jobs.
+        :return: A CommandHandler instance that can be used to monitor the outcome of the attempt to start a write job.
+        """
+        self._job_id = job.job_id
+        return self.worker_finder.try_start_job(job)
+
+    def get_state(self) -> JobState:
+        """
+        Get the state of the job.
+        """
+        return self.worker_finder.get_job_state(self._job_id)
+
+    def is_done(self) -> bool:
+        """
+        :return: True if job was completed without errors. False otherwise.
+        .. note:: If the job was completed with errors, this call will return False.
+        """
+        current_job_state = self.worker_finder.get_job_state(self._job_id)
+        if current_job_state == JobState.ERROR:
+            raise RuntimeError(f'Job failed with error message "{self.get_message()}".')
+        if current_job_state == JobState.TIMEOUT:
+            raise RuntimeError("Timed out while trying to start write job.")
+        return current_job_state == JobState.DONE
+
+    def get_message(self) -> str:
+        """
+        Get a string describing the error that was encountered when running the job. (If there was an error.)
+        """
+        current_status = self.worker_finder.get_job_status(self._job_id)
+        if current_status is None:
+            return ""
+        return current_status.message
+
+    def set_stop_time(self, stop_time: datetime) -> CommandHandler:
+        """
+        Set a new stop time for the file-writing job. There is no guarantee that the stop time will actually be changed.
+        This call is not blocking. Calling this member function will have no effect on the stop-time before the write
+        job has started.
+
+        :param stop_time: The new stop time of the job.
+        :return: A CommandHandler instance that can be used to monitor the outcome of the attempt to set a new stop time.
+        """
+        current_status = self.worker_finder.get_job_status(self._job_id)
+        return self.worker_finder.try_send_stop_time(
+            current_status.service_id if current_status else None,
+            self._job_id,
+            stop_time,
+        )
+
+    def stop_now(self) -> CommandHandler:
+        """
+        See the documentation for abort_write_job().
+        """
+        return self.abort_write_job()
+
+    def abort_write_job(self) -> CommandHandler:
+        """
+        Tell the file-writing to abort writing. There is no guarantee that will actually happen though.
+        This call is not blocking. Calling this member function will have no effect if done before a write job has
+        actually started.
+
+        :return: A CommandHandler instance that can be used to monitor the outcome of the attempt to set a new stop time.
+        """
+        current_status = self.worker_finder.get_job_status(self._job_id)
+        return self.worker_finder.try_send_abort(
+            current_status.service_id if current_status else None, self._job_id
+        )
+
+    @property
+    def job_id(self) -> str:
+        """
+        The job identifier of the job that this instance of the JobHandler class represent.
+        """
+        return self._job_id

mccode_plumber/file_writer_control/JobStatus.py

@@ -0,0 +1,147 @@
+from datetime import datetime, timedelta
+from enum import Enum, auto
+from typing import Dict, Optional
+
+DEFAULT_TIMEOUT = timedelta(seconds=15)
+
+
+class JobState(Enum):
+    """
+    The state of a job.
+    """
+
+    NO_JOB = auto()
+    WAITING = auto()
+    WRITING = auto()
+    TIMEOUT = auto()
+    ERROR = auto()
+    DONE = auto()
+    UNAVAILABLE = auto()
+
+
+class JobStatus:
+    """
+    Contains general information about the (execution) of a job.
+    """
+
+    def __init__(self, job_id: str, timeout: Optional[timedelta] = DEFAULT_TIMEOUT):
+        self._job_id = job_id
+        self._timeout = timeout
+        self._service_id = ""
+        self._file_name = ""
+        self._last_update = datetime.now()
+        self._state = JobState.WAITING
+        self._metadata: Optional[Dict] = None
+        self._message = ""
+
+    def update_status(self, new_status: "JobStatus") -> None:
+        """
+        Updates the status/state of a this instance of the JobStatus class, using another instance.
+        .. note:: The job identifier of this instance and the other must be identical.
+        :param new_status: The other instance of the JobStatus class.
+        """
+        if new_status.job_id != self.job_id:
+            raise RuntimeError(
+                f"Job id of status update is not correct ({self.job_id} vs {new_status.job_id})"
+            )
+        self._state = new_status.state
+        if new_status.message:
+            self._message = new_status.message
+        self._service_id = new_status.service_id
+        self._file_name = new_status.file_name
+        self._last_update = new_status.last_update
+        self._metadata = new_status.metadata
+
+    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 != JobState.DONE
+            and self.state != JobState.ERROR
+            and self.state != JobState.TIMEOUT
+            and current_time - self.last_update > self._timeout
+        ):
+            self._state = JobState.TIMEOUT
+            self._last_update = current_time
+
+    @property
+    def job_id(self) -> str:
+        """
+        The (unique) job identifier.
+        """
+        return self._job_id
+
+    @property
+    def service_id(self) -> str:
+        """
+        The (unique) service identifier of the instance of the file-writer that executes the current job.
+        """
+        return self._service_id
+
+    @service_id.setter
+    def service_id(self, new_service_id: str) -> None:
+        if not self._service_id:
+            self._service_id = new_service_id
+            self._last_update = datetime.now()
+        elif self._service_id == new_service_id:
+            return
+        else:
+            raise RuntimeError(
+                f'Can not set service_id of job with id "{self._job_id}" to "{new_service_id}" as it has already been set to "{self._service_id}".'
+            )
+
+    @property
+    def last_update(self) -> datetime:
+        """
+        The local time stamp of the last update of the status of the job.
+        """
+        return self._last_update
+
+    @property
+    def state(self) -> JobState:
+        """
+        The current state of the job.
+        """
+        return self._state
+
+    @property
+    def file_name(self) -> str:
+        """
+        The file name of the job. None if the file name is not known.
+        """
+        if self._file_name == "":
+            return None
+        return self._file_name
+
+    @file_name.setter
+    def file_name(self, new_file_name: str) -> None:
+        self._file_name = new_file_name
+        self._last_update = datetime.now()
+
+    @state.setter
+    def state(self, new_state: JobState) -> None:
+        self._state = new_state
+        self._last_update = datetime.now()
+
+    @property
+    def message(self) -> str:
+        """
+        Status/state message of the job as received from the file-writer.
+        """
+        return self._message
+
+    @property
+    def metadata(self) -> Optional[Dict]:
+        return self._metadata
+
+    @metadata.setter
+    def metadata(self, metadata: Dict) -> None:
+        self._metadata = metadata
+
+    @message.setter
+    def message(self, new_message: str) -> None:
+        if new_message:
+            self._message = new_message
+            self._last_update = datetime.now()

mccode_plumber/file_writer_control/KafkaTopicUrl.py

@@ -0,0 +1,22 @@
+import re
+
+
+class KafkaTopicUrl:
+    """
+    Class for extracting address, port and topic name from a Kafka topic url.
+    """
+
+    test_regexp = re.compile(
+        r"^\s*(?:kafka://)?(?:(?P<host>[^/?#:]+)(?::(?P<port>\d+){1,5})?)/(?P<topic>[a-zA-Z0-9._-]+)\s*$"
+    )
+
+    def __init__(self, url: str):
+        result = re.match(KafkaTopicUrl.test_regexp, url)
+        if result is None:
+            raise RuntimeError("Unable to match kafka url.")
+        self.port = 9092  # Default Kafka broker port
+        if result.group("port") is not None:
+            self.port = int(result.group("port"))
+        self.host = result.group("host")
+        self.host_port = f"{self.host}:{self.port}"
+        self.topic = result.group("topic")

mccode_plumber/file_writer_control/StateExtractor.py

@@ -0,0 +1,58 @@
+from json import loads
+from typing import Optional
+
+from streaming_data_types.action_response_answ import (
+    ActionOutcome,
+    ActionResponse,
+    ActionType,
+)
+from streaming_data_types.status_x5f2 import StatusMessage
+
+from file_writer_control.CommandStatus import CommandState
+from file_writer_control.JobStatus import JobState
+from file_writer_control.WorkerStatus import WorkerState
+
+
+def extract_worker_state_from_status(status: StatusMessage) -> WorkerState:
+    """
+    Determine the worker state (i.e. file-writer state) based on a file-writer status message.
+    :param status: A status update message from a file-writer.
+    :return: The extracted worker state.
+    """
+    json_struct = loads(status.status_json)
+    status_map = {"writing": WorkerState.WRITING, "idle": WorkerState.IDLE}
+    try:
+        status_string = json_struct["state"]
+        return status_map[status_string]
+    except KeyError:
+        return WorkerState.UNKNOWN
+
+
+def extract_state_from_command_answer(answer: ActionResponse) -> CommandState:
+    """
+    Determine the command state from a action response message.
+    :param answer: The action (either "start a job" or "set top time") response from a file-writer.
+    :return: The extracted command state/response.
+    """
+    status_map = {
+        ActionOutcome.Failure: CommandState.ERROR,
+        ActionOutcome.Success: CommandState.SUCCESS,
+    }
+    try:
+        return status_map[answer.outcome]
+    except KeyError:
+        return CommandState.ERROR
+
+
+def extract_job_state_from_answer(answer: ActionResponse) -> Optional[JobState]:
+    """
+    Determine the file writing job state from a action response message.
+    :param answer: The action (either "start a job" or "set top time") response from a file-writer.
+    :return: The extracted job state, None if job state can not be determined from this answer.
+    """
+    if answer.action == ActionType.StartJob:
+        if answer.outcome == ActionOutcome.Success:
+            return JobState.WRITING
+        else:
+            return JobState.ERROR
+    return None