#!/bin/bash
#
# Script to automate suspend / resume
#
# Copyright (C) 2008-2009 Canonical Ltd.
#
# Authors:
#  Michael Frey <michael.frey@canonical.com>
#  Andy Whitcroft <apw@canonical.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2,
# as published by the Free Software Foundation.
#
# This program 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

#
# Script to automate suspend / resume
#
# We set a RTC alarm that wakes the system back up and then sleep
# for  seconds before we go back to sleep.
#
# Changelog:
#
# V8:
#  - add a new suspend battery drain test
#  - track batteries disabling tests which require them automatically
#  - disable dbus tests when we have no primary user
#  - include the new power drain test in --full
#  - handle AC transitions better
#  - use minutes in messages where appropriate
#  - report AC transition failures
#  - only mention AC when we have batteries
#  - report results at the bottom for easy posting
#
# V7:
#  - add a --dry-run mode to simplify developement
#  - add a automation mode for checkbox integration
#  - add a new pm-suspend test
#  - record and restore timer_delay around the variable time test.
#
# V6:
#  - move an --enable/--disable interface for tests
#  - add --set to allow setting of approved parameters
#  - fix up prompting for interactive and non-interactive tests
#  - supply a sensible default for testing on servers (apw, kirkland)
#
# V5:
#  - send dbus messages as the original user
#  - stop clearing the dmesg as we go
#  - stop using trace generally as this affects the wakeups
#  - do a single dbus test then move to pm-suspend to avoid screensaver
#  - timeout waiting for a suspend to complete catching failure to go down
#
# V4:
#  - update the help output
#  - add --comprehensive to do AC related tests
#  - add --extensive to do a range of time related tests
#  - add --full to enable all harder tests
#  - add fallback to pm-suspend for Kbuntu
#  - collect dmesg output
#  - remove hwclock update
#
# V3:
#  - fix typo in fallback acpi interface
#  - when recording the RTC clock do not go direct
#  - pmi is now deprecated suspend using dbus
#
# V2:
#  - support newer rtc sysfs wakealarm interface
#  - move to using pmi action suspend
#  - allow the user to specify the number of iterations
#  - ensure we are running as root
#  - report the iterations to the user
#  - clean up the output and put it in a standard logfile
#  - add a descriptive warning and allow user cancel
#  - add tracing enable/disable
#  - fix logfile location
#  - add a failure cleanup mode
#  - make time sleep time and delay time configurable
#  - ensure the log directory exists
#  - clock will be fixed automatically on network connect
#  - default sleep before wakeup to 20s
#  - do not use dates after we have corrupted the clock
#  - sort out the copyright information
#  - we do not have any failure cleanup currently
#
# V1:
#  - add the suspend test scripts
#
P="test-suspend"

LOGDIR='/var/lib/pm-utils'
LOGFILE="$LOGDIR/stress.log"

setup_wakeup_timer ()
{
	timeout="$1"

	#
	# Request wakeup from the RTC or ACPI alarm timers.  Set the timeout
	# at 'now' + $timeout seconds.
	#
	ctl='/sys/class/rtc/rtc0/wakealarm'
	if [ -f "$ctl" ]; then
		# Cancel any outstanding timers.
		echo "0" >"$ctl"
		# rtcN/wakealarm can use relative time in seconds
		echo "+$timeout" >"$ctl"
		return 0
	fi
	ctl='/proc/acpi/alarm'
	if [ -f "$ctl" ]; then
		echo `date '+%F %H:%M:%S' -d '+ '$timeout' seconds'` >"$ctl"
		return 0
	fi

	echo "no method to awaken machine automatically" 1>&2
	exit 1
}

