KARP

KARP: Kafka Auto-Rebalancing & Partitioning

A story about a lean, powerful, and truly proactive solution to an old problem.

Start the Journey

The Unspoken Challenge of Kafka

"Everything looks fine." — a famous last sentence in a data-driven world. We've all been there. You've got a great Kafka cluster humming along, but under the hood, a single broker is starting to sweat. Its CPU is climbing, a disk is filling up, and it's handling a disproportionate share of the load. Manual rebalancing is a tedious game of "whack-a-mole." You fix one broker, and another one inevitably gets imbalanced a week later. There had to be a better way.

The solution wasn't a bigger hammer, it was a smarter one.

Why KARP? The Story of Proactive Scoring

We needed a solution that was more than a glorified script. We needed a **mindset shift**. Our old system was purely reactive. It would only intervene if a broker's metrics hit a theoretical maximum—a value of `1.00`. But in a real-world, high-traffic environment, waiting for a system to be at 100% capacity is like waiting for your car's engine to explode before you get an oil change. It's too late.

Our solution was simple: be proactive.

We decided that an `OVERLOAD_SCORE` of just **0.75** on a scale of 0 to 1 is a clear indicator of a problem. But a single metric can be misleading. A disk could be full but not used. A CPU could spike momentarily. So we built a comprehensive, weighted scoring system that combines every critical metric:

  • CPU Usage: `5%`
  • Disk Usage: `25%`
  • Messages In: `10%`
  • Bytes In/Out: `20%` each
  • Requests/sec: `20%`

This new **75-25** weighting approach, with a heavy bias towards disk and network I/O, ensures that our rebalancing logic is a direct reflection of our real-world bottlenecks.

Start
KARP Initiated
Poll Metrics
from Prometheus
Compute Score
Weighted scoring (CPU, Disk, etc.)
List Over/Underloaded
Thresholds: 0.75 / 0.25
Are brokers imbalanced?
NO
Wait & Poll Again
YES
Generate Rebalance Plan
Using Kafka Reassignment Tools
Apply Throttle
Manual "Human-in-the-Loop"
Execute Rebalance
Verify & Clear Throttle
End

KARP vs. The Other Guys: The Practical Choice

KARP: Our Solution

  • Built for Purpose: A lean, focused bash script designed for our specific production needs. No bloat, no unnecessary features.
  • Human-in-the-Loop: We control the process. A simple `Y`/`N` variable for `DRY_RUN` and other environment variables like `NEED_THROTTLE` gives us total control over when and how rebalances happen.
  • Proactive & Agile: Our custom scoring model ensures we rebalance at the first sign of a problem, not when a crisis hits.
  • Transparent & Auditable: The entire process is a simple, readable script. No black box, no hidden logic. We know exactly why a rebalance happened.
  • Zero Operational Overhead: It's a job on our existing Docker stack, not a new service to maintain, monitor, and scale.

The Other Guy: Cruise Control

  • Overly Complex: A massive tool with a multitude of configurations that often requires a dedicated operator to understand and maintain.
  • "Black Box" Operations: While powerful, its Java-based logic and complex data collection model can be a black box. Understanding why a specific rebalance plan was generated can be difficult and time-consuming.
  • Reactive by Nature: Without extensive tuning and custom metrics, it often reacts to problems after they've occurred, rather than preventing them.
  • High Operational Overhead: Requires a separate service to run, which adds to the operational complexity and resource footprint of your Kafka cluster.

How It Works: Every Variable Tells a Story

Every line of our script is a part of this story. From the `PROMETHEUS_URL` that tells us where our cluster's heartbeat is, to the `WEIGHT_*` variables that define our priorities, the code is a narrative of our solution.

1. Data from Prometheus

Our `get_metrics` function, with its `prom_query` and `jq` parsing, is the data gatherer. It's a small but mighty team that reports back every **60 seconds** on the health of our brokers.

2. The Grand Score Calculation

The `compute_scores` function is where the magic happens. It takes our raw data, normalizes it, and applies the strategic weights we've discussed. This gives us the single, holistic `OVERLOAD_SCORE` for each broker.

3. A Rebalancing Plan is Born

`list_over_under` identifies the problem brokers. `gen_plan` then uses the native Kafka rebalancing tools to propose a safe, effective solution.

4. Execution & Throttling

