#!/bin/bash
lockfile=/var/lock/subsys/os-cgroups

# import openshift node configuration
if [ -f /etc/openshift/node.conf ]
then
    . /etc/openshift/node.conf
fi

# import resource limit tuning values for cgroups
if [ -f /etc/openshift/resource_limits.conf ]
then
    . /etc/openshift/resource_limits.conf
fi

RETVAL=0
GROUP_RETVAL=0

#
# Set defaults if not provided
#
GEAR_GECOS=${GEAR_GECOS:="OpenShift guest"}

OPENSHIFT_CGROUP_ROOT=${OPENSHIFT_CGROUP_ROOT:="/openshift"}
OPENSHIFT_CGROUP_SUBSYSTEMS=${OPENSHIFT_CGROUP_SUBSYSTEMS:="cpu,cpuacct,memory,net_cls,freezer"}

CGROUP_RULES_FILE=${CGROUP_RULES_FILE:="/etc/cgrules.conf"}

CPU_VARS="cfs_period_us cfs_quota_us rt_period_us rt_runtime_us shares"
MEM_VARS="limit_in_bytes memsw_limit_in_bytes soft_limit_in_bytes swappiness"

# Get a user's UID
function uid() {
    # USERNAME=$1
    getent passwd | grep -e "^$1:" | cut -d: -f3
}

# ============================================================================
#  Functions for setting the net class
# ============================================================================

#
# Convert an MCS pair into a cgroup net class id
#
function classid() {
    # major: 1, minor UID
    printf "0x1%04x" $1
}

function set_net_cls() {
    # USERNAME=$1
    CGPATH=openshift/$1
    USERID=`uid $1`
    USERCLASSID=`classid $USERID`
    cgset -r net_cls.classid=$USERCLASSID $CGPATH
}

# ==========================================================================
#  Functions for tuning the user's CPU limits in cgroups
# ==========================================================================
CPUVARS="cfs_period_us cfs_quota_us rt_period_us rt_runtime_us shares"
function set_cpu() {
    # USERNAME=$1
    CGPATH=openshift/$1

    for VARNAME in $CPUVARS
    do
	# cgroups names can have periods(.)  shell varnames can't
	SAFENAME=`echo $VARNAME | tr . _`
	VALUE=`eval echo \\$cpu_$SAFENAME`
	if [ -n "${VALUE}" ]
	then
	    # TODO: get per-app increments
	    cgset -r "cpu.$VARNAME=$VALUE" $CGPATH
	fi
    done
}

# ==========================================================================
#  Functions for tuning the user's memory limits in cgroups
# ==========================================================================
MEMVARS="limit_in_bytes memsw.limit_in_bytes soft_limit_in_bytes swappiness"
function set_memory() {
    # USERNAME=$1
    CGPATH=openshift/$1

    # for each var get and set the value
    for VARNAME in $MEMVARS
    do
	# cgroups names can have periods(.)  shell varnames can't
	SAFENAME=`echo $VARNAME | tr . _`
	VALUE=`eval echo \\$memory_$SAFENAME`
	if [ -n "${VALUE}" ]
	then
	    # TODO: get per-app increments
	    cgset -r "memory.$VARNAME=$VALUE" $CGPATH
	fi
    done
}

# ==========================================================================
#  Functions for tuning the user's memory limits in cgroups
# ==========================================================================
BLKIOVARS="weight weight_device"
function set_blkio() {
    # USERNAME=$1
    CGPATH=/$1

    # for each var get and set the value
    for VARNAME in $BLKIOVARS
    do
	# cgroups names can have periods(.)  shell varnames can't
	SAFENAME=`echo $VARNAME | tr . _`
	VALUE=`eval echo \\$blkio_$SAFENAME`
	if [ -n "${VALUE}" ]
	then
	    # TODO: get per-app increments
	    # TODO: weight_device should really use the user's home device
	    #       and set the rest (if any) to 0
	    # cgset -r "blkio.$VARNAME=$VALUE" $CGPATH
	    echo nothing >>/dev/null
	fi
    done
}

# List the openshift guest users
#
openshift_users() {
    getent passwd | grep "${GEAR_GECOS}" | cut -d: -f1
}

valid_user() {
    # check if the user name exists and is tagged as a openshift guest user
    getent passwd | grep ":${GEAR_GECOS}:" | cut -d: -f1 | grep -e "^$1\$" >/dev/null 2>&1
}

