#!/bin/sh

#
# lxc: linux Container library

# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.

# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.

# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA

usage()
{
    echo "usage: $(basename $0) [--lxc | --host | --name NAME] [--] [PS_OPTIONS...]" >&2
}

help() {
    usage
    echo >&2
    echo "List current processes with container names." >&2
    echo >&2
    echo "  --lxc         show processes in all containers" >&2
    echo "  --host        show processes not related to any container, i.e. to the host" >&2
    echo "  --name NAME   show processes in the specified container" >&2
    echo "                 (multiple containers can be separated by commas)" >&2
    echo "  PS_OPTIONS    ps command options (see \`ps --help')" >&2
}

get_parent_cgroup()
{
    local hierarchies hierarchy fields subsystems init_cgroup mountpoint

    parent_cgroup=""

    # Obtain a list of hierarchies that contain one or more subsystems
    hierarchies=$(tail -n +2 /proc/cgroups | cut -f 2)

    # Iterate through the list until a suitable hierarchy is found
    for hierarchy in $hierarchies; do
        # Obtain information about the init process in the hierarchy
        fields=$(grep -E "^$hierarchy:" /proc/1/cgroup | head -n 1)
        if [ -z "$fields" ]; then continue; fi
        fields=${fields#*:}

        # Get a comma-separated list of the hierarchy's subsystems
        subsystems=${fields%:*}

        # Get the cgroup of the init process in the hierarchy
        init_cgroup=${fields#*:}

        # Get the filesystem mountpoint of the hierarchy
        mountpoint=$(awk -v subsysregex="(^|,)$subsystems(,|\$)" \
            '$3 == "cgroup" && $4 ~ subsysregex {print $2}' /proc/self/mounts)
        if [ -z "$mountpoint" ]; then continue; fi

        # Return the absolute path to the containers' parent cgroup
        # (do not append '/lxc' if the hierarchy contains the 'ns' subsystem)
        case ",$subsystems," in
            *,ns,*) parent_cgroup="${mountpoint}${init_cgroup%/}";;
            *) parent_cgroup="${mountpoint}${init_cgroup%/}/lxc";;
        esac
        break
    done
}

containers=""
list_container_processes=0
while true; do
    case $1 in
        -h|--help)
            help; exit 1;;
        -n|--name)
            containers=$2; list_container_processes=1; shift 2;;
        --lxc)
            list_container_processes=1; shift;;
        --host)
            list_container_processes=-1; shift;;
        --)
            shift; break;;
        *)
            break;;
        esac
done

if [ "$list_container_processes" -eq "1" ]; then
    set -- -e $@
fi

get_parent_cgroup
if [ ! -d "$parent_cgroup" ]; then
    echo "$(basename $0): no cgroup mount point found" >&2
    exit 1
fi

if [ -z "$containers" ]; then
    containers="$(find $parent_cgroup -mindepth 1 -maxdepth 1 -type d 2>/dev/null | sed 's:.*/::')"
fi

container_field_width=9
tasks_files=
for container in ${containers}; do
    if [ "${#container}" -gt "$container_field_width" ]; then
        container_field_width=${#container}
    fi

    if [ -f "$parent_cgroup/$container/tasks" ]; then
        tasks_files="$tasks_files $parent_cgroup/$container/tasks"
    fi
done

# first file is stdin, the rest are the container tasks
ps "$@" | awk -v container_field_width="$container_field_width" \
    -v list_container_processes="$list_container_processes" '
# first line is PS header
NR == 1 {
    header = $0
    # find pid field index
    for (i = 1; i<=NF; i++)
        if ($i == "PID") {
            pididx = i
            break
        }
    if (pididx == "") {
        print("No PID field found") > "/dev/stderr"
        header = ""  # to signal error condition to the END rule
        exit 1
    }
    next
}

# store lines from ps with pid as index
NR == FNR {
    ps_line[NR] = $0
    pid_of_line[NR] = $pididx
    next
}

# find container name from filename on first line
FNR == 1 {
    container = FILENAME
    sub(/\/tasks/, "", container)
    sub(/.*\//, "", container)
}

# container tasks
{
    container_of_pid[$0] = container
}

END {
    if (!header) exit 1 # quit due to internal error
    printf("%-" container_field_width "s %s\n", "CONTAINER", header)
    for (i in ps_line) {
        container = container_of_pid[pid_of_line[i]]
        if (list_container_processes == 0 || (container != "") == (list_container_processes > 0) )
            printf("%-" container_field_width "s %s\n", container, ps_line[i])
    }
}

' - $tasks_files