suspend_system ()
{
	if [ "$dry" -eq 1 ]; then
		echo "DRY-RUN: suspend machine for $timer_sleep"
		sleep 1
		return
	fi

	setup_wakeup_timer "$timer_sleep"

	dmesg >"$LOGFILE.dmesg.A"

	# Send a dbus message to initiate Suspend.
	if [ "$suspend_dbus" -eq 1 ]; then
		sudo -u $SUDO_USER dbus-send --session --type=method_call \
			--dest=org.freedesktop.PowerManagement \
			/org/freedesktop/PowerManagement \
			org.freedesktop.PowerManagement.Suspend \
			>> "$LOGFILE" || {
				ECHO "FAILED: dbus suspend failed"
				return
			}
	else
		pm-suspend >> "$LOGFILE"
	fi

	# Wait on the machine coming back up -- pulling the dmesg over.
	echo "v---" >>"$LOGFILE"
	retry=30
	while [ "$retry" -gt 0 ]; do
		let "retry=$retry-1"

		# Accumulate the dmesg delta.
		dmesg >"$LOGFILE.dmesg.B"
		diff "$LOGFILE.dmesg.A" "$LOGFILE.dmesg.B" | \
			grep '^>' >"$LOGFILE.dmesg"
		mv "$LOGFILE.dmesg.B" "$LOGFILE.dmesg.A"

		echo "Waiting for suspend to complete $retry to go ..." \
							>> "$LOGFILE"
		cat "$LOGFILE.dmesg" >> "$LOGFILE"

		if [ "`grep -c 'Back to C!' $LOGFILE.dmesg`" -ne 0 ]; then
			break;
		fi
		sleep 1
	done
	echo "^---" >>"$LOGFILE"
	rm -f "$LOGFILE.dmesg"*
	if [ "$retry" -eq 0 ]; then
		ECHO "SUSPEND FAILED, did not go to sleep"
	fi
}

delay_system ()
{
	if [ "$dry" -eq 1 ]; then
		echo "DRY-RUN: stay awake for $timer_delay"
		sleep 1
		return
	fi

	#
	# wait for $timer_delay seconds after system resume from S3
	#
	ECHO "wait for $timer_delay seconds"
	sleep $timer_delay
}

ECHO ()
{
	echo "$@" | tee -a "$LOGFILE"
}

run_suspend ()
{
	CNT=1
	TOTAL=$1
	ECHO "Suspend Test starting on $(date '+%F %H:%M:%S') ($TOTAL cycles)"
	while [ "$CNT" -le "$TOTAL" ]
	do
		ECHO "Suspend iteration $CNT of $TOTAL"

		suspend_system "$START"
		delay_system

		(( CNT++ ))
	done
	ECHO "Suspend Test completed"
}

enable_trace()
{
    echo 1 > '/sys/power/pm_trace'
}

disable_trace()
{
    echo 0 > '/sys/power/pm_trace'
}

# Battery
battery_count()
{
	cat /proc/acpi/battery/*/state 2>/dev/null | \
	awk '
		BEGIN			{ total = 0 }
		/present:.*yes/		{ total += 1 }
		END			{ print total }
	'
}
battery_capacity()
{
	cat /proc/acpi/battery/*/state 2>/dev/null | \
	awk '
		BEGIN			{ total = 0 }
		/remaining capacity:/	{ total += $3 }
		END			{ print total }
	'
}

#
# MAIN
#
usage() {
	cat - 1>&2 <<EOM
Usage: $P [<options>]
Options:
 --sleep <seconds>         - how long the machine wait before waking
 --delay <seconds>         - default delay between iterations

 --enable <test>           - enable a specific test
 --disable <test>          - disable a specific test
 --set <test>.<var>=<val>  - set a test specific variable
     dbus                    - perform a suspend via dbus
     ac                      - perform tests involving removing ac power
     timed                   - perform a variable timing test
     repeat                  - perform a longer repeat test
       .iterations             - the number of iterations in the repeat
     power                   - perform a battery consumption test
       .sleep                  - how long to sleep
 
 --full                    - run a basic set of tests
 --server                  - run those test appropriate for a server
EOM
}

# We need TEMP as the `eval set --' would nuke the return value of getopt.
TEMP=`getopt -o '' -l dry-run,auto,,sleep:,delay:,enable:,disable:,set:,full,desktop,server -n "$P" -- "$@"`
if [ $? != 0 ] ; then
	usage
        exit 1
fi

# Note the quotes around `$TEMP': they are essential!
eval set -- "$TEMP"

# Options helpers.
chk_test ()
{
	if ! declare -p "test_$1" 2>/dev/null 1>&2; then
		echo "$P: $1: test unknown" 1>&2
		exit 1
	fi
}
handle_set ()
{
	stmt=`echo "$1" | sed -e 's/\./_/g'`

	test="${stmt%%_*}"
	var="${stmt%%=*}"

	chk_test "$test"
	if ! declare -p "args_$var" 2>/dev/null 1>&2; then
		echo "$P: $var: test variable unknown" 1>&2
		exit 1
	fi
	
	RET="args_$stmt"
}
chk_number() {
	eval "val=\"\$$1\""
	let num="0+$val"
	if [ "$val" != "$num" ]; then
		name=`echo "$1" | sed -e 's/args_//' -e 's/_/./'`
		echo "$P: $name: $val: non-numeric value" 1>&2
		exit 1
	fi
}