The `rebalance` function takes over, applying the plan. With `NEED_THROTTLE` and `THROTTLE_VAL_*` variables, we ensure the rebalancing happens smoothly, without causing network congestion or performance issues.

A diagram of the Docker deployment and script execution.

Ready to Implement KARP?

Get started in minutes by configuring and running the provided shell script as a Docker container.

InitiateKARP.sh 

#!/usr/bin/env bash
set -euo pipefail

echo "==> Starting KARP (Kafka Auto-Rebalancing & Partitioning)"

#
# ===== Required env =====
#
: "${PROMETHEUS_URL:?PROMETHEUS_URL is required, e.g. http://prometheus:9090}"
: "${KAFKA_BOOTSTRAP_SERVERS:?KAFKA_BOOTSTRAP_SERVERS is required, e.g. ccf23broker6:13679,...}"
: "${ELIGIBLE_BROKERS:?ELIGIBLE_BROKERS is required, e.g. ccf23broker6:12956,ccf23broker7:13680,...}"
: "${BROKER_ID_MAP:?BROKER_ID_MAP is required, e.g. ccf23broker6:12956=6,ccf23broker7:13680=7,...}"
: "${KAFKA_CLUSTER:?KAFKA_CLUSTER is required, e.g. kafka-cluster-1}"
: "${PROM_METRIC_PREFIX:?PROM_METRIC_PREFIX is required, e.g. karp_metrics_}"
: "${METRICS_FILE:?METRICS_FILE is required, e.g. /opt/karp/metrics.json}"

#
# ===== Optional env (with defaults) =====
#
PROMETHEUS_POLL_INTERVAL="${PROMETHEUS_POLL_INTERVAL:-60}"
LOG_LEVEL="${LOG_LEVEL:-INFO}"
MSG_TOPIC_IGNORE_REGEX="${MSG_TOPIC_IGNORE_REGEX:-^(_.*|__.*|_confluent.*|_schemas|connect-.*|ksqldb.*)$}"

# weights (all configurable)
WEIGHT_CPU="${WEIGHT_CPU:-0.05}"
WEIGHT_DISK="${WEIGHT_DISK:-0.25}"
WEIGHT_MSGS="${WEIGHT_MSGS:-0.10}"
WEIGHT_BYTES_IN="${WEIGHT_BYTES_IN:-0.20}"
WEIGHT_BYTES_OUT="${WEIGHT_BYTES_OUT:-0.20}"
WEIGHT_REQUESTS="${WEIGHT_REQUESTS:-0.20}"

# thresholds (configurable)
OVERLOAD_SCORE_THRESHOLD="${OVERLOAD_SCORE_THRESHOLD:-0.75}"
UNDERLOAD_SCORE_THRESHOLD="${UNDERLOAD_SCORE_THRESHOLD:-0.25}"

MAX_PARTITIONS_PER_BROKER="${MAX_PARTITIONS_PER_BROKER:-0}" # 0 disables
DRY_RUN="${DRY_RUN:-Y}" # Y, N
NEED_THROTTLE="${NEED_THROTTLE:-N}"
THROTTLE_VAL_1="${THROTTLE_VAL_1:-1048576}"
THROTTLE_VAL_2="${THROTTLE_VAL_2:-1048576}"
THROTTLE_VAL_3="${THROTTLE_VAL_3:-11}"
THROTTLE_VAL_4="${THROTTLE_VAL_4:-1}"



#!/usr/bin/env bash
set -euo pipefail

echo "==> Starting KARP (Kafka Auto-Rebalancing & Partitioning)"
echo "came here 1"
COUNTER=0
#
# ===== Required env =====
#
: "${PROMETHEUS_URL:?PROMETHEUS_URL is required, e.g. http://prometheus:9090}"
: "${KAFKA_BOOTSTRAP_SERVERS:?KAFKA_BOOTSTRAP_SERVERS is required, e.g. ccf23broker6:13679,...}"
: "${ELIGIBLE_BROKERS:?ELIGIBLE_BROKERS is required, e.g. ccf23broker6:12956,ccf23broker7:12797,ccf23broker8:13134,ccf23broker9:12798}"
: "${BROKER_ID_MAP:?BROKER_ID_MAP is required, e.g. ccf23broker6:12956=6,ccf23broker7:12797=7,...}"
: "${KAFKA_CLUSTER:?KAFKA_CLUSTER is required, e.g. kafka-cluster-1}"
: "${PROM_METRIC_PREFIX:?PROM_METRIC_PREFIX is required, e.g. karp_metrics_}"
: "${METRICS_FILE:?METRICS_FILE is required, e.g. /opt/karp/metrics.json}"

