Use VmHWM to Catch “Memory-Heavy” YSQL Backends (and Help Prevent OOMs)

				
					#!/usr/bin/env bash
#
# yb_top_backend_mem.sh
#
# Show the Top-N YSQL processes ranked by peak resident memory (VmHWM),
# with current resident memory (VmRSS), DELTA (VmHWM - VmRSS),
# and a best-effort process classification.
#
# Enhancements:
#   --sort-delta           Sort by DELTA_MB instead of VmHWM_MB
#   Summary line           Total VmRSS MB across displayed rows + (optional) all scanned rows
#   Color highlighting     DELTA_MB warning/critical thresholds
#
# Columns:
#   PID | TYPE | VmHWM_MB | VmRSS_MB | DELTA_MB | DETAIL
#

set -euo pipefail

TOP_N=20
ONLY_CLIENT=0

HOST="${HOST:-127.0.0.1}"
PORT="${PORT:-5433}"
USER_NAME="${USER_NAME:-yugabyte}"
DB_NAME="${DB_NAME:-yugabyte}"
YSQLSH="${YSQLSH:-ysqlsh}"
TIMEOUT_S="${TIMEOUT_S:-5}"

SORT_MODE="hwm"                 # hwm | delta
DELTA_WARN_MB="${DELTA_WARN_MB:-50}"
DELTA_CRIT_MB="${DELTA_CRIT_MB:-200}"

COLOR_MODE="auto"               # auto | on | off

usage() {
  cat <<EOF
Usage: $(basename "$0") [options]
  -n, --top N                 Show top N rows (default: ${TOP_N})
  -h, --host HOST             YSQL host (default: ${HOST})
  -p, --port PORT             YSQL port (default: ${PORT})
  -U, --user USER             YSQL user (default: ${USER_NAME})
  -d, --db DB                 Database name (default: ${DB_NAME})
  --timeout SEC               statement_timeout in seconds (default: ${TIMEOUT_S})
  --only-client               Show only client backends (from pg_stat_activity)
  --sort-delta                Sort by DELTA_MB (spike hunting mode)
  --delta-warn-mb N           Warn threshold for DELTA_MB (default: ${DELTA_WARN_MB})
  --delta-crit-mb N           Crit threshold for DELTA_MB (default: ${DELTA_CRIT_MB})
  --color                     Force color output
  --no-color                  Disable color output
  --help                      Show help

Environment:
  YSQLSH                       Path to ysqlsh (default: ysqlsh)
  HOST, PORT, USER_NAME, DB_NAME, TIMEOUT_S
  DELTA_WARN_MB, DELTA_CRIT_MB
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    -n|--top) TOP_N="$2"; shift 2;;
    -h|--host) HOST="$2"; shift 2;;
    -p|--port) PORT="$2"; shift 2;;
    -U|--user) USER_NAME="$2"; shift 2;;
    -d|--db) DB_NAME="$2"; shift 2;;
    --timeout) TIMEOUT_S="$2"; shift 2;;
    --only-client) ONLY_CLIENT=1; shift;;
    --sort-delta) SORT_MODE="delta"; shift;;
    --delta-warn-mb) DELTA_WARN_MB="$2"; shift 2;;
    --delta-crit-mb) DELTA_CRIT_MB="$2"; shift 2;;
    --color) COLOR_MODE="on"; shift;;
    --no-color) COLOR_MODE="off"; shift;;
    --help) usage; exit 0;;
    *) echo "Unknown option: $1" >&2; usage; exit 2;;
  esac
done

command -v "$YSQLSH" >/dev/null 2>&1 || {
  echo "ERROR: '$YSQLSH' not found in PATH." >&2
  exit 1
}

timeout_ms=$(( TIMEOUT_S * 1000 ))

# Determine whether to use color
use_color=0
if [[ "$COLOR_MODE" == "on" ]]; then
  use_color=1
elif [[ "$COLOR_MODE" == "auto" && -t 1 ]]; then
  use_color=1
