#!/bin/bash
# @Function
## Find out the highest cpu consumed threads of java, and print the stack of these threads.
#
# @Usage
#   $ ./show-busy-java-threads
#
# @online-doc https://github.com/oldratlee/useful-scripts/blob/master/docs/java.md#beer-show-busy-java-threads
# @author Jerry Lee (oldratlee at gmail dot com)
# @author superhj1987 (superhj1987 at 126 dot com)

readonly PROG="`basename $0`"
readonly -a COMMAND_LINE=("$0" "$@")
# Get current user name via whoami command
#   See https://www.lifewire.com/current-linux-user-whoami-command-3867579
# Because if run command by `sudo -u`, env var $USER is not rewritten/correct, just inherited from outside!
readonly USER="`whoami`"

################################################################################
# util functions
################################################################################

# NOTE: $'foo' is the escape sequence syntax of bash
readonly ec=$'\033' # escape char
readonly eend=$'\033[0m' # escape end

colorEcho() {
    local color=$1
    shift

    # if stdout is console, turn on color output.
    [ -t 1 ] && echo "$ec[1;${color}m$@$eend" || echo "$@"
}

colorPrint() {
    local color=$1
    shift

    colorEcho "$color" "$@"
    # use "{} &> /dev/null" to discard bash error message:
    #    line 42: xxx.log: Permission denied
    { [ -n "$append_file" ] && echo "$@" >> "$append_file"; } &> /dev/null
}

redPrint() {
    colorPrint 31 "$@"
}

greenPrint() {
    colorPrint 32 "$@"
}

yellowPrint() {
    colorPrint 33 "$@"
}

bluePrint() {
    colorPrint 36 "$@"
}

normalPrint() {
    echo "$@"
    [ -n "$append_file" ] && echo "$@" >> "$append_file"
}

fatal() {
    redPrint "$@" 1>&2
    exit 1
}

usage() {
    local -r exit_code="$1"
    shift
    [ -n "$exit_code" -a "$exit_code" != 0 ] && local -r out=/dev/stderr || local -r out=/dev/stdout

    (( $# > 0 )) && { echo "$@"; echo; } > $out

    > $out cat <<EOF
Usage: ${PROG} [OPTION] [delay [ucount]]
Find out the highest cpu consumed threads of java, and print the stack of these threads.

Example: 
    ${PROG}       # show busy java threads info
    ${PROG} 1     # update every 1 second, (stop by eg: CTRL+C)
    ${PROG} 3 10  # update every 3 seconds, update 10 times

Options:
    -p, --pid       find out the highest cpu consumed threads from the specifed java process,
                    default from all java process.
    -c, --count     set the thread count to show, default is 5
    -a, --append-file <file>  specify the file to append output as log
    -s, --jstack-path <path>  specify the path of jstack command
    -F, --force               set jstack to force a thread dump(use jstack -F option)
    -S, --jstack-file-dir <path>  specifies the directory for storing jstack output files, and keep files.
                            default store jstack output files at tmp dir, and auto remove after run.
                            use this option to keep files so as to review jstack later.
    -m, --mix-native-frames   set jstack to print both java and native frames (mixed mode).
    -l, --lock-info           set jstack with long listing. Prints additional information about locks.
    -d, --top-delay  <deplay> specifies the delay between top samples, default is 0.5 (second).
                              get thread cpu percentage during this delay interval.
                              more info see top -d option. eg: -d 1 (1 second).
    -P, --use-ps              use ps command to find busy thread(cpu usage) instead of top command,
                              default use top command, because cpu usage of ps command is expressed as
                              the percentage of time spent running during the entire lifetime of a process,
                              this is not ideal.
    -h, --help      display this help and exit
    delay is the delay between updates in seconds. when this is not set, it will output once.
    ucount is the number of updates. when delay is set, ucount is not set, it will output in unstop mode.
EOF

    exit $exit_code
}

################################################################################
# Check os support
################################################################################

uname | grep '^Linux' -q || fatal "Error: $PROG only support Linux, not support `uname` yet!"

################################################################################
# parse options
################################################################################

# NOTE: ARGS can not be declared as readonly!!
# readonly declaration make exit code of assignment to be always 0, aka. the exit code of `getopt` in subshell is discarded.
#   tested on bash 4.2.46

ARGS=`getopt -n "$PROG" -a -o p:c:a:s:S:Pd:Fmlh -l pid:,count:,append-file:,jstack-path:,jstack-file-dir:,use-ps,top-delay:,force,mix-native-frames,lock-info,help -- "$@"`
[ $? -ne 0 ] && { echo; usage 1; }
eval set -- "${ARGS}"

while true; do
    case "$1" in
    -c|--count)
        count="$2"
        shift 2
        ;;
    -p|--pid)
        pid="$2"
        shift 2
        ;;
    -a|--append-file)
        append_file="$2"
        shift 2
        ;;
    -s|--jstack-path)
        jstack_path="$2"
        shift 2
        ;;
    -S|--jstack-file-dir)
        jstack_file_dir="$2"
        shift 2
        ;;
    -P|--use-ps)
        use_ps=true
        shift
        ;;
    -d|--top-delay)
        top_delay="$2"
        shift 2
        ;;
    -F|--force)
        force=-F
        shift
        ;;
    -m|--mix-native-frames)
        mix_native_frames=-m
        shift
        ;;
    -l|--lock-info)
        more_lock_info=-l
        shift
        ;;
    -h|--help)
        usage
        ;;
    --)
        shift
        break
        ;;
    esac