#
# ===== Optional env (with defaults) =====
#
PROMETHEUS_POLL_INTERVAL="${PROMETHEUS_POLL_INTERVAL:-60}"
LOG_LEVEL="${LOG_LEVEL:-INFO}"
MSG_TOPIC_IGNORE_REGEX="${MSG_TOPIC_IGNORE_REGEX:-^(_.*|__.*|_confluent.*|_schemas|connect-.*|ksqldb.*)$}"

# weights (all configurable)
WEIGHT_CPU="${WEIGHT_CPU:-0.05}"
WEIGHT_DISK="${WEIGHT_DISK:-0.25}"
WEIGHT_MSGS="${WEIGHT_MSGS:-0.10}"
WEIGHT_BYTES_IN="${WEIGHT_BYTES_IN:-0.20}"
WEIGHT_BYTES_OUT="${WEIGHT_BYTES_OUT:-0.20}"
WEIGHT_REQUESTS="${WEIGHT_REQUESTS:-0.20}"

# thresholds (configurable)
OVERLOAD_SCORE_THRESHOLD="${OVERLOAD_SCORE_THRESHOLD:-0.75}"
UNDERLOAD_SCORE_THRESHOLD="${UNDERLOAD_SCORE_THRESHOLD:-0.25}"

MAX_PARTITIONS_PER_BROKER="${MAX_PARTITIONS_PER_BROKER:-0}" # 0 disables
DRY_RUN="${DRY_RUN:-Y}" # Y, N
NEED_THROTTLE="${NEED_THROTTLE:-N}"
THROTTLE_VAL_1="${THROTTLE_VAL_1:-1048576}"
THROTTLE_VAL_2="${THROTTLE_VAL_2:-1048576}"
THROTTLE_VAL_3="${THROTTLE_VAL_3:-11}"
THROTTLE_VAL_4="${THROTTLE_VAL_4:-1}"

#
# ===== Functions =====
#

log() {
  local level="$1"
  local msg="$2"
  if [[ "$LOG_LEVEL" == "$level" ]]; then
    echo "HAHAHA $level $msg"
  fi
}

prom_query() {
  local query="$1"
  log INFO "Prometheus Query: $query"
  curl --fail --silent --show-error "$PROMETHEUS_URL/api/v1/query?query=$(jq -sRr @uri <<< "$query")"
}

normalize_to01() {
  local -n MAP="$1"
  local max=0
  for val in "${MAP[@]}"; do
    if (( $(echo "$val > $max" | bc -l) )); then
      max="$val"
    fi
  done

  if (( $(echo "$max == 0" | bc -l) )); then
    for key in "${!MAP[@]}"; do
      echo "$key|0"
    done
  else
    for key in "${!MAP[@]}"; do
      local normalized_val
      normalized_val=$(echo "scale=4; ${MAP[$key]} / $max" | bc)
      echo "$key|$normalized_val"
    done
  fi
}

declare -A BROKER_ID_MAP_ARRAY
load_broker_id_map() {
    IFS=',' read -r -a pairs <<< "$BROKER_ID_MAP"
    for pair in "${pairs[@]}"; do
        IFS='=' read -r host_port broker_id <<< "$pair"
        BROKER_ID_MAP_ARRAY["$host_port"]="$broker_id"
    done
}

