统计进程大页内存使用量

#!/bin/bash

# 文件名: show_hugepage.sh
# 作者: wanghongwei
# 日期: 2024年9月29日
# 版本: 1.0
# 描述: 统计进程大页内存使用量
# 使用方式: ./show_hugepage.sh [-k/-K] [-m/-M] [-g/-G]

# 初始化所有进程的统计变量
declare -A TOTAL_PRIVATE_HUGETLB
declare -A TOTAL_SHARED_HUGETLB
declare -A TOTAL_HUGE_PAGE_SIZE

TOTAL_PRIVATE_HUGETLB[KB]=0
TOTAL_SHARED_HUGETLB[KB]=0
TOTAL_HUGE_PAGE_SIZE[KB]=0

Usage() {
	echo "提示：请指定您想要显示的单位"
	echo "Usage: $0 [-k/-K] [-m/-M] [-g/-G]"
	exit 1
}

Title() {
	echo "+----------+----------------------+----------------------+----------------------+"
	printf "| %-8s | %-20s | %-20s | %-20s |\n" \
	"PID" \
	"Private_Hugetlb ($1)" \
	"Shared_Hugetlb ($1)" \
	"Total HugePages ($1)"
	echo "+----------+----------------------+----------------------+----------------------+"
}

ShowHugePage() {
	for PID in /proc/[0-9]*; do
		PID=${PID##*/}

		# 初始化单个进程的统计变量
		declare -A PROCESS_PRIVATE_HUGETLB
		declare -A PROCESS_SHARED_HUGETLB
		declare -A PROCESS_HUGE_PAGE_TOTAL

		PROCESS_PRIVATE_HUGETLB[KB]=0
		PROCESS_SHARED_HUGETLB[KB]=0
		PROCESS_HUGE_PAGE_TOTAL[KB]=0

		# 获取进程的大页内存信息
		PROCESS_HUGE_PAGE_INFO=$(cat /proc/${PID}/smaps 2>/dev/null | grep -v " 0 kB" | grep "Hugetlb")

		while read line; do
			if [[ $line =~ ^Private_Hugetlb:[[:space:]]+([0-9]+) ]]; then
				PAGE_SIZE=${BASH_REMATCH[1]} # BASH_REMATCH[1] 包含第一个括号中匹配的内容
				PROCESS_PRIVATE_HUGETLB[KB]=$((PROCESS_PRIVATE_HUGETLB[KB] + PAGE_SIZE))
				PROCESS_HUGE_PAGE_TOTAL[KB]=$((PROCESS_HUGE_PAGE_TOTAL[KB] + PAGE_SIZE))

			elif [[ $line =~ ^Shared_Hugetlb:[[:space:]]+([0-9]+) ]]; then
				PAGE_SIZE=${BASH_REMATCH[1]} # BASH_REMATCH[1] 包含第一个括号中匹配的内容
				PROCESS_SHARED_HUGETLB[KB]=$((PROCESS_SHARED_HUGETLB[KB] + PAGE_SIZE))
				PROCESS_HUGE_PAGE_TOTAL[KB]=$((PROCESS_HUGE_PAGE_TOTAL[KB] + PAGE_SIZE))
			else
				:
			fi
		done <<<"$PROCESS_HUGE_PAGE_INFO"

		# 累加全局统计
		TOTAL_PRIVATE_HUGETLB[KB]=$((TOTAL_PRIVATE_HUGETLB[KB] + PROCESS_PRIVATE_HUGETLB[KB]))
		TOTAL_SHARED_HUGETLB[KB]=$((TOTAL_SHARED_HUGETLB[KB] + PROCESS_SHARED_HUGETLB[KB]))
		TOTAL_HUGE_PAGE_SIZE[KB]=$((TOTAL_HUGE_PAGE_SIZE[KB] + PROCESS_HUGE_PAGE_TOTAL[KB]))

		# 单位换算
		case "$1" in
		"KB") : ;;
		"MB")
			PROCESS_PRIVATE_HUGETLB[MB]=$(echo "scale=3; ${PROCESS_PRIVATE_HUGETLB[KB]} / 1024" | bc)
			PROCESS_SHARED_HUGETLB[MB]=$(echo "scale=3; ${PROCESS_SHARED_HUGETLB[KB]} / 1024" | bc)
			PROCESS_HUGE_PAGE_TOTAL[MB]=$(echo "scale=3; ${PROCESS_HUGE_PAGE_TOTAL[KB]} / 1024" | bc)
			;;
		"GB")
			PROCESS_PRIVATE_HUGETLB[GB]=$(echo "scale=6; ${PROCESS_PRIVATE_HUGETLB[KB]} / 1024 / 1024" | bc)
			PROCESS_SHARED_HUGETLB[GB]=$(echo "scale=6; ${PROCESS_SHARED_HUGETLB[KB]} / 1024 / 1024" | bc)
			PROCESS_HUGE_PAGE_TOTAL[GB]=$(echo "scale=6; ${PROCESS_HUGE_PAGE_TOTAL[KB]} / 1024 / 1024" | bc)
			;;
		*) : ;;
		esac

		# 打印单个进程的大页内存统计
		if [[ ${PROCESS_HUGE_PAGE_TOTAL[KB]} -gt 0 ]]; then
			printf "| %-8s | %-20s | %-20s | %-20s |\n" \
			"${PID}" \
			"${PROCESS_PRIVATE_HUGETLB[$1]} $1" \
			"${PROCESS_SHARED_HUGETLB[$1]} $1" \
			"${PROCESS_HUGE_PAGE_TOTAL[$1]} $1"
		fi
	done
}

