# ex:ts=4
#
# forkfarm
#   -- idle accept preforking daemon framework
#
# $LinuxKorea: forkfarm.py,v 2.19 2001/11/01 09:29:00 perky Exp $
#
# Copyright 2001 Hye-Shik Chang. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without 
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. Neither the name of author nor the names of its contributors
#    may be used to endorse or promote products derived from this software
#    without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 
# ARE DISCLAIMED.  IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY 
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE. 
#
# code by Hye-Shik Chang <perky@linuxkorea.co.kr>
#

from __future__ import nested_scopes
import socket, select
import sys, os, signal
import string, errno
import time

# constants for arguments
SC_SIGHUP  = 1
SC_SIGTERM = 2
SC_SIGINT  = 3
TERMSIGS = {
	signal.SIGHUP:		SC_SIGHUP,
	signal.SIGTERM:		SC_SIGTERM,
	signal.SIGINT:		SC_SIGINT
}
SOCK_DGRAM, SOCK_STREAM = socket.SOCK_DGRAM, socket.SOCK_STREAM

class forkfarm_carrot:
	""" the forkfarm connection controller thread """

	def __init__(self, sock, addr, side):
		self.sock = sock
		self.addr = addr
		self.side = side
	
	def start(self):
		if self.access(self.addr):
			self.sock.close()
		else:
			return self.run()
	
	def access(self, addr):
		# override me
		return None
		
	def run(self):
		# return 1, if you want not to run anymore.
		raise NotImplemented
	
	def process(self, data, addr):
		# use this instead of start, access, run methods for dgram sockets
		raise NotImplemented

class forkfarm:
	""" prefork daemon framework """

	def __init__(self, address, controller=forkfarm_carrot, maxprocess=4, \
					socktype=SOCK_STREAM, runlimit=0, side={}):
		if type(address) is type(()):
			family = socket.AF_INET
		else:
			family = socket.AF_UNIX
		
		self.address    = address		# binding address (tuple or string)
		self.maxprocess = maxprocess	# n# of service processes
		self.sock		= socket.socket(family, socktype)
										# main binding socket
		self.controller = controller	# service controller class
		self.socktype	= socktype		# socket type
		self.stopcode	= 0				# return value of loop() method
		self.serial		= 0				# serial# of forked childs
		self.side		= side			# passing side data to controller
		self.children	= []			# children's PIDs
		self.runlimit	= runlimit		# limit of running time for each child
		try:
			self.sock.setsockopt (
				socket.SOL_SOCKET, socket.SO_REUSEADDR,
				self.sock.getsockopt (socket.SOL_SOCKET, socket.SO_REUSEADDR) | 1
			)
		except: 
			pass
		self.sock.bind(address)
		if socktype is SOCK_STREAM:
			self.childclass = forkfarm_child
			self.sock.listen(128)
		elif socktype is SOCK_DGRAM:
			self.childclass = forkfarm_child_dgram
		self.oninit()
	
	def __del__(self):
		self.atexit()
		self.sock.close()
	
	def preserve_children(self):
		while self.maxprocess > len(self.children):
			self.children.append(self.fork_child(self.serial))
			self.serial = self.serial + 1

	def SIGCHLD(self, signum, frame):
		try:	
			while 1:
				pid = os.waitpid(0, os.WNOHANG)[0]
				if not pid:
					raise OSError 
				self.children.remove(pid)
		except OSError:
			pass
	
	def signal_terminator(self, signum, frame):
		self.stopcode = TERMSIGS[signum]
		self.killall()

	def fork_child(self, childn):
		pid = os.fork()
		if not pid:
			try:
				try:
					c = self.childclass(self.sock, self.controller, 
											(childn, self.runlimit), self.side)
					self.childinit(c)
					c.loop()
					del c
				except KeyboardInterrupt:
					pass
				except socket.error, v:
					if v[0] != errno.EINTR:
						import traceback, sys
						traceback.print_exc(file=sys.stdout)
				except:
					import traceback, sys
					traceback.print_exc(file=sys.stdout)
			finally:
				self.childexit()
				os._exit(0)
		return pid
	
	def killall(self, signum=signal.SIGTERM):
		for i in self.children[:]:
			os.kill(i, signum)
	
	def wait(self):
		while self.children:
			signal.pause()
		
	def activewait(self):
		def KILLALL_SIG(sn, fr):
			self.killall()
		
		signal.signal(signal.SIGALRM, KILLALL_SIG)
		signal.alarm(7)
		
		while self.children:
			signal.pause()
		
		signal.signal(signal.SIGALRM, signal.SIG_DFL)
	
	def oninit(self):
		# on initializing server
		pass

	def atexit(self):
		# on terminating server
		pass
	
	def childinit(self, childinst):
		# on initializing each child process
		pass
	
	def childexit(self):
		# on terminating each child process
		pass
	
	def loop(self):
		signal.signal(signal.SIGCHLD, self.SIGCHLD)
		signal.signal(signal.SIGHUP,  self.signal_terminator)
		signal.signal(signal.SIGTERM, self.signal_terminator)
		signal.signal(signal.SIGINT,  self.signal_terminator)
		
		while not self.stopcode: # derived classes can stop this loop
			self.preserve_children()
			signal.pause()
		
		signal.signal(signal.SIGINT,  signal.SIG_DFL)
		signal.signal(signal.SIGTERM, signal.SIG_DFL)
		signal.signal(signal.SIGHUP,  signal.SIG_DFL)
		
		self.activewait()
		# in order to clean dying children
		signal.signal(signal.SIGCHLD, signal.SIG_DFL)
		
		return self.stopcode
	

