abots/abots/net/socket_client.py

"""

Socket Client
=============



"""

from abots.helpers import eprint, cast, jots, jsto, utc_now_timestamp, coroutine

from struct import pack, unpack
from socket import socket, timeout as sock_timeout
from socket import AF_INET, SOCK_STREAM, SOL_SOCKET, SO_REUSEADDR
from ssl import wrap_socket
from threading import Thread, Event
from queue import Queue, Empty
from time import sleep
from random import randint

class SocketClient(Thread):
    def __init__(self, host, port, buffer_size=4096, secure=False, 
        timeout=None, daemon=False, reconnects=10):
        super().__init__()
        self.setDaemon(daemon)

        self.host = host
        self.port = port
        self.buffer_size = buffer_size
        self.secure = secure
        self.timeout = timeout
        self.reconnects = reconnects
        self.sock = socket(AF_INET, SOCK_STREAM)
        if self.secure:
            self.sock = wrap_socket(self.sock)

        self.connection = (self.host, self.port)
        self.running = True

        self.kill_switch = Event()
        self.ready = Event()
        self.stopped = Event()
        self.broken = Event()
        self.reconnecting = Event()

        self._inbox = Queue()
        self._events = Queue()
        self._outbox = Queue()
        self.queues = dict()
        self.queues["inbox"] = self._inbox
        self.queues["outbox"] = self._outbox
        self.queues["events"] = self._events

        self._bridge = None

    def _send_event(self, message):
        self._events.put(jots(message))
        cast(self._bridge, "send", ("events", message))
    
    def _prepare(self):
        self.sock.setblocking(False)
        self.sock.settimeout(1)
        try:
            self.sock.connect(self.connection)
        except Exception as e:
            return True, e
        return False, None
    
    def _obtain(self, queue_name, timeout=False):
        queue = self.queues[queue_name]
        if timeout is False:
            timeout = self.timeout
        while True:
            try:
                if timeout is not None:
                    message = queue.get(timeout=timeout)
                else:
                    message = queue.get_nowait()
                yield message
                cast(self._bridge, "send", (queue_name, message))
                queue.task_done()
            except Empty:
                break

    def _queue_thread(self, inbox, timeout):
        while not self.kill_switch.is_set():
            for message in self._obtain("inbox", timeout):
                if self.broken.is_set():
                    self.reconnecting.wait()
                self.send_message(message)

    def _get_message(self, decode=True):
        try:
            packet = self.sock.recv(self.buffer_size)
            result = packet.decode() if decode else packet
            self._outbox.put(result)
            cast(self._bridge, "send", ("outbox", result))
        except (BrokenPipeError, OSError) as e:
            pass

    def _format_message(self, message, *args):
        if len(args) > 0:
            formatted = message.format(*args)
        else:
            formatted = message
        return formatted.encode()

    def _attempt_reconnect(self):
        if self.kill_switch.is_set():
            return
        print("BROKEN!")
        self.reconnecting.clear()
        self.broken.set()
        event = dict()
        event["name"] = "socket-down"
        event["data"] = dict()
        event["data"]["when"] = utc_now_timestamp()
        self._send_event(event)
        attempts = 0
        while attempts <= self.reconnects or not self.kill_switch.is_set():
            # Need to be run to prevent ConnectionAbortedError
            self.sock.__init__()
            err, report = self._prepare()
            if not err:
                self.reconnecting.set()
                self.broken.clear()
                event = dict()
                event["name"] = "socket-up"
                event["data"] = dict()
                event["data"]["when"] = utc_now_timestamp()
                self._send_event(event)
                return
            # Exponential backoff
            attempts = attempts + 1
            max_delay = (2**attempts) - 1
            delay = randint(0, max_delay)
            sleep(delay)
        self.stop()

    def send_message(self, message, *args):
        formatted = self._format_message(message, *args)
        try:
            self.sock.send(formatted)
        except (BrokenPipeError, OSError) as e:
            if not isinstance(e, sock_timeout):
                self._attempt_reconnect()
            self._attempt_reconnect()

    def recv(self):
        yield from self._obtain("outbox")

    def check(self):
        for letter in self.recv():
            if letter is not None and len(letter) > 0:
                print(letter)

    def send(self, message):
        self._inbox.put(message)
        cast(self._bridge, "send", ("inbox", message))

    def connect(self, bridge):
        self._bridge = bridge()

    @coroutine
    def bridge(self, inbox, outbox, events):
        try:
            while True:
                task = (yield)
                source, message = task
                if source == "inbox":
                    outbox.puts(message) 
                elif source == "outbox":
                    inbox.puts(message)
                elif source == "events":
                    events.put(message)
        except GeneratorExit:
            pass 

    def run(self):
        err, report = self._prepare()
        if err:
            eprint(report)
            return report
        queue_args = self._inbox, self.timeout
        Thread(target=self._queue_thread, args=queue_args).start()
        print("Client ready!")
        self.ready.set()
        while self.running:
            if self.broken.is_set():
                self.reconnecting.wait()
            self._get_message()

    def stop(self, done=None):
        # print("Stopping client!")
        self.kill_switch.set()
        event = dict()
        event["name"] = "closing"
        event["data"] = dict()
        event["data"]["when"] = utc_now_timestamp()
        self._send_event(event)
        self.running = False
        self.sock.close()
        self.stopped.set()
        cast(done, "set")
        # print("Stopped client!")