#!/bin/sh
#
# ktrace
#	Reset and enable kernel tracing
#	Run command in background, if specified
#	Sleep for duration, if specified
#	Wait for command completion, timeout, or keyboard interrupt
#	(Note: keyboard interrupt kills command)
#	Disable kernel tracing
#	Extract kernel trace buffer
#	Reset kernel tracing, if specified
#

TMPDIR=${TMPDIR:-/tmp}

output_file=$TMPDIR/ktrace.tnf
buffer_size=512k
duration=
events=
command=
do_reset=y
child=
alarm=

# usage message
usage() {
	echo ""
	echo $1
	echo "
Usage: ktrace	[-o <output_file>]	# default /tmp/ktrace.tnf
		[-s <buffer_size>]	# default 512k
		-e <events>		# kernel probes (keys) to enable
		[-r]			# don't reset kernel tracing
					# default is to reset after command
		<cmd> | -t <seconds>

Eg,
	# ktrace -e `thread vm io' -t 10
	# ktrace -e `thread' -s 256k myapp ...
"
	exit 1
}
# failure message
fail() {
	while [ $# -gt 0 ]
	do
		echo $1
		shift
	done
	echo "ktrace failed"
	exit 2
}

# Reset kernel tracing
ktrace_reset() {
	if [ $1 = "y" ]; then
		echo "Resetting kernel tracing"
		prex -k >/dev/null 2>&1 <<-EOF
			ktrace off
			untrace /.*/
			disable /.*/
			buffer dealloc
			EOF
		test $? -ne 0 && fail "Could not reset kernel tracing" \
			"`su root' and retry"
	fi
}

# Enable kernel tracing
ktrace_on() {
	echo "Enabling kernel tracing"
	prex -k >/dev/null 2>&1 <<-EOF
		buffer alloc $buffer_size
		trace $events
		enable $events
		ktrace on
		EOF
	test $? -ne 0 && fail "Could not enable kernel tracing" \
		"Check syntax of `-e' argument" \
		"Check buffer size is not too high"
}

# Disable kernel tracing
ktrace_off() {
	prex -k >/dev/null 2>&1 <<-EOF
		ktrace off
		EOF
	test $? -ne 0 && fail "Could not disable kernel tracing"
	echo "Kernel tracing disabled"
}

# Extract kernel trace buffer
ktrace_xtract() {
	echo "Extracting kernel trace buffer"
	tnfxtract $output_file || fail "Could not extract kernel trace
buffer"
	ls -l $output_file
}

# Run command, sleep for duration, or wait for command to complete
run_command() {
	trap `interrupt' 0 1 2 3 15
	if [ "$command" ]; then
		$command &
		child=$!
		echo "`$command' pid is $child"
	fi

	if [ "$duration" ]; then
		sleep $duration &
		alarm=$!
		wait $alarm
		# XXX test -z "$child" || kill -15 $child
	else
		wait $child
	fi
	trap 0 1 2 3 15
}

# Keyboard interrupt
interrupt() {
	test -z "$alarm" || kill -15 $alarm
	test -z "$child" || kill -15 $child
}

# Parse options
while getopts o:s:t:e:r opt
do
	case $opt in
	 o)	output_file="$OPTARG";;
	 s)	buffer_size="$OPTARG";;
	 t)	duration="$OPTARG";;
	 e)	events="$OPTARG";;
	 r)	do_reset="n";;
	 \?)	usage;;
	esac
done

shift `expr $OPTIND - 1`

# Get command to run
test $# -gt 0 && command="$*"

# Check and normalize options
test -z "$events" && usage "No kernel events specified"
test -z "$command" && test -z "$duration" && \
	usage "No command or time duration specified"

# Perform experiment
ktrace_reset y				# Reset kernel tracing
ktrace_on				# Enable kernel tracing
run_command				# Run command, wait, or sleep
ktrace_off				# Disable kernel tracing
ktrace_xtract				# Extract buffer
ktrace_reset $do_reset			# Reset kernel tracing again
exit 0