#!/bin/sh
#########################################################################################
#
# This script is used to collect memory monitoring data for WebSphere or Liberty on Linux
#
# Usage:
# ./linmem.sh [options]                       # System-wide monitoring
# ./linmem.sh <pid> [options]                 # PID + system-wide monitoring
#
# Options:
# <pid>                   Process ID to monitor (numeric argument)
# -d <threshold>          Auto-dump threshold (repeatable, e.g., -d 120m -d 200m)
#                         Units: m/M (megabytes), g/G (gigabytes)
# -i <interval>           Collection interval in seconds (default: 1800)
# --output-dir=<path>     Custom output directory (default: current directory)
# --no-archive            Don't create tar.gz archive
# -h, --help              Show this help message
#
# Examples:
# ./linmem.sh -i 1800                         # System-wide only
# ./linmem.sh 12345 -d 120m -d 200m -i 60     # PID with thresholds
# ./linmem.sh 12345 -i 300                    # PID without thresholds
#
# Maintainer: IBM WebSphere Support, WASSDK team, Piotr Zalewski
#
#########################################################################################

version=2026.06.11

# Variables
pid=""
interval=1800
system_wide=0
include_system=0
no_archive=0
output_dir="."
dump_at=""
dump_arr=""
current_dump_index=0
dir_name=""
has_data_produced=0
all_thresholds_reached=0

# Trap cleanup
trap cleanup INT TERM

cleanup() {
 exit 0
}

# Log function - writes timestamped message to screen.out
log() {
 message="$1"
 if [ -n "$message" ]; then
  printf "%s    %s\n" "$(date '+%Y-%m-%d %H:%M:%S')" "$message" | tee -a screen.out
 fi
}

# Help output
help_output() {
 echo "linmem.sh (version: $version)"
 echo ""
 echo "Usage:"
 echo " ./linmem.sh [options]                       # System-wide monitoring"
 echo " ./linmem.sh <pid> [options]                 # PID + system-wide monitoring"
 echo ""
 echo "Options:"
 echo " <pid>                   Process ID to monitor (numeric argument)"
 echo " -d <threshold>          Auto-dump threshold (repeatable, e.g., -d 120m -d 200m)"
 echo "                         Units: m/M (megabytes), g/G (gigabytes)"
 echo " -i <interval>           Collection interval in seconds (default: 1800)"
 echo " --output-dir=<path>     Custom output directory (default: current directory)"
 echo " --no-archive            Don't create tar.gz archive"
 echo " -h, --help              Show this help message"
 echo ""
 echo "Examples:"
 echo " # System-wide monitoring only"
 echo " ./linmem.sh -i 1800"
 echo ""
 echo " # Monitor PID with auto-dump thresholds"
 echo " ./linmem.sh 12345 -d 120m -d 200m -i 60"
 echo ""
 echo " # Monitor PID without thresholds"
 echo " ./linmem.sh 12345 -i 300"
 echo ""
 echo "Output:"
 echo " Creates a timestamped directory with multiple output files:"
 echo " - screen.out: Navigation log with timestamps"
 echo " - ps_vsz_rss.csv: Memory metrics (Excel-ready, PID mode)"
 echo " - meminfo.out: System memory information"
 echo " - ps_threads.out: Thread details (PID mode)"
 echo " - proc_maps.out: Process memory maps (PID mode)"
 echo " - ps_all.out: All processes snapshot (system-wide mode)"
 echo " - vmstat.out: Virtual memory statistics (system-wide mode)"
 echo ""
}

# Parse command-line arguments
while [ $# -gt 0 ]; do
 case "$1" in
  -d)
   dump_at="$2"
   if [ -z "$dump_arr" ]; then
    dump_arr="$2"
   else
    dump_arr="$dump_arr $2"
   fi
   shift 2
   ;;
  -i)
   interval="$2"
   shift 2
   ;;
  --output-dir=*)
   output_dir="${1#*=}"
   shift
   ;;
  --no-archive)
   no_archive=1
   shift
   ;;
  -*)
   echo "Unknown option: $1"
   help_output
   exit 1
   ;;
  *)
   # If it's a number, treat it as PID
   if echo "$1" | grep -qE '^[0-9]+$'; then
    pid="$1"
    shift
   else
    echo "Invalid argument: $1"
    help_output
    exit 1
   fi
   ;;
 esac
done


# Create output directory
create_output_directory() {
 dir_name="linmem_results.$(date '+%Y%m%d.%H%M%S')"
 mkdir -p "$output_dir/$dir_name"
 cd "$output_dir/$dir_name" || exit 1
}

# Initialize CSV file
init_csv() {
 echo "timestamp,pid,vsz_kb,rss_kb" > ps_vsz_rss.csv
}