done
count=${count:-5}
use_ps=${use_ps:-false}
top_delay=${top_delay:-0.5}
update_delay=${1:-0}
update_limit=${2:-0}

# check the directory of append-file(-a) mode, create if not exsit.
if [ -n "$append_file" ]; then
    if [ -e "$append_file" ]; then
        [ ! -f "$append_file" ] && fatal "Error: $append_file(specified by option -a, for storing run output files) exists but is not a file!"
        [ ! -w "$append_file" ] && fatal "Error: file $append_file(specified by option -a, for storing run output files) exists but is not writable!"
    else
        append_file_dir="$(dirname "$append_file")"
        if [ -e "$append_file_dir" ]; then
            [ ! -d "$append_file_dir" ] && fatal "Error: directory $append_file_dir(specified by option -a, for storing run output files) exists but is not a directory!"
            [ ! -w "$append_file_dir" ] && fatal "Error: directory $append_file_dir(specified by option -a, for storing run output files) exists but is not writable!"
        else
            mkdir -p "$append_file_dir" && fatal "Error: fail to create directory $append_file_dir(specified by option -a, for storing run output files)!"
        fi
    fi
fi

# check jstack-file directory(-S) mode, create directory if not exsit.
if [ -n "$jstack_file_dir" ]; then
    if [ -e "$jstack_file_dir" ]; then
        [ ! -d "$jstack_file_dir" ] && fatal "Error: $jstack_file_dir(specified by option -S, for storing jstack output files) exists but is not a directory!"
        [ ! -w "$jstack_file_dir" ] && fatal "Error: directory $jstack_file_dir(specified by option -S, for storing jstack output files) exists but is not writable!"
    else
        mkdir -p "$jstack_file_dir" && fatal "Error: fail to create directory $jstack_file_dir(specified by option -S, for storing jstack output files)!"
    fi
fi

################################################################################
# check the existence of jstack command
################################################################################

if [ -n "$jstack_path" ]; then
    [ -f "$jstack_path" ] || fatal "Error: $jstack_path is NOT found!"
    [ -x "$jstack_path" ] || fatal "Error: $jstack_path is NOT executalbe!"
elif which jstack &> /dev/null; then
    jstack_path="`which jstack`"