get_metrics() {
    echo "came here 3"
    echo "elig : $ELIGIBLE_BROKERS"

    # Declare associative arrays as global so they can be accessed by other functions
    declare -gA CPU_RATE MAP_DISK_BYTES MAP_MSGS_RATE MAP_BYTES_IN MAP_BYTES_OUT MAP_REQUESTS

    echo "came here p1"
    echo "ELIGIBLE_BROKERS : $ELIGIBLE_BROKERS"

    # Split ELIGIBLE_BROKERS into an array
    IFS=',' read -r -a ELIGIBLE_ARRAY <<< "$ELIGIBLE_BROKERS"

    # Prometheus queries
    local q_cpu='sum by (instance) (rate(process_cpu_seconds_total[1m]))'
    local q_disk='sum by (instance) (kafka_log_log_size)'
    local q_msgs="sum by (instance) (kafka_server_brokertopicmetrics_messagesinpersec{topic"'!'"~\"${MSG_TOPIC_IGNORE_REGEX}\"})"
    local q_leaders='sum by (instance) (kafka_server_replicamanager_leadercount)'
    local q_bytesin='sum by (instance) (rate(kafka_server_brokertopicmetrics_bytesinpersec[1m]))'
    local q_bytesout='sum by (instance) (rate(kafka_server_brokertopicmetrics_bytesoutpersec[1m]))'
    local q_requests='sum by (instance) (rate(kafka_network_requestmetrics_requestspersec[1m]))'

    echo "came here p3 $q_cpu $q_disk $q_msgs $q_leaders $q_bytesin $q_bytesout $q_requests"

    # Fetch all metrics at once
    local j_cpu j_disk j_msgs j_leaders j_bytesin j_bytesout j_requests
    j_cpu="$(prom_query "$q_cpu" || true)"
    echo "came here p4 $j_cpu"
    j_disk="$(prom_query "$q_disk" || true)"
    echo "came here p5 $j_disk"
    j_msgs="$(prom_query "$q_msgs" || true)"
    echo "came here p6 $j_msgs"
    j_leaders="$(prom_query "$q_leaders" || true)"
    echo "came here p7 $j_leaders"
    j_bytesin="$(prom_query "$q_bytesin" || true)"
    echo "came here p8 $j_bytesin"
    j_bytesout="$(prom_query "$q_bytesout" || true)"
    echo "came here p9 $j_bytesout"
    j_requests="$(prom_query "$q_requests" || true)"
    echo "came here p10 $j_requests"                

    echo "all done metrics 1"

    # Populate maps with a filter for eligible brokers using hostname prefix match
    
    # CPU metrics
    for broker in "${ELIGIBLE_ARRAY[@]}"; do
        local broker_host="${broker%%:*}"
        local cpu_val
        cpu_val=$(echo "$j_cpu" | jq -r --arg host "$broker_host" '.data.result[] | select(.metric.instance | startswith($host)) | .value[1]')
        CPU_RATE["$broker"]="${cpu_val:-0}"
    done

    # Disk metrics
    for broker in "${ELIGIBLE_ARRAY[@]}"; do
        local broker_host="${broker%%:*}"
        local disk_val
        disk_val=$(echo "$j_disk" | jq -r --arg host "$broker_host" '.data.result[] | select(.metric.instance | startswith($host)) | .value[1]')
        MAP_DISK_BYTES["$broker"]="${disk_val:-0}"
    done

    # Message rate metrics
    for broker in "${ELIGIBLE_ARRAY[@]}"; do
        local broker_host="${broker%%:*}"
        local msg_val
        msg_val=$(echo "$j_msgs" | jq -r --arg host "$broker_host" '.data.result[] | select(.metric.instance | startswith($host)) | .value[1]')
        MAP_MSGS_RATE["$broker"]="${msg_val:-0}"
    done

    # Bytes In metrics
    for broker in "${ELIGIBLE_ARRAY[@]}"; do
        local broker_host="${broker%%:*}"
        local bytesin_val
        bytesin_val=$(echo "$j_bytesin" | jq -r --arg host "$broker_host" '.data.result[] | select(.metric.instance | startswith($host)) | .value[1]')
        MAP_BYTES_IN["$broker"]="${bytesin_val:-0}"
    done

    # Bytes Out metrics
    for broker in "${ELIGIBLE_ARRAY[@]}"; do
        local broker_host="${broker%%:*}"
        local bytesout_val
        bytesout_val=$(echo "$j_bytesout" | jq -r --arg host "$broker_host" '.data.result[] | select(.metric.instance | startswith($host)) | .value[1]')
        MAP_BYTES_OUT["$broker"]="${bytesout_val:-0}"
    done

    # Requests Per Second metrics
    for broker in "${ELIGIBLE_ARRAY[@]}"; do
        local broker_host="${broker%%:*}"
        local requests_val
        requests_val=$(echo "$j_requests" | jq -r --arg host "$broker_host" '.data.result[] | select(.metric.instance | startswith($host)) | .value[1]')
        MAP_REQUESTS["$broker"]="${requests_val:-0}"
    done

    echo "HAHAHA DEBUG CPU_RATE: $(for k in "${!CPU_RATE[@]}"; do echo -n "$k=${CPU_RATE[$k]} "; done)"
    echo "HAHAHA DEBUG DISK_BYTES: $(for k in "${!MAP_DISK_BYTES[@]}"; do echo -n "$k=${MAP_DISK_BYTES[$k]} "; done)"
    echo "HAHAHA DEBUG MSGS_RATE: $(for k in "${!MAP_MSGS_RATE[@]}"; do echo -n "$k=${MAP_MSGS_RATE[$k]} "; done)"
    echo "HAHAHA DEBUG BYTES_IN: $(for k in "${!MAP_BYTES_IN[@]}"; do echo -n "$k=${MAP_BYTES_IN[$k]} "; done)"
    echo "HAHAHA DEBUG BYTES_OUT: $(for k in "${!MAP_BYTES_OUT[@]}"; do echo -n "$k=${MAP_BYTES_OUT[$k]} "; done)"
    echo "HAHAHA DEBUG REQUESTS: $(for k in "${!MAP_REQUESTS[@]}"; do echo -n "$k=${MAP_REQUESTS[$k]} "; done)"
    echo "came here 4"
}

