Benchmarking YSQL Connection Time in YugabyteDB with ysqlsh

				
					#!/bin/bash

set -euo pipefail

# --- Defaults ---
YSQLSH="${YSQLSH:-/home/yugabyte/yugabyte-db/bin/ysqlsh}"
USERNAME="yugabyte"
DBNAME="yugabyte"
COUNT=10
MODE="socket"
HOST=""
PORT="5433"

# --- Parse CLI args ---
while [[ $# -gt 0 ]]; do
  case "$1" in
    --tcp)
      MODE="tcp"
      shift
      ;;
    --socket)
      MODE="socket"
      shift
      ;;
    --host)
      HOST="$2"
      shift 2
      ;;
    -U|--user)
      USERNAME="$2"
      shift 2
      ;;
    -d|--dbname)
      DBNAME="$2"
      shift 2
      ;;
    -n|--count)
      COUNT="$2"
      shift 2
      ;;
    *)
      echo "Unknown option: $1"
      echo "Usage: $0 [--tcp|--socket] [--host <hostname>] [-U user] [-d dbname] [-n count]"
      exit 1
      ;;
  esac
done

# --- Validate ysqlsh path ---
if [[ ! -x "$YSQLSH" ]]; then
  echo "❌ Error: ysqlsh not found or not executable at: $YSQLSH"
  echo "👉 Please set the correct path at the top of the script or use:"
  echo "   YSQLSH=/path/to/ysqlsh ./yb_ysqlsh_connect_benchmark.sh ..."
  exit 1
fi

# --- Determine connection method ---
if [[ "$MODE" == "tcp" ]]; then
  HOST="${HOST:-127.0.0.1}"
  echo "✅ Using TCP connection to $HOST:$PORT"
else
  HOST=$(ps -ef | grep '[t]server' | awk '{for(i=1;i<=NF;i++)if($i~/-k/)print $(i+1)}')
  if [[ -z "$HOST" ]]; then
    echo "❌ Could not detect YSQL socket path. Is the tserver running?"
    exit 1
  fi
  echo "✅ Detected YSQL socket path: $HOST"
fi

echo "▶ Running $COUNT connections to database '$DBNAME' as user '$USERNAME' using ysqlsh ($MODE mode)..."

# --- Run benchmark ---
TIMINGS=()

for i in $(seq 1 "$COUNT"); do
  start_ns=$(date +%s%N)

  if [[ "$MODE" == "tcp" ]]; then
    output=$(
      PGPASSWORD="" "$YSQLSH" --no-password -h "$HOST" -p "$PORT" -U "$USERNAME" -d "$DBNAME" -c "SELECT 1;" 2>&1
    )
  else
    output=$(
      PGPASSWORD="" "$YSQLSH" --no-password -h "$HOST" -U "$USERNAME" -d "$DBNAME" -c "SELECT 1;" 2>&1
    )
  fi

  end_ns=$(date +%s%N)
  duration=$(awk -v s="$start_ns" -v e="$end_ns" 'BEGIN { printf "%.4f", (e - s)/1000000000 }')

  if echo "$output" | grep -q "could not connect"; then
    echo "⚠️  [$i] Connection failed:"
    echo "$output"
    continue
  fi

  TIMINGS+=("$duration")
  printf "  [%02d] Connection time: %.4f sec\n" "$i" "$duration"
done

# --- Summary stats ---
if [[ ${#TIMINGS[@]} -eq 0 ]]; then
  echo "❌ No successful connections."
  exit 1
fi

MIN=$(printf '%s\n' "${TIMINGS[@]}" | sort -n | head -n1)
MAX=$(printf '%s\n' "${TIMINGS[@]}" | sort -n | tail -n1)
AVG=$(printf '%s\n' "${TIMINGS[@]}" | awk '{sum+=$1} END{printf "%.4f", sum/NR}')
MEDIAN=$(printf '%s\n' "${TIMINGS[@]}" | sort -n | awk '
  { a[NR]=$1 }
  END {
    if (NR % 2 == 1)
      printf "%.4f", a[(NR+1)/2]
    else
      printf "%.4f", (a[NR/2] + a[NR/2+1]) / 2
  }')

echo ""
echo "📊 Connection Time Summary:"
echo "  Total:  ${#TIMINGS[@]} successful connections"
echo "  Min:    ${MIN} sec"
echo "  Max:    ${MAX} sec"
echo "  Avg:    ${AVG} sec"
echo "  Median: ${MEDIAN} sec"