#!/bin/bash

set -eEu -o pipefail

NAME=io-scheduler
CONF="/etc/default/${NAME:-}"  # support customization via /etc/default/io-scheduler

## configuration example to force execution inside virtual environment:
# FORCE_VIRT_RUN=true

## configuration example to avoid disk detection, but configure specific disks:
# DISKS="sda sdb"

## configuration example to use specific schedulers for different disks:
# declare -A SCHEDULER
# SCHEDULER[sda]='none'
# SCHEDULER[sdb]='deadline'

# Default settings
DEFAULT_SCHEDULER='none'  # SSDs, NVME,...
DEFAULT_ROTATIONAL_DISK_SCHEDULER='deadline'  # spinning disks

# helper functions
identify_disks() {
  local boot_disk
  boot_disk="$(pvs --noheadings --separator : | awk -F: '/:ngcp:/ {print $1; exit}' | tr -d ' ')"  # e.g. /dev/nvme0n1p3 or /dev/md0 or /dev/sda3

  case "${boot_disk}" in
    /dev/md*)
      # we need to report all disks present in SW-RAID
      local swraid_devices
      swraid_devices=$(mdadm --detail --export "${boot_disk}" | awk -F= '/MD_DEVICE_.*ev_.*_DEV/ {print $2}')
       for dev in ${swraid_devices} ; do
         local parent_device
         parent_device=$(dirname "$(udevadm info -n "${dev}" -q path)")  # e.g. /devices/pci0000:00/0000:00:05.0/0000:01:02.0/virtio4/host3/target3:0:0/3:0:0:1/block/sdb
         if [[ -L /sys/block/"${parent_device##*/}" ]] ; then
           echo "${parent_device##*/}"  # e.g. "sdb"
         fi
       done
       ;;
     *)
       local parent_device
       parent_device=$(dirname "$(udevadm info -n "${boot_disk}" -q path)")  # e.g. /devices/pci0000:00/0000:00:1c.6/0000:03:00.0/nvme/nvme0/nvme0n1
                                                                             # or   /devices/pci0000:00/0000:00:0d.0/ata3/host2/target2:0:0/2:0:0:0/block/sda
       if [[ -L /sys/block/"${parent_device##*/}" ]] ; then
         echo "${parent_device##*/}"  # e.g. "sdb"
       fi
       ;;
  esac
}

is_rotational_disk() {
  local disk="$1"

  # note: information can be in either `ATTR{queue/rotational}` or `ATTRS{queue/rotational}`
  if udevadm info -n "${disk}" -a | awk -F== '/ATTRS?{queue\/rotational}/ {print $2}' | grep -q '1' ; then
    return 0
  fi

  return 1
}

# Support overriding defaults
if [[ -r "${CONF}" ]] ; then
  # shellcheck disable=SC1090
  . "${CONF}" || true
fi

# detect disks, when not configured via /etc/default/io-scheduler
if [[ -z "${DISKS:-}" ]] ; then
  DISKS=$(identify_disks | sort -u)
fi

setup_default_scheduler() {
  if ! command -v ngcp-virt-identify >/dev/null; then
    printf "Warning: ngcp-virt-identify not available, skipping virtual environment detection.\n"
  else
    if ngcp-virt-identify ; then
      if [[ "${FORCE_VIRT_RUN:-}" == "true" ]] ; then
        printf "Virtualized environment detected, forcing execution due to FORCE_VIRT_RUN.\n"
      else
        printf "Virtualized environment detected. Skipping ngcp-io-scheduler execution.\n"
        return 0
      fi
    fi
  fi

  if [[ -z "${DISKS:-}" ]] ; then
    echo "Warning: no default disks found, nothing to do."
    return 0
  fi

  for disk in ${DISKS} ; do
    local disk_scheduler="${DEFAULT_SCHEDULER:-}"

    if is_rotational_disk "${disk}" ; then
      if [[ -n "${DEFAULT_ROTATIONAL_DISK_SCHEDULER:-}" ]] ; then
        printf "NOTE: Disk '%s' seems to be a spinning drive, considering default '%s' scheduler.\n" "${disk}" "${DEFAULT_ROTATIONAL_DISK_SCHEDULER}"
        disk_scheduler="${DEFAULT_ROTATIONAL_DISK_SCHEDULER}"
      fi
    fi

    if [[ -n "${SCHEDULER[*]}" ]] && [[ -n "${SCHEDULER[${disk}]}" ]] ; then
      disk_scheduler="${SCHEDULER[$disk]}"
      printf "NOTE: Disk specific configuration for '%s' present, considering '%s' scheduler.\n" "${disk}" "${disk_scheduler}"
    fi

    if ! [[ -r "/sys/block/${disk}/queue/scheduler" ]] ; then
      printf "Warning: /sys/block/%s/queue/scheduler can't be read, not modifying scheduler setting\n" "${disk}"
      continue
    fi

    if ! grep -q --word-regexp "${disk_scheduler}" "/sys/block/${disk}/queue/scheduler" ; then
      printf "Warning: '%s' scheduler not supported by disk %s, not modifying scheduler setting\n" "${disk_scheduler}" "${disk}"
      continue
    fi

    if ! [[ -w "/sys/block/${disk}/queue/scheduler" ]] ; then
      echo "Error: can not write to I/O scheduler settings file for ${disk}" >&2
      exit 1
    fi

    printf "Setting I/O scheduler for disk '%s' to '%s': " "${disk}" "${disk_scheduler}"
    if echo "${disk_scheduler}" > "/sys/block/${disk}/queue/scheduler" ; then
      echo "done"
    else
      echo "failed" >&2
      exit 1
    fi
  done
}

display_scheduler_settings() {
  if [[ -z "${DISKS:-}" ]] ; then
    echo "Warning: no default disks found, nothing to do."
    return 0
  fi

  local rc=0
  for disk in ${DISKS} ; do
    printf "Checking state of %s for '%s':\n" "${NAME:-}" "${disk:-}"

    if ! [[ -r "/sys/block/${disk}/queue/scheduler" ]] ; then
      echo "Error reading scheduler setting for disk ${disk}" >&2
      rc=1
      continue
    fi

    printf "* I/O scheduler setting for disk %s: " "${disk}"
    cat "/sys/block/${disk}/queue/scheduler"
  done

  exit "${rc}"
}

case "${1:-}" in
  start)
    setup_default_scheduler
    ;;
  status)
    display_scheduler_settings
    ;;
  *)
    echo "Usage: $0 {start|status}" >&2
    exit 1
    ;;
esac

exit 0