# Compute weighted score map
declare -A SCORE
compute_scores() {
  local -A NCPU NDISK NMSG NBYTES_IN NBYTES_OUT NREQUESTS
  while IFS='|' read -r inst n; do NCPU["$inst"]="$n"; done < <(normalize_to01 CPU_RATE)
  while IFS='|' read -r inst n; do NDISK["$inst"]="$n"; done < <(normalize_to01 MAP_DISK_BYTES)
  while IFS='|' read -r inst n; do NMSG["$inst"]="$n"; done < <(normalize_to01 MAP_MSGS_RATE)
  while IFS='|' read -r inst n; do NBYTES_IN["$inst"]="$n"; done < <(normalize_to01 MAP_BYTES_IN)
  while IFS='|' read -r inst n; do NBYTES_OUT["$inst"]="$n"; done < <(normalize_to01 MAP_BYTES_OUT)
  while IFS='|' read -r inst n; do NREQUESTS["$inst"]="$n"; done < <(normalize_to01 MAP_REQUESTS)

  SCORE=()
  for inst in "${ELIGIBLE_ARRAY[@]}"; do
    local s
    s=$(echo "${WEIGHT_CPU}*${NCPU[$inst]:-0} + ${WEIGHT_DISK}*${NDISK[$inst]:-0} + ${WEIGHT_MSGS}*${NMSG[$inst]:-0} + ${WEIGHT_BYTES_IN}*${NBYTES_IN[$inst]:-0} + ${WEIGHT_BYTES_OUT}*${NBYTES_OUT[$inst]:-0} + ${WEIGHT_REQUESTS}*${NREQUESTS[$inst]:-0}" | bc -l)
    SCORE["$inst"]="$s"
  done
}

declare -A OVER
declare -A UNDER
list_over_under() {
  OVER=()
  UNDER=()
  for inst in "${ELIGIBLE_ARRAY[@]}"; do
    if (( $(echo "${SCORE[$inst]} > ${OVERLOAD_SCORE_THRESHOLD}" | bc -l) )); then
      OVER["$inst"]="${SCORE[$inst]}"
    elif (( $(echo "${SCORE[$inst]} < ${UNDERLOAD_SCORE_THRESHOLD}" | bc -l) )); then
      UNDER["$inst"]="${SCORE[$inst]}"
    fi
  done
  log INFO "OVER brokers: $(for k in "${!OVER[@]}"; do echo -n "$k "; done)"
  log INFO "UNDER brokers: $(for k in "${!UNDER[@]}"; do echo -n "$k "; done)"
}

get_overloaded_broker() {
  # find max score in OVER
  local max=0
  local max_broker=""
  for broker in "${!OVER[@]}"; do
    local score="${OVER[$broker]}"
    if (( $(echo "$score > $max" | bc -l) )); then
      max="$score"
      max_broker="$broker"
    fi
  done
  echo "$max_broker"
}

