Rebalancing a Kafka Cluster in AWS MSK using CLI Commands

Apache Kafka is a widely used distributed streaming platform known for its fault tolerance, scalability, and performance. Amazon Managed Streaming for Apache Kafka (MSK) simplifies the setup, management, and monitoring of Kafka clusters in the AWS environment. In this guide, we’ll walk through the process of rebalancing a Kafka cluster hosted on AWS MSK using CLI commands.

Prerequisites

  1. Access to an AWS account with permissions to manage MSK clusters.
  2. Installed Kafka CLI tools.
  3. Basic understanding of Kafka architecture and concepts.

Steps to Rebalance a Kafka Cluster

Step 1: Navigate to Kafka Bin Directory

cd /path/to/kafka/bin

Step 2: Set Broker URL as a Variable

export bs={Replace_this_with_broker_url}

Step 3: Ensure that you have the necessary permissions to view the topic’s configuration and current statistics. Below is the output displaying this information.

./kafka-topics.sh --bootstrap-server $bs --describe --topic test-topic

Topic: test-topic Partition: 0    Leader: 6       Replicas: 6,4,5 Isr: 6,5,4
Topic: test-topic Partition: 1    Leader: 6       Replicas: 6,5,4 Isr: 6,5,4
Topic: test-topic Partition: 2    Leader: 6       Replicas: 6,4,5 Isr: 6,5,4
Topic: test-topic Partition: 3    Leader: 6       Replicas: 6,5,4 Isr: 6,5,4
Topic: test-topic Partition: 4    Leader: 6       Replicas: 6,5,4 Isr: 6,5,4
Topic: test-topic Partition: 5    Leader: 6       Replicas: 6,4,5 Isr: 6,5,4

If access errors occur (common in AWS MSK), set credentials:

export AWS_ACCESS_KEY_ID="replace_access_key"
export AWS_SECRET_ACCESS_KEY="replace_secret_key"
export AWS_SESSION_TOKEN="replace_session_token"

If you’re attempting to rebalance a topic authenticated by IAM, an additional configuration is required. Create a file and insert the provided content, then specify the location of this file using the –command-config option as demonstrated below.

security.protocol=SASL_SSL
sasl.mechanism=AWS_MSK_IAM
sasl.jaas.config=software.amazon.msk.auth.iam.IAMLoginModule required;
sasl.client.callback.handler.class=software.amazon.msk.auth.iam.IAMClientCallbackHandler
./kafka-topics.sh --bootstrap-server $bs --command-config /etc/kafka/iam-client.config --describe --topic test-topic

Step 4 : Generate a JSON configuration file to specify the list of topics you wish to rebalance as follows:

touch /tmp/test-topic.json

cat <<EOF >> /tmp/test-topic.json
{
  "topics": [
    {
      "topic": "test-topic"
    }
  ],
  "version": 1
}
EOF

Execute the following command, providing the configured file, to retrieve the existing and proposed partition and replica mappings. For instance, if your cluster comprises 6 brokers and the current partition distribution is across 3, 4, and 5, and you intend to shift it to 1,2 and 6 utilize the command below to obtain the Current and Proposed partition replica assignment:

./kafka-reassign-partitions.sh --bootstrap-server $bs --command-config /etc/kafka/iam-client.config --broker-list "1,2,6" --topics-to-move-json-file /tmp/test-topic.json --generate

The result of executing the aforementioned command appears as follows:

Current partition replica assignment

{
  "version":1,
  "partitions":[
     {
        "topic":"test-topic",
        "partition":0,
        "replicas":[
           4,
           5,
           3
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":1,
        "replicas":[
           3,
           4,
           5
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":2,
        "replicas":[
           5,
           3,
           4
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":3,
        "replicas":[
           4,
           3,
           5
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":4,
        "replicas":[
           3,
           5,
           4
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":5,
        "replicas":[
           5,
           4,
           3
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     }
  ]
}

Proposed partition reassignment configuration

{
  "version":1,
  "partitions":[
     {
        "topic":"test-topic",
        "partition":0,
        "replicas":[
           1,
           2,
           6
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":1,
        "replicas":[
           2,
           6,
           1
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":2,
        "replicas":[
           6,
           1,
           2
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":3,
        "replicas":[
           1,
           6,
           2
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":4,
        "replicas":[
           2,
           1,
           6
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":5,
        "replicas":[
           6,
           2,
           1
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     }
  ]
}

Step 6: Transfer the data from the Proposed partition reassignment configuration into a file as below

cat <<EOF >> /tmp/reassignment-test-topic.json
{
  "version":1,
  "partitions":[
     {
        "topic":"test-topic",
        "partition":0,
        "replicas":[
           1,
           2,
           6
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":1,
        "replicas":[
           2,
           6,
           1
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":2,
        "replicas":[
           6,
           1,
           2
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":3,
        "replicas":[
           1,
           6,
           2
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":4,
        "replicas":[
           2,
           1,
           6
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     },
     {
        "topic":"test-topic",
        "partition":5,
        "replicas":[
           6,
           2,
           1
        ],
        "log_dirs":[
           "any",
           "any",
           "any"
        ]
     }
  ]
}
EOF

Step 7 : The last step involves utilizing the aforementioned reassignment configuration file and executing the following command, initiating the reassignment of partitions and replicas across brokers.

./kafka-reassign-partitions.sh --bootstrap-server $bs --command-config /etc/kafka/iam-client.config --reassignment-json-file /tmp/reassignment-test-topic.json --execute

Conclusion

Rebalancing a Kafka cluster ensures even distribution of data across brokers, optimizing performance and fault tolerance. With AWS MSK and CLI commands, administrators can efficiently manage and rebalance Kafka clusters to meet changing workload demands. Remember to carefully review proposed changes before executing reassignment to avoid unintended consequences.