path: root/worker.py

#!/usr/bin/env python


"""
A module consisting of a zeromq worker receiving video stream
via router from clients. Able to send control messages to clients.
Needs libzmq and pyzmq with 'drafts' support.
"""

import os
import sys
from threading import Thread, Event
import time
import signal
import logging
from queue import Queue
from collections import defaultdict
import tempfile
import json
import zmq
import cv2
import numpy

from helpers import CustomLoggingFormatter

__author__ = "Franoosh Corporation"

# Various constants:  # TODO: put them in a config

HOST = "127.0.0.1"
ZMQPORT = "9979"
# WSPORT = "8000"
WSPORT = "8008"
ZMQ_BACKEND_ADDR = f"tcp://{HOST}:{ZMQPORT}"
WEB_BACKEND_ADDR = f"tcp://{HOST}:{WSPORT}"
# LOGLEVEL = logging.DEBUG
LOGLEVEL = logging.INFO

#   File paths:

LOGFILE = f"{os.path.splitext(os.path.basename(__file__))[0]}.log"
CWD = os.getcwd()
TMP_DIR = os.path.join(CWD, "tmp")
CLIENTS_JSON_FILE = os.path.join(CWD, 'clients.json')

# Other:

CLIENTS_DICT = {}

stop_event = Event()

log_formatter = CustomLoggingFormatter()
handler = logging.FileHandler(LOGFILE, encoding='utf-8', mode='a')
handler.setFormatter(log_formatter)
logging.root.addHandler(handler)
logging.root.setLevel(LOGLEVEL)
logger = logging.getLogger(__name__)
logging.basicConfig(
    filename=LOGFILE,
    datefmt='%Y-%m-%d %I:%M:%S',
    level=LOGLEVEL,
)


if not os.path.exists(TMP_DIR):
    try:
        os.makedirs(TMP_DIR)
    except Exception as exc:
        logger.error("Could not create temporary directory: %r", exc)
        sys.exit()

class MonitorTask(Thread):
    """Monitor task"""
    def __init__(self, socket):
        super().__init__(daemon=True)
        self.socket = socket
        self.running = True

    def run(self):
        """Monitor connection on the initial socket.
        This is heartbeat monitoring"""
        monitor = self.socket.get_monitor_socket()
        monitor.setsockopt(zmq.RCVTIMEO, 5000)
        logger.debug("Monitor socket started.")
        while self.running:
            try:
                event, _ = monitor.recv_multipart()
                # Resolve the received event type:
                event_type = int.from_bytes(event[:2], "little")
                if event_type in (zmq.EVENT_CLOSED, zmq.EVENT_DISCONNECTED):
                    logger.warning("Monitor socket: closed | disconnected")
                    stop_event.set()
                elif event_type == zmq.EVENT_CONNECT_DELAYED:
                    logger.debug("Monitor socket: event connect delayed")
                elif event_type == zmq.EVENT_CONNECT_RETRIED:
                    logger.debug("Monitor socket: event connect retried")
                elif event_type in (zmq.EVENT_CONNECTED, zmq.EVENT_HANDSHAKE_SUCCEEDED):
                    logger.debug("Monitor socket: client connected to router, handshake OK.")
                    stop_event.clear()
                else:
                    logger.warning("Monitor socket: other event: '%r'", event_type)
            except zmq.Again:
                logger.debug("Timeout on monitoring socket.")
            except Exception as exc: # W: Catching too general exception Exception
                logger.error("Other exception on monitoring socket: %r", exc)

        # monitor.close()

    def stop(self):
        """Stop thread"""
        logger.info("Stopping monitor thread ...")
        self.running = False


