#!/usr/local/bin/cbsd
#v11.2.1
CBSDMODULE="jail"
MYARG=""
. ${distsharedir}/rctl.conf
MYOPTARG="alljails display emulator header human jname mode node prometheus shownode ${RCTL} ${RCTL_EXTRA} quiet"
MYDESC="Set or flush resource limit for jail"
ADDHELP="\
alljails=1 - get jaillist from remote node\n\
display= list by comma for column. Default: jname,memoryuse,maxproc,openfiles,vmemoryuse,swapuse,pcpu,fsquota,nice,bw\n\
emulator=[jail,bhyve,xen] - show only emulator engine, default - all\n\
header=0 don't print header\n\
human=0 don't convert bytes to human readable form\n\
jname= show only specified jail\n\
mode= set (apply), unset, show (show current state) or get (show limits)\n\
node= only for current node\n\
prometheus=0 or 1 = prometheus metric mode output\n\
shownode=1 - show nodename for jails\n\
rctl params: ${RCTL} ${RCTL_EXTRA}\n\
quiet= 0,1: be quiet, dont output verbose message\n"
EXTHELP="wf_jrctl"

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

. ${cbsdinit}

[ -z "${quiet}" ] && 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

	local _emulator

	_emulator=$( cbsdsqlro local SELECT emulator FROM jails WHERE jname=\"${jname}\" )

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

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

		_val="0"

		if [ "${_i}" = "jname" ]; then
			_val="${jname}"
		else
			if [ -r "${_formfile}" ]; then
				_val=$( cbsdsqlro ${_formfile} "SELECT cur FROM forms WHERE param=\"${_i}\" LIMIT 1" 2>/dev/null )
				[ -z "${_val}" ] && _val="0"
			fi

			if rctl_humanize ${_i}; then
				if conv2human ${_val}; then
					_val="${convval}"
				fi
			fi
		fi

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

}

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

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

	if [ -n "${jname}" ]; then
		_sqlq="cbsdsqlro ${sqlfile} SELECT jname FROM jails WHERE jname=\"${jname}\" AND emulator = \"bhyve\" OR emulator = \"jail\""
	else
		if [ -n "${emulator}" ]; then
			_sqlq="cbsdsqlro ${sqlfile} SELECT jname FROM jails WHERE emulator = \"${emulator}\""
		else
			_sqlq="cbsdsqlro ${sqlfile} SELECT jname FROM jails WHERE emulator = \"bhyve\" OR emulator = \"jail\""
		fi
	fi

	${_sqlq}| while read jname; 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
}


show_local()
{
	local _errcode _status
	show_header
	show_jaildata_from_sql local
}


show_remote()
{
	show_header

	if [ -z "${node}" ]; then
		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}${CBSD_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}

	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
		fi

		_valnew=
		_valnew=$( cbsdsqlro ${_formfile} "SELECT new FROM forms WHERE param=\"${i}\" LIMIT 1" 2>/dev/null )
		if [ -z "${_valnew}" ]; then
			_val=$( cbsdsqlro ${_formfile} "SELECT cur FROM forms WHERE param=\"${i}\" LIMIT 1" 2>/dev/null )
		else
			# apply 'cur' values and truncate 'new'
			_val="${_valnew}"
			cbsdsqlrw ${_formfile} "UPDATE forms SET cur=\"${_val}\",new='' WHERE param=\"${i}\"" 2>/dev/null
		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
		fi
		_val=
		_valnew=
		_valnew=$( cbsdsqlro ${_formfile} "SELECT new FROM forms WHERE param=\"${i}\" LIMIT 1" 2>/dev/null )
		if [ -z "${_valnew}" ]; then
			_val=$( cbsdsqlro ${_formfile} "SELECT cur FROM forms WHERE param=\"${i}\" LIMIT 1" 2>/dev/null )
		else
			# apply 'cur' values and truncate 'new'
			_val="${_valnew}"
			cbsdsqlrw ${_formfile} "UPDATE forms SET cur=\"${_val}\",new='' WHERE param=\"${i}\"" 2>/dev/null
		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)
				;;
			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 in $( ${JOT_CMD} ${_val} ); 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=\"${_cpuset_args}\"" > ${CPU}
				else
					[ -r ${CPU} ] && ${RM_CMD} -f ${CPU}
				fi
				;;
		esac
	done

	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}${CBSD_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}\" )
				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}\" )
							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}\" )
							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}"
						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}"
						  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}\" )
						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}"
							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
	. ${subrdir}/rcconf.subr
	[ $? -eq 1 ] && err 1 "${N1_COLOR}No such jail: ${N2_COLOR}${jname}${N0_COLOR}"
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