# Write CSV entry
write_csv_entry() {
 timestamp=$(date '+%Y-%m-%d %H:%M:%S')
 ps_output=$(ps -p "$pid" -o vsz=,rss= 2>/dev/null)
 if [ $? -eq 0 ]; then
  vsz=$(echo "$ps_output" | awk '{print $1}')
  rss=$(echo "$ps_output" | awk '{print $2}')
  echo "$timestamp,$pid,$vsz,$rss" >> ps_vsz_rss.csv
 fi
}

# Check memory threshold
check_memory_threshold() {
 if [ -n "$dump_at" ]; then
  current_rss=$(ps -p "$pid" -o rss= 2>/dev/null | awk '{print $1}')
  if [ -z "$current_rss" ]; then
   return
  fi
  
  # Convert current RSS from KB to bytes
  current_rss_bytes=$((current_rss * 1024))
  
  # Get current threshold from array
  threshold=$(echo "$dump_arr" | awk -v idx=$((current_dump_index + 1)) '{print $idx}')
  if [ -z "$threshold" ]; then
   return
  fi
  
  # Convert threshold to bytes (simple conversion for g/G and m/M)
  threshold_upper=$(echo "$threshold" | tr '[:lower:]' '[:upper:]')
  threshold_value=$(echo "$threshold_upper" | sed 's/[GM]$//')
  
  if echo "$threshold_upper" | grep -q 'G$'; then
   threshold_bytes=$((threshold_value * 1024 * 1024 * 1024))
   threshold_display="${threshold_value} GB"
  elif echo "$threshold_upper" | grep -q 'M$'; then
   threshold_bytes=$((threshold_value * 1024 * 1024))
   threshold_display="${threshold_value} MB"
  else
   threshold_bytes=$threshold_value
   threshold_display="${threshold_value} bytes"
  fi
  
  # Display current status
  current_gb=$(awk "BEGIN {printf \"%.1f\", $current_rss_bytes / 1073741824}")
  
  log "Current RSS: ${current_gb} GB"
  
  # Check if threshold reached
  if [ "$current_rss_bytes" -ge "$threshold_bytes" ]; then
   log ""
   log "Threshold reached: ${threshold_display}"
   kill -3 "$pid"
   log "Issued kill -3 for PID $pid"
   
   # Move to next threshold
   num_thresholds=$(echo "$dump_arr" | wc -w)
   if [ $((current_dump_index + 1)) -lt "$num_thresholds" ]; then
    current_dump_index=$((current_dump_index + 1))
    next_threshold=$(echo "$dump_arr" | awk -v idx=$((current_dump_index + 1)) '{print $idx}')
    next_threshold_upper=$(echo "$next_threshold" | tr '[:lower:]' '[:upper:]')
    next_value=$(echo "$next_threshold_upper" | sed 's/[GM]$//')
    if echo "$next_threshold_upper" | grep -q 'G$'; then
     next_display="${next_value} GB"
    elif echo "$next_threshold_upper" | grep -q 'M$'; then
     next_display="${next_value} MB"
    else
     next_display="${next_value} bytes"
    fi
    log "Next threshold: ${next_display}"
   else
    log "All thresholds reached. Stopping data collection."
    dump_at=""
    all_thresholds_reached="1"
   fi
  fi
 fi
}

# Collect PID-specific data
collect_pid_data() {
 if [ ! -d "/proc/$pid" ]; then
  log "ERROR: PID $pid does not exist. Stopping."
  return 1
 fi
  
 # Write CSV entry
 write_csv_entry
 
 # Collect meminfo
 echo "$(date '+%Y-%m-%d %H:%M:%S')" >> meminfo.out
 echo "*** meminfo start" >> meminfo.out
 cat /proc/meminfo >> meminfo.out
 echo "*** meminfo stop" >> meminfo.out
 echo "--------------------" >> meminfo.out
 
 # Collect ps threads
 echo "$(date '+%Y-%m-%d %H:%M:%S')" >> ps_threads.out
 echo "*** ps thread start" >> ps_threads.out
 ps -mp "$pid" -o THREAD >> ps_threads.out 2>/dev/null
 echo "*** ps thread stop" >> ps_threads.out
 echo "--------------------" >> ps_threads.out
 
 # Collect proc maps
 echo "$(date '+%Y-%m-%d %H:%M:%S')" >> "proc_maps.out"
 echo "*** proc/maps start" >> "proc_maps.out"
 cat "/proc/$pid/maps" >> "proc_maps.out" 2>/dev/null
 echo "*** proc/maps stop" >> "proc_maps.out"
 echo "--------------------" >> "proc_maps.out"
 
 has_data_produced=1
 
 # Check memory threshold
 check_memory_threshold
 
 return 0
}

