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 kill(self, name, signum):
  90         assert name in self.namemap
  91         self.namemap[name].kill(name, stop)
  92
  93     def sigchild_handler(self, signum, stack_frame=None):
  94         try:
  95             (pid, status) = os.waitpid(-1, os.WNOHANG)
  96         except OSError, e:
  97             if e.errno is errno.ECHILD:
  98                 return
  99             raise
 100         while pid:
 101             if pid in self.pidmap:
 102                 p = self.pidmap[pid]
 103                 p.process.returncode = status
 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         raise 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