#!/bin/bash
#
LANG="en_US.UTF-8"
#
# NVME LOW POWER STATE SCRIPT v0.4 2023_12_22
# (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.

#
# 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 'minimum_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 -i current_power_state
declare -i lowest_power_state
declare -i minimum_power_state
declare -i delay_seconds



########## 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.
run_continuously="false"	# This when 'true' will ignore 'exit_after' and run the script and never exit.
				# - This is good if you run the scrip upon startup and never again.
output_to_stdout="true"		# This will list all actions occurring
output_to_file="true"		# This will output to a file any actions occurring (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 drive in $smartdrivesNVM; do
		### Get Power States
		nvme_drive=$(echo "nvme"$(echo $smartdrivesnvme | sed -r 's#^nvme##' | cut -d 'n' -f 1)" ")
		lowest_power_state=$(($(nvmecontrol power -l $drive | tail -1 | cut -c1-2)))
		current_power_state=$(($(nvmecontrol power $drive | rev | cut -c1-2 | rev)))
	done
	}


########## 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 changeable 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 response 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 continuously 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 continuous 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 "Note that 'run_continuously' must be 'false' to enable this feature."
	echo " "
	echo " run_continuously - This value is used to either allow the script to run continuously"
	echo "or exit based on the 'exit_after' time period.  The default value is 'false'."
	echo "Use this feature if you plan to run the script at bootstrap time vice in a reoccuring"
	echo "CRON job."
	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 volatile 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), then runs the script."
	echo " "
	exit 0
	}

# START THE PROGRAM

# Delete file option
# Use the -d switch to delete the output file upon running 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"
fi
if [[ $output_to_stdout != "true" && $output_to_file == "true" ]]; then
	echo $(date)" Recording Only to output file "$output_file >> $output_file
fi

purge_output_file

get_smartNVM_listings

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

	# Cycle through each drive
	for drive in $smartdrivesNVM; do
		change="false"
		# Check $drive for lowest and current power state
		lowest_power_state=$(($(nvmecontrol power -l $drive | tail -1 | cut -c1-2)))
		current_power_state=$(($(nvmecontrol power $drive | rev | cut -c1-2 | rev)))

################# ADD WATTS TO POWER STATE OUTPUT AND START TO ADD UP SECONDS BETWEEN LAST CHANGE AND NEW CHANGE, IF POSSIBLE ########

		# Compare to $minimum_power_state, if current is lower then set power level to minimum power state
		if [[ "$minimum_power_state" == "99" ]]; then
			if [[ $current_power_state -ne $lowest_power_state ]]; then
				# Set to lowest power state
			#	nvmecontrol power -p $lowest_power_state $drive
				power_state=$lowest_power_state
				change="true"
			fi
		else
			if [[ "$current_power_state" -le "$minimum_power_state" ]]; then
				# Set to $minimum_power_state
			#	nvmecontrol power -p $mimimum_power_state $drive
				power_state=$minimum_power_state
				change="true"
			fi
		fi
		
		#new_power_state=$(($(nvmecontrol power $drive | rev | cut -c1-2 | rev)))

		if [[ $change == "true" ]]; then
			if [[ $output_to_stdout == "true" ]] && [[ $output_on_change_only == "true" ]]; then
					echo $(date)" -> Drive "$drive" detected in power state "$current_power_state"."
					echo "        Delaying "$delay_seconds" seconds."
			fi
			if [[ $output_to_file == "true" ]] && [[ $output_on_change_only == "true" ]]; then
					echo $(date)" -> Drive "$drive" detected in power state "$current_power_state"." >> $output_file
					echo "        Delaying "$delay_seconds" seconds." >> $output_file
			fi

		fi

		sleep $delay_seconds

		if [[ $change == "true" ]]; then
			if [[ $output_to_stdout == "true" ]] && [[ $output_on_change_only == "true" ]]; then
					echo $(date)" -> Attempting to set power state "$power_state
					nvmecontrol power -p $power_state $drive
					new_power_state=$(($(nvmecontrol power $drive | rev | cut -c1-2 | rev)))
					echo $(date)" -> Actual power state set is "$new_power_state
					echo "-------------------------------------"
			fi
			if [[ $output_to_file == "true" ]] && [[ $output_on_change_only == "true" ]]; then
					echo $(date)" -> Attempting to set power state "$power_state >> $output_file
					nvmecontrol power -p $power_state $drive
					new_power_state=$(($(nvmecontrol power $drive | rev | cut -c1-2 | rev)))
					echo $(date)" -> Actual power state set is "$new_power_state >> $output_file
					echo "-------------------------------------" >> $output_file
			fi

		fi




			duration=$SECONDS
			if [[ $run_continuously == "false" ]]; then		# Set user changable variable to "true" to run all the time.
				runtime_hours=0
				runtime_minutes=0
				runtime_seconds=0
				if [[ $duration -ge $exit_after ]]; then
					if [[ $duration -gt 3599 ]]; then	# Subtract hours
						runtime_hours=$(( $duration/3600 ))
						duration=$(( $duration-$runtime_hours*3600 ))
					fi
					if [[ $duration -lt 3600 ]]; then	# Subtract minutes
						runtime_minutes=$(( $duration/60 ))
						duration=$(( $duration-$runtime_minutes*60 ))
					fi
					#if [[ $duration -lt 60 ]]; then	# Subtract minutes
						runtime_seconds=$duration
					#	duration=$(( $duration-$runtime_seconds*60 ))
					#fi
					
					if [[ $output_to_stdout == "true" ]]; then echo "Exiting Normally after exit timer completed after "$runtime_hours":"$runtime_minutes":"$runtime_seconds; fi
					if [[ $output_to_stdout == "true" ]]; then echo "Exiting Normally after exit timer completed after "$runtime_hours":"$runtime_minutes":"$runtime_seconds" minutes." >> $output_file; fi
					exit 0
				fi
			fi
	done

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

done