# Options handling.
dry=0
auto=0
timer_sleep=20
timer_delay=10

test_dbus=0
test_pmsuspend=0
test_ac=0
test_timed=0
test_repeat=0
args_repeat_iterations=10
test_power=0
args_power_sleep=1200

while :
do
        case "$1" in
	--dry-run)		dry=1; shift 1 ;;
	--auto)			auto=1; shift 1 ;;
	--sleep)		timer_sleep="$2"; shift 2 ;;
	--delay)		timer_delay="$2"; shift 2 ;;
	--disable)		chk_test "$2"; declare "test_$1=0"; shift 2 ;;
	--enable)		chk_test "$2"; declare "test_$2=1"; shift 2 ;;
	--set)			handle_set "$2"; declare "$RET"; shift 2 ;;
	--desktop|--full)	test_dbus=1; test_ac=1; test_timed=1;
				test_power=1; shift 1 ;;
	--server)		test_timed=1; shift 1 ;;
        --)                     shift; break ;;
	*)			echo "$1: ERROR"; exit 1 ;;
        esac
done

chk_number "args_repeat_iterations"
chk_number "args_power_sleep"

tests=`set | grep ^test_ | grep -c =1`

if [ "$#" -gt 1 ]; then
	usage
	exit 1
fi
if [ "$tests" -eq 0 ]; then
	usage
	echo "$P: no tests selected" 1>&2
	exit 1
fi

battery_count=`battery_count`

report_battery=''

suspend_dbus=0

# Check we are running as root as we are going to fiddle with the clock
# and use the rtc wakeups.
id=`id -u`
if [ "$id" -ne 0 ]; then
	echo "ERROR: must be run as root to perform this test, use sudo:" 1>&2
	echo "       sudo $0 $@" 1>&2
	exit 1
fi