fi

# ANSI colors (only used if use_color=1)
RED=$'\033[31m'
YELLOW=$'\033[33m'
RESET=$'\033[0m'

CLIENT_PIDS=""
if [[ "$ONLY_CLIENT" -eq 1 ]]; then
  SQL_PIDS="SET statement_timeout = ${timeout_ms};
SELECT pid
FROM pg_stat_activity
WHERE backend_type = 'client backend'
ORDER BY pid;"

  if ! CLIENT_PIDS=$("$YSQLSH" -h "$HOST" -p "$PORT" -U "$USER_NAME" -d "$DB_NAME" \
    -A -t -F $'\n' -c "$SQL_PIDS"); then
    cat >&2 <<EOF
ERROR: Failed to run ysqlsh query.

Try:
  $YSQLSH -h $HOST -p $PORT -U $USER_NAME -d $DB_NAME -c "SELECT 1;"
EOF
    exit 1
  fi
fi

printf "%10s  %-24s  %12s  %12s  %12s  %s\n" \
  "PID" "TYPE" "VmHWM_MB" "VmRSS_MB" "DELTA_MB" "DETAIL"
printf "%10s  %-24s  %12s  %12s  %12s  %s\n" \
  "----------" "------------------------" "------------" "------------" "------------" "------------------------------"

# Choose sort key (3=VmHWM_MB, 5=DELTA_MB in tab output)
sort_key_col=3
if [[ "$SORT_MODE" == "delta" ]]; then
  sort_key_col=5
fi

# We'll compute totals over:
#  - all scanned rows (ALL_RSS_SUM_MB)
#  - displayed top-N rows (TOP_RSS_SUM_MB)
ALL_RSS_SUM_MB=0

# Collect rows (tab-separated), while also accumulating ALL_RSS_SUM_MB
rows_tmp="$(mktemp)"
trap 'rm -f "$rows_tmp"' EXIT

{
  if [[ "$ONLY_CLIENT" -eq 1 ]]; then
    while read -r pid; do
      [[ -n "${pid:-}" ]] || continue

      cmd=$(ps -p "$pid" -o args= 2>/dev/null || true)
      [[ -n "$cmd" ]] || continue

      status_file="/proc/$pid/status"
      [[ -r "$status_file" ]] || continue

      hwm_kb=$(awk '/^VmHWM:/ {print $2}' "$status_file")
      rss_kb=$(awk '/^VmRSS:/ {print $2}' "$status_file")
      [[ -n "${hwm_kb:-}" ]] || continue
      rss_kb="${rss_kb:-0}"

      hwm_mb=$(awk "BEGIN {printf \"%.1f\", $hwm_kb/1024}")
      rss_mb=$(awk "BEGIN {printf \"%.1f\", $rss_kb/1024}")
      delta_mb=$(awk "BEGIN {printf \"%.1f\", ($hwm_kb-$rss_kb)/1024}")

      # Accumulate ALL_RSS_SUM_MB
      ALL_RSS_SUM_MB=$(awk "BEGIN {printf \"%.1f\", ${ALL_RSS_SUM_MB}+${rss_mb}}")

      detail="$cmd"
      if [[ "$cmd" == *"postgres:"* ]]; then
        detail="${cmd#*postgres: }"
      fi

      printf "%s\t%s\t%s\t%s\t%s\t%s\n" \
        "$pid" "Client backend" "$hwm_mb" "$rss_mb" "$delta_mb" "${detail:0:140}"
    done <<< "$CLIENT_PIDS"
  else
    for pid in $(ps -eo pid,comm | awk '$2=="postgres"{print $1}'); do
      cmd=$(ps -p "$pid" -o args= 2>/dev/null || true)
      [[ -n "$cmd" ]] || continue

      role=""
      if [[ "$cmd" == *"postgres:"* ]]; then
        role="${cmd#*postgres: }"
      fi

      type="Other"
      detail="${cmd:0:140}"

      if [[ "$cmd" != *"postgres:"* ]] && [[ "$cmd" == *"/postgres/bin/postgres"* || "$cmd" == *"/postgres" || "$cmd" == postgres* ]]; then
        type="YSQL postmaster parent"
      elif [[ "$role" == YSQL\ webserver* ]]; then
        type="YSQL webserver"; detail="$role"
      elif [[ "$role" == yb_ash\ collector* ]]; then
        type="yb_ash collector"; detail="$role"
      elif [[ "$role" =~ ^(logger|checkpointer|background\ writer|walwriter|autovacuum\ launcher|stats\ collector|logical\ replication\ launcher) ]]; then
        type="Background worker"; detail="$role"
      elif [[ -n "$role" ]]; then
        type="Client backend"; detail="$role"
      fi

      status_file="/proc/$pid/status"
      [[ -r "$status_file" ]] || continue

      hwm_kb=$(awk '/^VmHWM:/ {print $2}' "$status_file")
      rss_kb=$(awk '/^VmRSS:/ {print $2}' "$status_file")
      [[ -n "${hwm_kb:-}" ]] || continue
      rss_kb="${rss_kb:-0}"

      hwm_mb=$(awk "BEGIN {printf \"%.1f\", $hwm_kb/1024}")
      rss_mb=$(awk "BEGIN {printf \"%.1f\", $rss_kb/1024}")
      delta_mb=$(awk "BEGIN {printf \"%.1f\", ($hwm_kb-$rss_kb)/1024}")


      ALL_RSS_SUM_MB=$(awk "BEGIN {printf \"%.1f\", ${ALL_RSS_SUM_MB}+${rss_mb}}")

      printf "%s\t%s\t%s\t%s\t%s\t%s\n" \
        "$pid" "$type" "$hwm_mb" "$rss_mb" "$delta_mb" "${detail:0:140}"
    done
  fi
} > "$rows_tmp"

