#!/usr/local/bin/cbsd
#v11.2.1
CBSDMODULE="jail"
CIXARG=""
. ${distsharedir}/rctl.conf
CIXOPTARG="alljails display emulator header human jname mode node prometheus shownode sort ${RCTL} ${RCTL_EXTRA} quiet"
MYDESC="Set or flush resource limit for jail"
ADDHELP="
${H3_COLOR}Description${N0_COLOR}:

The 'jrctl' script used to view and manage jail limits.
By default, 'cbsd jrctl' shows the current usage and limits per jail (if installed).
Most limits depend on RACCT features, make sure the host is booted with

kern.racct.enable=1

settings in /boot/loader.conf.

${H3_COLOR}Options${N0_COLOR}:

 ${N2_COLOR}alljails=${N0_COLOR}     - when '1' - get jaillist from remote node, default: '0';
 ${N2_COLOR}display=${N0_COLOR}      - list by comma for column. Default: jname,memoryuse,maxproc,openfiles,vmemoryuse,swapuse,pcpu,fsquota,nice,bw;
 ${N2_COLOR}emulator=${N0_COLOR}     - [jail,bhyve,xen] - show only emulator engine, default - all;
 ${N2_COLOR}header=${N0_COLOR}       - 0 - don't print header;
 ${N2_COLOR}human=${N0_COLOR}        - 0 - don't convert bytes to human readable form, default: '1';
 ${N2_COLOR}jname=${N0_COLOR}        - show only specified jail;
 ${N2_COLOR}mode=${N0_COLOR}         - set (apply), unset, show (show current state) or get (show limits);
 ${N2_COLOR}node=${N0_COLOR}         - only for current node;
 ${N2_COLOR}prometheus=${N0_COLOR}   - 1 = prometheus metric mode output;
 ${N2_COLOR}shownode=${N0_COLOR}     - show nodename for jails;
 ${N2_COLOR}sort=${N0_COLOR}         - sort by column name (from display list);
 ${N2_COLOR}rctl params:${N0_COLOR}  - ${RCTL} ${RCTL_EXTRA};
 ${N2_COLOR}quiet=${N0_COLOR}        - 0,1: be quiet, dont output verbose message;

${H3_COLOR}Examples${N0_COLOR}:

 # cbsd jrctl
 # cbsd jrctl human=0
 # cbsd jrctl display=jname,openfiles,memoryuse,fsquota sort=memoryuse
 # cbsd jrctl mode=show jname=xxx prometheus=1
 # cbsd jrctl jname=jail1

${H3_COLOR}See also${N0_COLOR}:

 cbsd jrctl-tui --help
 cbsd jget --help

"
EXTHELP="wf_jrctl"

. ${subrdir}/nc.subr
. ${strings}
. ${tools}
. ${distsharedir}/rctl.conf
. ${subrdir}/jrctl.subr

sort=
quiet=
jname=
ojname=
cixinit

: "${quiet:=1}"

show_header()
{
	local _header="${H1_COLOR}${BOLD}${myheader}${N0_COLOR}"
	[ ${header} -ne 0 ] && ${ECHO} "${_header}"
}