# Collect system-wide data
collect_system_data() {
 log "Collecting system-wide data"
 
 # Collect meminfo
 echo "$(date '+%Y-%m-%d %H:%M:%S')" >> meminfo.out
 echo "*** meminfo start" >> meminfo.out
 cat /proc/meminfo >> meminfo.out
 echo "*** meminfo stop" >> meminfo.out
 echo "--------------------" >> meminfo.out
 
 # Collect vmstat
 echo "$(date '+%Y-%m-%d %H:%M:%S')" >> vmstat.out
 echo "*** vmstat start" >> vmstat.out
 vmstat 1 3 >> vmstat.out
 echo "*** vmstat stop" >> vmstat.out
 echo "--------------------" >> vmstat.out
 
 # Collect ps all
 echo "$(date '+%Y-%m-%d %H:%M:%S')" >> ps_all.out
 echo "*** ps auxwww start" >> ps_all.out
 ps auxwww >> ps_all.out
 echo "*** ps auxwww stop" >> ps_all.out
 echo "--------------------" >> ps_all.out
 
 has_data_produced=1
}



# Create archive
create_archive() {
 if [ "$no_archive" != "1" ]; then
  log " "
  log "Compressing files into ${dir_name}.tar.gz"
  
  # Build list of files to archive
  files_to_archive="screen.out meminfo.out"
  
  if [ -n "$pid" ]; then
   files_to_archive="$files_to_archive ps_vsz_rss.csv ps_threads.out proc_maps.out"
  fi
  
  # Always collect system-wide data
  files_to_archive="$files_to_archive ps_all.out vmstat.out"
  
  # Create tar archive
  tar -cf ../${dir_name}.tar $files_to_archive 2>/dev/null
  
  if [ $? -ne 0 ]; then
   log "ERROR: Failed to create tar archive"
   log "Files remain in directory: $dir_name"
  else
   # Gzip the tar file
   gzip ../${dir_name}.tar
   
   if [ $? -eq 0 ]; then
    # Remove individual files
    rm -f $files_to_archive
    
    # Check if directory is empty, then remove it
    if [ -z "$(ls -A)" ]; then
     cd ..
     rm -r "$dir_name"
     log "Archive created: ${dir_name}.tar.gz"
     log " "
    else
     log "Archive created: ${dir_name}.tar.gz"
     log "Note: Directory $dir_name not removed (contains additional files)"
     log " "
    fi
    printf "\n\tTo share with IBM support, upload:\n"
    printf "\t* ${dir_name}\n"
    printf "\t* javacores (and system cores)\n\n"
   else
    log "ERROR: Failed to gzip archive"
   fi
  fi
 fi
}

# End message
end_message() {
 log "========================================"
 log "Data collection complete"
 create_archive
 cleanup
}

# Trap Ctrl-C
trap_ctrlc() {
 echo ""
 log "Ctrl-C detected"
 log "Exiting without creating archive"
 log " "
 log "Output files saved in: $dir_name"
 cleanup
}
trap trap_ctrlc INT

# Main execution
main() {
 # Create output directory
 create_output_directory
 
 # Print header
 log "========================================"
 log "linmem.sh - Linux Memory Monitoring"
 log "Version: $version"
 log "========================================"
 log ""
 log "Configuration:"
 if [ -n "$pid" ]; then
  log "- Mode: PID + system-wide monitoring"
  log "- PID: $pid"
 else
  log "- Mode: System-wide monitoring only"
 fi
 log "- Interval: $interval seconds"
 if [ -n "$dump_at" ]; then
  log "- Dump thresholds: $dump_arr"
 fi
 log "- Output directory: $dir_name"
 log ""
 
 # List output files
 log "Output Files:"
 log "- screen.out - This navigation log"
 log "- meminfo.out - System memory information"
 
 if [ -n "$pid" ]; then
  init_csv
  log "- ps_vsz_rss.csv - Memory metrics"
  log "- ps_threads.out - Thread details (ps -mp)"
  log "- proc_maps.out - Process memory maps"
 fi
 
 log "- ps_all.out - All processes snapshot"
 log "- vmstat.out - Virtual memory statistics"
 
 log ""
 log "========================================"
 log "Starting data collection..."
 log "========================================"
 log ""
 
 # Main collection loop
 iteration=1
 while true; do
  log "==== Iteration $iteration"
  
  if [ -n "$pid" ]; then
   # PID monitoring: collect PID data + system-wide
   if ! collect_pid_data; then
    if [ "$has_data_produced" = "1" ]; then
     end_message
    else
     exit 1
    fi
   fi
   collect_system_data
   # Check if all thresholds reached and should stop
   if [ "$all_thresholds_reached" = "1" ]; then
    end_message
   fi
  else
   # System-wide only
   collect_system_data
  fi
  
  log ""
  iteration=$((iteration + 1))
  sleep "$interval"
 done
}

# Run main
main