Summary() {
	# 单位换算
	case "$1" in
	"KB") : ;;
	"MB")
		TOTAL_PRIVATE_HUGETLB[MB]=$(echo "scale=3; ${TOTAL_PRIVATE_HUGETLB[KB]} / 1024" | bc)
		TOTAL_SHARED_HUGETLB[MB]=$(echo "scale=3; ${TOTAL_SHARED_HUGETLB[KB]} / 1024" | bc)
		TOTAL_HUGE_PAGE_SIZE[MB]=$(echo "scale=3; ${TOTAL_HUGE_PAGE_SIZE[KB]} / 1024" | bc)
		;;
	"GB")
		TOTAL_PRIVATE_HUGETLB[GB]=$(echo "scale=6; ${TOTAL_PRIVATE_HUGETLB[KB]} / 1024 / 1024" | bc)
		TOTAL_SHARED_HUGETLB[GB]=$(echo "scale=6; ${TOTAL_SHARED_HUGETLB[KB]} / 1024 / 1024" | bc)
		TOTAL_HUGE_PAGE_SIZE[GB]=$(echo "scale=6; ${TOTAL_HUGE_PAGE_SIZE[KB]} / 1024 / 1024" | bc)
		;;
	*) : ;;
	esac

	# 打印总统计
	echo "+----------+----------------------+----------------------+----------------------+"
	printf "| %-8s | %-20s | %-20s | %-20s |\n" \
	"Total:" \
	"${TOTAL_PRIVATE_HUGETLB[$1]} $1" \
	"${TOTAL_SHARED_HUGETLB[$1]} $1" \
	"${TOTAL_HUGE_PAGE_SIZE[$1]} $1"
	echo "+----------+----------------------+----------------------+----------------------+"
}

main() {
	if [ -z "$1" ]; then
		Usage
	fi

	case "$1" in
	-k | -K)
		Title KB
		ShowHugePage KB
		Summary KB
		;;
	-m | -M)
		Title MB
		ShowHugePage MB
		Summary MB
		;;
	-g | -G)
		Title GB
		ShowHugePage GB
		Summary GB
		;;
	*) Usage ;;
	esac
	exit 0
}

main $1