[rustc.git] / src / llvm / utils / Misc / zkill

#!/usr/bin/env python

import os
import re
import sys

def _write_message(kind, message):
    import inspect, os, sys

    # Get the file/line where this message was generated.
    f = inspect.currentframe()
    # Step out of _write_message, and then out of wrapper.
    f = f.f_back.f_back
    file,line,_,_,_ = inspect.getframeinfo(f)
    location = '%s:%d' % (os.path.basename(file), line)

    print >>sys.stderr, '%s: %s: %s' % (location, kind, message)

note = lambda message: _write_message('note', message)
warning = lambda message: _write_message('warning', message)
error = lambda message: (_write_message('error', message), sys.exit(1))

def re_full_match(pattern, str):
    m = re.match(pattern, str)
    if m and m.end() != len(str):
        m = None
    return m

def parse_time(value):
    minutes,value = value.split(':',1)
    if '.' in value:
        seconds,fseconds = value.split('.',1)
    else:
        seconds = value
    return int(minutes) * 60 + int(seconds) + float('.'+fseconds)

def extractExecutable(command):
    """extractExecutable - Given a string representing a command line, attempt
    to extract the executable path, even if it includes spaces."""

    # Split into potential arguments.
    args = command.split(' ')

    # Scanning from the beginning, try to see if the first N args, when joined,
    # exist. If so that's probably the executable.
    for i in range(1,len(args)):
        cmd = ' '.join(args[:i])
        if os.path.exists(cmd):
            return cmd

    # Otherwise give up and return the first "argument".
    return args[0]

class Struct:
    def __init__(self, **kwargs):
        self.fields = kwargs.keys()
        self.__dict__.update(kwargs)

    def __repr__(self):
        return 'Struct(%s)' % ', '.join(['%s=%r' % (k,getattr(self,k))
                                         for k in self.fields])

kExpectedPSFields = [('PID', int, 'pid'),
                     ('USER', str, 'user'),
                     ('COMMAND', str, 'command'),
                     ('%CPU', float, 'cpu_percent'),
                     ('TIME', parse_time, 'cpu_time'),
                     ('VSZ', int, 'vmem_size'),
                     ('RSS', int, 'rss')]
def getProcessTable():
    import subprocess
    p = subprocess.Popen(['ps', 'aux'], stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE)
    out,err = p.communicate()
    res = p.wait()
    if p.wait():
        error('unable to get process table')
    elif err.strip():
        error('unable to get process table: %s' % err)

    lns = out.split('\n')
    it = iter(lns)
    header = it.next().split()
    numRows = len(header)

    # Make sure we have the expected fields.
    indexes = []
    for field in kExpectedPSFields:
        try:
            indexes.append(header.index(field[0]))
        except:
            if opts.debug:
                raise
            error('unable to get process table, no %r field.' % field[0])

    table = []
    for i,ln in enumerate(it):
        if not ln.strip():
            continue

        fields = ln.split(None, numRows - 1)
        if len(fields) != numRows:
            warning('unable to process row: %r' % ln)
            continue

        record = {}
        for field,idx in zip(kExpectedPSFields, indexes):
            value = fields[idx]
            try:
                record[field[2]] = field[1](value)
            except:
                if opts.debug:
                    raise
                warning('unable to process %r in row: %r' % (field[0], ln))
                break
        else:
            # Add our best guess at the executable.
            record['executable'] = extractExecutable(record['command'])
            table.append(Struct(**record))

    return table

def getSignalValue(name):
    import signal
    if name.startswith('SIG'):
        value = getattr(signal, name)
        if value and isinstance(value, int):
            return value
    error('unknown signal: %r' % name)

import signal
kSignals = {}
for name in dir(signal):
    if name.startswith('SIG') and name == name.upper() and name.isalpha():
        kSignals[name[3:]] = getattr(signal, name)

