统计进程内存使用量

统计进程内存使用量

#!/bin/bash

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

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

TOTAL_MEM_SIZE[KB]=0
TOTAL_MEM_SIZE[MB]=0
TOTAL_MEM_SIZE[GB]=0

Usage() {
  echo "Usage: $0 [-k/-K] [-m/-M] [-g/-G] <参数错误，请指定您想要显示的单位>"
  exit 1
}

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

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

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

  # 累加全局统计
  TOTAL_MEM_SIZE[KB]=$((TOTAL_MEM_SIZE[KB] + PROCESS_MEM_TOTAL[KB]))
}

ShowMemory() {
  for PID in /proc/[0-9]*; do
    PID=${PID##*/}
    # 初始化单个进程的统计变量
    declare -A PROCESS_MEM_TOTAL

    PROCESS_MEM_TOTAL[KB]=0
    PROCESS_MEM_TOTAL[MB]=0
    PROCESS_MEM_TOTAL[GB]=0

    # 获取单个进程的内存总量
    GetMemory

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

    # 如果进程使用了内存，则打印
    if [[ ${PROCESS_MEM_TOTAL[KB]} -gt 0 ]]; then
      printf "| %-8s | %-20s |\n" \
        "${PID}" \
        "${PROCESS_MEM_TOTAL[$1]} $1"
    fi
  done
}

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

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

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

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

main $1