#
# Create a new openshift user cgroup
#
add_cgroup() {
    # USERNAME=$1
    cgcreate -t $1:$1 -g ${OPENSHIFT_CGROUP_SUBSYSTEMS}:${OPENSHIFT_CGROUP_ROOT}/$1
}

#
# Delete a openshift user cgroup
#
delete_cgroup() {
    # USERNAME=$1
    cgdelete ${OPENSHIFT_CGROUP_SUBSYSTEMS}:${OPENSHIFT_CGROUP_ROOT}/$1
}


#
# check which user cgroups exist
#
cgroup_user_subsystems() {
    # USERNAME=$1
    lscgroup | grep ":${OPENSHIFT_CGROUP_ROOT}/$1\$" | cut -d: -f1
}

#
# Check that a group binding rule exists for a user
#
cgroup_rule_exists() {
    #USERNAME=$1
    # remove comments, get first field, match exactly, quiet
    grep -v '^#' ${CGROUP_RULES_FILE} | cut -f1 | grep -q -x $1
}


#
# Bind the user to the cgroup: update /etc/cgrules.conf and kick cgred
#
add_cgroup_rule() {
    # USERNAME=$1
    cat <<EOF >>${CGROUP_RULES_FILE}
$1	$OPENSHIFT_CGROUP_SUBSYSTEMS	$OPENSHIFT_CGROUP_ROOT/$1
EOF
}

#
# Unbind the user from any cgroup
#
delete_cgroup_rule() {
    # USERNAME=$1
    sed -i -e "/^$1\s/d" ${CGROUP_RULES_FILE}
}

#
# Add the user's processes to the new group
#
collect_tasks() {
    # USERNAME=$1

    # add existing processes to the group
    for PID in $(ps -opid= -u $1) ; do
	echo $PID > /cgroup/all/${OPENSHIFT_CGROUP_ROOT}/$1/tasks
    done
}

startuser() {
    NEWUSER=$1

    echo -n "starting cgroups for $NEWUSER..."

    add_cgroup $NEWUSER
    if [ $? != 0 ]
    then
        RETVAL=$?
    fi

    set_cpu $NEWUSER
    set_memory $NEWUSER
    #set_blkio $NEWUSER
    set_net_cls $NEWUSER

    # CHECK: don't trust old rules
    if ( cgroup_rule_exists $NEWUSER )
    then
        delete_cgroup_rule $NEWUSER
    fi
    add_cgroup_rule $NEWUSER
    if [ $? != 0 ]
    then
        RETVAL=$?
    fi

    collect_tasks $NEWUSER

    if [ $RETVAL -eq 0 ]
    then
        echo -n " [OK] "
    else
        GROUP_RETVAL=$(($GROUP_RETVAL+1))
        echo -n " [FAILED] "
    fi
    echo
}

startall() {
    echo "Initializing Openshift guest control groups: "

    if !(service cgconfig status >/dev/null)
    then
        RETVAL=1
        GROUP_RETVAL=3
        echo "cgconfig service not running. attempting to start it"
        service cgconfig start
        return $GROUP_RETVAL
    fi

    if !(service cgconfig status >/dev/null)
    then
        RETVAL=1
        GROUP_RETVAL=3
        echo "cgconfig service not running."

        return $GROUP_RETVAL
    fi

    # don't start if not configured for openshift
    if [ ! -d /cgroup/all ]
    then
        echo "cgconfig not set for Openshift: /cgconfig/all does not exist"
        RETVAL=1
        GROUP_RETVAL=3
        return $GROUP_RETVAL
    fi

    # create the root of the openshift user control group
    add_cgroup # defaults to creating the root group
    RETVAL=$?

    # This won't scale forever, but works fine in the '100 or so' range
    for USERNAME in `openshift_users`
    do
        startuser $USERNAME
    done

    # kick the Cgroups rules daemon
    #service cgred reload
    pkill -USR2 cgrulesengd

    [ $GROUP_RETVAL -eq 0 ] && touch ${lockfile}
    [ $GROUP_RETVAL -eq 0 ] && (echo -n "[ OK ]") || (echo -n "[ FAILED ]")

    echo -n $"Openshift cgroups initialized"
    echo
    return $GROUP_RETVAL
    echo
    echo "WARNING !!! WARNING !!! WARNING !!!"
    echo "Cgroups may have just restarted.  It's important to confirm all the openshift apps are actively running."
    echo "It's suggested you run service openshift restart now"
    echo "WARNING !!! WARNING !!! WARNING !!!"
    echo
}

