import logging import socket import struct import json import sys from typing import Tuple from backend import Backend from protocol import PacketType, PKT_CHANGE_TYPES, change_from_packet, packet_from_change, send_packet, recv_packet class SystemdHandler(logging.Handler): PREFIX = { # EMERG <0> # ALERT <1> logging.CRITICAL: "<2>", logging.ERROR: "<3>", logging.WARNING: "<4>", # NOTICE <5> logging.INFO: "<6>", logging.DEBUG: "<7>", logging.NOTSET: "<7>" } def __init__(self, stream=sys.stdout): self.stream = stream logging.Handler.__init__(self) def emit(self, record): try: msg = self.PREFIX[record.levelno] + self.format(record) + "\n" self.stream.write(msg) self.stream.flush() except Exception: self.handleError(record) def setup_server_logging(mode, level): root_logger = logging.getLogger() root_logger.setLevel(level.upper()) mode = mode.lower() if mode == 'systemd': # replace handler with systemd one root_logger.handlers = [] root_logger.addHandler(SystemdHandler()) else: assert mode == 'plain' class SocketServer: def __init__(self, addr: Tuple[str, int], backend: Backend) -> None: self.backend = backend self.addr = addr self.bind = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self.bind.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) self.bind.bind(addr) def _send_packet(self, typ: int, payload: bytes) -> None: send_packet(self.sock, typ, payload) def _recv_packet(self) -> Tuple[int, bytes]: return recv_packet(self.sock) def _handle_conn(self, conn) -> None: # Can only handle one connection at a time logging.info('Servicing incoming connection') self.sock = conn while True: try: (typ, payload) = self._recv_packet() except IOError: logging.info('Connection closed') break if typ in PKT_CHANGE_TYPES: change = change_from_packet(typ, payload) if typ == PacketType.CHANGE: logging.debug('Received CHANGE {}'.format(change.version)) else: logging.info('Received SNAPSHOT {}'.format(change.version)) self.backend.add_change(change) self._send_packet(PacketType.ACK, struct.pack("!I", self.backend.version)) elif typ == PacketType.REWIND: logging.info('Received REWIND') to_version, = struct.unpack('!I', payload) if to_version != self.backend.prev_version: logging.info('Cannot rewind to version {}'.format(to_version)) self._send_packet(PacketType.NACK, struct.pack("!I", self.backend.version)) else: self.backend.rewind() self._send_packet(PacketType.ACK, struct.pack("!I", self.backend.version)) elif typ == PacketType.REQ_METADATA: logging.debug('Received REQ_METADATA') blob = struct.pack("!IIIQ", 0x01, self.backend.version, self.backend.prev_version, self.backend.version_count) self._send_packet(PacketType.METADATA, blob) elif typ == PacketType.RESTORE: logging.info('Received RESTORE') for change in self.backend.stream_changes(): (typ, payload) = packet_from_change(change) self._send_packet(typ, payload) self._send_packet(PacketType.DONE, b'') elif typ == PacketType.COMPACT: logging.info('Received COMPACT') stats = self.backend.compact() self._send_packet(PacketType.COMPACT_RES, json.dumps(stats).encode()) elif typ == PacketType.ACK: logging.debug('Received ACK') elif typ == PacketType.NACK: logging.debug('Received NACK') elif typ == PacketType.METADATA: logging.debug('Received METADATA') elif typ == PacketType.COMPACT_RES: logging.debug('Received COMPACT_RES') else: raise Exception('Unknown or unexpected packet type {}'.format(typ)) self.conn = None def run(self) -> None: self.bind.listen(1) logging.info('Waiting for connection on {}'.format(self.addr)) while True: conn, _ = self.bind.accept() try: self._handle_conn(conn) except Exception: logging.exception('Got exception') finally: conn.close()