Daniel Dunbar | 7e3089d | 2009-11-08 21:51:53 +0000 | [diff] [blame] | 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() |