leader-rebalance.sh

leader-rebalance.sh

#!/bin/bash

set -e

# 🎨 Cores e emojis
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

WARN="⚠️"
INFO="ℹ️"
SUCCESS="✅"
FAIL="❌"

TMP_FILE="/tmp/kafka-topic-description.txt"
PARTITIONS_FILE="/tmp/partitions-to-reassign.json"

function show_help() {
  echo -e "${CYAN}Uso:${NC}"
  echo -e "  $0 <bootstrap-server> <topic-name>"
  echo -e "  $0 <bootstrap-server> --all-topics"
  echo -e "  $0 <bootstrap-server> --topic-pattern <regex>"
  echo -e "  $0 --help"
  echo ""
  echo -e "${CYAN}Exemplos:${NC}"
  echo -e "  ./leader-rebalance.sh kafka-1:9092 lab-1"
  echo -e "  ./leader-rebalance.sh kafka-1:9092 --all-topics"
  echo -e "  ./leader-rebalance.sh kafka-1:9092 --topic-pattern lab-"
}

function describe_topic() {
  local topic=$1
  kafka-topics --bootstrap-server "$BOOTSTRAP_SERVER" --describe --topic "$topic" > "$TMP_FILE"
}

function parse_and_check() {
  declare -A leaders_set
  declare -A isr_set
  declare -A replicas_set

  local topic=$1
  describe_topic "$topic"

  while IFS= read -r line; do
    partition=$(echo "$line" | sed -E 's/.*Partition: ([0-9]+).*/\1/')
    leader=$(echo "$line" | sed -E 's/.*Leader: ([0-9]+).*/\1/')
    replicas=$(echo "$line" | sed -E 's/.*Replicas: ([0-9,]+).*/\1/' | tr -d '[:space:]')
    isr=$(echo "$line" | sed -E 's/.*Isr: ([0-9,]+).*/\1/' | tr -d '[:space:]')

    leaders_set[$leader]=1

    IFS=',' read -ra rep_arr <<< "$replicas"
    for r in "${rep_arr[@]}"; do
      [[ $r =~ ^[0-9]+$ ]] && replicas_set[$r]=1
    done

    IFS=',' read -ra isr_arr <<< "$isr"
    for i in "${isr_arr[@]}"; do
      [[ $i =~ ^[0-9]+$ ]] && isr_set[$i]=1
    done

  done < <(grep "Partition:" "$TMP_FILE")

  for r in "${!replicas_set[@]}"; do
    if [[ -z "${isr_set[$r]}" ]]; then
      echo -e "${YELLOW}${WARN} Broker $r ainda não está no ISR. Aguardando sincronização...${NC}"
      echo -e "${INFO} ${YELLOW}Abortando rebalanceamento do tópico '$topic'.${NC}"
      return 1
    fi
  done

  local rebalance_needed=false
  for b in "${!isr_set[@]}"; do
    if [[ -z "${leaders_set[$b]}" ]]; then
      echo -e "${YELLOW}${WARN} Broker $b está no ISR, mas não lidera nenhuma partição do tópico '$topic'.${NC}"
      rebalance_needed=true
    fi
  done

  if $rebalance_needed; then
    echo "{" > "$PARTITIONS_FILE"
    echo "  \"partitions\": [" >> "$PARTITIONS_FILE"
    grep "Partition:" "$TMP_FILE" | while read -r line; do
      partition=$(echo "$line" | sed -E 's/.*Partition: ([0-9]+).*/\1/')
      echo "    {\"topic\": \"$topic\", \"partition\": $partition}," >> "$PARTITIONS_FILE"
    done
    sed -i '$ s/,$//' "$PARTITIONS_FILE"
    echo "  ]" >> "$PARTITIONS_FILE"
    echo "}" >> "$PARTITIONS_FILE"

    echo -e "${INFO} ${CYAN}Executando rebalanceamento de líderes no tópico '${topic}'...${NC}"
    kafka-leader-election \
      --bootstrap-server "$BOOTSTRAP_SERVER" \
      --election-type PREFERRED \
      --path-to-json-file "$PARTITIONS_FILE"
    echo -e "${SUCCESS} ${GREEN}Rebalanceamento concluído para '$topic'!${NC}"
  else
    echo -e "${GREEN}${SUCCESS} Tópico '$topic' já está balanceado.${NC}"
  fi
}

function process_topics() {
  for topic in "${TOPICS[@]}"; do
    echo -e "\n🔍 ${CYAN}Processando tópico: $topic${NC}"
    parse_and_check "$topic"
  done
}

# ==== MAIN ====

# --help
if [[ "$1" == "--help" || "$1" == "-h" ]]; then
  show_help
  exit 0
fi

# Verificação mínima
if [ $# -lt 2 ]; then
  echo -e "${RED}${FAIL} Uso incorreto.${NC}"
  show_help
  exit 1
fi

BOOTSTRAP_SERVER=$1
shift

# Coleta lista de tópicos conforme a flag
if [ "$1" == "--all-topics" ]; then
  TOPICS=($(kafka-topics --bootstrap-server "$BOOTSTRAP_SERVER" --list))
elif [ "$1" == "--topic-pattern" ]; then
  if [ -z "$2" ]; then
    echo -e "${RED}${FAIL} Você precisa informar um padrão após --topic-pattern${NC}"
    exit 1
  fi
  PATTERN=$2
  TOPICS=($(kafka-topics --bootstrap-server "$BOOTSTRAP_SERVER" --list | grep "$PATTERN"))
else
  TOPICS=("$1")
fi

# Executa
process_topics

# Limpeza
rm -f "$TMP_FILE" "$PARTITIONS_FILE"