added enable_sighandler for scgi.forkpool

1 files changed, 50 insertions, 2 deletions

diff --git a/wsgitools/scgi/forkpool.py b/wsgitools/scgi/forkpool.py
index e6b65fd..dd8b42d 100644
--- a/wsgitools/scgi/forkpool.py
+++ b/wsgitools/scgi/forkpool.py
@@ -10,6 +10,7 @@ import os
 import select
 import sys
 import errno
+import signal
 
 class SocketFileWrapper:
     """Wraps a socket to a wsgi-compliant file-like object."""
@@ -201,6 +202,18 @@ class SCGIServer:
         self.server = None # becomes a socket
         # maps filedescriptors to WorkerStates
         self.workers = {}
+        self.running = False
+        self.ischild = False
+
+    def enable_sighandler(self, sig=signal.SIGTERM):
+        """
+        Changes the signal handler for the given signal to terminate the run()
+        loop.
+        @param sig: is the signal to handle
+        @returns: self
+        """
+        signal.signal(sig, self.shutdownhandler)
+        return self
 
     def run(self):
         """
@@ -210,13 +223,19 @@ class SCGIServer:
         self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         self.server.bind((self.interface, self.port))
         self.server.listen(5)
-        while True:
+        self.running = True
+        while self.running:
             while (len(self.workers) < self.minworkers or # less than min
                     (len(self.workers) < self.maxworkers and # less than max
                         not len([w for w in # no inactive
                                  self.workers.values() if w.state == 0]))):
                 self.spawnworker()
-            rs, _, _ = select.select(self.workers.keys(), [], [])
+            try:
+                rs, _, _ = select.select(self.workers.keys(), [], [])
+            except select.error, e:
+                if e[0] != errno.EINTR:
+                    raise
+                rs = []
             for s in rs:
                 try:
                     data = self.workers[s].sock.recv(1)
@@ -236,6 +255,34 @@ class SCGIServer:
                     pid, _ = os.waitpid(0, os.WNOHANG)
             except OSError:
                 pass
+        self.server.close()
+        self.server = None
+        self.killworkers()
+
+    def killworkers(self, sig=signal.SIGTERM):
+        """
+        Kills all worker children.
+        @param sig: is the signal used to kill the children
+        """
+        while self.workers:
+            _, state = self.workers.popitem()
+            state.sock.close()
+            os.kill(state.pid, sig)
+            # TODO: handle working children with a timeout
+
+    def shutdownhandler(self, sig=None, stackframe=None):
+        """
+        Signal handler function for stopping the run() loop. It works by
+        setting a variable that run() evaluates in each loop. As a signal
+        interrupts accept the loop is terminated, the accepting socket is
+        closed and the workers are killed.
+        @param sig: ignored for usage with signal.signal
+        @param stackframe: ignored for usage with signal.signal
+        """
+        if self.ischild:
+            sys.exit()
+        else:
+            self.running = False
 
     def spawnworker(self):
         """
@@ -245,6 +292,7 @@ class SCGIServer:
 
         pid = os.fork()
         if pid == 0:
+            self.ischild = True
             # close unneeded sockets
             srvsock.close()
             for worker in self.workers.values():