1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
|
#!/usr/bin/env python3
from __future__ import print_function, absolute_import, unicode_literals
import os
import argparse
# Define MYPY as False and use it as a conditional for typing import. Despite
# this declaration mypy will really treat MYPY as True when type-checking.
# This is required so that we can import typing on Python 2.x without the
# typing module installed. For more details see:
# https://mypy.readthedocs.io/en/latest/common_issues.html#import-cycles
MYPY = False
if MYPY:
from typing import List, Tuple, Text
def invoke(cmd):
# type: (List[Text]) -> Tuple[int, int]
"""
invoke invokes command and returns the pair (result, max_rss)
For normal termination result is >= 0 and conveys the exit code.
For abnormal termination result is < 0 and conveys the signal number.
"""
pid = os.fork()
if pid == 0:
# child
os.execvpe(cmd[0], cmd, os.environ)
os.exit(127)
else:
# parent
_, wait_status, rusage = os.wait4(pid, 0)
if os.WIFEXITED(wait_status):
exit_code = os.WEXITSTATUS(wait_status)
return exit_code, rusage.ru_maxrss
elif os.WIFSIGNALED(wait_status):
sig_num = os.WTERMSIG(wait_status)
return -sig_num, rusage.ru_maxrss
else:
raise OSError("unexpected wait status {}".format(wait_status))
def _make_parser():
# type: () -> argparse.ArgumentParser
parser = argparse.ArgumentParser(
description="""
This program executes COMMAND and stores the maximum size of RSS
into the specified file. The exit status of COMMAND is preserved.
"""
)
parser.add_argument(
"-o", dest="output",
metavar="RSS-FILE",
type=argparse.FileType(mode="w"),
help="write maximum RSS size to given file",
)
parser.add_argument(
"cmd", metavar="COMMAND", nargs="...", help="command to execute"
)
return parser
def main():
# type: () -> None
parser = _make_parser()
ns = parser.parse_args()
# The command cannot be empty but it is difficult to express in argparse
# itself.
if len(ns.cmd) == 0:
parser.print_usage()
parser.exit(0)
try:
result, max_rss = invoke(ns.cmd)
if ns.output is not None:
ns.output.write("{}\n".format(max_rss))
ns.output.close()
raise SystemExit(result)
except OSError as exc:
raise SystemExit(exc)
if __name__ == "__main__":
main()
|
