#! /bin/bash
#  Copyright (C) 2019 Cumulus Networks, Inc. All rights reserved
#  This script is intended *only* for use by the cumulus-tools postinst script
#  Enable discard option in fstab on disks which require it so ATA TRIM is used
#
#  This software is subject to the Cumulus Networks End User License Agreement
#  available  at the following locations:
#
#  Internet: https://cumulusnetworks.com/downloads/eula/latest/view/
#  Cumulus Linux systems: /usr/share/cumulus/EULA.txt

# Array which contains one entry for each disk device, i.e. /dev/sda (or
# partition, i.e. /dev/sda2) which requires TRIM support.
declare -a disk_devs

# Array which contains one entry for each mounted file system on a disk which
# requires TRIM.
declare -a trim_req_mounts

# Array which contains one entry for each mounted file system on a disk which
# requires TRIM but which is not currently mounted with the discard option.
declare -a no_discard_mounts

isElement()
{
    local e match="$1"
    shift
    for e; do [[ "${e}" == "${match}" ]] && return 0; done
    return 1
}

# Populate the disk_devs array with a list of disk drives, or partitions thereof,
# which require TRIM support. This is done by examining the file names in the
# /dev/disk/by-id directory for disks which are known to require TRIM.
getTrimRequiredDiskDevs()
{
    local file
    local dev

    # Does the file system have a directory with disks listed by ID?
    disk_devs=()
    [ -d /dev/disk/by-id ] || return 0

    # Examine each file name in the directory
    for file in /dev/disk/by-id/*
    do
        # Does the file name contain one of the strings identifying TRIM-required disks?
        if [[ ${file} == *"_3ME3_"* ]] ||
           [[ ${file} == *"_3IE3_"* ]] ||
           [[ ${file} == *"_3IE4_"* ]]
        then
            # Get the name of the disk device, e.g. /dev/sda, and add it to the
            # disk_devs array if it is not already there.
            dev=$(realpath -e ${file})
            isElement ${dev} "${disk_devs[@]}" || disk_devs+=(${dev})
        fi
    done
    return ${#disk_devs[@]}
}

# Examine the /etc/mtab for mounts of disks which require TRIM but were not
# mounted with the discard option. All such mount points are added to the
# no_discard_mounts array and the number of elements in the array is returned.
getMountsWithoutRequiredDiscard()
{
    local line
    local words

    no_discard_mounts=()
    trim_req_mounts=()
    # Does the /etc/mtab file exist and is it readable?
    [ -r /etc/mtab ] || return 0

    # Process each line the the /etc/mtab file. If the first token is a device
    # which requires TRIM and the discard option isn't in the options, then add
    # the mount point to the array.
    while read -r line
    do
        # Break the line into tokens and make sure there are enough of them
        words=( ${line} )
        [ ${#words[@]} -ge 4 ] || continue
        isElement ${words[0]} "${disk_devs[@]}" || continue
        # Add to list of mounts of trim required devices if not already there
        isElement ${words[1]} "${trim_req_mounts[@]}" ||
            trim_req_mounts+=(${words[1]})
        # Is the discard option set? If not, is this mount point already in the
        # array? If not, add it.
        [[ ${words[3]} == *"discard"* ]] ||
            isElement ${words[1]} "${no_discard_mounts[@]}" ||
                no_discard_mounts+=(${words[1]})
    done < /etc/mtab

    return ${#no_discard_mounts[@]}
}


# Add the discard option to entries in /etc/fstab which require it. In order to
# correlate an entry in /etc/fstab with a disk which requires TRIM the /dev/disk
# directory structure is examined.
# An array of file system specifications is created for all possible ways to
# specify filesystems of disks which require TRIM. If the first token in a line
# of /etc/fstab matches an entry in the array, then that line in /etc/fstab
# will have the discard option added, if it is not already there.
addDiscardOptionToFstab()
{
    local trim_req_fs disk_id prefix file dev mount sed_cmds

    # Is there an /etc/fstab file?
    [ -e /etc/fstab ] || return

    # Build the trim_req_fs array. Start with the TRIM required device names.
    trim_req_fs=("${disk_devs[@]}")
    # Go through each of the directories in /dev/disk, except by-id
    for disk_id in label partlabel partuuid uuid
    do
        # Convert the prefix to uppercase then look at each file in the directory
        prefix=$(echo ${disk_id} | awk '{print toupper($0)}')
        for file in /dev/disk/by-${disk_id}/*
        do
            # Get the device to which the file points. Is it a device which
            # requires trim? If so, is the file system spec already in the
            # trim_req_fs array? If not, add it.
            dev=$(realpath -e ${file})
            mount=${prefix}"="$(basename ${file})
            isElement ${dev} "${disk_devs[@]}" &&
                isElement ${mount} "${trim_req_fs[@]}" ||
                    trim_req_fs+=(${mount})
        done
    done

    sed_cmds=()
    for mount in "${trim_req_fs[@]}"
    do
        # If the discard option is not present, add it to the options if the first
        # token matches an entry in trim_req_fs.
        sed_cmds+=("-e /\(\S\+\s\+\)\{3\}\S*discard\S*\s\+/!s ^\(${mount}\(\s\+\S\+\)\{3\}\) \1,discard ")
    done

    cp /etc/fstab /tmp/fstab.orig || return
    sed "${sed_cmds[@]}" /tmp/fstab.orig > /tmp/fstab.new || return
    cmp -s /tmp/fstab.new /etc/fstab || cp /tmp/fstab.new /etc/fstab
}

# Run the fstrim command on all file system which require TRIM but are not
# currently mounted with the discard option (mount points in no_discard_mounts
# array).
runFstrim()
{
    local mounts
    local mount

    mounts=("${no_discard_mounts[@]}")

    # Run fstrim on each mount point
    echo "Running fstrim on filesystems.  May take a minute or two"
    for mount in "${mounts[@]}"
    do
        /sbin/fstrim ${mount}
    done
}

# Populate the disk_devs array with devices which require TRIM.
getTrimRequiredDiskDevs && exit 0

# Add the discard option to all entries in /etc/fstab which require it
addDiscardOptionToFstab

# Prepare to run fstrim by making a list of all mounts which require TRIM
getMountsWithoutRequiredDiscard
[ $? -eq 0 ] && exit 0
runFstrim
