from concurrent import futures
from forecaster.prophet import Forecaster as ProphetForecaster
from multiprocessing import Event, Process, cpu_count
from pythonjsonlogger import jsonlogger
import contextlib
import grpc
import logging
import model.api.forecast_pb2_grpc as grpc_lib
import os
import signal
import socket
import sys
import time

class ForecastServicer(ProphetForecaster):
    def __init__(self, logger):
        self.logger = logger

    def pretty_timedelta(self, seconds):
        seconds = int(seconds)
        days, seconds = divmod(seconds, 86400)
        hours, seconds = divmod(seconds, 3600)
        minutes, seconds = divmod(seconds, 60)
        if days > 0:
            return '{:d}d{:d}h{:d}m{:d}s'.format(days, hours, minutes, seconds)
        elif hours > 0:
            return '{:d}h{:d}m{:d}s'.format(hours, minutes, seconds)
        elif minutes > 0:
            return '{:d}m{:d}s'.format(minutes, seconds)
        else:
            return '{:d}s'.format(seconds)

class GracefulShutdown:
    def __init__(self, logger):
        self.logger = logger
        self.event = Event()
        signal.signal(signal.SIGINT, self.handler('SIGINT'))
        signal.signal(signal.SIGTERM, self.handler('SIGTERM'))
        signal.signal(signal.SIGHUP, self.handler('SIGHUP'))

    def handler(self, signal_name):
        def fn(signal_received, frame):
            self.logger.info('signal received', extra={'signal': signal_name})
            self.event.set()
        return fn

class Config(object):
    def __init__(self):
        self.grpc_server_address = os.getenv('GRPC_SERVER_ADDRESS', '')
        self.grpc_server_key = str.encode(os.getenv('GRPC_SERVER_KEY', ''))
        self.grpc_server_cert = str.encode(os.getenv('GRPC_SERVER_CERT', ''))
        self.grpc_root_ca = str.encode(os.getenv('GRPC_ROOT_CA', ''))
        self.gprc_server_process_num = int(os.getenv('GPRC_SERVER_PROCESS_NUM', cpu_count()))
        self.grpc_server_thread_num = int(os.getenv('GRPC_SERVER_THREAD_NUM', 1))
        self.grpc_server_grace_period_in_secs = int(os.getenv('GRPC_SERVER_GRACE_PERIOD_IN_SECS', 2))
        self.grpc_server_kill_period_in_secs = int(os.getenv('GRPC_SERVER_KILL_PERIOD_IN_SECS', 5))

class Server(object):
    def __init__(self, config, logger):
        self.config = config
        self.logger = logger

    @contextlib.contextmanager
    def _reserve_port(self):
        """Find and reserve a port for all subprocesses to use"""
        sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
        if sock.getsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT) == 0:
            raise RuntimeError('failed to set SO_REUSEPORT.')
        _, port = self.config.grpc_server_address.split(':')
        sock.bind(('', int(port)))
        try:
            yield sock.getsockname()[1]
        finally:
            sock.close()

    def _run_server(self, shutdown_event):
        server_credentials = grpc.ssl_server_credentials(
            [(self.config.grpc_server_key, self.config.grpc_server_cert)],
            root_certificates=self.config.grpc_root_ca,
            require_client_auth=True
        )
        server = grpc.server(
            futures.ThreadPoolExecutor(max_workers=self.config.grpc_server_thread_num),
            options=[
                ("grpc.so_reuseport", 1),
                ("grpc.use_local_subchannel_pool", 1),
            ],
        )
        grpc_lib.add_ForecastServicer_to_server(ForecastServicer(self.logger), server)
        server.add_secure_port(self.config.grpc_server_address, server_credentials)
        self.logger.info('starting python gRPC server...')
        server.start()

        while not shutdown_event.is_set():
            time.sleep(1)

        server.stop(5).wait()
        self.logger.info('python gRPC server stopped')

    def serve(self):
        with self._reserve_port():
            procs = []
            shutdown = GracefulShutdown(self.logger)
            for _ in range(self.config.gprc_server_process_num):
                proc = Process(target=self._run_server, args=(shutdown.event,))
                procs.append(proc)
                proc.start()
            while not shutdown.event.is_set():
                time.sleep(1)

            t = time.time()
            grace_period = self.config.grpc_server_grace_period_in_secs
            kill_period = self.config.grpc_server_kill_period_in_secs
            while True:
                # Send SIGINT if process doesn't exit quickly enough, and kill it as last resort
                # .is_alive() also implicitly joins the process (good practice in linux)
                alive_procs = [proc for proc in procs if proc.is_alive()]
                if len(alive_procs) == 0:
                    break
                elapsed = time.time() - t
                if elapsed >= grace_period and elapsed < kill_period:
                    for proc in alive_procs:
                        proc.terminate()
                        self.logger.info("sending SIGTERM to subprocess", extra={'proc': proc})
                elif elapsed >= kill_period:
                    for proc in alive_procs:
                        self.logger.warning("sending SIGKILL to subprocess", extra={'proc': proc})
                        # Queues and other inter-process communication primitives can break when
                        # process is killed, but we don't care here
                        proc.kill()
                time.sleep(1)

            time.sleep(1)
            for proc in procs:
                self.logger.info("subprocess terminated", extra={'proc': proc})