def main():
    global opts
    from optparse import OptionParser, OptionGroup
    parser = OptionParser("usage: %prog [options] {pid}*")

    # FIXME: Add -NNN and -SIGNAME options.

    parser.add_option("-s", "", dest="signalName",
                      help="Name of the signal to use (default=%default)",
                      action="store", default='INT',
                      choices=kSignals.keys())
    parser.add_option("-l", "", dest="listSignals",
                      help="List known signal names",
                      action="store_true", default=False)

    parser.add_option("-n", "--dry-run", dest="dryRun",
                      help="Only print the actions that would be taken",
                      action="store_true", default=False)
    parser.add_option("-v", "--verbose", dest="verbose",
                      help="Print more verbose output",
                      action="store_true", default=False)
    parser.add_option("", "--debug", dest="debug",
                      help="Enable debugging output",
                      action="store_true", default=False)
    parser.add_option("", "--force", dest="force",
                      help="Perform the specified commands, even if it seems like a bad idea",
                      action="store_true", default=False)

    inf = float('inf')
    group = OptionGroup(parser, "Process Filters")
    group.add_option("", "--name", dest="execName", metavar="REGEX",
                      help="Kill processes whose name matches the given regexp",
                      action="store", default=None)
    group.add_option("", "--exec", dest="execPath", metavar="REGEX",
                      help="Kill processes whose executable matches the given regexp",
                      action="store", default=None)
    group.add_option("", "--user", dest="userName", metavar="REGEX",
                      help="Kill processes whose user matches the given regexp",
                      action="store", default=None)
    group.add_option("", "--min-cpu", dest="minCPU", metavar="PCT",
                      help="Kill processes with CPU usage >= PCT",
                      action="store", type=float, default=None)
    group.add_option("", "--max-cpu", dest="maxCPU", metavar="PCT",
                      help="Kill processes with CPU usage <= PCT",
                      action="store", type=float, default=inf)
    group.add_option("", "--min-mem", dest="minMem", metavar="N",
                      help="Kill processes with virtual size >= N (MB)",
                      action="store", type=float, default=None)
    group.add_option("", "--max-mem", dest="maxMem", metavar="N",
                      help="Kill processes with virtual size <= N (MB)",
                      action="store", type=float, default=inf)
    group.add_option("", "--min-rss", dest="minRSS", metavar="N",
                      help="Kill processes with RSS >= N",
                      action="store", type=float, default=None)
    group.add_option("", "--max-rss", dest="maxRSS", metavar="N",
                      help="Kill processes with RSS <= N",
                      action="store", type=float, default=inf)
    group.add_option("", "--min-time", dest="minTime", metavar="N",
                      help="Kill processes with CPU time >= N (seconds)",
                      action="store", type=float, default=None)
    group.add_option("", "--max-time", dest="maxTime", metavar="N",
                      help="Kill processes with CPU time <= N (seconds)",
                      action="store", type=float, default=inf)
    parser.add_option_group(group)

    (opts, args) = parser.parse_args()

    if opts.listSignals:
        items = [(v,k) for k,v in kSignals.items()]
        items.sort()
        for i in range(0, len(items), 4):
            print '\t'.join(['%2d) SIG%s' % (k,v)
                             for k,v in items[i:i+4]])
        sys.exit(0)

    # Figure out the signal to use.
    signal = kSignals[opts.signalName]
    signalValueName = str(signal)
    if opts.verbose:
        name = dict((v,k) for k,v in kSignals.items()).get(signal,None)
        if name:
            signalValueName = name
            note('using signal %d (SIG%s)' % (signal, name))
        else:
            note('using signal %d' % signal)

    # Get the pid list to consider.
    pids = set()
    for arg in args:
        try:
            pids.add(int(arg))
        except:
            parser.error('invalid positional argument: %r' % arg)

    filtered = ps = getProcessTable()

    # Apply filters.
    if pids:
        filtered = [p for p in filtered
                    if p.pid in pids]
    if opts.execName is not None:
        filtered = [p for p in filtered
                    if re_full_match(opts.execName,
                                     os.path.basename(p.executable))]
    if opts.execPath is not None:
        filtered = [p for p in filtered
                    if re_full_match(opts.execPath, p.executable)]
    if opts.userName is not None:
        filtered = [p for p in filtered
                    if re_full_match(opts.userName, p.user)]
    filtered = [p for p in filtered
                if opts.minCPU <= p.cpu_percent <= opts.maxCPU]
    filtered = [p for p in filtered
                if opts.minMem <= float(p.vmem_size) / (1<<20) <= opts.maxMem]
    filtered = [p for p in filtered
                if opts.minRSS <= p.rss <= opts.maxRSS]
    filtered = [p for p in filtered
                if opts.minTime <= p.cpu_time <= opts.maxTime]

    if len(filtered) == len(ps):
        if not opts.force and not opts.dryRun:
            error('refusing to kill all processes without --force')

    if not filtered:
        warning('no processes selected')

    for p in filtered:
        if opts.verbose:
            note('kill(%r, %s) # (user=%r, executable=%r, CPU=%2.2f%%, time=%r, vmem=%r, rss=%r)' %
                 (p.pid, signalValueName, p.user, p.executable, p.cpu_percent, p.cpu_time, p.vmem_size, p.rss))
        if not opts.dryRun:
            try:
                os.kill(p.pid, signal)
            except OSError:
                if opts.debug:
                    raise
                warning('unable to kill PID: %r' % p.pid)

if __name__ == '__main__':
    main()
Commit	Line	Data
223e47cc LB	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()