# if $1 = "Unregister" then overwrite status to "Unregister"
populate_output_data()
{
	local _formfile= _cur_val= _val= _emulator=

	cbsdsqlro_vars local "SELECT emulator FROM jails WHERE jname='${jname}'" _emulator

	case "${_emulator}" in
		jail)
			_formfile="${jailsysdir}/${jname}/helpers/jrctl.sqlite"
			;;
		bhyve)
			_formfile="${jailsysdir}/${jname}/helpers/brctl.sqlite"
			;;
	esac

	[ ! -r "${_formfile}" ] && return 0

	# RACCT
	eval $( ${RCTL_CMD} -u jail:${jname} )

	#populate values for in output string
	for _i in ${mydisplay}; do

		_val="0"

		case "${_i}" in
			jname)
				_val="${jname}"
				if [ -z "${_status}" ]; then
					_status="${_val}"
				else
					_status="${_status} ${_val}"
				fi
				;;
			*)
				if [ -r "${_formfile}" ]; then
					cbsdsqlro_vars ${_formfile} "SELECT cur FROM forms WHERE param='${_i}' LIMIT 1" _val
					[ -z "${_val}" ] && _val="0"
				fi

				if rctl_humanize ${_i}; then
					if conv2human ${_val}; then
						_val="${convval}"
					fi
				fi
				if [ "${jid}" = "0" ]; then
					_cur_val="0"
				else
					if [ "${_i}" = "memoryuse" ]; then
						_cur_val="${memoryuse}"
					elif [ "${_i}" = "maxproc" ]; then
						_cur_val="${maxproc}"
					elif [ "${_i}" = "openfiles" ]; then
						_cur_val="${openfiles}"
					elif [ "${_i}" = "vmemoryuse" ]; then
						_cur_val="${vmemoryuse}"
					elif [ "${_i}" = "swapuse" ]; then
						_cur_val="${swapuse}"
					elif [ "${_i}" = "pcpu" ]; then
						_cur_val="${pcpu}"
					elif [ "${_i}" = "fsquota" ]; then
						if [ "${zfsfeat}" = "1" ]; then
							. ${system}			# is_mounted func
							data="${jaildatadir}/${jname}-${jaildatapref}"
							. ${subrdir}/jfs.subr
							if ! zfsmnt ${data}; then
								capture ZPOOL ${ZFS_CMD} get -Ho value name ${data}
								if [ -n "${ZPOOL}" ]; then
									capture _cur_val ${ZFS_CMD} get -Hp -o value refer ${ZPOOL}
									[ -z "${_cur_val}" ] && cur_val="0"
								else
									_cur_val="0"
								fi
							else
								_cur_val="0"
							fi
						else
							_cur_val="0"
						fi
					elif [ "${_i}" = "nice" ]; then
						_cur_val="${nice}"
					elif [ "${_i}" = "bw" ]; then
						_cur_val="${bw}"
					else
						_cur_val="0"
					fi
				fi

				if [ -z "${_status}" ]; then
					_status="${_cur_val}/${_val}"
				else
					_status="${_status} ${_cur_val}/${_val}"
				fi
				;;
		esac
	done
}

# $1 - which file from. Eg: local
show_jaildata_from_sql()
{
	local _i= _sqlq=
	local _res=
	local jname=

	# set sqlfile for ". rcconf" including
	if [ -n "${1}" ]; then
		sqlfile="${1}"
	else
		sqlfile="local"
	fi

	if [ -n "${ojname}" ]; then
		cbsdsqlro_vars ${sqlfile} "SELECT jname FROM jails WHERE jname='${ojname}'" _res
	else
		if [ -n "${emulator}" ]; then
			cbsdsqlro_vars ${sqlfile} "SELECT jname FROM jails WHERE emulator = '${emulator}' ORDER BY jname ASC" _res
		else
			cbsdsqlro_vars ${sqlfile} "SELECT jname FROM jails WHERE emulator = 'bhyve' OR emulator = 'jail' ORDER BY jname ASC" _res
		fi
	fi

	for jname in ${_res}; do
		_status=""
		. ${subrdir}/rctl.subr
		. ${subrdir}/rcconf.subr
		populate_output_data
		printf "${N0_COLOR}" # for column sort
		if [ "${jid}" != "0" ]; then
			printf "${N2_COLOR}"
		else
			printf "${N4_COLOR}"
		fi
		printf "${_status}"
		printf "${N0_COLOR}\n"
	done > /tmp/jrctl.$$

	if [ -n "${sort}" ]; then
		sort=$( echo ${sort} | ${TR_CMD} '[:upper:]' '[:lower:]' )
	else
		reverse=
	fi
	sort_column=1
	column_index=0

	# get column index for sort
	for _i in ${mydisplay}; do
		column_index=$(( column_index + 1 ))
		if [ "${_i}" = "${sort}" ]; then
			sort_column="${column_index}"
			reverse="-r"
			break
		fi
	done

	${SORT_CMD} -n -k${sort_column} ${reverse} /tmp/jrctl.$$ | while read ${mydisplay}; do
		for _i in ${mydisplay}; do
			_val="0"
			eval _val="\$${_i}"

			if [ "${human}" = "1" ]; then
				case "${_i}" in
					memoryuse|vmemoryuse|swapuse|fsquota)
						p1=${_val%%/*}
						p2=${_val##*/}
						if conv2human ${p1}; then
							p1="${convval}"
						fi
						if conv2human ${p2}; then
							p2="${convval}"
						fi
						[ -z "${p1}" ] && p1="0"
						[ -z "${p2}" ] && p2="0"
						printf "${p1}/${p2} "
						;;
					*)
						printf "${_val} "
						;;
				esac
			else
				printf "${_val} "
			fi
		done
		printf "${N0_COLOR}\n"
	done

	${RM_CMD} -f /tmp/jrctl.$$
}


