utils/Misc/zkill

   1 #!/usr/bin/env python
   2
   3 import os
   4 import re
   5 import sys
   6
   7 def _write_message(kind, message):
   8     import inspect, os, sys
   9
  10     # Get the file/line where this message was generated.
  11     f = inspect.currentframe()
  12     # Step out of _write_message, and then out of wrapper.
  13     f = f.f_back.f_back
  14     file,line,_,_,_ = inspect.getframeinfo(f)
  15     location = '%s:%d' % (os.path.basename(file), line)
  16
  17     print >>sys.stderr, '%s: %s: %s' % (location, kind, message)
  18
  19 note = lambda message: _write_message('note', message)
  20 warning = lambda message: _write_message('warning', message)
  21 error = lambda message: (_write_message('error', message), sys.exit(1))
  22
  23 def re_full_match(pattern, str):
  24     m = re.match(pattern, str)
  25     if m and m.end() != len(str):
  26         m = None
  27     return m
  28
  29 def parse_time(value):
  30     minutes,value = value.split(':',1)
  31     if '.' in value:
  32         seconds,fseconds = value.split('.',1)
  33     else:
  34         seconds = value
  35     return int(minutes) * 60 + int(seconds) + float('.'+fseconds)
  36
  37 def extractExecutable(command):
  38     """extractExecutable - Given a string representing a command line, attempt
  39     to extract the executable path, even if it includes spaces."""
  40
  41     # Split into potential arguments.
  42     args = command.split(' ')
  43
  44     # Scanning from the beginning, try to see if the first N args, when joined,
  45     # exist. If so that's probably the executable.
  46     for i in range(1,len(args)):
  47         cmd = ' '.join(args[:i])
  48         if os.path.exists(cmd):
  49             return cmd
  50
  51     # Otherwise give up and return the first "argument".
  52     return args[0]
  53
  54 class Struct:
  55     def __init__(self, **kwargs):
  56         self.fields = kwargs.keys()
  57         self.__dict__.update(kwargs)
  58
  59     def __repr__(self):
  60         return 'Struct(%s)' % ', '.join(['%s=%r' % (k,getattr(self,k))
  61                                          for k in self.fields])
  62
  63 kExpectedPSFields = [('PID', int, 'pid'),
  64                      ('USER', str, 'user'),
  65                      ('COMMAND', str, 'command'),
  66                      ('%CPU', float, 'cpu_percent'),
  67                      ('TIME', parse_time, 'cpu_time'),
  68                      ('VSZ', int, 'vmem_size'),
  69                      ('RSS', int, 'rss')]
  70 def getProcessTable():
  71     import subprocess
  72     p = subprocess.Popen(['ps', 'aux'], stdout=subprocess.PIPE,
  73                          stderr=subprocess.PIPE)
  74     out,err = p.communicate()
  75     res = p.wait()
  76     if p.wait():
  77         error('unable to get process table')
  78     elif err.strip():
  79         error('unable to get process table: %s' % err)
  80
  81     lns = out.split('\n')
  82     it = iter(lns)
  83     header = it.next().split()
  84     numRows = len(header)
  85
  86     # Make sure we have the expected fields.
  87     indexes = []
  88     for field in kExpectedPSFields:
  89         try:
  90             indexes.append(header.index(field[0]))
  91         except:
  92             if opts.debug:
  93                 raise
  94             error('unable to get process table, no %r field.' % field[0])
  95
  96     table = []
  97     for i,ln in enumerate(it):
  98         if not ln.strip():
  99             continue
 100
 101         fields = ln.split(None, numRows - 1)
 102         if len(fields) != numRows:
 103             warning('unable to process row: %r' % ln)
 104             continue
 105
 106         record = {}
 107         for field,idx in zip(kExpectedPSFields, indexes):
 108             value = fields[idx]
 109             try:
 110                 record[field[2]] = field[1](value)
 111             except:
 112                 if opts.debug:
 113                     raise
 114                 warning('unable to process %r in row: %r' % (field[0], ln))
 115                 break
 116         else:
 117             # Add our best guess at the executable.
 118             record['executable'] = extractExecutable(record['command'])
 119             table.append(Struct(**record))
 120
 121     return table
 122
 123 def getSignalValue(name):
 124     import signal
 125     if name.startswith('SIG'):
 126         value = getattr(signal, name)
 127         if value and isinstance(value, int):
 128             return value
 129     error('unknown signal: %r' % name)
 130
 131 import signal
 132 kSignals = {}
 133 for name in dir(signal):
 134     if name.startswith('SIG') and name == name.upper() and name.isalpha():
 135         kSignals[name[3:]] = getattr(signal, name)
 136
 137 def main():
 138     global opts
 139     from optparse import OptionParser, OptionGroup
 140     parser = OptionParser("usage: %prog [options] {pid}*")
 141
 142     # FIXME: Add -NNN and -SIGNAME options.
 143
 144     parser.add_option("-s", "", dest="signalName",
 145                       help="Name of the signal to use (default=%default)",
 146                       action="store", default='INT',
 147                       choices=kSignals.keys())
 148     parser.add_option("-l", "", dest="listSignals",
 149                       help="List known signal names",
 150                       action="store_true", default=False)
 151
 152     parser.add_option("-n", "--dry-run", dest="dryRun",
 153                       help="Only print the actions that would be taken",
 154                       action="store_true", default=False)
 155     parser.add_option("-v", "--verbose", dest="verbose",
 156                       help="Print more verbose output",
 157                       action="store_true", default=False)
 158     parser.add_option("", "--debug", dest="debug",
 159                       help="Enable debugging output",
 160                       action="store_true", default=False)
 161     parser.add_option("", "--force", dest="force",
 162                       help="Perform the specified commands, even if it seems like a bad idea",
 163                       action="store_true", default=False)
 164
 165     inf = float('inf')
 166     group = OptionGroup(parser, "Process Filters")
 167     group.add_option("", "--name", dest="execName", metavar="REGEX",
 168                       help="Kill processes whose name matches the given regexp",
 169                       action="store", default=None)
 170     group.add_option("", "--exec", dest="execPath", metavar="REGEX",
 171                       help="Kill processes whose executable matches the given regexp",
 172                       action="store", default=None)
 173     group.add_option("", "--user", dest="userName", metavar="REGEX",
 174                       help="Kill processes whose user matches the given regexp",
 175                       action="store", default=None)
 176     group.add_option("", "--min-cpu", dest="minCPU", metavar="PCT",
 177                       help="Kill processes with CPU usage >= PCT",
 178                       action="store", type=float, default=None)
 179     group.add_option("", "--max-cpu", dest="maxCPU", metavar="PCT",
 180                       help="Kill processes with CPU usage <= PCT",
 181                       action="store", type=float, default=inf)
 182     group.add_option("", "--min-mem", dest="minMem", metavar="N",
 183                       help="Kill processes with virtual size >= N (MB)",
 184                       action="store", type=float, default=None)
 185     group.add_option("", "--max-mem", dest="maxMem", metavar="N",
 186                       help="Kill processes with virtual size <= N (MB)",
 187                       action="store", type=float, default=inf)
 188     group.add_option("", "--min-rss", dest="minRSS", metavar="N",
 189                       help="Kill processes with RSS >= N",
 190                       action="store", type=float, default=None)
 191     group.add_option("", "--max-rss", dest="maxRSS", metavar="N",
 192                       help="Kill processes with RSS <= N",
 193                       action="store", type=float, default=inf)
 194     group.add_option("", "--min-time", dest="minTime", metavar="N",
 195                       help="Kill processes with CPU time >= N (seconds)",
 196                       action="store", type=float, default=None)
 197     group.add_option("", "--max-time", dest="maxTime", metavar="N",
 198                       help="Kill processes with CPU time <= N (seconds)",
 199                       action="store", type=float, default=inf)
 200     parser.add_option_group(group)
 201
 202     (opts, args) = parser.parse_args()
 203
 204     if opts.listSignals:
 205         items = [(v,k) for k,v in kSignals.items()]
 206         items.sort()
 207         for i in range(0, len(items), 4):
 208             print '\t'.join(['%2d) SIG%s' % (k,v)
 209                              for k,v in items[i:i+4]])
 210         sys.exit(0)
 211
 212     # Figure out the signal to use.
 213     signal = kSignals[opts.signalName]
 214     signalValueName = str(signal)
 215     if opts.verbose:
 216         name = dict((v,k) for k,v in kSignals.items()).get(signal,None)
 217         if name:
 218             signalValueName = name
 219             note('using signal %d (SIG%s)' % (signal, name))
 220         else:
 221             note('using signal %d' % signal)
 222
 223     # Get the pid list to consider.
 224     pids = set()
 225     for arg in args:
 226         try:
 227             pids.add(int(arg))
 228         except:
 229             parser.error('invalid positional argument: %r' % arg)
 230
 231     filtered = ps = getProcessTable()
 232
 233     # Apply filters.
 234     if pids:
 235         filtered = [p for p in filtered
 236                     if p.pid in pids]
 237     if opts.execName is not None:
 238         filtered = [p for p in filtered
 239                     if re_full_match(opts.execName,
 240                                      os.path.basename(p.executable))]
 241     if opts.execPath is not None:
 242         filtered = [p for p in filtered
 243                     if re_full_match(opts.execPath, p.executable)]
 244     if opts.userName is not None:
 245         filtered = [p for p in filtered
 246                     if re_full_match(opts.userName, p.user)]
 247     filtered = [p for p in filtered
 248                 if opts.minCPU <= p.cpu_percent <= opts.maxCPU]
 249     filtered = [p for p in filtered
 250                 if opts.minMem <= float(p.vmem_size) / (1<<20) <= opts.maxMem]
 251     filtered = [p for p in filtered
 252                 if opts.minRSS <= p.rss <= opts.maxRSS]
 253     filtered = [p for p in filtered
 254                 if opts.minTime <= p.cpu_time <= opts.maxTime]
 255
 256     if len(filtered) == len(ps):
 257         if not opts.force and not opts.dryRun:
 258             error('refusing to kill all processes without --force')
 259
 260     if not filtered:
 261         warning('no processes selected')
 262
 263     for p in filtered:
 264         if opts.verbose:
 265             note('kill(%r, %s) # (user=%r, executable=%r, CPU=%2.2f%%, time=%r, vmem=%r, rss=%r)' %
 266                  (p.pid, signalValueName, p.user, p.executable, p.cpu_percent, p.cpu_time, p.vmem_size, p.rss))
 267         if not opts.dryRun:
 268             try:
 269                 os.kill(p.pid, signal)
 270             except OSError:
 271                 if opts.debug:
 272                     raise
 273                 warning('unable to kill PID: %r' % p.pid)
 274
 275 if __name__ == '__main__':
 276     main()