ac_needed=-1
ac_is=-1
ac_becomes=-1
ac_required()
{
	ac_check

	ac_needed="$1"
	ac_becomes="$1"
}
ac_transitions()
{
	ac_check

	ac_needed="$1"
	ac_becomes="$2"
}
ac_online()
{
	cat /proc/acpi/ac_adapter/*/state 2>/dev/null | \
	awk '
		BEGIN			{ online = 0; offline = 0 }
		/on-line/		{ online = 1 }
		/off-line/		{ offline = 1 }
		END			{
						if (online) {
							print "1"
						} else if (offline) {
							print "0"
						} else {
							print "-1"
						}
					}
	'
}
ac_check()
{
	typeset ac_current=`ac_online`

	if [ "$ac_becomes" -ne -1 -a "$ac_current" -ne -1 -a \
			"$ac_current" -ne "$ac_becomes" ]; then
		ECHO "*** WARNING: AC power not in expected state" \
			"($ac_becomes) after test"
	fi
	ac_is="$ac_becomes"
}

phase=0
phase_first=1
phase_interactive=1
phase()
{
	typeset sleep

	let phase="$phase+1"

	echo ""
	echo "*** TEST $phase -- $1"
	shift 1
	for line in "$@"
	do
		echo "*** $line"
	done
	if [ "$battery_count" -ne 0 -a "$ac_needed" -ne "$ac_is" ]; then
		case "$ac_needed" in
		0) echo "*** please ensure your AC cord is detached" ;;
		1) echo "*** please ensure your AC cord is attached" ;;
		esac
		ac_is="$ac_needed"
	fi
	
	if [ "$timer_sleep" -gt 60 ]; then
		let sleep="$timer_sleep / 60"
		sleep="$sleep minutes"
	else
		sleep="$timer_sleep seconds"
	fi
	echo "*** machine will suspend for $sleep"

	if [ "$auto" -eq 1 ]; then
		:

	elif [ "$phase_interactive" -eq 1 ]; then
		echo "*** press return when ready"
		read x

	elif [ "$phase_first" -eq 1 ]; then
		echo "*** NOTE: there will be no further user interaction from this point"
		echo "*** press return when ready"
		phase_first=0
		read x
	fi
	echo ""
}

[ "$auto" -eq 0 ] && cat - <<EOM
This script will attempt to suspend and resume your computer a number of times.
Should the machine fail to resume, first attempt to manually resume it.  If
that fails power your system off and on which will generate an apport bug
report automatically.

Press CTRL-C now to abort testing ...
EOM

# Ensure the log directory exists.
mkdir -p "$LOGDIR"

phase_interactive=1
if [ "$test_dbus" -eq 1 -a \
		\( "$SUDO_USER" = "" -o "$SUDO_USER" = "root" \) ]; then
	ECHO "*** no primary user (via sudo) dbus tests skipped ..."
elif [ "$test_dbus" -eq 1 ]; then
	ac_required 1
	suspend_dbus=1
	phase "suspend triggered via dbus message"
	suspend_system
	suspend_dbus=0
fi
if [ "$test_pmsuspend" -eq 1 ]; then
	ac_required 1
	phase "suspend triggered via pm-suspend"
	suspend_system
fi
if [ "$test_ac" -eq 1 -a "$battery_count" -eq 0 ]; then
	ECHO "*** no BATTERY detected ac tests skipped ..."
elif [ "$test_ac" -eq 1 ]; then
	ac_required 0
	phase "suspend with AC disconnected"
	suspend_system

	ac_required 1
	phase "suspend with AC connected"
	suspend_system
	
	ac_transitions 1 0
	phase "loss of AC while suspended" \
	      "please remove the AC cord while the machine is suspended"
	suspend_system

	ac_transitions 0 1
	phase "return of AC while suspended" \
	      "please insert the AC cord while the machine is suspended"
	suspend_system
fi
if [ "$test_power" -eq 1 -a "$battery_count" -eq 0 ]; then
	ECHO "*** no BATTERY detected power test skipped ..."
elif [ "$test_power" -eq 1 ]; then
	save_timer_sleep="$timer_sleep"
	let timer_sleep="$args_power_sleep"

	ac_required 0
	phase "battery drain during suspend" \
	      "calculates overall power drain during a long-term suspend"

	# get start values
	date_before=`date +%s`
	bat_before=`battery_capacity`
  
	# Suspend
	suspend_system
  
	# get end values
	date_after=`date +%s`
	bat_after=`battery_capacity`

	# do the calculations 
	let consumed="$bat_before - $bat_after"
	let elapsed="$date_after - $date_before"
	let usage="($consumed * 60*60) / $elapsed"

	# output the results
	ECHO "before: $bat_before mWh"
	ECHO "after: $bat_after mWh"
	ECHO "consumed: $consumed mW"
	ECHO "sleep seconds: $elapsed sec"
	ECHO "overall usage: $usage mW"
	 
	report_battery="$usage mW"

	if [ $elapsed -lt 1200 ]
	then
		ECHO "WARNING: the suspend was less than 20 minutes"
		ECHO "         to get reliable numbers increase the sleep time"
		report_battery="$report_battery (unreliable)"
	fi

	timer_sleep="$save_timer_sleep"
fi

phase_interactive=0
if [ "$test_timed" -eq 1 ]; then
	save_timer_delay="$timer_delay"
	timer_delay=60

	ac_required 1
	phase "30 iteration variable delay suspend/resume stress test"
	while [ "$timer_delay" -gt 0 ]; do
		echo "delay $timer_delay ..."
		suspend_system
		delay_system
		let timer_delay="$timer_delay - 2"
	done
	timer_delay="$save_timer_delay"
fi
if [ "$test_repeat" -eq 1 ]; then
	ac_required 1
	phase "basic $args_repeat_iterations iteration suspend/resume stress test"
	run_suspend "$args_repeat_iterations"
fi

ac_check

#
# REPORT: final report stage.
#
report_this()
{	
	if [ "$2" != "" ]; then
		echo "      $1 $2"
	fi
}
if [ "$auto" -eq 0 ]; then
	echo ""
	echo "*** Please report your results on the Ubuntu WIKI:"
	echo "    https://wiki.ubuntu.com/KernelTeam/SuspendResumeTesting"
	echo ""
	report_this "Battery Consumption:" "$report_battery"
fi

# All suceessful, clean up.
rm -f "$LOGFILE"
