pl.allegro.tech.hermes.common.kafka.KafkaConsumerPool Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of hermes-common Show documentation
Fast and reliable message broker built on top of Kafka.
The newest version!
package pl.allegro.tech.hermes.common.kafka;

import static org.apache.kafka.clients.CommonClientConfigs.SECURITY_PROTOCOL_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.FETCH_MIN_BYTES_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.GROUP_ID_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.RECEIVE_BUFFER_CONFIG;
import static org.apache.kafka.common.config.SaslConfigs.SASL_JAAS_CONFIG;
import static org.apache.kafka.common.config.SaslConfigs.SASL_MECHANISM;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.cache.RemovalListener;
import com.google.common.cache.RemovalNotification;
import com.google.common.util.concurrent.UncheckedExecutionException;
import java.util.Properties;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.common.TopicPartition;
import pl.allegro.tech.hermes.common.broker.BrokerStorage;
import pl.allegro.tech.hermes.common.exception.BrokerNotFoundForPartitionException;

/**
 * This class help us to avoid unnecessarily creating new kafka consumers for the same broker
 * instance, mainly in case of retransmission and migrating topic from json to avro. We map topic
 * and its partitions to specific kafka broker which without caching would lead to creating exactly
 * the same consumer multiple times in short period of time. Additionaly consumers created here are
 * rarely used, so we don't bother about concurrency and thread safety.
 */
public class KafkaConsumerPool {

  private final LoadingCache> kafkaConsumers;
  private final BrokerStorage brokerStorage;

  public KafkaConsumerPool(
      KafkaConsumerPoolConfig poolConfig,
      BrokerStorage brokerStorage,
      String configuredBootstrapServers) {
    this.brokerStorage = brokerStorage;
    this.kafkaConsumers =
        CacheBuilder.newBuilder()
            .expireAfterAccess(poolConfig.getCacheExpirationSeconds(), TimeUnit.SECONDS)
            .removalListener(new KafkaConsumerRemoveListener())
            .build(new KafkaConsumerSupplier(poolConfig, configuredBootstrapServers));
  }

  public KafkaConsumer get(KafkaTopic topic, int partition) {
    return get(topic.name().asString(), partition);
  }

  public KafkaConsumer get(String topicName, int partition) {
    try {
      int leaderId = brokerStorage.readLeaderForPartition(new TopicPartition(topicName, partition));
      return kafkaConsumers.get(leaderId);

    } catch (ExecutionException e) {
      String message =
          String.format(
              "Cannot get KafkaConsumer for topic %s and partition %d", topicName, partition);
      throw new KafkaConsumerPoolException(message, e);
    } catch (UncheckedExecutionException e) {
      if (e.getCause() instanceof BrokerNotFoundForPartitionException) {
        throw (BrokerNotFoundForPartitionException) e.getCause();
      }
      throw e;
    }
  }

  private static class KafkaConsumerSupplier
      extends CacheLoader> {

    private final KafkaConsumerPoolConfig poolConfig;
    private final String configuredBootstrapServers;

    KafkaConsumerSupplier(KafkaConsumerPoolConfig poolConfig, String configuredBootstrapServers) {
      this.poolConfig = poolConfig;
      this.configuredBootstrapServers = configuredBootstrapServers;
    }

    @Override
    public KafkaConsumer load(Integer leaderId) throws Exception {
      return createKafkaConsumer();
    }

    private KafkaConsumer createKafkaConsumer() {

      Properties props = new Properties();
      props.put(BOOTSTRAP_SERVERS_CONFIG, configuredBootstrapServers);
      props.put(
          GROUP_ID_CONFIG, poolConfig.getIdPrefix() + "_" + poolConfig.getConsumerGroupName());
      props.put(RECEIVE_BUFFER_CONFIG, poolConfig.getBufferSizeBytes());
      props.put(ENABLE_AUTO_COMMIT_CONFIG, false);
      props.put(FETCH_MAX_WAIT_MS_CONFIG, poolConfig.getFetchMaxWaitMillis());
      props.put(FETCH_MIN_BYTES_CONFIG, poolConfig.getFetchMinBytes());
      props.put("key.deserializer", "org.apache.kafka.common.serialization.ByteArrayDeserializer");
      props.put(
          "value.deserializer", "org.apache.kafka.common.serialization.ByteArrayDeserializer");

      if (poolConfig.isSaslEnabled()) {
        props.put(SASL_MECHANISM, poolConfig.getSecurityMechanism());
        props.put(SECURITY_PROTOCOL_CONFIG, poolConfig.getSecurityProtocol());
        props.put(SASL_JAAS_CONFIG, poolConfig.getSaslJaasConfig());
      }
      return new KafkaConsumer<>(props);
    }
  }

  private static class KafkaConsumerRemoveListener
      implements RemovalListener> {
    @Override
    public void onRemoval(
        RemovalNotification> notification) {
      notification.getValue().close();
    }
  }
}