show_local()
{
	local _errcode _status
	show_header
	show_jaildata_from_sql local
}


show_remote()
{
	show_header

	if [ -z "${node}" ]; then
		capture node node mode=list header=0 allinfo=0
	fi

	for _n in $node; do
		nodename="${_n}"
		show_jaildata_from_sql ${_n}
	done
}

show_jails()
{
	if [ -n "${node}" ]; then
		show_remote
		exit
	fi

	if [ "${alljails}" = "1" ]; then
		show_local
		header=0
		show_remote
	else
		show_local
	fi
}

set_limit()
{
	local LIMITS="${jailsysdir}/${jname}/jail.limits"
	local CPU="${jailsysdir}/${jname}/cpu"
	local _formfile= _mymodule=
	local _val= _valnew=
	local i= DATA= _ret= _cpuset_args=

	case "${emulator}" in
		jail|qemu*)
			_formfile="${jailsysdir}/${jname}/helpers/jrctl.sqlite"
			_mymodule="jrctl"
			;;
		bhyve)
			_formfile="${jailsysdir}/${jname}/helpers/brctl.sqlite"
			_mymodule="brctl"
			;;
	esac

	if [ ! -r "${_formfile}" ]; then
		${ECHO} "${N1_COLOR}${CIX_APP}: no such formfile: ${_formfile}${N0_COLOR}"
		forms module=${_mymodule} jname=${jname} inter=0
		if [ ! -r "${_formfile}" ]; then
			${ECHO} "${N1_COLOR}No such ${_formfile}${N0_COLOR}"
			return 1
		fi
	fi

	${TRUNCATE_CMD} -s0 ${LIMITS}

	SQL_FORMFILE_QUERY=

	for i in ${RCTL}; do
		# rctl params update if args sets
		_val=

		eval _val="\$rctl_${i}"
		if [ -n "${_val}" ]; then
			${ECHO} "${N1_COLOR}update rctl for: ${N2_COLOR}${i}${N0_COLOR}"
			#cbsdsqlrw ${_formfile} "UPDATE forms SET new='${_val}',cur='${_val}' WHERE param='${i}'" 2>/dev/null
			SQL_FORMFILE_QUERY="${SQL_FORMFILE_QUERY}; UPDATE forms SET new='${_val}',cur='${_val}' WHERE param='${i}'"
		fi

		#_valnew=$( cbsdsqlro ${_formfile} "SELECT new FROM forms WHERE param='${i}' LIMIT 1" 2>/dev/null )
		cbsdsqlro_vars ${_formfile} "SELECT new FROM forms WHERE param='${i}' LIMIT 1" _valnew

		if [ -z "${_valnew}" ]; then
			#_val=$( cbsdsqlro ${_formfile} "SELECT cur FROM forms WHERE param='${i}' LIMIT 1" 2>/dev/null )
			cbsdsqlro_vars ${_formfile} "SELECT cur FROM forms WHERE param='${i}' LIMIT 1" _val
		else
			# apply 'cur' values and truncate 'new'
			_val="${_valnew}"
			#cbsdsqlrw ${_formfile} "UPDATE forms SET cur='${_val}',new='' WHERE param='${i}'" 2>/dev/null
			SQL_FORMFILE_QUERY="${SQL_FORMFILE_QUERY}; UPDATE forms SET cur='${_val}',new='' WHERE param='${i}'"
		fi
		[ -z "${_val}" ] && _val=0
		[ "${_val}" = "0" ] && continue
		echo "${i}:deny=${_val}" >> ${LIMITS}
	done

	for i in ${RCTL_EXTRA}; do
		# rctl params update if args sets
		_val=
		eval _val="\$rctl_${i}"
		if [ -n "${_val}" ]; then
			${ECHO} "${N1_COLOR}update rctl-extra for: ${N2_COLOR}${i}${N0_COLOR}"
			#cbsdsqlrw ${_formfile} "UPDATE forms SET new='${_val}',cur='${_val}' WHERE param='${i}'" 2>/dev/null
			SQL_FORMFILE_QUERY="${SQL_FORMFILE_QUERY}; UPDATE forms SET new='${_val}',cur='${_val}' WHERE param='${i}'"
		fi
		_val=
		#_valnew=$( cbsdsqlro ${_formfile} "SELECT new FROM forms WHERE param='${i}' LIMIT 1" 2>/dev/null )
		cbsdsqlro_vars ${_formfile} "SELECT new FROM forms WHERE param='${i}' LIMIT 1" _valnew

		if [ -z "${_valnew}" ]; then
			#_val=$( cbsdsqlro ${_formfile} "SELECT cur FROM forms WHERE param='${i}' LIMIT 1" 2>/dev/null )
			cbsdsqlro_vars ${_formfile} "SELECT cur FROM forms WHERE param='${i}' LIMIT 1" _val
		else
			# apply 'cur' values and truncate 'new'
			_val="${_valnew}"
			#cbsdsqlrw ${_formfile} "UPDATE forms SET cur='${_val}',new='' WHERE param='${i}'" 2>/dev/null
			SQL_FORMFILE_QUERY="${SQL_FORMFILE_QUERY}; UPDATE forms SET cur='${_val}',new='' WHERE param='${i}'"
		fi

		[ "${_val}" = "0" -o -z "${_val}" ] && continue

		case "${i}" in
			nice)
				case "${emulator}" in
					jail)
						jrenice jname=${jname}
						;;
					bhyve)
						brenice jname=${jname}
						;;
				esac
				;;
			fsquota)
				. ${subrdir}/zfs.subr
				if [ ${zfsfeat} -ne 1 ]; then
					${ECHO} "${N1_COLOR}fsquota supported only on ZFS. Currently zfsfeat is: ${N2_COLOR}${zfsfeat}${N1_COLOR}. Skip set fsquota${N0_COLOR}"
					break
				fi

				zfsmnt ${data}
				if [ $? -eq 0 ]; then
					${ECHO} "${N1_COLOR}${data} is not ZFS separated fileset. Skip set fsquota${N0_COLOR}"
					break
				fi

				DATA=$( ${ZFS_CMD} get -Ho value name ${data} )
				if [ -z "${DATA}" ]; then
					${ECHO} "${N1_COLOR}Unable to get dataset name for: ${N2_COLOR}${jaildatadir}${N1_COLOR}. Skip set fsquota${N0_COLOR}"
					break
				fi

				${ECHO} "${N1_COLOR}Set new quota for ${jname}:${data}: ${N2_COLOR}${_val}${N0_COLOR}"
				${ZFS_CMD} set quota=${_val} ${DATA}
				;;
			bw)
				true
				;;
			cpu)
				if [ "${_val}" != "0" ]; then
					if [ ${_val} -ge ${ncpu} ]; then
						${ECHO} "${N1_COLOR}Invalid CPU limit settings: number of env CPU >= HOST CPU (${_val} >= ${ncpu})${N0_COLOR}"
					else
						${ECHO} "${N1_COLOR}Set CPU limit ${jname}: ${N2_COLOR}${_val}${N0_COLOR}"
					fi
					_cpuset_args=
					for ((i=1; i<=${_val}; i++)); do
						jrctl_get_next_cpuset
						jrctl_increment_cpuset ${JRCTL_CPUSET_CORE}
						if [ -z "${_cpuset_args}" ]; then
							_cpuset_args="${JRCTL_CPUSET_CORE}"
						else
							_cpuset_args="${_cpuset_args},${JRCTL_CPUSET_CORE}"
						fi
					done
					echo "cpuset=\"-c -l ${_cpuset_args}\"" > ${CPU}
				else
					[ -r ${CPU} ] && ${RM_CMD} -f ${CPU}
				fi
				;;
		esac
	done

	[ -n "${SQL_FORMFILE_QUERY}" ] && cbsdsqlrw ${_formfile} ${SQL_FORMFILE_QUERY}

	if [ -f "${LIMITS}"  ]; then
		[ ${quiet} -ne 1 ] && printf "${N1_COLOR}${mode} resource limit: [ ${N2_COLOR}"
		${CAT_CMD} ${LIMITS} | while read _p; do
			case ":${_p}" in
				:#* | :)
					continue
				;;
			esac
			_str="${RCTL_CMD} -a jail:${jname}:$_p"
			_out=$( eval ${_str} )
			_ret=$?
			if [ ${_ret} -eq 0 ]; then
				[ ${quiet} -ne 1 ] && printf "${_p} "
			fi
		done
		_ret=$?
		[ ${quiet} -ne 1 ] && printf "${N1_COLOR}]${N0_COLOR}\n"
	fi

	return ${_ret}
}


