#!/bin/bash
#
LANG="en_US.UTF-8"
#
# NVME LOW POWER STATE SCRIPT v0.3 2023_11_10
# (by Joe Schmuck)
#
# TRUENAS CORE (FREEBSD) ONLY !

# WORK IN PROGRESS --- Calculate the wattages.  Each power state duration and number of time it changed to a specific state.
#  How many watts were used in each power state and total, and what the theoretical total would be if the nvme drives were
#  all left at power state 0.  Right now I can't see straight so I will finish this later.  The script should work fine.

#
# Update the instructions too.

#
# This script has two purposes:
#  1) Check the lowest power state and if the state is higher (lower numerically)
#     then it will change the power state to the minimum power state immediately and then
#     check periodically (1 hour by default) and if reset the power state if needed.
#     The minimum power state limit can be set 'minumum_power_setting', this feature is
#     available due to possible minimum power stability problems.
#
#  2) This script will list the date/time each time the power state changes to either
#     the stdout, a file, or both.
#
# TrueNAS SCALE (Debian) properly supports ATSP so this script is not required.

# Variables:
declare -a nvme_array
declare -a nvme_current_power_state
declare -a nvme_lowest_power_state
declare -i minimum_power_state
declare -a nvme_old_power_state
declare -i delay_seconds
declare -a nvme_maximum_current	# used to calculate total watts
declare -a nvme_total_current	# Used to calculate total watts
declare -a nvme_total_time	# Used to calculate total watts

########## USER DEFINABLE VARIABLES ##########
minimum_power_state=99		# The minimum power setting script is allowed to be changed to. 99=Ignore
check_periodicity=3		# 3600 seconds = 1 hour, 900 = 15 minutes.  How often to verify the power state of the nvme's.
delay_seconds=5			# Number of seconds to delay before forcing back to minimum power state.
exit_after=82800		# 82800 seconds = 23 hours.  How long to remain in a loop before exiting.
				# - If run from CRON once every 24 hours, set to exit before the CRON executes again.
output_to_stdout="true"		# This will list all actions occuring
output_to_file="true"		# This will output to a file any actions occuring (same as to the stdout)
output_on_change_only="true"	# Output message ONLY on a power status change.
max_output_file_size=100	# Maximum number of lines to retain in the text file before deleting the older entries.
output_file="/tmp/nvme_power_state_log.txt"   # The file to output to. Change to a drive path to retain the data through reboot.
#############################################

SECONDS=0
# Let's check and exit if this is not FreeBSD.
softver=$(uname -s)
if [[ $softver != "FreeBSD" ]]; then
	echo "Not FreeBSD/TrueNAS CORE... Exiting"
	exit 1
fi

# Define Functions

########## GET NVMe DRIVES ##########
# Get listing of the SMART NVMe's and put into $smartdrivesNVM variable

get_smartNVM_listings () {

	smartdrivesNVM=$(for drive in $(sysctl -n kern.disks); do
		if [ "$(smartctl -i /dev/"${drive}" | grep "NVM")" ]; then printf "%s " "${drive}"; fi
	done | awk '{for (i=NF; i!=0 ; i--) print $i }' | tr ' ' '\n' | sort | tr '\n' ' ')

	### Convert nvdx to nvmexx in smartdrivesNVM ###
	smartdrivesNVM=$( echo "$smartdrivesNVM" | sed 's/nvd/nvme/g' )

	### Convert nvme0n1 to nvme0, or nvme34n1 to nvme34 to make compatible with Smartmontools 7.4
	### Maybe this can go away with smartmontools v7.5?
	for smartdrivesnvme in $smartdrivesNVM; do
		nvme_drive=$nvme_drive$(echo "nvme"$(echo $smartdrivesnvme | sed -r 's#^nvme##' | cut -d 'n' -f 1)" ")
		# Let's assign the nvme number to an array variable. [0]=0, [1]=1, only if the drives exists will it increment
		# There is an assumption that each nvme drive will increment by +1.  Ex: 0, 1, 2, 3, 4, 5...
		nvme_number=$(echo $smartdrivesnvme | sed -r 's#^nvme##' | cut -d 'n' -f 1)
		nvme_array[$nvme_number]+=$nvme_number
	done
	smartdrivesNVM=$nvme_drive

	# Call Sort Routine with the drive string.
	if [[ "$smartdrivesNVM" != "" ]]; then
		sort_list=$smartdrivesNVM
		sort_drives
		smartdrivesNVM=$sort_list
	fi
	if [[ $output_to_stdout == "true" ]]; then echo "List of drives="$smartdrivesNVM; fi
	if [[ $output_to_file == "true" ]]; then echo "List of drives="$smartdrivesNVM >> $output_file; fi
	}


