#!/usr/bin/python

#

# Tool to set up Linux large page support with minimal effort

#

# by Jarod Wilson <jarod@redhat.com>

# (c) Red Hat, Inc., 2009

#

# Requires hugeadm from libhugetlbfs 2.7 (or backported support)

#

import os

debug = False

# must be executed under the root to operate

if os.geteuid() != 0:

    print "You must be root to setup hugepages!"

    os._exit(1)

# config files we need access to

sysctlConf = "/etc/sysctl.conf"

if not os.access(sysctlConf, os.W_OK):

    print "Cannot access %s" % sysctlConf

    if debug == False:

        os._exit(1)

# This file will be created if it doesn't exist

limitsConf = "/etc/security/limits.d/hugepages.conf"

# Figure out what we've got in the way of memory

memTotal = 0

hugePageSize = 0

hugePages = 0

hugeadmexplain = os.popen("/usr/bin/hugeadm --explain 2>/dev/null").readlines()

for line in hugeadmexplain:

    if line.startswith("Total System Memory:"):

        memTotal = int(line.split()[3])

        break

if memTotal == 0:

    print "Your version of libhugetlbfs' hugeadm utility is too old!"

    os._exit(1)

# Pick the default huge page size and see how many pages are allocated

poolList = os.popen("/usr/bin/hugeadm --pool-list").readlines()

for line in poolList:

    if '*' in line:

        hugePageSize = int(line.split()[0])

        hugePages = int(line.split()[2])

        break

if hugePageSize == 0:

    print "Aborting, cannot determine system huge page size!"

    os._exit(1)

# Get initial sysctl settings

shmmax = 0

hugeGID = 0

for line in hugeadmexplain:

    if line.startswith("A /proc/sys/kernel/shmmax value of"):

        shmmax = int(line.split()[4])

        break

for line in hugeadmexplain:

    if line.strip().startswith("vm.hugetlb_shm_group = "):

        hugeGID = int(line.split()[2])

        break

# translate group into textual version

hugeGIDName = "null"

groupNames = os.popen("/usr/bin/getent group").readlines()

for line in groupNames:

    curGID = int(line.split(":")[2])

    if curGID == hugeGID:

        hugeGIDName = line.split(":")[0]

        break

# dump system config as we see it before we start tweaking it

print "Current configuration:"

print " * Total System Memory......: %6d MB" % memTotal

print " * Shared Mem Max Mapping...: %6d MB" % (shmmax / (1024 * 1024))

print " * System Huge Page Size....: %6d MB" % (hugePageSize / (1024 * 1024))

print " * Number of Huge Pages.....: %6d"    % hugePages

print " * Total size of Huge Pages.: %6d MB" % (hugePages * hugePageSize / (1024 * 1024))

print " * Remaining System Memory..: %6d MB" % (memTotal - (hugePages * hugePageSize / (1024 * 1024)))

print " * Huge Page User Group.....:  %s (%d)" % (hugeGIDName, hugeGID)

print

# ask how memory they want to allocate for huge pages

userIn = None

while not userIn:

    try:

        userIn = raw_input("How much memory would you like to allocate for huge pages? "

                           "(input in MB, unless postfixed with GB): ")

        if userIn[-2:] == "GB":

            userHugePageReqMB = int(userIn[0:-2]) * 1024

        elif userIn[-1:] == "G":

            userHugePageReqMB = int(userIn[0:-1]) * 1024

        elif userIn[-2:] == "MB":

            userHugePageReqMB = int(userIn[0:-2])

        elif userIn[-1:] == "M":

            userHugePageReqMB = int(userIn[0:-1])

        else:

            userHugePageReqMB = int(userIn)

        # As a sanity safeguard, require at least 128M not be allocated to huge pages

        if userHugePageReqMB > (memTotal - 128):

            userIn = None

            print "Refusing to allocate %d, you must leave at least 128MB for the system" % userHugePageReqMB

        elif userHugePageReqMB < (hugePageSize / (1024 * 1024)):

            userIn = None

            print "Sorry, allocation must be at least a page's worth!"

        else:

            break

    except ValueError:

        userIn = None

        print "Input must be an integer, please try again!"

userHugePageReqKB = userHugePageReqMB * 1024

userHugePagesReq = userHugePageReqKB / (hugePageSize / 1024)

