#!/bin/sh
# Copyright 2015 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.

# Check if the stateful partition needs to be trimmed.
# Trim is triggered if:
# - Previous run was interrupted, in case of reboot or heavy disk usage.
# - The amount of write since the last trim is greater the X% of remaining
#   space in stateful.
# - Not trimmed for a long time.
#
# When trim is triggered, we constantly check that the underlying device is
# not heavily used.

. /usr/share/misc/shflags

DEFINE_boolean 'test' ${FLAGS_FALSE} "For unit testing."

# Parse command line.
FLAGS "$@" || exit 1
eval set -- "${FLAGS_ARGV}"

set -e

if [ "${FLAGS_test}" -eq "${FLAGS_FALSE}" ]; then
  . /usr/share/misc/chromeos-common.sh
  . /usr/share/cros/disk_utils.sh
  . /usr/share/cros/trim_utils.sh

  # Names of device and sysfs entries the script will work with.
  # These variables are mocked by the test harness.
  PARTITION_DEV="$(get_stateful_df_data | awk '{print $1}' | sed -e 's#/dev/##')"
  DEVICE="$(get_block_dev_from_partition_dev "${PARTITION_DEV}")"
  DEVICE_STATS="/sys/block/${DEVICE}/stat"
  PARTITION_DEV_STATS="/sys/block/${DEVICE}/${PARTITION_DEV}/stat"
else
  . ./share/trim_utils.sh
fi

# Maximal amount of time before forcing a trim.
MAX_DELAY_DAYS=7
MAX_DELAY=$(( MAX_DELAY_DAYS * 24 * 3600 ))

# Ratio of write/free space to trigger a trim.
MAX_TOTAL_WRITE_PERCENT=75

# Slice of disk processed by each fstrim call.
FSTRIM_SLICE_MB=128

# Maximum amount of time to spend in this script before rescheduling.
MAX_TRIM_WORKER_TIME_SEC=3000

# Time between fstrim calls when the disk is in steady state.
DELAY_BTW_FSTRIM_SEC=5

# Time to wait if disk is heavily used.
DELAY_HIGH_USAGE_SEC=15

# Maximal disk utilization for fstrim to proceed.
MAX_DISK_UTIL_PERCENT=20

# Maximal amount in milliseconds of disk utilization to still trigger trim.
MAX_DISK_UTIL_MS=$(( DELAY_BTW_FSTRIM_SEC * 1000 * \
                     MAX_DISK_UTIL_PERCENT / 100 ))

# Global variables:
# current_offset : current_offset to use when TRIM process resume,
# duration : Amount of time spent trimming so far,
# write_last_trim : Amount of data written when last TRIM completed,
# date_last_trim : When the last TRIM completed,
# trimmed_amount : Amount of data trimmed by the last TRIM process,
# trim_status : Status of the current/last TRIM process.

log_msg() {
  logger -t "$0[${PPID}]" "$*"
  echo "$*"
}

load_int() {
  local var="$1"
  local val

  if [ -f "${CHROMEOS_TRIM_DATA}" ]; then
    # Make sure we only permit the value to be set to an integer.
    val=$(sed -n "/^${var}=/{s:.*=::;s:[^0-9]::g;p;q}" "${CHROMEOS_TRIM_DATA}")
  fi
  eval "${var}=${val:-0}"
}

# Read status file, to know current process, Populate global variables.
read_status_file() {
  load_int current_offset
  load_int duration
  load_int write_last_trim
  load_int date_last_trim
  load_int trimmed_amount

  trim_status="$(cat "${CHROMEOS_TRIM_STATE}" 2> /dev/null || : )"
}

write_status_file() {
  cat > "${CHROMEOS_TRIM_DATA}" << EOF
# Automatically generated by $0
current_offset=${current_offset}
duration=${duration}
write_last_trim=${write_last_trim}
date_last_trim=${date_last_trim}
trimmed_amount=${trimmed_amount}
EOF
  echo "${trim_status}" > "${CHROMEOS_TRIM_STATE}"
}

cleanup() {
  write_status_file
  exit 0
}

# Return the number of ms spend doing disk io since boot.
get_disk_time_ms() {
  awk '{print $10}' "${DEVICE_STATS}"
}