########## SORT DRIVES ROUTINE ##########
# Sort drives into alphabetical order.
sort_drives () {
	sort_list=$(for i in `echo $sort_list`; do
	echo "$i"
	done | sort -V)
	}

######### Purge Output File ##########
purge_output_file () {

	line_count=$(wc -l < $output_file)	# Get the number of lines
	lines_to_delete=$(( $line_count-$max_output_file_size ))
	if [[ "$lines_to_delete" -gt "10" ]]; then $(echo sed -i -e 1,${lines_to_delete}d $output_file); fi
	}

########## Display Help ##########
display_help () {
	clear
	echo "NVMe Power Adjustment Script"
	echo " "
	echo "Due to FreeBSD (TrueNAS CORE) not supporting Autonomous Power State Transition (APST)"
	echo "this script will identify the minimum power state of each nvme and set the minimum"
	echo "power state value."
	echo " "
	echo "If the nvme does require a higher power state for the work that it is conducting,"
	echo "the nvme device will change itself automatically apart from this script however"
	echo "it will not return itself to the previous lower power state."
	echo " "
	echo "How to use this script:"
	echo "There are 9 user changable variables, they are:"
	echo " "
	echo " minimum_power_state - This defines the minimum power state you want the nvme drives to"
	echo "enter. 99=Disable and the minimum state of the nvme will be used. Use a higher state"
	echo "to limit how low the power state is allowed to go. If you set this value to '3',"
	echo "then the nvme drives will not be able to reach power state 4.  The purpose of this"
	echo "value is to improve responce times from the nvme drives.  The default value should"
	echo "be fine for most people."
	echo " "
	echo " check_periodicity - This value tells the script how long to delay before checking"
	echo "the nvme drives if they are at a new power state.  You may not desire for the script"
	echo "to continuiously keep the NAS busy so I'd recommend setting this value to '60' or"
	echo "higher.  For testing this value is nice to change to a low number."
	echo " "
	echo " delay_seconds - This value is used when a power status change has been detected."
	echo "It is the number of seconds to delay returning the power state back to low power state."
	echo "The purpose is to leave the nvme drive in the higher power state so it may complete"
	echo "whatever task caused the power state to increase.  '5' seconds it the default however"
	echo "larger values like '60' may be more appropriate."
	echo " "
	echo " exit_after - This value is used to force the script to exit after a specific period of"
	echo "time.  The script must run in a continuios loop in order to check the nvme power state."
	echo "There is an assumption that this script would be launched from a CRON job once every"
	echo "24 hours.  The default value of '82800' represents 23 hours. If you increase this value"
	echo "ensure that you do not have a running instance of the script when another instant launches."
	echo " "
	echo " output_to_stdout / output_to_file - These two settings when 'true' will write to the"
	echo "screen and/or log file.  Set these to 'false' and the script will be silent."
	echo " "
	echo " output_on_change_only - This when 'true' will output only changed data on the screen"
	echo "or log file."
	echo " "
	echo " max_output_file_size - This value will set the maximum log file line length."
	echo "This action is only performed when the script starts and will purge the older lines."
	echo " "
	echo " output_file - This is the location of the log file. The default is the '/tmp/'"
	echo "directory.  This is volitle and is gone during a reboot of TrueNAS.  It is recommended"
	echo "that if you desire to retain the log file, change this value to a location on a drive."
	echo "The log file will contain a date/time stamp when the power state changed to include"
	echo "the previous power state and the new power state.  There will also be a delay message."
	echo "And then another date/time entry and message changing the power state back to the"
	echo "lower power state."
	echo " "
	echo " "
	echo "Command Line Switches:"
	echo "-h	Displays this information"
	echo "-d	Deletes the log file (clean slate)"
	echo " "
	exit 0
	}


# START THE PROGRAM

# Delete file option
# Use the -d switch to delete the output file upon runnign the script.
if [[ $1 == "-d" ]]; then
	if test -e "$output_file"; then
		if [[ $output_to_stdout == "true" ]]; then echo "Deleting Output File"; fi
		rm $output_file
	else
		if [[ $output_to_stdout == "true" ]]; then echo "Output File - "$output_file" does not exist."; fi
	fi
fi

if [[ $1 == "-h" ]]; then
	display_help
fi

if [[ $1 != "-d" && $1 != "" && $1 != "-h" ]]; then echo "Invalid switch.  Use -h for help."; exit 1; fi

if [[ $output_to_stdout == "true" && $output_to_file == "true" ]]; then
	echo $(date)" Recording to both stdout and "$output_file
	echo $(date)" Recording to both stdout and "$output_file >> $output_file
fi
if [[ $output_to_stdout == "true" && $output_to_file != "true" ]]; then
	echo $(date)" Recording Only to stdout"
	#echo "Recording Only to stdout" >> $output_file
