#!/bin/bash
#***********************************************************************************************************
#
# Starfish Storage Corporation ("Starfish") CONFIDENTIAL
# Unpublished Copyright (c) 2011 - present Starfish Storage Corporation, All Rights Reserved.
#
# NOTICE: This file and its contents (1) constitute Starfish's "External Code" under Starfish's most-recent
# Limited Software End-User License Agreement, and (2) is and remains the property of Starfish. The
# intellectual and technical concepts contained herein are proprietary to Starfish and may be covered by
# U.S. and/or foreign patents or patents in process, and are protected by trade secret or copyright law.
# Dissemination of this information or reproduction of this material is strictly forbidden unless prior
# written permission is obtained from Starfish. Access to the source code contained herein is hereby
# forbidden to anyone except (A) current Starfish employees, managers, or contractors who have executed
# confidentiality or nondisclosure agreements explicitly covering such access, and (B) licensees of
# Starfish's software.
#
# ANY REPRODUCTION, COPYING, MODIFICATION, DISTRIBUTION, PUBLIC PERFORMANCE, OR PUBLIC DISPLAY OF OR
# THROUGH USE OF THIS SOURCE CODE WITHOUT THE EXPRESS WRITTEN CONSENT OF STARFISH IS STRICTLY PROHIBITED
# AND IS IN VIOLATION OF APPLICABLE LAWS AND INTERNATIONAL TREATIES. THE RECEIPT OR POSSESSION OF THIS
# FILE OR ITS CONTENTS AND/OR RELATED INFORMATION DOES NOT CONVEY OR IMPLY ANY RIGHTS TO REPRODUCE,
# DISCLOSE, OR DISTRIBUTE ITS CONTENTS, OR TO MANUFACTURE, USE, OR SELL ANYTHING THAT IT MAY DESCRIBE, IN
# WHOLE OR IN PART.
#
# FOR U.S. GOVERNMENT CUSTOMERS REGARDING THIS DOCUMENTATION/SOFTWARE
#   These notices shall be marked on any reproduction of this data, in whole or in part.
#   NOTICE: Notwithstanding any other lease or license that may pertain to, or accompany the delivery of,
#   this computer software, the rights of the Government regarding its use, reproduction and disclosure are
#   as set forth in Section 52.227-19 of the FARS Computer Software-Restricted Rights clause.
#   RESTRICTED RIGHTS NOTICE: Use, duplication, or disclosure by the Government is subject to the
#   restrictions as set forth in subparagraph (c)(1)(ii) of the Rights in Technical Data and Computer
#   Software clause at DFARS 52.227-7013.
#
#***********************************************************************************************************

set -euo pipefail

# This script returns 0 if there is enough space, 1 otherwise. Any other returncode indicates failure.

readonly AWK="$(command -v gawk 2>&1 || command -v awk 2>&1)"
readonly DF="$(command -v gdf 2>&1 || command -v df 2>&1)"

readonly STATUS_NOT_THROTTLED=0
readonly STATUS_THROTTLED=1
readonly STATUS_ARGPARSE_ERROR=2
readonly STATUS_GENERAL_FAILURE=99

LOW_WATERMARK=0
HIGH_WATERMARK=0
JOB_ID=""
MOUNT=""

usage() {

    cat <<EOF
This command should be used with '--throttle-cmd' option of 'sf job start'.
When volume usage goes over --high-watermark, this command returns 1 which
indicates that the job should pause. Job will be paused until volume usage
falls below --low-watermark.

    --help    - display this help
    --mount   - mount point of volume
    --bytes-to-process - number of bytes about to be processed by dispatcher
    --low-watermark - volume usage level at which processing should
                be resumed
    --high-watermark - volume usage level at which processing should
                be paused
    --job-id - id of job being processed
EOF
    exit 0
}

parse_cmdline() {
    while [[ $# -gt 0 ]]; do
        key="$1"

        case "${key}" in
            --help)
                usage
                exit 0
                ;;
            --mount)
                MOUNT="$2"
                shift
                ;;
            --bytes-to-process)
                BYTES_TO_PROCESS="$2"
                shift
                ;;
            --low-watermark)
                LOW_WATERMARK="$2"
                shift
                ;;
            --high-watermark)
                HIGH_WATERMARK="$2"
                shift
                ;;
            --job-id)
                JOB_ID="$2"
                shift
                ;;
            *)
                echo "Unknown option ${key}"
                exit "${STATUS_ARGPARSE_ERROR}"
                ;;
        esac
        shift
    done

    if [ -z "${JOB_ID}" ]; then
        echo "Job id is an empty string" >&2
        exit "${STATUS_ARGPARSE_ERROR}"
    fi
    if [ -z "${MOUNT}" ]; then
        echo "Option --mount is required" >&2
        exit "${STATUS_ARGPARSE_ERROR}"
    fi
}

file_exists() {
    [ -f "$1" ]
}

count_avail() {
    local mount="$1"

    # Columns of df --portability --block-size=1: Filesystem, 1-blocks, Used, Available, Capacity, Mounted on
    # shellcheck disable=SC2016
    ${DF} --portability --block-size=1 "${mount}" | ${AWK} '{ print $4 }' | tail -1
}

read_last_status() {
    local status_file="$1"

    if file_exists "${status_file}"; then
        head -n1 "${status_file}"
    else
        echo "${STATUS_NOT_THROTTLED}"
    fi
}

should_throttle() {
    local avail avail_after last_status_file last_status status

    avail="$(count_avail "${MOUNT}")"
    avail_after="$((avail - BYTES_TO_PROCESS))"
    last_status_file="${TMPDIR:-/tmp}/df_throttle_${JOB_ID}"
    last_status="$(read_last_status "${last_status_file}")"
    status="${STATUS_GENERAL_FAILURE}"

    if [ "${last_status}" -eq "${STATUS_NOT_THROTTLED}" ]; then
        if [ "${avail_after}" -ge "${LOW_WATERMARK}" ]; then
            status="${STATUS_NOT_THROTTLED}"
        else
            status="${STATUS_THROTTLED}"
        fi
    elif [ "${last_status}" -eq "${STATUS_THROTTLED}" ]; then
        # if a single batch is bigger than gap between LOW_WATERMARK and HIGH_WATERMARK
        # the job will be throttled forever
        if [ "${avail}" -ge "${HIGH_WATERMARK}" ] && [ "${avail_after}" -ge "${LOW_WATERMARK}" ]; then
            status="${STATUS_NOT_THROTTLED}"
        else
            status="${STATUS_THROTTLED}"
        fi
    else
        echo "Unexpected last status: ${last_status}" >&2
    fi

    echo -e "${status}\nmount=${MOUNT};low=${LOW_WATERMARK};high=${HIGH_WATERMARK};avail=${avail};bytes=${BYTES_TO_PROCESS};avail_after=${avail_after}" > "${last_status_file}"
    exit "${status}"
}

main() {
    parse_cmdline "$@"
    should_throttle
}

main "$@"