# Trim the stateful partition.
#
# Trim a slice of the partition, sleep a little to be sure
# we are not in the middle of heavy disk traffic.
# Wait a little longer if we are, and bails if the process takes
# too long.
do_trim() {
  local partition_size
  partition_size="$(get_stateful_total_space_blocks "1M")"
  local start_time
  start_time="$(date +%s)"
  local current_duration=0
  local before_time_spent_io_ms
  local after_time_spent_io_ms
  local output
  local trimmed
  local trim_length=${FSTRIM_SLICE_MB}

  read_status_file
  trim_status="${TRIM_IN_PROGRESS}"

  if [ "${current_offset}" -eq 0 ]; then
    trimmed_amount=0
    duration=0
  fi

  while [ "${current_offset}" -lt "${partition_size}" ] && \
          [ "${current_duration}" -lt "${MAX_TRIM_WORKER_TIME_SEC}" ]; do
    before_time_spent_io_ms=$(get_disk_time_ms)
    sleep "${DELAY_BTW_FSTRIM_SEC}"
    after_time_spent_io_ms=$(get_disk_time_ms)
    if [ $(( after_time_spent_io_ms - before_time_spent_io_ms )) -gt \
         ${MAX_DISK_UTIL_MS} ]; then
      sleep "${DELAY_HIGH_USAGE_SEC}"
    else
      if [ $(( partition_size - current_offset )) -le ${FSTRIM_SLICE_MB} ]; then
        : $(( trim_length = partition_size - current_offset ))
      fi
      if ! output=$(fstrim --verbose --offset "${current_offset}M" \
                    --length "${trim_length}M" "${STATEFUL}"); then
        log_msg "Error Trimming ${STATEFUL} at" \
                "current_offset ${current_offset}M: ${output}"
        trim_status="${TRIM_FAILED}"
        break
      fi
      trimmed=$(echo "${output}" | sed -ne 's/.*(\([0-9]\+\) bytes) trimmed.*/\1/p')
      if [ -z "${trimmed}" ]; then
        log_msg "Unexpected fstrim output at" \
                "current_offset ${current_offset}M : ${output}"
        trim_status="${TRIM_FAILED}"
        break
      fi
      : $(( trimmed_amount += trimmed ))
      : $(( current_offset += trim_length ))
    fi
    date_last_trim=$(date +%s)
    : $(( current_duration = date_last_trim - start_time ))
  done

  : $(( duration += current_duration ))

  # Prepare to save status in status file.
  if [ "${current_offset}" -eq "${partition_size}" ]; then
    current_offset=0
    write_last_trim="$(awk '{print $7}' "${PARTITION_DEV_STATS}")"
    trim_status="${TRIM_COMPLETED}"
    log_msg "${STATEFUL} fully trimmed in ${duration} seconds"
  else
    log_msg "${STATEFUL} trim interrupted at ${current_offset}M out of" \
            "${partition_size}M, ${duration} seconds so far"
  fi
}

get_device_block_size() {
  cat "/sys/block/${DEVICE}/queue/logical_block_size"
}

# Checks if the SSD needs to be trimmed.
main() {
  local trigger_trim=false
  local block_size
  local free_blocks
  local written_blocks

  if [ -f "${CHROMEOS_TRIM_DATA}" ]; then
    read_status_file

    if [ "${current_offset}" -ne 0 ]; then
      # Trim was interrupted due to intense disk activity, resume.
      trigger_trim=true
    fi

    # Last trim was completed a long time ago.
    if [ "$(( date_last_trim + MAX_DELAY ))" -lt "$(date +%s)" ]; then
      trigger_trim=true
    fi

  else
    # Status file not present, trim.
    trigger_trim=true
  fi

  if ! ${trigger_trim}; then
    block_size="$(get_device_block_size)"
    free_blocks="$(get_stateful_free_space_blocks "${block_size}")"
    written_blocks=$(awk '{print $7}' "${PARTITION_DEV_STATS}")
    if [ "$(date +%s -d "$(uptime -s)")" -gt "${date_last_trim}" ]; then
      # Machine has rebooted since last trim, the amount of writes since
      # last trim is irrelevant.
      write_last_trim=0
    fi
    : $(( written_blocks -= write_last_trim ))

    # Trigger Trim if at least the amount of write since boot is more than
    # 75% of the free blocks.
    if [ $(( written_blocks * 100 )) -gt \
         $(( free_blocks * MAX_TOTAL_WRITE_PERCENT )) ]; then
      trigger_trim=true
    fi
  fi

  if "${trigger_trim}"; then
      do_trim
  fi
}

# Invoke main if not in test mode, otherwise let the test code call.
if [ "${FLAGS_test}" -eq "${FLAGS_FALSE}" ]; then
  (
  if ! flock -n 9; then
    log_msg "Already running; quitting."
    exit 1
  fi

  # In case trim is interrupt by regular shutdown, save the current status.
  trap cleanup EXIT INT TERM

  main "$@"
  ) 9>"${CHROMEOS_TRIM_LOCK}"
fi