#!/usr/bin/env python3

import sys, re, argparse
sys.dont_write_bytecode = True
from subprocess import run

"""
hugepages - Configure hugepages on Linux EL7/EL8/EL9
Copyright (c) 2023 - Bart Sjerps <bart@dirty-cache.com>
License: GPLv3+
"""

# Max ratio of memory that can be used for hugepages
maxhuge = 0.8

# Location of sysctl file
sysctlpath = '/etc/sysctl.d/99-hp.conf'

def regsearch(exp, data):
    r = re.search(exp, data, re.M)
    if r:
        return r.group(1)

def getinfo():
    "Get the current memory values"
    with open('/proc/meminfo') as f:
        data = f.read()
    memtotal  = regsearch(r'MemTotal:\s+(\d+)\s+.*', data)
    hugetotal = regsearch(r'HugePages_Total:\s+(\d+)\s+.*', data)
    hugefree  = regsearch(r'HugePages_Free:\s+(\d+)\s+.*', data)
    return int(memtotal)//1024, int(hugetotal), int(hugefree)

def hugelist():
    "Show the current memory values"
    memory_mb, hugtotal, hugefree = getinfo()
    huglimit = round(maxhuge * memory_mb) // 2

    print(f"Memory Total:    {memory_mb:>10}  (           MiB)")
    print(f"Hugepages Total: {hugtotal:>10}  ({hugtotal*2:>10.0f} MiB)")
    print(f"Hugepages Free:  {hugefree:>10}  ({hugefree*2:>10.0f} MiB)")
    print(f"Hugepages Limit: {huglimit:>10}  ({huglimit*2:>10} MiB)")

def update(args):
    memory_mb, hugtotal, hugefree = getinfo()
    huglimit = round(maxhuge * memory_mb) // 2
    if args.update > huglimit:
        raise ValueError(f"Too high, max={huglimit}")

    with open(sysctlpath, 'w') as f:
        f.write(f"vm.nr_hugepages = {args.update}\n")
    
    run(['/usr/sbin/sysctl', '--system'])
    hugelist()

if __name__ == '__main__':
    formatter = lambda prog: argparse.HelpFormatter(prog, max_help_position=40)
    parser = argparse.ArgumentParser(formatter_class=formatter)

    parser.add_argument('-u', '--update', type=int, metavar='pages', help='Set nr of hugepages')
    args = parser.parse_args()
    try:
        if args.update:
            update(args)
        else:
            hugelist()
    except ValueError as e:
        print(e)
        sys.exit(10)
