#!/bin/sh

CONFIG="/etc/sysconfig/libvirt-hook-config.xml"

function addhugepages { # pagesize(2M|1G) count [node]
  # There's a recommendation to use 'virsh allocpages' for manipulating
  # hugepage pools, but a) that's a fairly recent addition to virsh (not
  # available in RHEL7.2) and b) I managed to get a hang using 'virsh
  # freepages', giving some credence to the warning in libvirt.org hook
  # documentation not to call back into libvirt.  Use raw sysfs instead.

  if [ -n "$3" ]; then
    SYSFSNODE="/sys/devices/system/node/node$3/hugepages/hugepages-"
  else
    SYSFSNODE="/sys/kernel/mm/hugepages/hugepages-"
  fi

  case "$1" in
    2M)
      SYSFSNODE+="2048kB"
      ;;
    1G)
      SYSFSNODE+="1048576kB"
      ;;
    *)
      echo "$0: Invalid hugepage size: $1" >&2
      exit 1
  esac

  if [ ! -d $SYSFSNODE ]; then
    echo "$0: Hugepage size $1 not supported by kernel" >&2
    exit 1
  fi

  echo $(( $(cat $SYSFSNODE/nr_hugepages) + $2 )) > $SYSFSNODE/nr_hugepages

  # We assume that we never need to retry releasing huge pages.  However,
  # allocating hugepages is not guaranteed, especially with 1G pages. It's
  # also racy, someone else can be adding or removing at the same time.
  # Therefore we try above to allocate for ourselves and below we switch to
  # the free_hugepages count to keep trying until we get what we need.

  if [ $2 -gt 0 ]; then
    for i in $(seq 1 10); do
      if [ $(cat $SYSFSNODE/free_hugepages) -lt $2 ]; then
        NEED=$(( $2 - $(cat $SYSFSNODE/free_hugepages) ))
        TOTAL=$(( $(cat $SYSFSNODE/nr_hugepages) + $NEED ))
        echo $TOTAL > $SYSFSNODE/nr_hugepages
      else
        break
      fi
    done

    if [ $(cat $SYSFSNODE/free_hugepages) -lt $2 ]; then
      echo "$0: Unable to allocate $1 hugepages for VM" >&2
      exit 1
    fi
  fi
}

XMLLINT=$(which xmllint 2>/dev/null)
if [ $? -ne 0 ]; then
  exit 0
fi

if [ ! -e $CONFIG ]; then
  exit 0
fi

BASE="//config/dynamic-hugepages/domain"

for DOMAIN in $(xmllint --xpath "$BASE[@name=\"$1\"]/@name" $CONFIG 2>/dev/null); do
  for SIZE in $(xmllint --xpath "$BASE[@$DOMAIN]/hugepages/@size" $CONFIG 2>/dev/null); do
    COUNT=$(xmllint --xpath "$BASE[@$DOMAIN]/hugepages[@$SIZE]/@count" $CONFIG)
    COUNT=$(echo $COUNT | cut -f 2 -d \")
    SIZE=$(echo $SIZE | cut -f 2 -d \")

    case "$2" in
      prepare)
        addhugepages $SIZE $COUNT
        ;;
      release)
        addhugepages $SIZE -$COUNT
        ;;
    esac
  done

  for NODE in $(xmllint --xpath "$BASE[@$DOMAIN]/node/@id" $CONFIG 2>/dev/null); do
    for SIZE in $(xmllint --xpath "$BASE[@$DOMAIN]/node[@$NODE]/hugepages/@size" $CONFIG 2>/dev/null); do
      COUNT=$(xmllint --xpath "$BASE[@$DOMAIN]/node[@$NODE]/hugepages[@$SIZE]/@count" $CONFIG)
      COUNT=$(echo $COUNT | cut -f 2 -d \")
      SIZE=$(echo $SIZE | cut -f 2 -d \")
      NODE=$(echo $NODE | cut -f 2 -d \")

      case "$2" in
        prepare)
          addhugepages $SIZE $COUNT $NODE
          ;;
        release)
          addhugepages $SIZE -$COUNT $NODE
          ;;
      esac
    done
  done
done

exit 0
