#!/bin/sh

# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

. /usr/share/misc/shflags

DEFINE_integer count 10000 "number of iterations" c
DEFINE_integer memory_check_size 0 \
               "Amount of memory to allocate (0 means as much as possible)"
DEFINE_integer suspend_max 10 "Max seconds to suspend"
DEFINE_integer suspend_min 5 "Min seconds to suspend"
DEFINE_integer wake_max 10 "Max seconds to stay awake for"
DEFINE_integer wake_min 5 "Min seconds to stay awake for"
DEFINE_boolean backup_rtc ${FLAGS_FALSE} "Use second rtc if present for backup"
DEFINE_boolean bug_fatal ${FLAGS_TRUE} "Abort on BUG dmesg lines"
DEFINE_boolean crc_fatal ${FLAGS_TRUE} "Abort on CRC error dmesg lines"
DEFINE_boolean warn_fatal ${FLAGS_FALSE} "Abort on WARNING dmesg lines"
DEFINE_boolean errors_fatal ${FLAGS_TRUE} "Abort on errors"
DEFINE_boolean fw_errors_fatal ${FLAGS_TRUE} \
               "Abort on firmware errors (subset of --errors_fatal)"
DEFINE_boolean memory_check ${FLAGS_FALSE} "Use memory_suspend_test to suspend"
DEFINE_boolean pm_print_times ${FLAGS_TRUE} \
               "Print the time taken by devices to suspend and resume"
DEFINE_string  ignore_wakeup_source "none" "Wakeup source to ignore" i
DEFINE_string  record_dmesg_dir "" "dir to record dmesgs for failed iterations"
DEFINE_string  post_resume_command "" "Command to run after each resume"

# Minimum value for --suspend_min when --memory_check is true.
readonly MEMORY_TEST_SUSPEND_MIN_DELAY=10

# Location of disable_dark_resume preference file.
readonly DISABLE_DARK_RESUME=/var/lib/power_manager/disable_dark_resume

# Scheduling priority of memory_suspend_test
readonly MEMORY_TEST_PRIORITY=-8

# Location of the temporary preference file. Set by disable_dark_resume.
MUTABLE_PREF_FILE=

# Record the dmesg for the last iteration.
# Accepts iteration number as an argument (to name the file dmesg_<iter_num>)
record_dmesg_for_iteration() {
  local lines=0
  local filename=${FLAGS_record_dmesg_dir}/dmesg_$1

  echo -n " (dmesg > ${filename})..."
  lines=$(dmesg | tac | grep -m1 -n "PM: Syncing filesystems" \
          | cut -f1 -d:)
  dmesg | tail -${lines} > ${filename}
}

get_success_count() {
  awk '$1 == "success:" { print $2 }' /sys/kernel/debug/suspend_stats
}

# Return the current event count for the wakeup source.
get_wakeup_source_event_count() {
  local wakeup_source=$1
  local event_count=0

  if [ "${wakeup_source}" -a "none" != "${wakeup_source}" ]; then
    # Get the event count field for the wakeup source.
    event_count=$(awk '$1 == "'"${wakeup_source}"'" { print $3 }' \
                  /sys/kernel/debug/wakeup_sources)
    if [ -z "$event_count" ]; then
      event_count=0
    fi
  fi
  echo "$event_count"
}

random() {
  hexdump -n 2 -e '/2 "%u"' /dev/urandom
}

boolean_value() {
  if [ $1 -eq ${FLAGS_TRUE} ]; then
    echo "true"
  else
    echo "false"
  fi
}

# Restarts powerd and blocks until it's re-registered its D-Bus name.
restart_powerd() {
  restart powerd
  while ! dbus-send --system --print-reply --dest=org.freedesktop.DBus \
      /org/freedesktop/DBus org.freedesktop.DBus.NameHasOwner \
      string:org.chromium.PowerManager | grep -q true; do
    sleep 0.1
  done
}