get_underloaded_broker() {
  # find min score in UNDER
  local min=1
  local min_broker=""
  for broker in "${!UNDER[@]}"; do
    local score="${UNDER[$broker]}"
    if (( $(echo "$score < $min" | bc -l) )); then
      min="$score"
      min_broker="$broker"
    fi
  done
  echo "$min_broker"
}

get_leader_partitions() {
  local broker_id="$1"
  local topic_name="${2:-}"
  local partitions_file="/tmp/partitions_${broker_id}.json"
  kafka-topics.sh --bootstrap-server "$KAFKA_BOOTSTRAP_SERVERS" --describe | \
    jq -r --argjson broker_id "$broker_id" '.[] | select(.partitions[] | .leader | select(. == $broker_id)) | {topic, partition}' > "$partitions_file"
  if [[ -n "$topic_name" ]]; then
    jq -r --arg topic "$topic_name" 'select(.topic == $topic)' "$partitions_file"
  else
    cat "$partitions_file"
  fi
}

get_partition_to_move() {
    local overloaded_broker_id=$1
    local underloaded_broker_id=$2
    local topic_to_move=$3
    local partitions_file_over="/tmp/partitions_over_${overloaded_broker_id}.json"
    local partitions_file_under="/tmp/partitions_under_${underloaded_broker_id}.json"
    
    # Get leader partitions for both brokers
    get_leader_partitions "$overloaded_broker_id" > "$partitions_file_over"
    get_leader_partitions "$underloaded_broker_id" > "$partitions_file_under"

    # Find a partition that is not a leader on the underloaded broker
    jq -r --argjson overloaded_id "$overloaded_broker_id" --argjson underloaded_id "$underloaded_broker_id" \
        '
        . as $over_partitions |
        ($under_partitions) |
        . as $under_partitions |
        ($over_partitions | .[] | select(.topic == $topic_to_move)) |
        select( .topic == $topic_to_move and (.partition | select(in($under_partitions[] | .partition) | not)) ) |
        .topic, .partition'
}


