102 lines
3.3 KiB
Python
102 lines
3.3 KiB
Python
import json
|
|
import os
|
|
import queue
|
|
import socket
|
|
import threading
|
|
from pathlib import Path
|
|
|
|
from colony.core.logger.logging import get_logger
|
|
from colony.transport.framing import send_frame, recv_frame
|
|
|
|
CONTROL_SOCKET = Path(__file__).resolve().parents[2] / "data" / "colony.sock"
|
|
|
|
|
|
def request(payload, path=CONTROL_SOCKET, timeout=10):
|
|
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
|
|
sock.settimeout(timeout)
|
|
sock.connect(str(path))
|
|
send_frame(sock, json.dumps(payload).encode())
|
|
return json.loads(recv_frame(sock).decode())
|
|
|
|
|
|
class ControlQueue:
|
|
def __init__(self, actions):
|
|
self.logger = get_logger(self.__class__.__name__)
|
|
self.actions = actions
|
|
self.tasks = queue.Queue()
|
|
|
|
def submit(self, payload, timeout):
|
|
reply = queue.Queue(maxsize=1)
|
|
self.tasks.put((payload, reply))
|
|
return reply.get(timeout=timeout)
|
|
|
|
def run_once(self):
|
|
payload, reply = self.tasks.get()
|
|
try:
|
|
body = json.dumps({"status": "ok", "result": self.actions.run(payload)})
|
|
except Exception as e:
|
|
self.logger.exception("Control command failed")
|
|
body = json.dumps({"status": "error", "error": f"{type(e).__name__}: {e}"})
|
|
reply.put(body.encode())
|
|
|
|
def run_forever(self):
|
|
while True:
|
|
self.run_once()
|
|
|
|
|
|
class ControlSocket:
|
|
CLIENT_TIMEOUT = 30
|
|
|
|
def __init__(self, control_queue, path=CONTROL_SOCKET):
|
|
self.logger = get_logger(self.__class__.__name__)
|
|
self.control_queue = control_queue
|
|
self.path = Path(path)
|
|
self._sock = None
|
|
|
|
def _remove_stale_socket(self):
|
|
if not self.path.exists():
|
|
return
|
|
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe:
|
|
try:
|
|
probe.connect(str(self.path))
|
|
except OSError:
|
|
self.path.unlink()
|
|
return
|
|
raise OSError(f"another node is already running on {self.path}")
|
|
|
|
def bind(self):
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._remove_stale_socket()
|
|
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
try:
|
|
sock.bind(str(self.path))
|
|
os.chmod(self.path, 0o600)
|
|
sock.listen()
|
|
except OSError:
|
|
sock.close()
|
|
raise
|
|
self._sock = sock
|
|
|
|
def start(self):
|
|
self.logger.info(f"Control socket on {self.path}")
|
|
while True:
|
|
conn, _ = self._sock.accept()
|
|
threading.Thread(target=self.handle, args=(conn,), daemon=True, name="HandleControl").start()
|
|
|
|
def handle(self, conn):
|
|
with conn:
|
|
try:
|
|
conn.settimeout(self.CLIENT_TIMEOUT)
|
|
payload = json.loads(recv_frame(conn).decode())
|
|
body = self.control_queue.submit(payload, self.CLIENT_TIMEOUT)
|
|
except queue.Empty:
|
|
body = json.dumps({"status": "error", "error": "timed out waiting for the node"}).encode()
|
|
except (OSError, ValueError) as e:
|
|
self.logger.warning(f"Bad control request: {e}")
|
|
return
|
|
|
|
try:
|
|
send_frame(conn, body)
|
|
except OSError as e:
|
|
self.logger.warning(f"Could not answer control request: {e}")
|