io.lenses.topology.client.kafka.metrics.TopologyKafkaStreamsClientSupplier Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of lenses-topology-client-kafka Show documentation
lenses-topology-client
The newest version!
package io.lenses.topology.client.kafka.metrics;

import io.lenses.topology.client.Topology;
import io.lenses.topology.client.TopologyClient;
import io.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());
  }

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