openrvdas.logger.readers.socket_reader

  1#!/usr/bin/env python3
  2import os
  3import socket
  4import tempfile
  5import atexit
  6import threading
  7from typing import Optional, Dict
  8
  9# Add parent directory to path
 10from logger.readers.reader import Reader  # noqa: E402
 11
 12# Global reference counter for channels
 13_channel_refs: Dict[str, int] = {}
 14_lock = threading.Lock()
 15
 16
 17class SocketReader(Reader):
 18    """Reader class for socket-based IPC mechanism.
 19
 20    Reads records from a Unix domain socket.
 21    """
 22
 23    def __init__(self, channel: str, timeout: Optional[float] = None,
 24                 buffer_size: int = 4096, keep_binary: bool = False, **kwargs):
 25        """Initialize a Reader for the specified channel.
 26
 27        Args:
 28            channel: A string identifier for the communication channel
 29
 30            timeout: Maximum time to wait for a record when reading (None means wait forever)
 31
 32            buffer_size: Max record size to expect
 33
 34            keep_binary: If true, don't convert received record to string
 35        """
 36        super().__init__(**kwargs)
 37
 38        self.timeout = timeout
 39        self.buffer_size = buffer_size
 40        self.keep_binary = keep_binary
 41
 42        # Create a unique socket path based on the channel name
 43        import hashlib
 44        channel_hash = hashlib.md5(channel.encode()).hexdigest()[:8]
 45        self.channel = channel
 46
 47        # Set socket path in temp directory to avoid permission issues
 48        temp_dir = tempfile.gettempdir()
 49        self.socket_path = os.path.join(temp_dir, f'ipc_socket_{channel_hash}')
 50
 51        # Socket for receiving data
 52        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
 53
 54        try:
 55            # Try to bind to the socket path
 56            self.sock.bind(self.socket_path)
 57        except OSError:
 58            # If socket already exists but no process is bound to it
 59            if os.path.exists(self.socket_path):
 60                os.unlink(self.socket_path)
 61                self.sock.bind(self.socket_path)
 62            else:
 63                raise
 64
 65        # Set socket timeout
 66        if self.timeout is not None:
 67            self.sock.settimeout(self.timeout)
 68
 69        # Update reference counter
 70        with _lock:
 71            _channel_refs[self.channel] = _channel_refs.get(self.channel, 0) + 1
 72
 73        # Register cleanup on exit
 74        atexit.register(self.close)
 75
 76    def read(self) -> str:
 77        """Read a record from the channel, blocking until one is available.
 78
 79        Returns:
 80            str: The data that was read
 81
 82        Raises:
 83            TimeoutError: If no data is available within the timeout period
 84        """
 85        try:
 86            record, _ = self.sock.recvfrom(self.buffer_size)
 87            if not self.keep_binary:
 88                record = record.decode('utf-8')
 89            return record
 90        except socket.timeout:
 91            raise TimeoutError('No data available within timeout period')
 92
 93    def close(self):
 94        """Clean up resources and potentially clean up the channel."""
 95        if not hasattr(self, 'sock') or self.sock is None:
 96            # Already closed
 97            return
 98
 99        self.sock.close()
100        self.sock = None
101
102        # Decrement reference counter and cleanup if this is the last reference
103        with _lock:
104            _channel_refs[self.channel] = _channel_refs.get(self.channel, 1) - 1
105            if _channel_refs[self.channel] <= 0:
106                if os.path.exists(self.socket_path):
107                    os.unlink(self.socket_path)
108
109                # Remove from reference counter
110                _channel_refs.pop(self.channel, None)
111
112    def __del__(self):
113        """Ensure cleanup happens."""
114        self.close()