prlimit.py

# 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()