print "Okay, we'll try to allocate %d MB for huge pages..." % userHugePageReqMB

print

# some basic user input validation

badchars = list(' \\\'":;~`!$^&*(){}[]?/><,')

inputIsValid = False

# ask for the name of the group allowed access to huge pages

while inputIsValid == False:

    foundbad = False

    userGroupReq = raw_input("What group should have access to the huge pages?"

                             "(The group will be created, if need be) [hugepages]: ")

    if userGroupReq is '':

        userGroupReq = 'hugepages'

    if userGroupReq[0].isdigit() or userGroupReq[0] == "-":

        foundbad = True

        print "Group names cannot start with a number or dash, please try again!"

    for char in badchars:

        if char in userGroupReq:

            foundbad = True

            print "Illegal characters in group name, please try again!"

            break

    if len(userGroupReq) > 16:

        foundbad = True

        print "Group names can't be more than 16 characaters, please try again!"

    if foundbad == False:

        inputIsValid = True

print "Okay, we'll give group %s access to the huge pages" % userGroupReq

# see if group already exists, use it if it does, if not, create it

userGIDReq = -1

for line in groupNames:

    curGroupName = line.split(":")[0]

    if curGroupName == userGroupReq:

        userGIDReq = int(line.split(":")[2])

        break

if userGIDReq > -1:

    print "Group %s (gid %d) already exists, we'll use it" % (userGroupReq, userGIDReq)

else:

    if debug == False:

        os.popen("/usr/sbin/groupadd %s" % userGroupReq)

    else:

        print "/usr/sbin/groupadd %s" % userGroupReq

    groupNames = os.popen("/usr/bin/getent group %s" % userGroupReq).readlines()

    for line in groupNames:

        curGroupName = line.split(":")[0]

        if curGroupName == userGroupReq:

            userGIDReq = int(line.split(":")[2])

            break

    print "Created group %s (gid %d) for huge page use" % (userGroupReq, userGIDReq)

print

# basic user input validation, take 2

# space is valid in this case, wasn't in the prior incarnation

badchars = list('\\\'":;~`!$^&*(){}[]?/><,')

inputIsValid = False

# ask for user(s) that should be in the huge page access group

while inputIsValid == False:

    foundbad = False

    userUsersReq = raw_input("What user(s) should have access to the huge pages (space-delimited list, users created as needed)? ")

    for char in badchars:

        if char in userUsersReq:

            foundbad = True

            print "Illegal characters in user name(s) or invalid list format, please try again!"

            break

    for n in userUsersReq.split():

        if len(n) > 32:

            foundbad = True

            print "User names can't be more than 32 characaters, please try again!"

            break

        if n[0] == "-":

            foundbad = True

            print "User names cannot start with a dash, please try again!"

            break

    if foundbad == False:

        inputIsValid = True

# see if user(s) already exist(s)

curUserList = os.popen("/usr/bin/getent passwd").readlines()

hugePageUserList = userUsersReq.split()

for hugeUser in hugePageUserList:

    userExists = False

    for line in curUserList:

        curUser = line.split(":")[0]

        if curUser == hugeUser:

            print "Adding user %s to huge page group" % hugeUser

            userExists = True

            if debug == False:

                os.popen("/usr/sbin/usermod -a -G %s %s" % (userGroupReq, hugeUser))

            else:

                print "/usr/sbin/usermod -a -G %s %s" % (userGroupReq, hugeUser)

        if userExists == True:

            break

    if userExists == False:

        print "Creating user %s with membership in huge page group" % hugeUser

        if debug == False:

            if hugeUser == userGroupReq:

                os.popen("/usr/sbin/useradd %s -g %s" % (hugeUser, userGroupReq))

            else:

                os.popen("/usr/sbin/useradd %s -G %s" % (hugeUser, userGroupReq))

        else:

            print "/usr/sbin/useradd %s -G %s" % (hugeUser, userGroupReq)

print

# set values for the current running environment

if debug == False:

    os.popen("/usr/bin/hugeadm --pool-pages-min DEFAULT:%sM" % userHugePageReqMB)

    os.popen("/usr/bin/hugeadm --pool-pages-max DEFAULT:%sM" % userHugePageReqMB)

    os.popen("/usr/bin/hugeadm --set-shm-group %d" % userGIDReq)

    os.popen("/usr/bin/hugeadm --set-recommended-shmmax")

