#!/bin/bash
# =============================================================================
# machine-setup — using the whole disk
# =============================================================================
#
# Definitions only, like the other lib/ files.
#
# ── The problem this exists for ──
#
# Ubuntu Server's installer, left on its defaults, creates an LVM logical volume
# at a fixed size and leaves the rest of the disk as free extents in the volume
# group. On a 2TB drive you get a root filesystem of around 100GB and no
# indication anything is wrong: `lsblk` shows the whole disk, `df` shows 100G,
# and the two are never seen side by side until the day it fills.
#
# The same shape turns up two other ways:
#
#   a virtual disk grown at the hypervisor or provider, where the partition still
#   ends where it used to
#
#   a partition that was resized without the filesystem inside it being told
#
# Three layers, and any one of them can be the short one:
#
#   disk        the physical or virtual device
#   container   the partition, or the logical volume
#   filesystem  what df reports
#
# So all three are measured and reported together. Seeing them in one place is
# most of the value; the fix is usually two commands once you know which layer is
# short.
#
# ── Only ever grows ──
#
# Nothing here shrinks anything, and nothing here creates or deletes a partition.
# ext4, xfs and btrfs all grow while mounted, so there is no unmount and no
# reboot, and a failure part-way leaves a smaller filesystem on a larger
# container — which is exactly the state it started in.

[[ -n "${MACHINE_SETUP_DISK_LOADED:-}" ]] && return 0
MACHINE_SETUP_DISK_LOADED=1

# -----------------------------------------------------------------------------
# What is where
# -----------------------------------------------------------------------------

root_device() { findmnt -no SOURCE / 2>/dev/null; }
root_fstype() { findmnt -no FSTYPE / 2>/dev/null; }

# Size of a block device in bytes.
dev_bytes() { lsblk -bndo SIZE "$1" 2>/dev/null || echo 0; }

# Bytes, formatted the way df and lsblk format them.
human_bytes() { numfmt --to=iec --suffix=B --format='%.1f' "$1" 2>/dev/null || echo "${1}B"; }

# Is the root filesystem on a logical volume?
root_is_lvm() { [[ "$(lsblk -ndo TYPE "$(root_device)" 2>/dev/null)" == "lvm" ]]; }

# The whole disk a device ultimately sits on: /dev/sda1 -> /dev/sda, and through
# LVM as well, since PKNAME walks one level at a time.
parent_disk() {
  local dev="$1" name
  while true; do
    name="$(lsblk -ndo PKNAME "$dev" 2>/dev/null)"
    [[ -z "$name" ]] && break
    dev="/dev/${name}"
  done
  echo "$dev"
}

# The partition immediately below a device — for LVM, the one holding the PV.
backing_partition() {
  local dev="$1" name
  while [[ "$(lsblk -ndo TYPE "$dev" 2>/dev/null)" != "part" ]]; do
    name="$(lsblk -ndo PKNAME "$dev" 2>/dev/null)"
    [[ -z "$name" ]] && return 1
    dev="/dev/${name}"
  done
  echo "$dev"
}

# Split /dev/sda1 into "/dev/sda 1" — growpart wants them as separate arguments.
# The digits come off the end because that is where a partition number is, on
# /dev/sda1 and /dev/nvme0n1p2 alike.
partition_parts() {
  local part="$1" num disk
  num="${part##*[!0-9]}"
  disk="${part%"$num"}"
  disk="${disk%p}" # nvme0n1p2 -> nvme0n1
  echo "$disk $num"
}

# -----------------------------------------------------------------------------
# Sizes of the three layers
# -----------------------------------------------------------------------------

# What the filesystem itself believes it is, which is the number df reports and
# the only one of the three that is asked of the filesystem rather than the
# kernel's block layer.
fs_bytes() {
  local dev="$1"
  case "$(root_fstype)" in
    ext2 | ext3 | ext4)
      local count size
      count="$(tune2fs -l "$dev" 2>/dev/null | awk -F: '/^Block count:/ { gsub(/ /, "", $2); print $2 }')"
      size="$(tune2fs -l "$dev" 2>/dev/null | awk -F: '/^Block size:/ { gsub(/ /, "", $2); print $2 }')"
      [[ -n "$count" && -n "$size" ]] && echo $((count * size)) || echo 0
      ;;
    xfs | btrfs)
      # Both report through the mount rather than the device.
      echo $(($(findmnt -bno SIZE / 2>/dev/null || echo 0)))
      ;;
    *) echo 0 ;;
  esac
}

# Unallocated extents in the volume group behind root. This is the Ubuntu
# installer case, and the one that is invisible without asking LVM directly.
vg_free_bytes() {
  local vg
  command -v vgs &>/dev/null || {
    echo 0
    return
  }
  vg="$(lvs --noheadings -o vg_name "$(root_device)" 2>/dev/null | tr -d ' ')"
  [[ -z "$vg" ]] && {
    echo 0
    return
  }
  vgs --noheadings --nosuffix --units b -o vg_free "$vg" 2>/dev/null | tr -d ' ' || echo 0
}

# -----------------------------------------------------------------------------
# Can anything be reclaimed?
# -----------------------------------------------------------------------------

# growpart answers this better than arithmetic on sector counts: it exits 0 when
# it would change something and 1 with NOCHANGE when the partition already
# reaches the end of the disk. Needs cloud-guest-utils, which is not installed by
# default on every image.
partition_can_grow() {
  local part="$1" disk num
  command -v growpart &>/dev/null || return 1
  read -r disk num <<<"$(partition_parts "$part")"
  growpart --dry-run "$disk" "$num" &>/dev/null
}

ensure_growpart() {
  command -v growpart &>/dev/null && return 0
  info "     installing cloud-guest-utils, which provides growpart"
  pkg_install_now cloud-guest-utils
}

# -----------------------------------------------------------------------------
# Growing
# -----------------------------------------------------------------------------

grow_partition() {
  local part="$1" disk num
  read -r disk num <<<"$(partition_parts "$part")"
  growpart "$disk" "$num"
}

# Tell LVM the partition under the physical volume got bigger.
grow_pv() { pvresize "$1"; }

# Take every free extent in the volume group.
grow_lv() { lvextend -l +100%FREE "$(root_device)"; }

# Grow the filesystem into whatever room it now has. All three do this online, so
# the root filesystem is grown while it is mounted and in use.
grow_fs() {
  case "$(root_fstype)" in
    ext2 | ext3 | ext4) resize2fs "$(root_device)" ;;
    xfs) xfs_growfs / ;;
    btrfs) btrfs filesystem resize max / ;;
    *)
      warn "do not know how to grow a $(root_fstype) filesystem"
      return 1
      ;;
  esac
}