disable_dark_resume() {
  # We can only bind-mount over the DISABLE_DARK_RESUME file if it
  # already exists. If it doesn't, we can just use it directly and clean
  # it up later the same way.
  if [ ! -e ${DISABLE_DARK_RESUME} ]; then
    MUTABLE_PREF_FILE=${DISABLE_DARK_RESUME}
  else
    MUTABLE_PREF_FILE=$(mktemp /tmp/powerd_disable_dark_resume.XXXXXXXXXXXX)
    mount --bind ${MUTABLE_PREF_FILE} ${DISABLE_DARK_RESUME}
  fi
  echo 1 > ${MUTABLE_PREF_FILE}
  restart_powerd
}

restore_dark_resume_state() {
  if [ ${MUTABLE_PREF_FILE} != ${DISABLE_DARK_RESUME} ]; then
    umount ${DISABLE_DARK_RESUME}
  fi
  rm -f ${MUTABLE_PREF_FILE}
  restart_powerd
}

dump_stats_and_exit() {
  echo "${preserved_pm_print_times}" > /sys/power/pm_print_times
  restore_dark_resume_state
  start tlsdated
  echo ""
  echo "Finished ${cur} iterations."
  echo "Suspend failures: $(( cur -
                              ($(get_success_count) - initial_successes) ))"
  echo "Wakealarm errors: ${wakealarm_errors}"
  echo "Firmware log errors: ${firmware_errors}"

  if [ "${check_s0ix_errors}" = 1 ]; then
    echo "s0ix errors: ${s0ix_errors}"
  fi
  exit 0
}

# Returns the S0ix residency counter as returned by the kernel.
# The driver that provides these counters depends on the platform.
get_s0ix_count() {
  if [ -e /sys/kernel/debug/pmc_core/slp_s0_residency_usec ]; then
    # Kabylake / Skylake systems.
    cat /sys/kernel/debug/pmc_core/slp_s0_residency_usec
  elif [ -e /sys/kernel/debug/telemetry/s0ix_residency_usec ]; then
    # Apollolake Systems.
    cat /sys/kernel/debug/telemetry/s0ix_residency_usec
  else
    echo 0
  fi
}

FLAGS "$@" || exit 1

if [ "${USER}" != root ]; then
  echo "This script must be run as root." 1>&2
  exit 1
fi

if [ ${FLAGS_backup_rtc} -eq ${FLAGS_TRUE} ] &&
   [ ! -e /sys/class/rtc/rtc1/wakealarm ]; then
  echo "rtc1 not present, not setting second wakealarm"
  FLAGS_backup_rtc=${FLAGS_FALSE}
fi

if [ ${FLAGS_memory_check} -eq ${FLAGS_TRUE} ]; then
  # Increase priority to make it less likely that memory_suspend_test takes
  # so long to free memory that the wake alarm fires prematurely.
  # Default size is set to 0, which tells memory_suspend_test to allocate as
  # much memory as possible.
  memory_test_cmd="memory_suspend_test --size=${FLAGS_memory_check_size}"
  suspend_cmd="nice -n ${MEMORY_TEST_PRIORITY} ${memory_test_cmd}"

  # Writing to memory can take a while. Make it less likely that the
  # wake alarm will fire before the system has suspended.
  if [ ${FLAGS_suspend_min} -lt ${MEMORY_TEST_SUSPEND_MIN_DELAY} ]; then
    orig_diff=$(( FLAGS_suspend_max - FLAGS_suspend_min ))
    FLAGS_suspend_min=${MEMORY_TEST_SUSPEND_MIN_DELAY}
    FLAGS_suspend_max=$(( FLAGS_suspend_min + orig_diff ))
  fi
else
  suspend_cmd="powerd_dbus_suspend --delay=0"
fi

if [ ! -z "${FLAGS_record_dmesg_dir}" ]; then
  mkdir -p ${FLAGS_record_dmesg_dir}
  touch ${FLAGS_record_dmesg_dir}/dmesg_test
  if [ $? -ne 0 ]; then
    echo "${FLAGS_record_dmesg_dir} not writable for recording dmesg(s)"
    exit 1
  fi
  rm -rf ${FLAGS_record_dmesg_dir}/dmesg_*
fi