jail_rctl()
{
	local _ret=

	local LIMITS="${jailsysdir}/${jname}/jail.limits"

	case "${mode}" in
		"set")
			set_limit
			_ret=$?
			[ ${_ret} -ne 0 ] && err ${_ret} "${N1_COLOR}${CIX_APP}: set_limit error${N0_COLOR}"
			;;
		"unset")
			[ -z "${jname}" ] && err 1 "${N1_COLOR}jname= must be set${N0_COLOR}"
			${RCTL_CMD} -r jail:${jname}
			return 0
			;;
		"show")
			if [ -n "${jname}" ]; then
				#emulator=$( cbsdsqlro local "SELECT emulator FROM jails WHERE jname='${jname}'" )
				cbsdsqlro_vars /usr/jails/var/db/local.sqlite "SELECT emulator FROM jails WHERE jname='${jname}'" emulator
				if [ "${prometheus}" = "0" ]; then
					case "${emulator}" in
						jail)
							${RCTL_CMD} -hu jail:${jname}
							;;
						bhyve)
							#pid=$( cbsdsqlro local "SELECT jid FROM jails WHERE jname='${jname}'" )
							#vm_cpus=$( cbsdsqlro local "SELECT vm_cpus FROM bhyve WHERE jname='${jname}'" )
							cbsdsqlro_vars local "SELECT jails.jid, bhyve.vm_cpus FROM jails JOIN bhyve ON jails.jname = bhyve.jname WHERE jails.jname = '${jname}'" pid vm_cpus
							eval $( ${RCTL_CMD} -hu process:${pid} )
							for res in ${RCTL}; do
								eval _val=\$${res}
								if [ "${res}" = "pcpu" -a -n "${_val}" ]; then
									# take into account multi-core guest in pcpu value
									# On multi-core guest we need to: pcpu / vm_cpus
									[ ${vm_cpus} -gt 1 ] && _val=$(( _val / vm_cpus ))
								fi
								[ -n "${_val}" ] && echo "${res}=${_val}"
							done
							;;
					esac
				else
					case "${emulator}" in
						jail)
							eval $( ${RCTL_CMD} -u jail:${jname} )
							;;
						bhyve)
							#pid=$( cbsdsqlro local "SELECT jid FROM jails WHERE jname='${jname}'" )
							#vm_cpus=$( cbsdsqlro local "SELECT vm_cpus FROM bhyve WHERE jname='${jname}'" )
							cbsdsqlro_vars local "SELECT jails.jid, bhyve.vm_cpus FROM jails JOIN bhyve ON jails.jname = bhyve.jname WHERE jails.jname = '${jname}'" pid vm_cpus
							eval $( ${RCTL_CMD} -u process:${pid} )
							# take into account multi-core guest in pcpu value
							# On multi-core guest we need to: pcpu / vm_cpus
							[ ${vm_cpus} -gt 1 ] && pcpu=$(( pcpu / vm_cpus ))
							;;
					esac

					for res in ${RCTL}; do
						eval _mydesc="\$${res}_desc"
						eval _val=\$${res}
						echo "# HELP ${emulator}_${res} ${_mydesc} ${jname}"
						echo "${emulator}_${res}{name=\"${jname}\"} ${_val:-0}"
					done
				fi
			else
				if [ "${prometheus}" = "0" ]; then
					for i in $( jorder ); do
						echo "--- ${i} ---"
						${RCTL_CMD} -hu jail:${i}
					done
					for i in $( border ); do
						echo "--- ${i} ---"
						pid=$( cbsdsqlro local "SELECT jid FROM jails WHERE jname='${i}'" )
						vm_cpus=$( cbsdsqlro local "SELECT vm_cpus FROM bhyve WHERE jname='${i}'" )
						eval $( ${RCTL_CMD} -hu process:${pid} )
						for res in ${RCTL}; do
							eval _val=\$${res}
							if [ "${res}" = "pcpu" -a -n "${_val}" ]; then
								# take into account multi-core guest in pcpu value
								# On multi-core guest we need to: pcpu / vm_cpus
								[ ${vm_cpus} -gt 1 ] && _val=$(( _val / vm_cpus ))
							fi
							[ -n "${_val}" ] && echo "${res}=${_val}"
						done
					done
				else
					for i in $( jorder ); do
						eval $( ${RCTL_CMD} -u jail:${i} )
						for res in ${RCTL}; do
							eval _mydesc="\$${res}_desc"
							eval _val=\$${res}
							echo "# HELP jail_${res} ${_mydesc} ${i}"
							echo "jail_${res}{name=\"${i}\"} ${_val:-0}"
						done
					done
					for i in $( border ); do
						#pid=$( cbsdsqlro local "SELECT jid FROM jails WHERE jname='${i}'" )
						#vm_cpus=$( cbsdsqlro local "SELECT vm_cpus FROM bhyve WHERE jname='${i}'" )
						cbsdsqlro_vars local "SELECT jails.jid, bhyve.vm_cpus FROM jails JOIN bhyve ON jails.jname = bhyve.jname WHERE jails.jname = '${jname}'" pid vm_cpus
						eval $( ${RCTL_CMD} -u process:${pid} )
						# take into account multi-core guest in pcpu value
						# On multi-core guest we need to: pcpu / vm_cpus
						[ ${vm_cpus} -gt 1 ] && pcpu=$(( pcpu / vm_cpus ))

						for res in ${RCTL}; do
							eval _mydesc="\$${res}_desc"
							eval _val=\$${res}
							echo "# HELP bhyve_${res} ${_mydesc} ${i}"
							echo "bhyve_${res}{name=\"${i}\"} ${_val:-0}"
						done
					done
				fi
			fi
			return 0
			;;
		"get")
			# Here must be
			# rctl -l jail:$jname::=/
			# or
			# rctl -u jail:$jname::=/
			# but it still unusable or crashed
			[ -f "${LIMITS}" ] && ${CAT_CMD} ${LIMITS}
			return 0
			;;
		*)
			show_jails | ${COLUMN_CMD} -t
			;;
	esac
}

