Version 1.3 - cleanup some stuff, better more reliable termination

kata198 · kata198 · commit dd56d895c53b · 2015-06-09T15:26:46.000-04:00
diff --git a/PumpkinLB.py b/PumpkinLB.py
@@ -1,11 +1,13 @@
 #!/usr/bin/python2
 
-import os
+import math
 import multiprocessing
+import os
 import platform
 import socket
 import sys
 import signal
+import threading
 import traceback
 import time
 
@@ -14,6 +16,7 @@
 from pumpkinlb.listener import PumpkinListener
 
 if __name__ == '__main__':
+
     configFilename = None
     for arg in sys.argv[1:]:
         if arg == '--help':
@@ -69,43 +72,70 @@ def handleSigTerm(*args):
                 pass
         sys.stderr.write('Sent signal to children, waiting up to 4 seconds then trying to clean up\n')
         time.sleep(1)
+        startTime = time.time()
         remainingListeners = listeners
-        for i in xrange(3):
-            remainingListeners2 = []
-            for listener in remainingListeners:
-                sys.stderr.write('Waiting on %d...\n' %(listener.pid,))
-                sys.stderr.flush()
-                listener.join(.005)
-                if listener.is_alive() is True:
-                    remainingListeners2.append(listener)
-            remainingListeners = remainingListeners2
-            sys.stderr.write('Remaining (%d) listeners are: %s\n' %(len(remainingListeners), [listener.pid for listener in remainingListeners]))
+        remainingListeners2 = []
+        for listener in remainingListeners:
+            sys.stderr.write('Waiting on %d...\n' %(listener.pid,))
             sys.stderr.flush()
-            if len(remainingListeners) == 0:
-                break
-            time.sleep(1)
-
-        if len(remainingListeners) > 0:
-            sys.stderr.write('After trying to clean up, %d listeners remain.\n' %(len(remainingListeners)))
-            for listener in remainingListeners:
-                try:
-                    os.kill(listener.pid, signal.SIGKILL)
-                except:
-                    pass
-            time.sleep(.1)
-            sys.stderr.write('Starting final join\n')
-            sys.stderr.flush()
-            for listener in remainingListeners:
-                listener.join()
-            sys.stderr.write('Done\n')
-            sys.stderr.flush()
-         
-        if '_NT' in platform.system():
-            # Some issue on windows, or at least cygwin on windows, causes an infinite loop in this signal handler when trying to kill the process. CThe only way out is to cause an exception between here and the next call, so provide a function with invalid arguments.
-            sys.stderr.write('You can ignore the below exception, it is used to quit on NT\n')
-            signal.signal(signal.SIGTERM, sys.exit)
-            signal.signal(signal.SIGINT, sys.exit)
+            listener.join(.05)
+            if listener.is_alive() is True:
+                remainingListeners2.append(listener)
+        remainingListeners = remainingListeners2
+        sys.stderr.write('Remaining (%d) listeners are: %s\n' %(len(remainingListeners), [listener.pid for listener in remainingListeners]))
+        sys.stderr.flush()
+
+        afterJoinTime = time.time()
+
+        if remainingListeners:
+            delta = afterJoinTime - startTime
+            remainingSleep = int(6 - math.floor(afterJoinTime - startTime))
+            if remainingSleep > 0:
+                anyAlive = False
+                # If we still have time left, see if we are just done or if there are children to clean up using remaining time allotment
+                if threading.activeCount() > 1 or len(multiprocessing.active_children()) > 0:
+                    sys.stderr.write('Listener closed in %1.2f seconds. Waiting up to %d seconds before terminating.\n' %(delta, remainingSleep))
+                    sys.stderr.flush()
+                    thisThread = threading.current_thread()
+                    for i in range(remainingSleep):
+                        allThreads = threading.enumerate()
+                        anyAlive = False
+                        for thread in allThreads:
+                            if thread is thisThread or thread.name == 'MainThread':
+                                continue
+                            thread.join(.05)
+                            if thread.is_alive() == True:
+                                anyAlive = True
+
+                        allChildren = multiprocessing.active_children()
+                        for child in allChildren:
+                            child.join(.05)
+                            if child.is_alive() == True:
+                                anyAlive = True
+                        if anyAlive is False:
+                            break
+                        time.sleep(1)
+
+                if anyAlive is True:
+                    sys.stderr.write('Could not kill in time.\n')
+                else:
+                    sys.stderr.write('Shutdown successful after %1.2f seconds.\n' %( time.time() - startTime))
+                sys.stderr.flush()
+                    
+            else:
+                sys.stderr.write('Listener timed out in closing, exiting uncleanly.\n')
+                sys.stderr.flush()
+                time.sleep(.05) # Why not? :P
+
+        sys.stdout.write('exiting...\n')
+        sys.stdout.flush()
+        signal.signal(signal.SIGTERM, signal.SIG_DFL)
+        signal.signal(signal.SIGINT, signal.SIG_DFL)
         sys.exit(0)