class ServerWorker(Thread):
    """ServerWorker"""
    def __init__(self, identity, context=None):
        super().__init__(daemon=True)
        self.id = identity
        self.context = context or zmq.Context.instance()
        self.socket = self.context.socket(zmq.DEALER)
        self.socket.identity = self.id.encode("utf-8")
        self.monitor = MonitorTask(self.socket)
        self.video_threads = defaultdict(lambda: defaultdict(dict))
        self.poller = zmq.Poller()
        self.web_sock = self.context.socket(zmq.DEALER)
        self.connected = False
        self.running = True

    def start_client(self, client_id, camera_id, filename):
        """
        Start a video thread for new client_id and camera_id.
        """
        if client_id not in self.video_threads or not self.video_threads[client_id].get(camera_id):  # New client or new camera
            q = Queue()
            logger.debug("Starting new video thread for client '%s', camera '%s'", client_id, camera_id)
            video_worker = VideoWorker(client_id, camera_id, filename, q)
            video_worker.start()
            self.video_threads[client_id][camera_id] = video_worker

    def stop_client(self, client_id, camera_id):
        """
        Stop video thread for a client_id and camera_id.
        """
        if client_id in self.video_threads and camera_id in self.video_threads[client_id]:
            logger.debug("Stopping video thread for client '%s', camera '%s'", client_id, camera_id)
            self.video_threads[client_id][camera_id].stop() # Stop the thread
            del self.video_threads[client_id][camera_id]
            logger.info("Stopped video thread for client '%s', camera '%s'", client_id, camera_id)

    def run(self):
        """
        Main loop of the worker.
        Full of wonders.
        """
        logger.debug("ServerWorker '%s' starting ...", self.id)
        self.socket.connect(ZMQ_BACKEND_ADDR)
        try:
           self.web_sock.bind(WEB_BACKEND_ADDR)
        except Exception as exc:
            logger.error("Connection to zmq websocket failed: %r", exc)
        self.poller.register(self.socket, zmq.POLLIN)
        self.poller.register(self.web_sock, zmq.POLLIN)
        self.monitor.start()

        while self.running:
            logger.debug("ServerWorker '%s' waiting for a message ...", self.id)
            if not self.connected or stop_event.is_set():
                self.socket.send_multipart([b"READY"])  # Router needs worker identity, hence this
                time.sleep(1)  # Wait a bit before trying to connect again

            sockets = dict(self.poller.poll(1000))
            if self.socket in sockets:
                self.connected = True
                msg = self.socket.recv_multipart()
                logger.debug("ServerWorker '%s' received message of length: %d.", self.id, len(msg))
                filename = None
                # At the moment we don't expect any other message than a start/stop message (of length 4) and a video message:
                if len(msg) == 4:  # This is a message with start/stop directive and no video data.
                    logger.debug("Received start/stop directive: (?)")
                    client_id, camera_id, metadata = msg[0], msg[1], msg[2]
                    # Convert bytes to str once and for all:
                    client_id = client_id.decode('utf-8')
                    camera_id = camera_id.decode('utf-8')
                    update_clients(client_id, camera_id)
                    try:
                        # Directive and data are fron now on converted to strings:
                        directive, data = metadata.decode("utf-8").split(":")
                        logger.info(
                            "Received directive '%s' with data: %r for client '%s', camera '%s'",
                            directive,
                            data,
                            client_id,
                            camera_id,
                        )
                    except ValueError:
                        logger.error("Invalid metadata format: %r", metadata)
                        directive = None
                        continue
                    if directive == 'rename':
                        old_name, new_name, timestamp = None, None, None
                        try:
                            old_name, new_name, timestamp = data.split(":")
                            logger.info("Renamed video thread from '%s' to '%s'.", old_name, new_name)
                        except ValueError:
                            logger.error("Invalid rename data format: %r", data)
                            continue
                        if old_name and new_name and timestamp:
                            # I think it's better to stop the old thread and start a new one,
                            # rather than reuse the old one as it's less mucking about.
                            self.stop_client(old_name, camera_id)
                            self.start_client(new_name, camera_id, f"{new_name}_{camera_id}-{timestamp}.mkv")
                            to_remove= b':'.join([client_id, camera_id])
                            try:
                                self.web_sock.send_multipart([to_remove, b'', b''])  # Notify webserver of rename
                            except Exception as exc:
                                logger.error("Sending rename notification to websocket failed: %r", exc)
                    else:
                        timestamp = data
                        filename = f"{client_id}_{camera_id}-{timestamp}.mkv"
                        logger.debug("Directive: '%s', Timestamp: '%s', Filename: '%s'", directive, timestamp, filename)
                        if directive == "start":
                            self.start_client(client_id, camera_id, filename)
                        elif directive == "stop":
                            self.stop_client(client_id, camera_id)
                        else:
                            logger.error("Unknown directive: %r", directive)

                elif len(msg) == 3:  # This is a video message with data
                    logger.debug("Received video message with data only.")
                    client_id, camera_id, content = msg[0], msg[1], msg[2]
                    client_id = client_id.decode('utf-8')
                    camera_id = camera_id.decode('utf-8')
                    if client_id in self.video_threads and camera_id in self.video_threads[client_id]:
                        self.video_threads[client_id][camera_id].queue.put(content)
                        # Send only [client_id, camera_id, jpeg_bytes] to the webserver:
                        # zmq subsciber can subscribe to a topic defined by the first
                        # part of the multipart message, so in order to allow for a 
                        # per camera subscription, we need to join client_id and camera_id
                        topic = ':'.join([client_id, camera_id]).encode('utf-8')
                        try:
                            self.web_sock.send_multipart([topic, content], flags=zmq.NOBLOCK)
                        except Exception as exc:
                            logger.error("Sending message to websocket failed: %r", exc)
                    else:
                        logger.error("No video thread found for client '%s', camera '%s'", client_id, camera_id)
                        logger.error("Available video threads keys: %r", self.video_threads.keys())
                        logger.error("Available video threads values: %r", self.video_threads.values())
                else:
                    logger.warning("Received a message of unexpected length from client. Message length: %d", len(msg))
            else:
                logger.debug("No message received, polling again ...")
                time.sleep(5)
            if self.web_sock in sockets:
                frontend_msg = self.web_sock.recv_multipart()
                logger.info("Received message from frontend: %r", frontend_msg)
                self.socket.send_multipart(frontend_msg)
                logger.info("Forwarded message to backend: %r", frontend_msg)

        self.monitor.stop()
        self.monitor.join()

        for camera_thread in self.video_threads.values():
            for thread in camera_thread.values():
                thread.queue.put(None)  # Sentinel to unblock queue.get()
                thread.stop()
                thread.join()

    def send_control_message(self, client_id, camera_id, message):
        """Send control message to a specific client and camera."""
        if client_id in self.video_threads and camera_id in self.video_threads[client_id]:
            self.socket.send_multipart([client_id.encode("utf-8"), camera_id.encode("utf-8"), message])
            logger.info("Sent control message to client '%s', camera '%s': %r", client_id, camera_id, message)
        else:
            logger.error("No video thread found for client '%s', camera '%s'", client_id, camera_id)

    def stop(self):
        logger.info("ServerWorker '%s' exiting ...", self.id)
        self.running = False


