#!/usr/bin/env python

"""Support pool based servers.

This model is good when you want to dedicate N processes to a job
without the overhead of continually forking.

The server forks N workers, and each then sits in an endless loop
accepting and servicing requests.
"""

"""
RCSid:
	$Id: PoolServer.py,v 1.22 2022/05/26 22:10:03 sjg Exp $

	@(#) Copyright (c) 2012 Simon J. Gerraty

	This file is provided in the hope that it will
	be of use.  There is absolutely NO WARRANTY.
	Permission to copy, redistribute or otherwise
	use this file is hereby granted provided that 
	the above copyright notice and this notice are
	left intact. 
      
	Please send copies of changes and bug-fixes to:
	sjg@crufty.net

"""

import os
import sys
import socket
import signal
import errno
from syslog import *

def daemonize():
    """ensure we have no controlling tty"""
    try:
        t = os.ttyname(0)
    except:
        return None
    os.close(0)
    p = os.fork()
    assert p >= 0
    if p > 0:
        os._exit(0)
    os.setsid()
    return os.getpid()

def sig_handler(sig, frame):
    syslog(LOG_DEBUG, 'caught {0}'.format(sig))
    if sig == signal.SIGALRM:
        raise IOError(errno.ETIMEDOUT, 'Alarm')
    if sig == signal.SIGHUP:
        PoolServer._shutdown_requested = True
    if sig == signal.SIGTERM:
        sys.exit(0)

class PoolServer(object):
    """base class for pool servers

    a pool server creates a number of children that accept connections
    from its listening socket.
    """
    _shutdown_requested = False

    def __init__(self, listen_address, nchild=4, conf={}):
        self.conf = conf
        self.nchild = nchild
        self.debug = 0
        self.bind_listen_sock(listen_address)
        self.server_address = self.listen_sock.getsockname()
        self.daemonize = True
        self._shutdown_requested = False
        self.children = []
        self.timeout = 5
        signal.signal(signal.SIGALRM, sig_handler)
        signal.signal(signal.SIGHUP, sig_handler)
        signal.signal(signal.SIGTERM, sig_handler)
        self.server_setup()
        if not self.daemonize or not daemonize():
            os.setpgrp()
        self.server_pid = os.getpid()

    def bind_listen_sock(self, listen_address):
        """bind the socket we listen on

        an AF_UNIX socket if conf['sock_af'] is 'unix'
        otherwise AF_INET.
        """
        sock_af = self.conf.get('sock_af', 'inet')

        if sock_af == 'unix':
            self.listen_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            if listen_address.find('/') < 0:
                sock_dir = self.conf.get('sock_dir',
                                         '/tmp/.{}-sock'.format(os.getenv('USER')))
                listen_address = os.path.join(sock_dir, listen_address)
            else:
                sock_dir = os.path.dirname(listen_address)
            try:
                os.makedirs(sock_dir, int(self.conf.get('sock_umask',0o700)))
            except:
                pass
            try:
                os.unlink(listen_address)
            except:
                pass
        else:
            self.listen_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.listen_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            if 'SO_REUSEPORT' in socket.__dict__:
                self.listen_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)

        self.listen_sock.bind(listen_address)

    def server_setup(self):
        pass

    def server_finish(self):
        pass

    def child_setup(self):
        pass

    def child_finish(self):
        pass

    def start_child(self):
        pid = os.fork()
        if pid:
            self.children.append(pid)
            syslog(LOG_INFO, 'started worker {0}'.format(pid))
            return
        else:
            self.pid = os.getpid()
            self.nchild = 0
            signal.signal(signal.SIGHUP, sig_handler)
            self.child_setup()
            self.request_loop()
            self.child_finish()
            os._exit(0)

    def shutdown_requested(self):
        """make this more complex if you like"""
        sr = PoolServer._shutdown_requested
        if self.nchild == 0:            # child
            if os.getppid() != self.server_pid or sr:
                return True
        elif sr:       # parent
            # tap the children - since we know who they are
            for pid in self.children:
                try:
                    syslog(LOG_DEBUG, "Hup'ing worker {}".format(pid))
                    os.kill(pid, signal.SIGHUP)
                except:
                    pass
            return False                # wait for them
        return sr
    
    def server_loop(self):
        """start listening, then start the workers"""
        self.listen_sock.listen(self.nchild + 5)

        syslog(LOG_INFO, 'Listending on {0}'.format(self.server_address))

        for i in range(self.nchild):
            self.start_child()
            
        while len(self.children) > 0:
            try:
                pid, status = os.wait()
                self.children.remove(pid)
                x = os.WEXITSTATUS(status)
                syslog(LOG_INFO, 'worker {} exited {}'.format(pid, x))
                if x == os.EX_UNAVAILABLE:
                    PoolServer._shutdown_requested = True
                    if self.shutdown_requested():
                        break
                if not self.debug and not PoolServer._shutdown_requested:
                    self.start_child()
            except OSError as e:
                if e.errno == errno.EINTR:
                    continue
                raise

        self.listen_sock.close()
        self.server_finish()
        syslog(LOG_INFO, 'Shutting down...')
        

    def request_shutdown(self, request, client_address):
        """called by child to start shutdown"""
        syslog(LOG_NOTICE, 'shutdown request from: {0}'.format(self.client_address))
        self.child_finish()
        sys.exit(os.EX_UNAVAILABLE) # let server know

    def verify_request(self, request, client_address):
        """return True if request should be processed"""
        return True

    def process_request(self, request, client_address):
        """override with something useful"""
        while True:
            try:
                signal.alarm(self.timeout)
                msg = request.recv(1024)
                signal.alarm(0)
            except Exception as e:
                if e.errno in [errno.EINTR,errno.EIO]:
                    return
                if e.errno == errno.ETIMEDOUT:
                    syslog(LOG_NOTICE, 'timeout waiting for {}'.format(self.client_address))
                    return
                if self.debug:
                    print('caught: {} {}'.format(e.errno, e.strerror))
                raise
            if not msg:
                return
            response = 'echo[{0}]: {1:.64}\n'.format(self.pid, msg)
            request.send(response)
            
    def request_loop(self):
        """This is where children spend their lives"""

        while not self.shutdown_requested():
            try:
                request, client_address = self.listen_sock.accept()
            except:
                continue

            if self.verify_request(request, client_address):
                syslog(LOG_DEBUG, 'connection from {0}'.format(client_address))
                self.process_request(request, client_address)
            else:
                syslog(LOG_NOTICE, 'rejected {0}'.format(client_address))

            request.close()

if __name__ == '__main__':
    import getopt
    
    addr = '0.0.0.0'
    port = 0
    
    opts,args = getopt.getopt(sys.argv[1:], 'dL:O:t:k:c:a:p:')
    for o,a in opts:
        if o == '-a':
            addr = a
        elif o == '-p':
            port = int(a)

    openlog('echo', LOG_PID|LOG_PERROR, LOG_USER)
    server = PoolServer((addr, port))
    server.server_loop()
    