else:

    print "/usr/bin/hugeadm --pool-pages-min DEFAULT:%sM" % userHugePageReqMB

    print "/usr/bin/hugeadm --pool-pages-max DEFAULT:%sM" % userHugePageReqMB

    print "/usr/bin/hugeadm --set-shm-group %d" % userGIDReq

    print "/usr/bin/hugeadm --set-recommended-shmmax"

    print

# figure out what that shmmax value we just set was

hugeadmexplain = os.popen("/usr/bin/hugeadm --explain 2>/dev/null").readlines()

for line in hugeadmexplain:

    if line.strip().startswith("kernel.shmmax = "):

        shmmax = int(line.split()[2])

        break

# write out sysctl config changes to persist across reboot

if debug == False:

    sysctlConfLines = "# sysctl configuration\n"

    if os.access(sysctlConf, os.W_OK):

        try:

            sysctlConfLines = open(sysctlConf).readlines()

            os.rename(sysctlConf, sysctlConf + ".backup")

            print("Saved original %s as %s.backup" % (sysctlConf, sysctlConf))

        except:

            pass

    fd = open(sysctlConf, "w")

    for line in sysctlConfLines:

        if line.startswith("kernel.shmmax"):

            continue

        elif line.startswith("vm.nr_hugepages"):

            continue

        elif line.startswith("vm.hugetlb_shm_group"):

            continue

        else:

            fd.write(line);

    fd.write("kernel.shmmax = %d\n" % shmmax)

    fd.write("vm.nr_hugepages = %d\n" % userHugePagesReq)

    fd.write("vm.hugetlb_shm_group = %d\n" % userGIDReq)

    fd.close()

else:

    print "Add to %s:" % sysctlConf

    print "kernel.shmmax = %d" % shmmax

    print "vm.nr_hugepages = %d" % userHugePagesReq

    print "vm.hugetlb_shm_group = %d" % userGIDReq

    print

# write out limits.conf changes to persist across reboot

if debug == False:

    limitsConfLines = "# Huge page access configuration\n"

    if os.access(limitsConf, os.W_OK):

        try:

            limitsConfLines = open(limitsConf).readlines()

            os.rename(limitsConf, limitsConf + ".backup")

            print("Saved original %s as %s.backup" % (limitsConf, limitsConf))

        except:

            pass

    fd = open(limitsConf, "w")

    for line in limitsConfLines:

        cfgExist = False

        for hugeUser in hugePageUserList:

            try:

                if line.split()[0] == hugeUser:

                    cfgExist = True

            except IndexError:

                # hit either white or comment line, it is safe not to take

                # any action and continue.

                pass

        if cfgExist == True:

            continue

        else:

            fd.write(line)

    for hugeUser in hugePageUserList:

        fd.write("%s		soft	memlock		%d\n" % (hugeUser, userHugePageReqKB))

        fd.write("%s		hard	memlock		%d\n" % (hugeUser, userHugePageReqKB))

    fd.close()

else:

    print "Add to %s:" % limitsConf

    for hugeUser in hugePageUserList:

        print "%s		soft	memlock		%d" % (hugeUser, userHugePageReqKB)

        print "%s		hard	memlock		%d" % (hugeUser, userHugePageReqKB)

# dump the final configuration of things now that we're done tweaking

print

print "Final configuration:"

print " * Total System Memory......: %6d MB" % memTotal

if debug == False:

    print " * Shared Mem Max Mapping...: %6d MB" % (shmmax / (1024 * 1024))

else:

    # This should be what we *would* have set it to, had we actually run hugeadm --set-recommended-shmmax

    print " * Shared Mem Max Mapping...: %6d MB" % (userHugePagesReq * hugePageSize / (1024 * 1024))

print " * System Huge Page Size....: %6d MB" % (hugePageSize / (1024 * 1024))

print " * Available Huge Pages.....: %6d"    % userHugePagesReq

print " * Total size of Huge Pages.: %6d MB" % (userHugePagesReq * hugePageSize / (1024 * 1024))

print " * Remaining System Memory..: %6d MB" % (memTotal - userHugePageReqMB)

print " * Huge Page User Group.....:  %s (%d)" % (userGroupReq, userGIDReq)

print