gen_plan() {
  local overloaded_broker="$1"
  local underloaded_broker="$2"
  local plan_file="/tmp/plan_${overloaded_broker_id}_${underloaded_broker_id}.json"

  local current_json
  current_json=$(kafka-reassign-partitions.sh --bootstrap-server "$KAFKA_BOOTSTRAP_SERVERS" --topics-to-move-json-file <(echo '{"topics": [], "version": 1}') --generate | grep 'Current partition replica assignment')

  local topics_to_move_list=$(jq -r '.partitions[] | .topic' <<< "$current_json" | sort | uniq)

  # Check if there are any topics
  if [[ -z "$topics_to_move_list" ]]; then
    log INFO "No topics found to generate a plan for."
    return 1
  fi
  
  local partitions_to_move=()
  local topic_to_move
  for topic_to_move in $topics_to_move_list; do
    local partitions
    partitions=$(get_partition_to_move "$(BROKER_ID_MAP_ARRAY["$overloaded_broker"])" "$(BROKER_ID_MAP_ARRAY["$underloaded_broker"])" "$topic_to_move")
    partitions_to_move+=("$partitions")
  done

  if [[ ${#partitions_to_move[@]} -eq 0 ]]; then
    log INFO "No suitable partitions found to move from $overloaded_broker to $underloaded_broker."
    return 1
  fi

  local partitions_json='[]'
  for p in "${partitions_to_move[@]}"; do
    local topic=$(echo "$p" | jq -r '.topic')
    local partition=$(echo "$p" | jq -r '.partition')
    partitions_json=$(echo "$partitions_json" | jq --arg topic "$topic" --argjson partition "$partition" '. + [{topic: $topic, partition: $partition}]')
  done

  echo "$partitions_json" > "$plan_file"
  log INFO "Generated plan file: $plan_file"

  return 0
}

apply_throttle() {
    if [[ "$NEED_THROTTLE" == "Y" ]]; then
        log INFO "Applying throttle..."
        kafka-configs.sh --bootstrap-server "$KAFKA_BOOTSTRAP_SERVERS" --entity-type brokers --entity-name 1 --alter --add-config "leader.replication.throttled.rate=$THROTTLE_VAL_1,follower.replication.throttled.rate=$THROTTLE_VAL_2,replica.throttled.rate=$THROTTLE_VAL_3"
    fi
}

clear_throttle() {
    if [[ "$NEED_THROTTLE" == "Y" ]]; then
        log INFO "Clearing throttle..."
        kafka-configs.sh --bootstrap-server "$KAFKA_BOOTSTRAP_SERVERS" --entity-type brokers --entity-name 1 --alter --delete-config "leader.replication.throttled.rate,follower.replication.throttled.rate,replica.throttled.rate"
    fi
}

rebalance() {
  log INFO "Checking for over/underloaded brokers..."
  local overloaded_broker
  overloaded_broker=$(get_overloaded_broker)
  
  local underloaded_broker
  underloaded_broker=$(get_underloaded_broker)

  if [[ -z "$overloaded_broker" || -z "$underloaded_broker" ]]; then
    log INFO "No suitable overloaded/underloaded brokers found. Skipping rebalance."
    return
  fi

  log INFO "Found overloaded broker: $overloaded_broker (score: ${OVER[$overloaded_broker]})"
  log INFO "Found underloaded broker: $underloaded_broker (score: ${UNDER[$underloaded_broker]})"
  
  local partitions_to_move_count=$(get_leader_partitions "$(BROKER_ID_MAP_ARRAY["$overloaded_broker"])" | jq '. | length')
  local partitions_under_count=$(get_leader_partitions "$(BROKER_ID_MAP_ARRAY["$underloaded_broker"])" | jq '. | length')
  
  if [[ "$MAX_PARTITIONS_PER_BROKER" -gt 0 && "$partitions_under_count" -ge "$MAX_PARTITIONS_PER_BROKER" ]]; then
    log INFO "Underloaded broker $underloaded_broker already has more than $MAX_PARTITIONS_PER_BROKER partitions. Skipping rebalance to this broker."
    return
  fi

  local topics_list=$(jq -r '.[] | .topic' <<< "$(get_leader_partitions "$(BROKER_ID_MAP_ARRAY["$overloaded_broker"])")" | sort | uniq)
  
  for topic in $topics_list; do
    local plan="$(gen_plan "$overloaded_broker" "$underloaded_broker" "$topic")"
    if [[ "$?" -ne 0 ]]; then
      log INFO "gen_plan failed for topic $topic. Skipping."
      continue
    fi
    
    local num_partitions_to_move=$(jq '.partitions | length' "$plan")
    if [[ "$num_partitions_to_move" -eq 0 ]]; then
      log INFO "Found no partitions to move for topic $topic."
      continue
    fi

    log INFO "Proposed plan for $topic:"
    jq '. | {count:(.partitions|length), sample:(.partitions[:10])}' "$plan"

    if [[ "$DRY_RUN" == "Y" ]]; then
      log INFO "DRY_RUN=Y -> not executing. (Plan saved at $plan)"
      continue
    fi
    
    apply_throttle

    log INFO "Executing reassignment for $topic…"
    kafka-reassign-partitions.sh --bootstrap-server "$KAFKA_BOOTSTRAP_SERVERS" \
      --reassignment-json-file "$plan" --execute

    # Verify loop for this topic
    log INFO "Verifying reassignment for $topic…"
    while kafka-reassign-partitions.sh --bootstrap-server "$KAFKA_BOOTSTRAP_SERVERS" \
        --reassignment-json-file "$plan" --verify | grep -q 'is still in progress'; do
      log INFO "Still in progress for $topic… sleeping 20s"
      sleep 20
    done
    log INFO "Reassignment finished for $topic."

    clear_throttle       
    
  done
}

main_loop() {
  while true; do
echo "came here 3"
    get_metrics
echo "came here 4"
    compute_scores
echo "came here 5"
    list_over_under
echo "came here 6"
    #if [[ ${#OVER[@]:-} -gt 0 && ${#UNDER[@]:-} -gt 0 ]]; then
if [[ ${#OVER[@]} -gt 0 && ${#UNDER[@]} -gt 0 ]]; then
echo "came here 7"
      rebalance
    fi
echo "came here 8"
    log INFO "Main loop finished. Sleeping for ${PROMETHEUS_POLL_INTERVAL}s."
    sleep "$PROMETHEUS_POLL_INTERVAL"
  done
}

# --- Main script starts here ---
load_broker_id_map
main_loop