com.landoop.lenses.topology.client.kafka.metrics.TopologyKafkaStreamsClientSupplier Maven / Gradle / Ivy

Go to download
package com.landoop.lenses.topology.client.kafka.metrics;

import com.landoop.lenses.topology.client.Topology;
import com.landoop.lenses.topology.client.TopologyClient;
import com.landoop.lenses.topology.client.metrics.MetricsBuilder;
import org.apache.kafka.clients.admin.AdminClient;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.apache.kafka.streams.KafkaClientSupplier;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * An implementation of a KafkaClientSupplier to be used with KStreams that will
 * register metrics with the topology client.
 */
public class TopologyKafkaStreamsClientSupplier implements KafkaClientSupplier {

    private final TopologyClient client;
    private final Topology topology;

    public TopologyKafkaStreamsClientSupplier(TopologyClient client, Topology topology) {
        this.client = client;
        this.topology = topology;
    }

    @Override
    public AdminClient getAdminClient(Map config) {
        return AdminClient.create(config);
    }

    @Override
    public Producer getProducer(Map config) {
        KafkaProducer producer = new KafkaProducer<>(config, new ByteArraySerializer(), new ByteArraySerializer());

        List builders = new ArrayList<>();
        for (String topic : topology.getTopics()) {
            KafkaMetricsBuilder builder = new KafkaMetricsBuilder(producer);
            builders.add(builder);
            client.register(topology.getAppName(), topic, builder);
        }
        Runnable callback = () -> {
            for (MetricsBuilder builder : builders) {
                client.unregister(builder);
            }
        };
        return new CallbackProducer<>(producer, callback);
    }

    @Override
    public Consumer getConsumer(Map config) {
        KafkaConsumer consumer = new KafkaConsumer<>(config, new ByteArrayDeserializer(), new ByteArrayDeserializer());

        List builders = new ArrayList<>();
        for (String topic : topology.getTopics()) {
            KafkaMetricsBuilder builder = new KafkaMetricsBuilder(consumer);
            builders.add(builder);
            client.register(topology.getAppName(), topic, builder);
        }
        Runnable callback = () -> {
            for (MetricsBuilder builder : builders) {
                client.unregister(builder);
            }
        };
        return new CallbackConsumer<>(consumer, callback);
    }

    @Override
    public Consumer getRestoreConsumer(Map config) {
        return new KafkaConsumer<>(config, new ByteArrayDeserializer(), new ByteArrayDeserializer());
    }
}