rpt

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# rpt is a Runtime performance tuner.

import os
import sys
import argparse
import shutil
import logging
import subprocess


def SplitArgv(argv):
    i = -1
    try:
        i = argv.index('--')
    except ValueError:
        pass
    if i == -1:
        return argv[:], []
    return argv[:i], argv[i + 1:]


ext_LD_PRELOAD = set()
ext_LD_LIBRARY_PATH = set()
# TODO: Support customizing environment variables.
ext_ENV = set()


def EnableProtonRTX():
    ext_ENV.add('PROTON_HIDE_NVIDIA_GPU=0')
    ext_ENV.add('VKD3D_CONFIG=dxr11,dxr')
    ext_ENV.add('PROTON_ENABLE_NVAPI=1')
    ext_ENV.add('PROTON_ENABLE_NGX_UPDATER=1')


def EnableProtonFSR(level='ultra'):
    ext_ENV.add('WINE_FULLSCREEN_FSR=1')
    ext_ENV.add(f'WINE_FULLSCREEN_FSR_MODE={level}')


def EnableMangoHud():
    ext_LD_LIBRARY_PATH.add('/usr/$LIB/mangohud')
    ext_LD_PRELOAD.add('libMangoHud.so')
    ext_ENV.add('MANGOHUD=1')


def EnableMangoHudDlsym():
    ext_LD_LIBRARY_PATH.add('/usr/$LIB/mangohud')
    ext_LD_PRELOAD.add('libMangoHud_dlsym.so')
    ext_ENV.add('MANGOHUD=1')


def UseMalloc(allocator):
    if allocator == '<default>':
        return
    elif allocator == 'jemalloc':
        ext_LD_PRELOAD.add('libjemalloc.so')
    elif allocator == 'mimalloc':
        ext_LD_PRELOAD.add('libmimalloc.so')
    elif allocator == 'tcmalloc':
        ext_LD_PRELOAD.add('libtcmalloc.so')
    else:
        logging.error('Unsupported allocator: {}'.format(allocator))


def RunProgram(command):
    argv = ['env']
    argv.extend(ext_ENV)

    def UpdatePath(key, extra):
        argv.append('{}={}'.format(
            key,
            os.environ.get(key, '') + os.path.pathsep +
            os.path.pathsep.join(extra)))

    UpdatePath('LD_LIBRARY_PATH', ext_LD_LIBRARY_PATH)
    UpdatePath('LD_PRELOAD', ext_LD_PRELOAD)
    argv.extend(command)
    return os.execv(shutil.which('env'), argv)


def main():
    parser = argparse.ArgumentParser(
        description="{} -- %command% in steam's launch options".format(
            sys.argv[0]))
    parser.add_argument('--mangohud',
                        action='store_true',
                        default=False,
                        help='Enable mangohud hooking')
    parser.add_argument('--mangohud_dlsym',
                        action='store_true',
                        default=False,
                        help='Enable mangohud dlsym hooking')
    parser.add_argument('--malloc',
                        default='<default>',
                        help='Memory allocator to use')
    parser.add_argument('-l',
                        '--preload',
                        action='append',
                        help='Append library to LD_PRELOAD')
    parser.add_argument('--fsr',
                        type=str,
                        default='ultra',
                        help='Try to enable FSR support in proton')
    # Use `lscpu --all --extended` to find out which core a CPU belongs to.
    parser.add_argument('--bind_cpus', help='Bind process to specific CPUs')
    parser.add_argument('--cpu_affinity', help='Set process affinity')
    parser.add_argument('--high_priority',
                        help='Set process at high priority',
                        action='store_true',
                        default=False)
    parser.add_argument('--rtx',
                        help='Tell proton to use NV RTX card features',
                        action='store_true',
                        default=False)
    my_argv, command = SplitArgv(sys.argv[1:])
    if not command:
        parser.print_help(sys.stderr)
        return 1
    config = parser.parse_args(my_argv)
    if config.mangohud:
        EnableMangoHud()
    if config.mangohud_dlsym:
        EnableMangoHudDlsym()
    if config.fsr:
        EnableProtonFSR(config.fsr)
    if config.rtx:
        EnableProtonRTX()
    UseMalloc(config.malloc)
    if config.preload:
        ext_LD_PRELOAD.update(config.preload)
    if config.bind_cpus:
        command_prefix = [
            'numactl',
            '-l',
        ]
        command_prefix.append(f'-C{config.bind_cpus}')
        command = command_prefix + command
    pid = os.getpid()
    if config.cpu_affinity:
        cmd = [
            'taskset',
            '-acp',
            config.cpu_affinity,
            f'{pid}',
        ]
        subprocess.run(cmd,
                       stdout=subprocess.DEVNULL,
                       stderr=subprocess.DEVNULL)
    if config.high_priority:
        # User has to modify /etc/security/limits.conf to allow negative nice
        # value.
        # For example:
        # @sudo            -       nice            -20
        cmd = [
            'renice',
            '-n',
            '-10',
            f'{pid}',
        ]
        subprocess.run(cmd,
                       stdout=subprocess.DEVNULL,
                       stderr=subprocess.DEVNULL)
    return RunProgram(command)


if __name__ == '__main__':
    sys.exit(main())