# Select Top-N
top_tmp="$(mktemp)"
trap 'rm -f "$rows_tmp" "$top_tmp"' EXIT

sort -t $'\t' -k${sort_key_col},${sort_key_col}nr "$rows_tmp" | head -n "$TOP_N" > "$top_tmp"

# Compute TOP_RSS_SUM_MB
TOP_RSS_SUM_MB=$(awk -F $'\t' '{sum+=$4} END {printf "%.1f", sum+0}' "$top_tmp")

# Print Top-N with optional color highlighting based on DELTA_MB
awk -F $'\t' -v use_color="$use_color" -v warn="$DELTA_WARN_MB" -v crit="$DELTA_CRIT_MB" \
  -v RED="$RED" -v YELLOW="$YELLOW" -v RESET="$RESET" '
{
  pid=$1; type=$2; hwm=$3; rss=$4; delta=$5; detail=$6;

  delta_out=delta;
  if (use_color==1) {
    if (delta+0 >= crit+0) delta_out=RED delta RESET;
    else if (delta+0 >= warn+0) delta_out=YELLOW delta RESET;
  }

  printf "%10s  %-24s  %12s  %12s  %12s  %s\n", pid, type, hwm, rss, delta_out, detail;
}' "$top_tmp"


# Summary
printf "\n%s\n" "Summary"
printf '%s\n' "-------"

printf "  %-34s %12s\n" \
  "Sort mode:" "$( [[ "$SORT_MODE" == "delta" ]] && echo "DELTA_MB" || echo "VmHWM_MB" )"

printf "  %-34s %12.1f MB\n" \
  "Total VmRSS (all scanned rows):" "$ALL_RSS_SUM_MB"

printf "  %-34s %12.1f MB\n" \
  "Total VmRSS (displayed Top-${TOP_N}):" "$TOP_RSS_SUM_MB"

printf "  %-34s %12.1f MB\n" \
  "DELTA warn threshold:" "$DELTA_WARN_MB"

printf "  %-34s %12.1f MB\n" \
  "DELTA crit threshold:" "$DELTA_CRIT_MB"