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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
|
# Copyright 2016 Red Hat.
# All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
import argparse
import os
import resource
import shutil
import sys
USAGE_PROGRAM = (
f'{os.path.basename(sys.executable)} -m oslo_concurrency.prlimit'
)
RESOURCES = (
# argparse argument => resource
('as', resource.RLIMIT_AS),
('core', resource.RLIMIT_CORE),
('cpu', resource.RLIMIT_CPU),
('data', resource.RLIMIT_DATA),
('fsize', resource.RLIMIT_FSIZE),
('memlock', resource.RLIMIT_MEMLOCK),
('nofile', resource.RLIMIT_NOFILE),
('nproc', resource.RLIMIT_NPROC),
('rss', resource.RLIMIT_RSS),
('stack', resource.RLIMIT_STACK),
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description='prlimit', prog=USAGE_PROGRAM)
parser.add_argument('--as', type=int, help='Address space limit in bytes')
parser.add_argument(
'--core', type=int, help='Core file size limit in bytes'
)
parser.add_argument('--cpu', type=int, help='CPU time limit in seconds')
parser.add_argument('--data', type=int, help='Data size limit in bytes')
parser.add_argument('--fsize', type=int, help='File size limit in bytes')
parser.add_argument(
'--memlock', type=int, help='Locked memory limit in bytes'
)
parser.add_argument(
'--nofile', type=int, help='Maximum number of open files'
)
parser.add_argument(
'--nproc', type=int, help='Maximum number of processes'
)
parser.add_argument(
'--rss', type=int, help='Maximum Resident Set Size (RSS) in bytes'
)
parser.add_argument('--stack', type=int, help='Stack size limit in bytes')
parser.add_argument('program', help='Program (absolute path)')
parser.add_argument(
'program_args', metavar="arg", nargs='...', help='Program parameters'
)
args = parser.parse_args()
return args
def main() -> None:
args = parse_args()
program = args.program
if not os.path.isabs(program):
# program uses a relative path: try to find the absolute path
# to the executable
program_abs = shutil.which(program)
if program_abs:
program = program_abs
for arg_name, rlimit in RESOURCES:
value = getattr(args, arg_name)
if value is None:
continue
try:
resource.setrlimit(rlimit, (value, value))
except ValueError as exc:
print(
f"{USAGE_PROGRAM}: failed to set the {arg_name.upper()} "
f"resource limit: {exc}",
file=sys.stderr,
)
sys.exit(1)
try:
# it's okay to start this user-provided program
os.execv(program, [program] + args.program_args) # noqa: S606
except Exception as exc:
print(
f"{USAGE_PROGRAM}: failed to execute {program}: {exc}",
file=sys.stderr,
)
sys.exit(1)
if __name__ == "__main__":
main()
|