+        os.kill(os.getpid(), signal.SIGTERM)
+        return 0
+    # END handleSigTerm
+        
 
     signal.signal(signal.SIGTERM, handleSigTerm)
     signal.signal(signal.SIGINT, handleSigTerm)
@@ -116,4 +146,4 @@ def handleSigTerm(*args):
         except:
             os.kill(os.getpid(), signal.SIGTERM)
 
-    # vim: ts=4 sw=4 expandtab
+# vim: set ts=4 sw=4 expandtab
diff --git a/pumpkinlb/listener.py b/pumpkinlb/listener.py
@@ -27,6 +27,8 @@ def __init__(self, localAddr, localPort, workers):
 
         self.listenSocket = None
 
+        self.cleanupThread = None
+
         self.keepGoing = True
 
     def cleanup(self):
@@ -40,9 +42,10 @@ def cleanup(self):
             time.sleep(1.5)
 
     def closeWorkers(self, *args):
-        sys.stdout.write("GOT SIGNAL on %s:%d\n" %(self.localAddr, self.localPort))
         self.keepGoing = False
 
+        time.sleep(1)
+
         try:
             self.listenSocket.shutdown(socket.SHUT_RDWR)
         except:
@@ -53,6 +56,8 @@ def closeWorkers(self, *args):
             pass
 
         if not self.pumpkinWorkers:
+            self.cleanupThread and self.cleanupThread.join(3)
+            signal.signal(signal.SIGTERM, signal.SIG_DFL)
             sys.exit(0)
 
         for pumpkinWorker in self.pumpkinWorkers:
@@ -64,6 +69,7 @@ def closeWorkers(self, *args):
 
         time.sleep(1)
 
+
         remainingWorkers = []
         for pumpkinWorker in self.pumpkinWorkers:
             pumpkinWorker.join(.03)
@@ -75,6 +81,10 @@ def closeWorkers(self, *args):
             for pumpkinWorker in remainingWorkers:
                 pumpkinWorker.join(.2)
 
+        self.cleanupThread and self.cleanupThread.join(2)
+
+        signal.signal(signal.SIGTERM, signal.SIG_DFL)
+
         sys.exit(0)
 
     def retryFailedWorkers(self, *args):
@@ -122,7 +132,7 @@ def run(self):
 
         listenSocket.listen(5)
 
-        cleanupThread = threading.Thread(target=self.cleanup)
+        self.cleanupThread = cleanupThread = threading.Thread(target=self.cleanup)
         cleanupThread.start()
         if len(self.workers) > 1:
             retryThread = threading.Thread(target=self.retryFailedWorkers)
@@ -148,10 +158,9 @@ def run(self):
         except Exception as e:
             sys.stderr.write('Got exception: %s, shutting down worker on %s:%d\n' %(str(e), self.localAddr, self.localPort))
             self.closeWorkers()
+            return
 
 
-        # If we got here, must have been signlaed to terminate. Stay open for a short bit to try to close children, then tank.
-        time.sleep(6)
-        sys.exit(0)
+        self.closeWorkers()
 
