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