class forkfarm_child:
	
	def __init__(self, sock, controller, cinfo, side):
		self.childn = cinfo[0]
		self.sock   = sock
		self.side   = side
		self.runlimit = cinfo[1] or -1
		self.controller = controller
	
	def loop(self):
		signal.signal(signal.SIGTERM, self.signal_terminator)
		
		while self.runlimit:
			conn, addr = self.sock.accept()
			self.controller(conn, addr, self.side).start() # blocks until ends
			if self.runlimit >= 0:
				self.runlimit = self.runlimit - 1

		signal.signal(signal.SIGTERM, signal.SIG_DFL)
	
	def signal_terminator(self, signum, frame):
		self.runlimit = 0


class forkfarm_child_dgram:

	def __init__(self, sock, controller, cinfo, side):
		self.childn = cinfo[0]
		self.sock   = sock
		self.side   = side
		self.runlimit = cinfo[1] or -1
		if side.has_key('recvbuff'):
			self.recvbuff = side['recvbuff']
		else:
			self.recvbuff = 8192
		self.controller = controller(self.sock, None, self.side)
	
	def loop(self):
		signal.signal(signal.SIGTERM, self.signal_terminator)
		
		while self.runlimit:
			data, addr = self.sock.recvfrom(self.recvbuff)
			self.controller.process(data, addr)
			if self.runlimit >= 0:
				self.runlimit = self.runlimit - 1

		signal.signal(signal.SIGTERM, signal.SIG_DFL)
	
	def signal_terminator(self, signum, frame):
		self.runlimit = 0


if __name__ == "__main__":
	
	# this is example for tcp

	class mycontrol(forkfarm_carrot):
		def run(self):
			self.sock.send("Hello!!")
			self.sock.close()

	forkfarm(('', 9988), mycontrol).loop()


	# this is DGRAM usage

	#import os, time
	#
	#class mycontrol(forkfarm_carrot):
	#	def process(self, data, addr):
	#		print os.getpid(), addr, ">>", data
	#		self.sock.sendto('reply:' + data, addr)
	#
	#forkfarm(('', 9988), mycontrol, socktype=SOCK_DGRAM, maxprocess=50).loop)