check_s0ix_errors=0
if check_powerd_config --suspend_to_idle; then
  if [ -e /sys/kernel/debug/pmc_core/slp_s0_residency_usec ] || \
     [ -e /sys/kernel/debug/telemetry/s0ix_residency_usec ]; then
    check_s0ix_errors=1
  else
    echo "Warning! platform uses s0ix, but we have no way to verify s0ix"
  fi
fi

echo "Running ${FLAGS_count} iterations with:"
echo "  suspend: ${FLAGS_suspend_min}-${FLAGS_suspend_max} seconds"
echo "  wake: ${FLAGS_wake_min}-${FLAGS_wake_max} seconds"
echo "  backup_rtc: $(boolean_value ${FLAGS_backup_rtc})"
echo "  errors_fatal: $(boolean_value ${FLAGS_errors_fatal})"
echo "  fw_errors_fatal: $(boolean_value ${FLAGS_fw_errors_fatal})"
echo "  bugs fatal:  $(boolean_value ${FLAGS_bug_fatal})"
echo "  warnings fatal:  $(boolean_value ${FLAGS_warn_fatal})"
echo "  crcs fatal:  $(boolean_value ${FLAGS_crc_fatal})"
echo "  suspend command: ${suspend_cmd}"
echo "  wakeup source to ignore: ${FLAGS_ignore_wakeup_source}"
echo "  record_dmesg_dir: ${FLAGS_record_dmesg_dir}"
echo "  pm_print_times:  $(boolean_value ${FLAGS_pm_print_times})"

initial_successes=$(get_success_count)
suspend_interval=$(( FLAGS_suspend_max - FLAGS_suspend_min + 1 ))
wake_interval=$(( FLAGS_wake_max - FLAGS_wake_min + 1 ))
preserved_pm_print_times=$(cat /sys/power/pm_print_times)

if [ ${FLAGS_pm_print_times} -eq ${FLAGS_TRUE} ]; then
  echo 1 > /sys/power/pm_print_times
else
  echo 0 > /sys/power/pm_print_times
fi

trap dump_stats_and_exit INT
disable_dark_resume

# Turn off network sync of time to prevent any spurious RTC wake events
stop tlsdated

cur=0
wakealarm_errors=0
firmware_errors=0
s0ix_errors=0
last_successes=${initial_successes}
exit_loop=0
cur_ignore_count=0
last_ignore_count=0

