Leveraging CountDownLatch to Control Kafka Producer in Asynchronous Messaging

Understanding Kafka Producer

Apache Kafka is a powerful tool for handling real-time data streams. The Kafka producer is responsible for sending data to the Kafka cluster. An essential feature of the Kafka producer is its ability to send messages asynchronously, which can significantly enhance performance in distributed systems.

Asynchronous Message Sending

When sending messages using a Kafka producer, you can choose to send them asynchronously. This means that your application can continue executing without waiting for the broker to acknowledge that the message has been sent.

To achieve this, the producer uses a callback mechanism that notifies you once the message is successfully delivered or if it encounters an error.

Following is a simple implementation that demonstrates the callback:

String topicName = "my-topic";
String message = "test message";

// create producer config
Properties config = new Properties();
config.put("bootstrap.servers", "localhost:9092");
config.put("key.serializer", StringSerializer.class.getName());
config.put("value.serializer", StringSerializer.class.getName());

// create producer object using the above config
KafkaProducer<String, String> producer = new KafkaProducer<>(config);

// create an object of ProducerRecord which represents 1 message
ProducerRecord<String, String> record = new ProducerRecord(topicName, message);

// send the message asynchronously using send method of the KafkaProducer class
producer.send(record, (metadata, exception) -> {
    // this code block is a callback function that gets called once 
    // each message is sent.. it also lets us know if there was 
    // an exception while sending the message

    if (exception != null) {
        // log the exception
    } else {
        System.out.println("Message sent successfully: " + metadata.toString());
    }
});

Implementing CountDownLatch

To wait for the completion of all messages being sent, you can utilize the CountDownLatch class from the Java concurrency package.

The idea is to create a CountDownLatch initialized with the number of messages you are sending. As each message is sent, decrement the latch. Once all messages are acknowledged, the latch reaches zero, allowing your application to proceed.

Here’s a simple implementation outline:

String topicName = "my-topic";

// total number of messages to send
int totalMessages = 1000;

// CountDownLatch to help us keep track of the messages sent so far
CountDownLatch latch = new CountDownLatch (totalMessages);

// send messages in a loop
for (int i = 0; i < totalMessages; i++) {
    String message = "test message - " + i;
    ProducerRecord<String, String> record = new ProducerRecord<>(topicName, message);
    
    producer.send (record, (metadata, exception) -> 
    // callback handler
    {
        // countDown the latch first
        latch.countDown();
        
        // check for exception
        if (exception != null) {
            // log the exception
        }
        else {
            System.out.println("Message sent successfully");        
        }
    }
}

// Now wait for the latch to count down to zero
try {
    // wait infinitely
    latch.await ();
    
    // or wait with a timeout
    // latch.await (1, TimeUnit.MINUTES);
}
catch (InterruptedException ex) {
    // log the exception
}

This approach not only provides the benefit of asynchronous sending but also ensures you can manage the program’s flow based on the sending status of the messages.