else
    [ -z "$JAVA_HOME" ] && fatal "Error: jstack not found on PATH and No JAVA_HOME setting! Use -s option set jstack path manually."
    [ -f "$JAVA_HOME/bin/jstack" ] || fatal "Error: jstack not found on PATH and \$JAVA_HOME/bin/jstack($JAVA_HOME/bin/jstack) file does NOT exists! Use -s option set jstack path manually."
    [ -x "$JAVA_HOME/bin/jstack" ] || fatal "Error: jstack not found on PATH and \$JAVA_HOME/bin/jstack($JAVA_HOME/bin/jstack) is NOT executalbe! Use -s option set jstack path manually."
    jstack_path="$JAVA_HOME/bin/jstack"
fi

################################################################################
# biz logic
################################################################################

readonly run_timestamp="`date "+%Y-%m-%d_%H:%M:%S.%N"`"
readonly uuid="${PROG}_${run_timestamp}_${RANDOM}_$$"
readonly tmp_store_dir="/tmp/${uuid}"
mkdir -p "$tmp_store_dir"

if [ -n "$jstack_file_dir" ]; then
    readonly jstack_file_path_prefix="$jstack_file_dir/jstack_${run_timestamp}_"
else
    readonly jstack_file_path_prefix="$tmp_store_dir/jstack_${run_timestamp}_"
fi

cleanupWhenExit() {
    rm /tmp/${uuid}_* &> /dev/null
}
trap "cleanupWhenExit" EXIT

headInfo() {
    colorEcho "0;34;42" ================================================================================
    echo "$(date "+%Y-%m-%d %H:%M:%S.%N") [$(( update_count + 1 ))/$update_limit]: ${COMMAND_LINE[@]}"
    colorEcho "0;34;42" ================================================================================
    echo
}

# output field: pid, thread id(lwp), pcpu, user
#   order by pcpu(percentage of cpu usage)，this value reprents the cpu running time for this process / total process time,no realtime cpu usage.
findBusyJavaThreadsByPs() {
    if [ -n "${pid}" ]; then
        local -r ps_options="-p $pid"
    else
        local -r ps_options="-C java"
    fi
    # 1. sort by %cpu by ps option `--sort -pcpu`
    # 2. use wide output(unlimited width) by ps option `-ww`
    #    avoid trunk user column to username_fo+ or $uid alike
    ps $ps_options -wwLo pid,lwp,pcpu,user --sort -pcpu --no-headers | head -n "${count}"
}

# top with output field: thread id, %cpu, thsi value is realtime cpu usage
__top_threadId_cpu() {
    # 1. sort by %cpu by top option `-o %CPU`
    #    unfortunately, top version 3.2 does not support -o option(supports from top version 3.3+),
    #    use
    #       HOME="$tmp_store_dir" top -H -b -n 1
    #    combined
    #       sort
    #    instead of
    #       HOME="$tmp_store_dir" top -H -b -n 1 -o '%CPU'
    # 2. change HOME env var when run top,
    #    so as to prevent top command output format being change by .toprc user config file unexpectedly
    # 3. use option `-d 0.5`(interval 0.5 second) and `-n 2`(show 2 times), and use second time update data
    #    to get cpu percentage of thread in 0.5 second interval
    HOME="$tmp_store_dir" top -H -b -d $top_delay -n 2 |
        awk '{
                if (idx == 4 && $NF == "java")    # $NF is command
                    # only print 4th text block(idx == 3), aka. process info of second top update
                    print $1 " " $9    # $1 is thread id, $9 is %cpu
                if ($0 == "")
                    idx++
        }' | sort -k2,2nr
}

# output format is same as function findBusyJavaThreadsByPs
findBusyJavaThreadsByTop() {
    if [ -n "${pid}" ]; then
        local -r ps_options="-p $pid"
    else
        local -r ps_options="-C java"
    fi
    # ps output field: pid, thread id(lwp), user
    local -r ps_out="$(ps $ps_options -wwLo pid,lwp,user --no-headers)"

    local idx=0
    local -a line
    while IFS=" " read -a line ; do
        (( idx < count )) || break

        local threadId="${line[0]}"
        local pcpu="${line[1]}"

        # output field: pid, threadId, pcpu, user
        local output_fields="$( echo "$ps_out" |
            awk -v "threadId=$threadId" -v "pcpu=$pcpu" '$2==threadId {
                print $1 " " threadId " " pcpu " " $3; exit
            }' )"
        if [ -n "$output_fields" ]; then
            (( idx++ ))
            echo "$output_fields"
        fi
    done < <( __top_threadId_cpu )
}