while true; do
  : $(( cur += 1 ))
  printf "Suspend %5d of ${FLAGS_count}: " "${cur}"

  # Flush file system buffer before setting wake alarm to make it
  # less likely that the kernel's sys_sync() call takes so long
  # that the alarm fires early.
  sync

  sus_time=$(( ( $(random) % suspend_interval ) + FLAGS_suspend_min ))
  printf "sleep for %2d seconds..." "${sus_time}"
  last_s0ix_count=$(get_s0ix_count)
  last_ignore_count=$(get_wakeup_source_event_count \
                      ${FLAGS_ignore_wakeup_source})
  wakeup_count=$(cat /sys/power/wakeup_count)
  echo 0 > /sys/class/rtc/rtc0/wakealarm
  echo "+${sus_time}" > /sys/class/rtc/rtc0/wakealarm
  if [ ${FLAGS_backup_rtc} -eq ${FLAGS_TRUE} ]; then
    echo 0 > /sys/class/rtc/rtc1/wakealarm
    echo "+$(( sus_time + 5 ))" > /sys/class/rtc/rtc1/wakealarm
  fi

  ${suspend_cmd} --wakeup_count=${wakeup_count}

  if [ -n "${FLAGS_post_resume_command}" ]; then
    # Leading space because this will appear on the same line as the
    # printf output from above.
    echo " running ${FLAGS_post_resume_command}"
    ${FLAGS_post_resume_command}
  fi

  # Wakealarm is expected to be cleared after the set wakeup time,
  # which is the resume time if the device does not resume due to
  # other sources. Although we always clear the wakealarm file before
  # setting it, we still report this error if it is not cleared.
  cur_time="$(cat /sys/class/rtc/rtc0/since_epoch)"
  main_wakealarm="$(cat /sys/class/rtc/rtc0/wakealarm)"
  if [ -n "${main_wakealarm}" ] && [ ${main_wakealarm} -lt ${cur_time} ]; then
    : $(( wakealarm_errors += 1 ))
    echo "Wakealarm is not cleared (ignoring)."
  fi
  if [ ${FLAGS_backup_rtc} -eq ${FLAGS_TRUE} ]; then
    backup_wakealarm="$(cat /sys/class/rtc/rtc1/wakealarm)"
    if [ -n "${backup_wakealarm}" ] &&
       [ ${backup_wakealarm} -lt ${cur_time} ]; then
      : $(( wakealarm_errors += 1 ))
      echo "Backup wakealarm is not cleared (ignoring)."
    fi
  fi

  # Look for errors in firmware log.
  if cbmem -1 | grep ERROR; then
    : $(( firmware_errors += 1 ))
    if [ ${FLAGS_errors_fatal} -eq ${FLAGS_TRUE} -a \
         ${FLAGS_fw_errors_fatal} -eq ${FLAGS_TRUE} ]; then
      exit_loop=1
      echo "Firmware error found."
    else
      echo "Firmware error found (ignoring)."
    fi
  fi
  # Make sure suspend succeeded
  cur_successes=$(get_success_count)
  cur_ignore_count=$(get_wakeup_source_event_count \
                     ${FLAGS_ignore_wakeup_source})
  cur_s0ix_count=$(get_s0ix_count)
  if [ ${cur_successes} -eq ${last_successes} -a \
       ${cur_ignore_count} -eq ${last_ignore_count} ]; then
    if [ ${FLAGS_errors_fatal} -eq ${FLAGS_TRUE} ]; then
      echo "Suspend failed."
      exit_loop=1
    else
      if [ ! -z "${FLAGS_record_dmesg_dir}" ]; then
        record_dmesg_for_iteration ${cur}
      fi
      printf "(suspend failed, ignoring)"
    fi
  elif [ "${check_s0ix_errors}" = 1 ] && \
       [ "${cur_s0ix_count}" = "${last_s0ix_count}" ]; then
    echo "Suspend failed, s0ix count did not increment from ${last_s0ix_count})"
    : $(( s0ix_errors += 1 ))
    if [ ${FLAGS_errors_fatal} -eq ${FLAGS_TRUE} ]; then
      exit_loop=1
    elif [ ! -z "${FLAGS_record_dmesg_dir}" ]; then
        record_dmesg_for_iteration ${cur}
    fi
  fi
  last_successes=${cur_successes}
  # For BUG and CRC errors counting existing occurrences in dmesg
  # is not that useful as dmesg will wrap so we would need to account
  # for the count shrinking over time.
  # Exit on BUG
  if [ ${FLAGS_bug_fatal} -eq ${FLAGS_TRUE} ] &&
        dmesg | grep -w BUG; then
    echo "BUG found."
    exit_loop=1
  fi
  # Exit on WARNING
  if [ ${FLAGS_warn_fatal} -eq ${FLAGS_TRUE} ] &&
        dmesg | grep -w WARNING; then
    echo "WARNING found."
    exit_loop=1
  fi
  # Exit on CRC error
  if [ ${FLAGS_crc_fatal} -eq ${FLAGS_TRUE} ] && dmesg | grep "CRC.*error"; then
    echo "CRC error found."
    exit_loop=1
  fi
  # Exit the loop if requested from errors or done with iterations
  if [ ${cur} -eq ${FLAGS_count} ] || [ ${exit_loop} -eq 1 ]; then
    if [ ${exit_loop} -eq 1 -a \
         ! -z "${FLAGS_record_dmesg_dir}" ]; then
      record_dmesg_for_iteration ${cur}
    fi
    echo ""
    break
  fi
  wake_time=$(( ( $(random) % wake_interval ) + FLAGS_wake_min ))
  printf " wake for %2d seconds..." "${wake_time}"
  sleep "${wake_time}"
  echo ""
done

dump_stats_and_exit