pymin/procman.py

   1 # vim: set encoding=utf-8 et sw=4 sts=4 :
   2
   3 import os
   4 import errno
   5 import signal
   6 import subprocess
   7
   8 __ALL__ = ('ProcessManager', 'manager', 'register', 'unregister', 'call',
   9            'start', 'stop', 'kill', 'get', 'has', 'sigchild_handler')
  10
  11 class ProcessInfo:
  12     def __init__(self, name, command, callback=None, persist=False,
  13                        args=None, kw=None, max_errors=3):
  14         self.name = name
  15         self.command = command
  16         self.callback = callback
  17         if args is None: args = list()
  18         self.args = args
  19         if kw is None: kw = dict()
  20         self.kw = kw
  21         self.persist = persist
  22         self.max_errors = max_errors
  23         self.clear()
  24     def clear(self):
  25         self.dont_run = False
  26         self.signal = None
  27         self.process = None
  28         self.error_count = 0
  29         self.last_return = None
  30         self.running = False
  31     def start(self):
  32         assert self.process is None
  33         self.restart()
  34     def restart(self):
  35         self.clear()
  36         self.process = subprocess.Popen(self.command, *self.args, **self.kw)
  37         self.running = True
  38     def stop(self):
  39         assert self.process is not None
  40         self.dont_run = True
  41         if self.signal == signal.SIGTERM or self.signal == signal.SIGKILL:
  42             # Allready stopped, kill it
  43             self.kill(signal.SIGKILL)
  44         else:
  45             # Stop it
  46             self.kill(signal.SIGTERM)
  47     def kill(self, signum):
  48         assert self.process is not None
  49         os.kill(pi.process.pid, signum)
  50         self.signal = signum
  51     def __repr__(self):
  52         pid = None
  53         if self.process is not None:
  54             pid = self.process.pid
  55         return 'ProcessInfo(name=%s, pid=%s command=%s, persist=%s, cb=%s)' % (
  56                     self.name, pid, self.command, self.persist,
  57                     self.callback.__name__)
  58
  59 class ProcessManager:
  60
  61     def __init__(self):
  62         self.services = dict()
  63         self.namemap = dict()
  64         self.pidmap = dict()
  65
  66     def register(self, name, command, callback=None, persist=False,
  67                        *args, **kw):
  68         self.services[name] = ProcessInfo(name, command, callback, persist,
  69                                           args, kw)
  70
  71     def unregister(self, name):
  72         del self.services[name]
  73
  74     def _call(self, pi):
  75         pi.start()
  76         self.namemap[pi.name] = self.pidmap[pi.process.pid] = pi
  77
  78     def call(self, name, command, callback=None, persist=False, *args, **kw):
  79         pi = ProcessInfo(name, command, callback, persist, args, kw)
  80         self._call(pi)
  81
  82     def start(self, name):
  83         assert name not in self.namemap
  84         self._call(self.services[name])
  85
  86     def stop(self, name):
  87         assert name in self.namemap
  88         self.namemap[name].stop(name)
  89
  90     def kill(self, name, signum):
  91         assert name in self.namemap
  92         self.namemap[name].kill(name, stop)
  93
  94     def sigchild_handler(self, signum, stack_frame=None):
  95         try:
  96             (pid, status) = os.waitpid(-1, os.WNOHANG)
  97         except OSError, e:
  98             if e.errno is e.ECHILD:
  99                 return
 100             raise
 101         while pid:
 102             if pid in self.pidmap:
 103                 p = self.pidmap[pid]
 104                 if p.callback is not None:
 105                     p.callback(self, p)
 106                 if p.dont_run or not p.persist or p.error_count >= p.max_errors:
 107                     del self.namemap[p.name]
 108                     del self.pidmap[pid]
 109                     p.clear()
 110                 else:
 111                     if p.process.returncode == 0:
 112                         p.error_count = 0
 113                     else:
 114                         p.error_count += 1
 115                     del self.pidmap[pid]
 116                     p.restart()
 117                     self.pidmap[p.process.pid] = p
 118             try:
 119                 (pid, status) = os.waitpid(-1, os.WNOHANG)
 120             except OSError, e:
 121                 if e.errno == errno.ECHILD:
 122                     return
 123                 raise
 124
 125     def get(self, name):
 126         if isinstance(name, basestring): # is a name
 127             if name in self.namemap:
 128                 return self.namemap[name]
 129             if name in self.services:
 130                 return self.services[name]
 131         else: # is a pid
 132             if name in self.pidmap:
 133                 return self.pidmap[name]
 134         return KeyError, name
 135
 136     def has(self, name):
 137         if isinstance(name, basestring): # is a name
 138             if name in self.namemap:
 139                 return True
 140             if name in self.services:
 141                 return True
 142         else: # is a pid
 143             if name in self.pidmap:
 144                 return True
 145         return False
 146
 147     def __getitem__(self, name):
 148         return self.get(name)
 149
 150     def __contains__(self, name):
 151         return self.has(name)
 152
 153 # Globals
 154 manager = ProcessManager()
 155 register = manager.register
 156 unregister = manager.unregister
 157 call = manager.call
 158 start = manager.start
 159 stop = manager.stop
 160 kill = manager.kill
 161 get = manager.get
 162 has = manager.has
 163 sigchild_handler = manager.sigchild_handler
 164
 165
 166 if __name__ == '__main__':
 167
 168     import signal
 169     import time
 170
 171     sig = None
 172     count = 0
 173
 174     def SIGCHLD_handler(signum, stacktrace):
 175         global sig
 176         sig = signum
 177         print 'SIGCHLD', signum
 178
 179     def notify(pm, pi):
 180         global count
 181         if pi.name == 'test-service':
 182             print 'test-service count =', count
 183             count += 1
 184             if count > 4:
 185                 print 'test-service not persistent anymore, start test2'
 186                 pi.persist = False
 187                 pm.start('test2')
 188         print 'died:', pi.name, pi.command
 189
 190     register('test-service', ('sleep', '2'), notify, True)
 191     register('test2', ('sleep', '3'), notify, False)
 192
 193     signal.signal(signal.SIGCHLD, SIGCHLD_handler)
 194
 195     call('test', ('sleep', '5'), notify)
 196     start('test-service')
 197
 198     print "Esperando...", [pi.name for pi in manager.namemap.values()]
 199     while manager.pidmap:
 200         signal.pause()
 201         if sig == signal.SIGCHLD:
 202             sig = None
 203             sigchild_handler(sig)
 204         print "Esperando...", [pi.name for pi in manager.namemap.values()]
 205