-# vim: ts=4 sw=4 expandtab
+ 
diff --git a/pumpkinlb/worker.py b/pumpkinlb/worker.py
@@ -42,6 +42,7 @@ def closeConnections(self):
             self.clientSocket.close()
         except:
             pass
+        signal.signal(signal.SIGTERM, signal.SIG_DFL)
 
     def closeConnectionsAndExit(self, *args):
         self.closeConnections()
@@ -61,44 +62,47 @@ def run(self):
 
         signal.signal(signal.SIGTERM, self.closeConnectionsAndExit)
 
-        dataToClient = ''
-        dataFromClient = ''
-        while True:
-            waitingToWrite = []
+        try:
+            dataToClient = ''
+            dataFromClient = ''
+            while True:
+                waitingToWrite = []
 
-            if dataToClient:
-                waitingToWrite.append(clientSocket)
-            if dataFromClient:
-                waitingToWrite.append(workerSocket)
-            
+                if dataToClient:
+                    waitingToWrite.append(clientSocket)
+                if dataFromClient:
+                    waitingToWrite.append(workerSocket)
+                
 
-            (hasDataForRead, readyForWrite, hasError) = select.select( [clientSocket, workerSocket], waitingToWrite, [clientSocket, workerSocket], .3)
+                (hasDataForRead, readyForWrite, hasError) = select.select( [clientSocket, workerSocket], waitingToWrite, [clientSocket, workerSocket], .3)
 
-            if hasError:
-                break
-        
-            if clientSocket in hasDataForRead:
-                nextData = clientSocket.recv(4096)
-                if not nextData:
+                if hasError:
                     break
-                dataFromClient += nextData
+            
+                if clientSocket in hasDataForRead:
+                    nextData = clientSocket.recv(4096)
+                    if not nextData:
+                        break
+                    dataFromClient += nextData
+
+                if workerSocket in hasDataForRead:
+                    nextData = workerSocket.recv(4096)
+                    if not nextData:
+                        break
+                    dataToClient += nextData
+            
+                if workerSocket in readyForWrite:
+                    while dataFromClient:
+                        workerSocket.send(dataFromClient[:4096])
+                        dataFromClient = dataFromClient[4096:]
 
-            if workerSocket in hasDataForRead:
-                nextData = workerSocket.recv(4096)
-                if not nextData:
-                    break
-                dataToClient += nextData
-        
-            if workerSocket in readyForWrite:
-                while dataFromClient:
-                    workerSocket.send(dataFromClient[:4096])
-                    dataFromClient = dataFromClient[4096:]
+                if clientSocket in readyForWrite:
+                    while dataToClient:
+                        clientSocket.send(dataToClient[:4096])
+                        dataToClient = dataToClient[4096:]
 
-            if clientSocket in readyForWrite:
-                while dataToClient:
-                    clientSocket.send(dataToClient[:4096])
-                    dataToClient = dataToClient[4096:]
+        except Exception as e:
+            sys.stderr.write('Error: ' + str(e) + '\n')
 
-        self.closeConnections()
+        self.closeConnectionsAndExit()
 
-# vim: ts=4 sw=4 expandtab
diff --git a/setup.py b/setup.py
@@ -37,7 +37,7 @@
 """
 
 setup(name='PumpkinLB',
-        version='1.2.1',
+        version='1.3',
         scripts=['PumpkinLB.py'],
         packages=['pumpkinlb'],
         author='Tim Savannah',
@@ -46,6 +46,7 @@
         maintainer_email='kata198@gmail.com',
         provides=['PumpkinLB'],
         description='A simple, fast, pure-python load balancer',
+        url='https://github.com/kata198/PumpkinLB',
         long_description=long_description,
         license='GPLv3',
         keywords=['load balancer', 'load balance', 'python', 'balance', 'lb', 'http', 'socket', 'port', 'forward', 'tcp', 'fast', 'server', 'network'],