def json_logger():
    logger = logging.getLogger()
    log_handler = logging.StreamHandler(sys.stdout)
    formatter = jsonlogger.JsonFormatter(fmt='%(asctime)s %(name)s %(levelname)s %(message)s')
    log_handler.setFormatter(formatter)
    log_handler.flush = sys.stdout.flush
    logger.setLevel(logging.INFO)
    logger.addHandler(log_handler)
    return logger

[((5703, 5722), 'logging.getLogger', 'logging.getLogger', ([], {}), '()\n', (5720, 5722), False, 'import logging\n'), ((5741, 5774), 'logging.StreamHandler', 'logging.StreamHandler', (['sys.stdout'], {}), '(sys.stdout)\n', (5762, 5774), False, 'import logging\n'), ((5791, 5869), 'pythonjsonlogger.jsonlogger.JsonFormatter', 'jsonlogger.JsonFormatter', ([], {'fmt': '"""%(asctime)s %(name)s %(levelname)s %(message)s"""'}), "(fmt='%(asctime)s %(name)s %(levelname)s %(message)s')\n", (5815, 5869), False, 'from pythonjsonlogger import jsonlogger\n'), ((1101, 1108), 'multiprocessing.Event', 'Event', ([], {}), '()\n', (1106, 1108), False, 'from multiprocessing import Event, Process, cpu_count\n'), ((1579, 1615), 'os.getenv', 'os.getenv', (['"""GRPC_SERVER_ADDRESS"""', '""""""'], {}), "('GRPC_SERVER_ADDRESS', '')\n", (1588, 1615), False, 'import os\n'), ((2481, 2531), 'socket.socket', 'socket.socket', (['socket.AF_INET6', 'socket.SOCK_STREAM'], {}), '(socket.AF_INET6, socket.SOCK_STREAM)\n', (2494, 2531), False, 'import socket\n'), ((2998, 3167), 'grpc.ssl_server_credentials', 'grpc.ssl_server_credentials', (['[(self.config.grpc_server_key, self.config.grpc_server_cert)]'], {'root_certificates': 'self.config.grpc_root_ca', 'require_client_auth': '(True)'}), '([(self.config.grpc_server_key, self.config.\n    grpc_server_cert)], root_certificates=self.config.grpc_root_ca,\n    require_client_auth=True)\n', (3025, 3167), False, 'import grpc\n'), ((1658, 1690), 'os.getenv', 'os.getenv', (['"""GRPC_SERVER_KEY"""', '""""""'], {}), "('GRPC_SERVER_KEY', '')\n", (1667, 1690), False, 'import os\n'), ((1735, 1768), 'os.getenv', 'os.getenv', (['"""GRPC_SERVER_CERT"""', '""""""'], {}), "('GRPC_SERVER_CERT', '')\n", (1744, 1768), False, 'import os\n'), ((1809, 1838), 'os.getenv', 'os.getenv', (['"""GRPC_ROOT_CA"""', '""""""'], {}), "('GRPC_ROOT_CA', '')\n", (1818, 1838), False, 'import os\n'), ((1976, 2014), 'os.getenv', 'os.getenv', (['"""GRPC_SERVER_THREAD_NUM"""', '(1)'], {}), "('GRPC_SERVER_THREAD_NUM', 1)\n", (1985, 2014), False, 'import os\n'), ((2068, 2116), 'os.getenv', 'os.getenv', (['"""GRPC_SERVER_GRACE_PERIOD_IN_SECS"""', '(2)'], {}), "('GRPC_SERVER_GRACE_PERIOD_IN_SECS', 2)\n", (2077, 2116), False, 'import os\n'), ((2169, 2216), 'os.getenv', 'os.getenv', (['"""GRPC_SERVER_KILL_PERIOD_IN_SECS"""', '(5)'], {}), "('GRPC_SERVER_KILL_PERIOD_IN_SECS', 5)\n", (2178, 2216), False, 'import os\n'), ((3247, 3321), 'concurrent.futures.ThreadPoolExecutor', 'futures.ThreadPoolExecutor', ([], {'max_workers': 'self.config.grpc_server_thread_num'}), '(max_workers=self.config.grpc_server_thread_num)\n', (3273, 3321), False, 'from concurrent import futures\n'), ((3776, 3789), 'time.sleep', 'time.sleep', (['(1)'], {}), '(1)\n', (3786, 3789), False, 'import time\n'), ((4312, 4323), 'time.time', 'time.time', ([], {}), '()\n', (4321, 4323), False, 'import time\n'), ((5545, 5558), 'time.sleep', 'time.sleep', (['(1)'], {}), '(1)\n', (5555, 5558), False, 'import time\n'), ((1920, 1931), 'multiprocessing.cpu_count', 'cpu_count', ([], {}), '()\n', (1929, 1931), False, 'from multiprocessing import Event, Process, cpu_count\n'), ((4097, 4153), 'multiprocessing.Process', 'Process', ([], {'target': 'self._run_server', 'args': '(shutdown.event,)'}), '(target=self._run_server, args=(shutdown.event,))\n', (4104, 4153), False, 'from multiprocessing import Event, Process, cpu_count\n'), ((4281, 4294), 'time.sleep', 'time.sleep', (['(1)'], {}), '(1)\n', (4291, 4294), False, 'import time\n'), ((5518, 5531), 'time.sleep', 'time.sleep', (['(1)'], {}), '(1)\n', (5528, 5531), False, 'import time\n'), ((4844, 4855), 'time.time', 'time.time', ([], {}), '()\n', (4853, 4855), False, 'import time\n')]