printStackOfThreads() {
    local update_round_num="$1"

    local -a line
    local idx=0
    while IFS=" " read -a line ; do
        local pid="${line[0]}"
        local threadId="${line[1]}"
        local threadId0x="0x`printf %x ${threadId}`"
        local pcpu="${line[2]}"
        local user="${line[3]}"

        (( idx++ ))
        local jstackFile="$jstack_file_path_prefix${update_round_num}_${pid}"
        [ -f "${jstackFile}" ] || {
            if [ "${user}" == "${USER}" ]; then
                # run without sudo, when java process user is current user
                "$jstack_path" ${force} $mix_native_frames $more_lock_info ${pid} > ${jstackFile}
            elif [ $UID == 0 ]; then
                # if java process user is not current user, must run jstack with sudo
                sudo -u "${user}" "$jstack_path" ${force} $mix_native_frames $more_lock_info ${pid} > ${jstackFile}
            else
                # current user is not root user, so can not run with sudo; print error message and rerun suggestion
                redPrint "[$idx] Fail to jstack busy(${pcpu}%) thread(${threadId}/${threadId0x}) stack of java process(${pid}) under user(${user})."
                redPrint "User of java process($user) is not current user($USER), need sudo to rerun:"
                yellowPrint "    sudo ${COMMAND_LINE[@]}"
                normalPrint
                continue
            fi || {
                redPrint "[$idx] Fail to jstack busy(${pcpu}%) thread(${threadId}/${threadId0x}) stack of java process(${pid}) under user(${user})."
                normalPrint
                rm "${jstackFile}" &> /dev/null
                continue
            }
        }

        bluePrint "[$idx] Busy(${pcpu}%) thread(${threadId}/${threadId0x}) stack of java process(${pid}) under user(${user}):"

        if [ -n "$mix_native_frames" ]; then
            local sed_script="/--------------- $threadId ---------------/,/^---------------/ {
                /--------------- $threadId ---------------/b # skip first seperator line
                /^---------------/s/.*// # replace sencond seperator line to empty line
                p
            }"
        elif [ -n "$force" ]; then
            local sed_script="/^Thread ${threadId}:/,/^$/p"
        else
            local sed_script="/nid=${threadId0x} /,/^$/p"
        fi
        sed "$sed_script" -n ${jstackFile} | tee ${append_file:+-a "$append_file"}
    done
}

update_count=0
while true
do
    [ -n "$append_file" ] && head_info >> "$append_file"
    headInfo
    
   if $use_ps; then
            findBusyJavaThreadsByPs
        else
            findBusyJavaThreadsByTop
        fi | printStackOfThreads $(( update_count + 1 ))
    
    if (( "$update_delay" <= 0 )); then #不需要定时刷新数据
        break
    fi
    
    ((update_count++))
    if (( "$update_limit" > 0 && "$update_count" == "$update_limit" )); then #刷新次数到达上限
        break #
    fi

    sleep "$update_delay" #休眠
done

# if update_limit <= 0, infinite loop till user interupted (eg: CTRL+C)
#for ((i = 0; update_limit <= 0 || i < update_limit; ++i)); do
#    [ "$i" -gt 0 ] && sleep "$update_delay"

#    [ -n "$append_file" ] && headInfo >> "$append_file"
#    [ "$update_count" -ne 1 ] && headInfo

#    ps -Leo pid,lwp,user,comm,pcpu --no-headers | {
#       [ -z "${pid}" ] &&
#        awk '$4=="java"{print $0}' ||
#        awk -v "pid=${pid}" '$1==pid,$4=="java"{print $0}'
#    } | sort -k5 -r -n | head -n "${count}" | printStackOfThreads
#done