class VideoWorker(Thread):
    """Class for video threads."""

    def __init__(self, client_id, camera_id, filename, queue):
        super().__init__(daemon=True)
        self.context = zmq.Context()
        self.context.setsockopt(zmq.LINGER, 0)
        self.client_id = client_id
        self.camera_id = camera_id
        self.filename = filename
        self.queue = queue
        self.live = True

    def stop(self):
        logger.info("VideoWorker %r exiting ...", self.camera_id)
        self.live = False

    def run(self):
        if os.path.exists(self.filename):
            logger.warning("File '%s' already exists, overwriting ...", self.filename)

        fourcc = cv2.VideoWriter_fourcc(*'VP80')
        out = cv2.VideoWriter(self.filename, fourcc, 30.0, (640, 480))  # Assuming 640x480 resolution
        logger.info("This is the first run, binding websocket ...")
        while self.live:
            logger.debug("VideoWorker writing to file: %s", self.filename)
            frame_bytes = self.queue.get()
            if frame_bytes is None:
                logger.debug("Received None, stopping video worker for camera: '%s'", self.camera_id)
                break
            frame = cv2.imdecode(numpy.frombuffer(frame_bytes, dtype=numpy.uint8), cv2.IMREAD_COLOR)
            logger.debug("Processing ('writing to a file') frame for camera: '%s'", self.camera_id)
            # Write frame to file:
            out.write(frame)

        # Release video writer
        out.release()
        logger.info("VideoWorker finished writing to file: %s", self.filename)

def signal_handler(sig, frame):
    worker.stop()

def update_clients(client_id, camera_id):
    """Update client and camera dictionary and write to a shared file."""
    global CLIENTS_DICT
    if client_id not in CLIENTS_DICT:
        logger.debug("Client_id not in CLIENTS_DICT, adding an empty list for it.")
        CLIENTS_DICT[client_id] = []
    if camera_id not in CLIENTS_DICT[client_id]:
        logger.debug("Camera_id not in CLIENTS_DICT[%s] list, adding", client_id)
        CLIENTS_DICT[client_id].append(camera_id)
    # Atomic write using tempfile. Works only when both files on the same filesystem
    with tempfile.NamedTemporaryFile('w', dir=TMP_DIR, delete=False) as tf:
        logger.debug("Dumping to file CLIENTS_DICT: %r", CLIENTS_DICT)
        json.dump(CLIENTS_DICT, tf)
        tempname = tf.name
    os.replace(tempname, CLIENTS_JSON_FILE)


if __name__ == "__main__":
    logger.info("Starting up ...")

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    worker = ServerWorker("worker-task")
    worker.start()

    worker.join()
    try:
        os.remove(CLIENTS_JSON_FILE)
    except FileNotFoundError:
        pass