fi
if [[ $output_to_stdout != "true" && $output_to_file == "true" ]]; then
	#echo "Recording Only to output file "$output_file
	echo $(date)" Recording Only to output file "$output_file >> $output_file
fi

purge_output_file

get_smartNVM_listings

# Set Initial Power State on Each nvme

for drive in $smartdrivesNVM; do
	# Check lowest power state and set it.
	if [[ $output_to_stdout == "true" ]]; then echo "Initial power states for each nvme drive:"; fi
	if [[ $output_to_file == "true" ]]; then echo "Initial power state for each nvme drive:" >> $output_file; fi
	lowest_power_state=$(($(nvmecontrol power -l $drive | tail -1 | cut -c1-2)))
	current_power_state=$(($(nvmecontrol power $drive | rev | cut -c1-2 | rev)))
	nvme_number=$(echo $smartdrivesnvme | sed -r 's#^nvme##' | cut -d 'n' -f 1)
	eval "nvme_lowest_power_state[$nvme_number]+=$lowest_power_state"
	eval "nvme_current_power_state[$nvme_number]+=$current_power_state"

# Lets grab the wattage for each mode
#	nvme_maximum_current[$nvme_number]+=$(($(nvmecontrol power -l $drive | grep '0:' | grep 'W' | cut -d ':' -f 2 | cut -d 'W' -f 1)))
# echo "max_curr="$nvme_maximum_cur
#	nvme_maximum_current[$nvme_number]+=$(echo $nvme_maximum_cur)
# nvme_maximum_current[$nvme_number]+="3"
# echo "max_power_state="$nvme_maximum_current[$nvme_number]

	if [[ $output_to_stdout == "true" ]]; then echo "Drive "$drive" current power state "${nvme_current_power_state[$nvme_number]}; fi
	if [[ $output_to_stdout == "true" ]]; then echo "Drive "$drive" lowest power state "${nvme_lowest_power_state[$nvme_number]}; fi
	if [[ $output_to_file == "true" ]]; then echo "Drive "$drive" current power state "${nvme_current_power_state[$nvme_number]} >> $output_file; fi
	if [[ $output_to_file == "true" ]]; then echo "Drive "$drive" lowest power state "${nvme_lowest_power_state[$nvme_number]} >> $output_file; fi
done

for (( ; ; )); do		# Infinate loop, Check Power State on all drives and Lower if Needed

	# Cycle through each drive
	for drive in $smartdrivesNVM; do

			nvme_old_power_state[$nvme_number]=$(($(nvmecontrol power $drive | rev | cut -c1-2 | rev)))  # Setup initial value so loop does not exit on first run.

			if [[ $output_to_stdout == "true" && $output_on_change_only != "true" ]]; then echo $(date)" "$drive" Power State = "${nvme_old_power_state[$nvme_number]}; fi
			if [[ $output_to_file == "true" && $output_on_change_only != "true" ]]; then echo $(date)" "$drive" Power State = "${nvme_old_power_state[$nvme_number]} >> $output_file; fi

			power_state=$(($(nvmecontrol power $drive | rev | cut -c1-2 | rev)))

			if [[ $minimum_power_state != 99 ]]; then
				new_power_state=$minimum_power_state
			else
				new_power_state=${nvme_lowest_power_state[$nvme_drive]}
			fi

			if [[ $power_state != $new_power_state ]]; then
				if [[ $output_to_stdout == "true" ]]; then
					echo $(date)" -> Drive "$drive" now in power state "${nvme_old_power_state[$nvme_number]}"."
					echo "Pausing "$delay_seconds" seconds to allow for work to be completed."
					sleep $delay_seconds
					echo $(date)" -> Setting power state "$new_power_state
					echo "-------------------------------------"
				fi

				if [[ $output_to_file == "true" ]]; then
					echo $(date)" -> Drive "$drive" now in power state "${nvme_old_power_state[$nvme_number]}"." >> $output_file
					echo "Pausing "$delay_seconds" seconds to allow for work to be completed." >> $output_file
					sleep $delay_seconds
					echo $(date)" -> Setting power state "$new_power_state >> $output_file
					echo "-------------------------------------" >> $output_file
				fi
				nvmecontrol power -p $new_power_state $drive
				power_state=$(($(nvmecontrol power $drive | rev | cut -c1-2 | rev))) # Read the power state again to be set to old power state.
			fi

			duration=$SECONDS

			if [[ $duration -ge $exit_after ]]; then
				if [[ $output_to_stdout == "true" ]]; then echo "Exiting Normally after exit timer completed."; fi
				if [[ $output_to_stdout == "true" ]]; then echo "Exiting Normally after exit timer completed." >> $output_file; fi
				exit 0
			fi

		eval "nvme_old_power_state[$nvme_number]=$power_state"
	done

	sleep $check_periodicity	# Wait for periodicity to expire.  Good to not always have the script looping eating up CPU cycles.
done