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