stopuser() {
    DELUSER=$1
    echo -n "stopping cgroups for $DELUSER..."

    # kill any processes owned by these users
    #pkill -u $DELUSER
    
    # remove the user's cgroup
    delete_cgroup $DELUSER
    if [ $? != 0 ]
    then
	RETVAL=$?
    fi
    
    # remove the user's cgroup binding rule
    delete_cgroup_rule $DELUSER
    if [ $? != 0 ]
    then
	RETVAL=$?
    fi

    if [ $RETVAL -eq 0 ]
    then
        echo -n "[ OK ]"
    else
        GROUP_RETVAL=$(($GROUP_RETVAL+1))
        echo -n "[ FAILED ]"
    fi
}

stopall() {
    echo "Removing Openshift guest control groups: "

    if !(service cgconfig status >/dev/null)
    then
       RETVAL=1
       GROUP_RETVAL=3
       echo "cgconfig service not running"

       return $GROUP_RETVAL
    fi

    # This won't scale forever, but works fine in the '100 or so' range
    for USERNAME in `openshift_users`
    do
	stopuser $USERNAME
    done

    # notify the cgroup rule daemon
    #service cgred reload
    pkill -USR2 cgrulesengd

    # remove the openshift root cgroup
    delete_cgroup

    if [ $RETVAL -eq 0 ]
    then
        echo -n "[ OK ]"
    else
        GROUP_RETVAL=$(($GROUP_RETVAL+1))
        echo -n "[ FAILED ]"
    fi

    [ $GROUP_RETVAL -eq 0 ] && touch ${lockfile}
    echo -n $"Openshift cgroups uninitialized"
    echo
    return $GROUP_RETVAL
}

restartall() {
    stopall
    startall
}

status() {
    echo "Checking Openshift Services: "

    # don't start if not configured for openshift
    if [ ! -d /cgroup/all ]
    then
        echo "Openshift cgroups not configured: /cgconfig/all does not exist"
        return 1
    fi

    lscgroup | grep -e  ":${OPENSHIFT_CGROUP_ROOT}\$" >/dev/null 2>&1
    if [ $? -ne 0 ]
    then
	echo "Openshift cgroups uninitialized"
	echo
	return 1
    else
	echo "Openshift cgroups initialized"
    fi
    
    if [ -z "$1" ]
    then
	USERLIST=`openshift_users`
    else
        USERLIST=$1
    fi

    # check that the /openshift cgroup exists

    # This won't scale forever, but works fine in the '100 or so' range
    #  would be easy to convert to a 'in `find...`'     jj
    for USERNAME in $USERLIST
    do
	# check that /openshift/<username> exists
	SUBSYSTEMS=`cgroup_user_subsystems`
	if ( cgroup_rule_exists $USERNAME )
        then
	    RETVAL=0
            BOUND="BOUND"
        else
	    RETVAL=1
            BOUND="UNBOUND"
        fi

	echo -n "${USERNAME}: $BOUND	" `echo $SUBSYSTEMS | tr ' ' ,`
	# check that cgrule exists

        if [ $RETVAL -eq 0 ]
        then
            echo -n "[ OK ]"
        else
            GROUP_RETVAL=$(($GROUP_RETVAL+1))
            echo -n "[ FAILED ]"
        fi
	echo
    done
    return $GROUP_RETVAL
}

case "$1" in
  startall)
    startall
    ;;

  stopall) 
    stopall
    ;;

  restartall)
    restartall
    ;;

  condrestartall)
    [ -f "$lockfile" ] && restartall
    ;;

  status)
    status $2
    ;;

  startuser)
    if (service cgconfig status >/dev/null)
    then
        startuser $2
        #service cgred reload
        pkill -USR2 cgrulesengd
    else
        RETVAL=1
        echo "cgconfig service not running"
    fi
    ;;

  stopuser)
    if (service cgconfig status >/dev/null)
    then
        stopuser $2
        #service cgred reload
        pkill -USR2 cgrulesengd
    else
        RETVAL=1
        echo "cgconfig service not running"
    fi
    ;;

  *)
    echo $"Usage: $0 {start|stop|status|restart|condrestart|startuser <username>|stopuser <username>}"
    exit 1
esac

exit $RETVAL
