#!/bin/bash

TRACE_DIR=/sys/kernel/debug/tracing/
BUFFER_SIZE=$TRACE_DIR/buffer_size_kb
CURRENT_TRACER=$TRACE_DIR/current_tracer
AVAILABLE_TRACERS=$TRACE_DIR/available_tracers
SET_EVENTS=$TRACE_DIR/set_event
TRACE_ON=$TRACE_DIR/tracing_on
TRACE_MARKER=$TRACE_DIR/trace_marker
SAVED_CMDLINES=$TRACE_DIR/saved_cmdlines
PER_CPU=$TRACE_DIR/per_cpu
TRACE=$TRACE_DIR/trace
TRACE_PIPE=$TRACE_DIR/trace_pipe
RENICEWRAP=""
DEFAULT_DUMP_DIR=/var/log/kernel_trace

DEFAULT_EVENTS=""
DEFAULT_BUFFER="1408"

TRACE_EVENTS="sched:*"
TRACE_BUFFER="5000"
TRACER="nop"

parse_options() {
    while getopts "e:b:t:n:" OPTION
    do
        case $OPTION in
        e) TRACE_EVENTS="$OPTARG" ;;
        b) TRACE_BUFFER="$OPTARG" ;;
        t) TRACER="$OPTARG" ;;
        *) echo "ERROR: Unknown option: $OPTION" >&2
        esac
    done
}

OPERATION="$1"
shift
case "$OPERATION" in
    enable|start)
      parse_options "$@"
      echo "${TRACE_BUFFER}" > $BUFFER_SIZE
      echo "${TRACER}" > $CURRENT_TRACER
      echo "${TRACE_EVENTS}" > $SET_EVENTS
      echo 1 > $TRACE_ON
      ;;
    disable|stop)
      echo 0 > $TRACE_ON
      echo $DEFAULT_BUFFER > $BUFFER_SIZE
      echo $DEFAULT_EVENTS > $SET_EVENTS
      ;;
    status)
      if [[ "$(<$TRACE_ON)" -eq 1 ]]; then
          echo "Tracing: on"
      else
          echo "Tracing: off"
      fi
      echo
      echo "Buffer size: $(<$BUFFER_SIZE) kB"
      echo "Current tracer: $(<$CURRENT_TRACER)"
      echo
      echo "Traced events: "
      if [[ "$(wc -l $SET_EVENTS | cut -d' ' -f 1)" -gt 10 ]]; then
          head -5 $SET_EVENTS
          echo ...
          tail -5 $SET_EVENTS
      else
          cat $SET_EVENTS
      fi
      ;;
    pause)
      echo 0 > $TRACE_ON ;;
    resume)
      echo 1 > $TRACE_ON ;;
    mark)
      echo "$*" > $TRACE_MARKER ;;
    follow)
      cat $TRACE_PIPE;;
    dump)
      echo "KERNEL TRACE DUMPED" > $TRACE_MARKER
      DUMP_DIR=${1:-$DEFAULT_DUMP_DIR}
      TRACE_DUMP=trace_dump_$(date +'%Y-%m-%d_%H-%M-%S')
      TRACE_ARCHIVE=$TRACE_DUMP.tar.gz
      mkdir -p $DUMP_DIR/$TRACE_DUMP
      pushd $DUMP_DIR/$TRACE_DUMP > /dev/null || { echo ERROR: Could not change directory. ; exit 0; }
      {
          echo -n "uptime:"
          cat /proc/uptime
          echo -n "date:"
          date +"%s.%N"
      } > uptime.txt
      cat $TRACE > trace.txt
      cat $SAVED_CMDLINES > saved_cmdlines.txt
      date +"%Y-%m-%d %H:%M:%S.%N" > ps.txt
      ps wwax -o stat,psr,nice,user,ppid,pid,%cpu,%mem,vsz,rss,start,time,cmd --sort=-%cpu >> ps.txt
      echo >> ps.txt
      ps wwax --forest -o user,ppid,pid,psr,%cpu,%mem,vsz,rss,stat,nice,start,time,cmd >> ps.txt
      dmesg > dmesg.txt
      echo Dump temporarily written to: $DUMP_DIR/$TRACE_DUMP
      popd > /dev/null
      if tar -C $DUMP_DIR -cf - $TRACE_DUMP | gzip - > $DUMP_DIR/$TRACE_ARCHIVE; then
          if [[ -n "$DUMP_DIR" ]] && [[ -n "$TRACE_DUMP" ]]; then
              rm -rf $DUMP_DIR/$TRACE_DUMP
          fi
          echo Dump created: $DUMP_DIR/$TRACE_ARCHIVE
      fi
      ;;
    convert)
      TRACE="${1:-trace.txt}"
      BASE_UPTIME="$(grep uptime: ${2:-uptime.txt} | cut -d: -f 2 | cut -d' ' -f 1)"
      BASE_DATE="$(grep date: ${2:-uptime.txt} | cut -d: -f 2)"
      echo "Base uptime: $BASE_UPTIME"
      echo "Base date: $BASE_DATE"
      ;;
    --help|help)
      PROG_NAME=$(basename $0)
      echo "Kernel trace - simple wrapper"
      echo ""
      echo "Usage:"
      echo "    $PROG_NAME start [-b BUFFER_SIZE] [-t TRACER] [-e TRACE_EVENTS]"
      echo "         -b - Buffer size in kilobytes."
      echo "         -t - Tracer to use - current: $(<$CURRENT_TRACER) available: $(<$AVAILABLE_TRACERS)"
      echo "         -e - Events to trace - default: $TRACE_EVENTS"
      echo ""
      echo "    $PROG_NAME stop"
      echo ""
      echo "    $PROG_NAME pause"
      echo ""
      echo "    $PROG_NAME resume"
      echo ""
      echo "    $PROG_NAME mark MARKER"
      echo "         MARKER - Some text marker to input in the trace at the point when it is executed."
      echo ""
      echo "    $PROG_NAME dump DIRECTORY"
      echo ""
      echo "    $PROG_NAME help"
      echo ""
      echo "Example usage:"
      echo " $PROG_NAME start"
      echo " <wait for/reproduce traced issue>"
      echo " $PROG_NAME pause"
      echo " $PROG_NAME dump"
      echo " $PROG_NAME resume"
      echo "  or"
      echo " $PROG_NAME stop"
      ;;
    *) echo "ERROR: Unknown operation: $OPERATION!" >&2
      exit 1
      ;;
esac