# store args settings
for i in ${RCTL} ${RCTL_EXTRA}; do
	_val=
	eval _val="\$${i}"
	[ -n "${_val}" ] && export rctl_${i}="${_val}"
done

#### MAIN
if [ -n "${jname}" ]; then
	# support for multi/node-sqlite?
	. ${subrdir}/rcconf.subr
	[ $? -eq 1 ] && err 1 "${N1_COLOR}No such jail: ${N2_COLOR}${jname}${N0_COLOR}"
	ojname="${jname}"
fi

case "${emulator}" in
	jail)
		readconf jrctl.conf
		;;
	bhyve)
		readconf brctl.conf
		;;
esac

[ -z "${display}" ] && display="jname,memoryuse,maxproc,openfiles,vmemoryuse,swapuse,pcpu,fsquota,nice,bw"
[ "${shownode}" = "1" ] && display="nodename,${display}"
[ -z "${prometheus}" ] && prometheus=0

#remove commas for loop action on header
mydisplay=$( echo ${display} | ${TR_CMD} ',' '  ' )

# upper for header
myheader=$( echo ${mydisplay} | ${TR_CMD} '[:lower:]' '[:upper:]' )
JLS=""

[ -z "${header}" ] && header=1
[ -z "${human}" ] && human=1

sqldelimer=" "
jail_rctl

exit 0
