src/main/java/org/apache/sling/distribution/journal/kafka/KafkaClientProvider.java - sling-org-apache-sling-distribution-journal-kafka - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *   http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */
 package org.apache.sling.distribution.journal.kafka;

 import static java.lang.String.format;
 import static java.util.Collections.singleton;
 import static java.util.Collections.unmodifiableMap;
 import static java.util.Objects.requireNonNull;
 import static org.apache.kafka.clients.CommonClientConfigs.SECURITY_PROTOCOL_CONFIG;
 import static org.apache.kafka.clients.consumer.ConsumerConfig.AUTO_OFFSET_RESET_CONFIG;
 import static org.apache.kafka.clients.consumer.ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG;
 import static org.apache.kafka.clients.consumer.ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG;
 import static org.apache.kafka.clients.consumer.ConsumerConfig.GROUP_ID_CONFIG;
 import static org.apache.kafka.clients.consumer.ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG;
 import static org.apache.kafka.clients.consumer.ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG;
 import static org.apache.kafka.clients.consumer.ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG;
 import static org.apache.kafka.clients.producer.ProducerConfig.ACKS_CONFIG;
 import static org.apache.kafka.clients.producer.ProducerConfig.BOOTSTRAP_SERVERS_CONFIG;
 import static org.apache.kafka.clients.producer.ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG;
 import static org.apache.kafka.clients.producer.ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG;
 import static org.apache.kafka.common.config.SaslConfigs.SASL_JAAS_CONFIG;
 import static org.apache.kafka.common.config.SaslConfigs.SASL_MECHANISM;

 import java.io.Closeable;
 import java.io.IOException;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;

 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;

 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
 import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.serialization.StringDeserializer;
 import org.apache.kafka.common.serialization.StringSerializer;
 import org.apache.sling.distribution.journal.HandlerAdapter;
 import org.apache.sling.distribution.journal.JsonMessageSender;
 import org.apache.sling.distribution.journal.MessageHandler;
 import org.apache.sling.distribution.journal.MessageSender;
 import org.apache.sling.distribution.journal.MessagingException;
 import org.apache.sling.distribution.journal.MessagingProvider;
 import org.apache.sling.distribution.journal.Reset;
 import org.osgi.service.component.annotations.Activate;
 import org.osgi.service.component.annotations.Component;
 import org.osgi.service.component.annotations.ConfigurationPolicy;
 import org.osgi.service.component.annotations.Deactivate;
 import org.osgi.service.metatype.annotations.Designate;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import com.google.protobuf.GeneratedMessage;

 @Component(service = MessagingProvider.class, configurationPolicy = ConfigurationPolicy.REQUIRE)
 @Designate(ocd = KafkaEndpoint.class)
 public class KafkaClientProvider implements MessagingProvider, Closeable {

     private static final Logger LOG = LoggerFactory.getLogger(KafkaClientProvider.class);

     public static final int PARTITION = 0;

     private volatile KafkaProducer<String, byte[]> rawProducer = null;

     private volatile KafkaProducer<String, String> jsonProducer = null;

     private String kafkaBootstrapServers;

     private int requestTimeout;

     private int defaultApiTimeout;

     private String securityProtocol;

     private String saslMechanism;

     private String saslJaasConfig;

     @Activate
     public void activate(KafkaEndpoint kafkaEndpoint) {
         kafkaBootstrapServers = requireNonNull(kafkaEndpoint.kafkaBootstrapServers());
         requestTimeout = kafkaEndpoint.kafkaRequestTimeout();
         defaultApiTimeout = kafkaEndpoint.kafkaDefaultApiTimeout();
         securityProtocol = kafkaEndpoint.securityProtocol();
         saslMechanism = kafkaEndpoint.saslMechanism();
         saslJaasConfig = kafkaEndpoint.saslJaasConfig();
     }

     @Deactivate
     public void close() {
         closeQuietly(rawProducer);
         closeQuietly(jsonProducer);
     }

     @Override
     public <T extends GeneratedMessage> MessageSender<T> createSender() {
         return new KafkaMessageSender<>(buildKafkaProducer());
     }

     @Override
     public <T> Closeable createPoller(String topicName, Reset reset, HandlerAdapter<?>... adapters) {
         return createPoller(topicName, reset, null, adapters);
     }

     @Override
     public Closeable createPoller(String topicName, Reset reset, @Nullable String assign, HandlerAdapter<?>... adapters) {
         KafkaConsumer<String, byte[]> consumer = createConsumer(ByteArrayDeserializer.class, reset);
         TopicPartition topicPartition = new TopicPartition(topicName, PARTITION);
         Collection<TopicPartition> topicPartitions = singleton(topicPartition);
         consumer.assign(topicPartitions);
         if (assign != null) {
             consumer.seek(topicPartition, offset(assign));
         } else if (reset == Reset.earliest) {
             consumer.seekToBeginning(topicPartitions);
         } else {
             consumer.seekToEnd(topicPartitions);
         }
         Closeable poller = new KafkaMessagePoller(consumer, adapters);
         LOG.info("Created poller for reset {}, topicName {}, assign {}", reset, topicName, assign);
         return poller;
     }

     @Override
     public <T> JsonMessageSender<T> createJsonSender() {
         return new KafkaJsonMessageSender<>(buildJsonKafkaProducer());
     }

     @Override
     public <T> Closeable createJsonPoller(String topicName, Reset reset, MessageHandler<T> handler, Class<T> type) {
         KafkaConsumer<String, String> consumer = createConsumer(StringDeserializer.class, reset);
         TopicPartition topicPartition = new TopicPartition(topicName, PARTITION);
         Collection<TopicPartition> topicPartitions = singleton(topicPartition);
         consumer.assign(topicPartitions);
         if (reset == Reset.earliest) {
             consumer.seekToBeginning(topicPartitions);
         } else {
             consumer.seekToEnd(topicPartitions);
         }
         return new KafkaJsonMessagePoller<>(consumer, handler, type);
     }

     @Override
     public void assertTopic(String topic) throws MessagingException {
         Map<String, List<PartitionInfo>> topics;
         try (KafkaConsumer<String, String> consumer = createConsumer(StringDeserializer.class, Reset.latest)) {
             topics = consumer.listTopics();
         } catch (Exception e) {
             throw new MessagingException(format("Unable to load topic stats for %s", topic), e);
         }
         if (! topics.containsKey(topic)) {
             throw new MessagingException(format("Topic %s does not exist", topic));
         }
     }

     @Override
     public long retrieveOffset(String topicName, Reset reset) {
         try (KafkaConsumer<String, String> consumer = createConsumer(StringDeserializer.class, reset)) {;
             TopicPartition topicPartition = new TopicPartition(topicName, PARTITION);
             Map<TopicPartition, Long> offsets = getOffsets(reset, consumer, topicPartition);
             Long offset = offsets.get(topicPartition);
             return offset;
         }
     }

     private Map<TopicPartition, Long> getOffsets(Reset reset, KafkaConsumer<String, String> consumer,
             TopicPartition topicPartition) {
         Collection<TopicPartition> topicPartitions = singleton(topicPartition);
         return (reset == Reset.earliest) //
                 ? consumer.beginningOffsets(topicPartitions)
                 : consumer.endOffsets(topicPartitions);
     }

     @Override
     public String assignTo(long offset) {
         return format("%s:%s", PARTITION, offset);
     }

     protected <T> KafkaConsumer<String, T> createConsumer(Class<? extends Deserializer<?>> deserializer, Reset reset) {
         String groupId = UUID.randomUUID().toString();
         ClassLoader oldClassloader = Thread.currentThread().getContextClassLoader();
         Thread.currentThread().setContextClassLoader(KafkaConsumer.class.getClassLoader());
         try {
             return new KafkaConsumer<>(consumerConfig(deserializer, groupId, reset));
         } finally {
             Thread.currentThread().setContextClassLoader(oldClassloader);
         }
     }

     private void closeQuietly(final Closeable closeable) {
         try {
             if (closeable != null) {
                 closeable.close();
             }
         } catch (final IOException ioe) {
             // ignore
         }
     }

     @Nonnull
     private synchronized KafkaProducer<String, byte[]> buildKafkaProducer() {
         if (rawProducer == null) {
             ClassLoader oldClassloader = Thread.currentThread().getContextClassLoader();
             Thread.currentThread().setContextClassLoader(KafkaProducer.class.getClassLoader());
             try {
                 rawProducer = new KafkaProducer<>(producerConfig(ByteArraySerializer.class));
             } finally {
                 Thread.currentThread().setContextClassLoader(oldClassloader);
             }

             rawProducer = new KafkaProducer<>(producerConfig(ByteArraySerializer.class));
         }
         return rawProducer;
     }

     @Nonnull
     private synchronized KafkaProducer<String, String> buildJsonKafkaProducer() {
         if (jsonProducer == null) {
             jsonProducer = new KafkaProducer<>(producerConfig(StringSerializer.class));
         }
         return jsonProducer;
     }

     private Map<String, Object> consumerConfig(Object deserializer, String consumerGroupId, Reset reset) {
         Map<String, Object> config = commonConfig();
         config.put(GROUP_ID_CONFIG, consumerGroupId);
         config.put(ENABLE_AUTO_COMMIT_CONFIG, false);
         config.put(DEFAULT_API_TIMEOUT_MS_CONFIG, defaultApiTimeout);
         config.put(KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
         config.put(VALUE_DESERIALIZER_CLASS_CONFIG, deserializer);
         config.put(AUTO_OFFSET_RESET_CONFIG, reset.name());
         return config;
     }

     private Map<String, Object> producerConfig(Object serializer) {
         Map<String, Object> config = commonConfig();
         config.put(REQUEST_TIMEOUT_MS_CONFIG, requestTimeout);
         config.put(KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class);
         config.put(VALUE_SERIALIZER_CLASS_CONFIG, serializer);
         config.put(ACKS_CONFIG, "all");
         return unmodifiableMap(config);
     }

     private Map<String, Object> commonConfig() {
         Map<String, Object> config = new HashMap<>();
         config.put(BOOTSTRAP_SERVERS_CONFIG, kafkaBootstrapServers);
         config.put(SASL_MECHANISM, saslMechanism);
         config.put(SECURITY_PROTOCOL_CONFIG, securityProtocol);
         config.put(SASL_JAAS_CONFIG, saslJaasConfig);
         return config;
     }

     private long offset(String assign) {
         String[] chunks = assign.split(":");
         if (chunks.length != 2) {
             throw new IllegalArgumentException(format("Illegal assign %s", assign));
         }
         return Long.parseLong(chunks[1]);
     }


 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing,
	* software distributed under the License is distributed on an
	* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	* KIND, either express or implied. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*/
	package org.apache.sling.distribution.journal.kafka;

	import static java.lang.String.format;
	import static java.util.Collections.singleton;
	import static java.util.Collections.unmodifiableMap;
	import static java.util.Objects.requireNonNull;
	import static org.apache.kafka.clients.CommonClientConfigs.SECURITY_PROTOCOL_CONFIG;
	import static org.apache.kafka.clients.consumer.ConsumerConfig.AUTO_OFFSET_RESET_CONFIG;
	import static org.apache.kafka.clients.consumer.ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG;
	import static org.apache.kafka.clients.consumer.ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG;
	import static org.apache.kafka.clients.consumer.ConsumerConfig.GROUP_ID_CONFIG;
	import static org.apache.kafka.clients.consumer.ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG;
	import static org.apache.kafka.clients.consumer.ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG;
	import static org.apache.kafka.clients.consumer.ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG;
	import static org.apache.kafka.clients.producer.ProducerConfig.ACKS_CONFIG;
	import static org.apache.kafka.clients.producer.ProducerConfig.BOOTSTRAP_SERVERS_CONFIG;
	import static org.apache.kafka.clients.producer.ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG;
	import static org.apache.kafka.clients.producer.ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG;
	import static org.apache.kafka.common.config.SaslConfigs.SASL_JAAS_CONFIG;
	import static org.apache.kafka.common.config.SaslConfigs.SASL_MECHANISM;

	import java.io.Closeable;
	import java.io.IOException;
	import java.util.Collection;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;
	import java.util.UUID;

	import javax.annotation.Nonnull;
	import javax.annotation.Nullable;

	import org.apache.kafka.clients.consumer.KafkaConsumer;
	import org.apache.kafka.clients.producer.KafkaProducer;
	import org.apache.kafka.common.PartitionInfo;
	import org.apache.kafka.common.TopicPartition;
	import org.apache.kafka.common.serialization.ByteArrayDeserializer;
	import org.apache.kafka.common.serialization.ByteArraySerializer;
	import org.apache.kafka.common.serialization.Deserializer;
	import org.apache.kafka.common.serialization.StringDeserializer;
	import org.apache.kafka.common.serialization.StringSerializer;
	import org.apache.sling.distribution.journal.HandlerAdapter;
	import org.apache.sling.distribution.journal.JsonMessageSender;
	import org.apache.sling.distribution.journal.MessageHandler;
	import org.apache.sling.distribution.journal.MessageSender;
	import org.apache.sling.distribution.journal.MessagingException;
	import org.apache.sling.distribution.journal.MessagingProvider;
	import org.apache.sling.distribution.journal.Reset;
	import org.osgi.service.component.annotations.Activate;
	import org.osgi.service.component.annotations.Component;
	import org.osgi.service.component.annotations.ConfigurationPolicy;
	import org.osgi.service.component.annotations.Deactivate;
	import org.osgi.service.metatype.annotations.Designate;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import com.google.protobuf.GeneratedMessage;

	@Component(service = MessagingProvider.class, configurationPolicy = ConfigurationPolicy.REQUIRE)
	@Designate(ocd = KafkaEndpoint.class)
	public class KafkaClientProvider implements MessagingProvider, Closeable {

	private static final Logger LOG = LoggerFactory.getLogger(KafkaClientProvider.class);

	public static final int PARTITION = 0;

	private volatile KafkaProducer<String, byte[]> rawProducer = null;

	private volatile KafkaProducer<String, String> jsonProducer = null;

	private String kafkaBootstrapServers;

	private int requestTimeout;

	private int defaultApiTimeout;

	private String securityProtocol;

	private String saslMechanism;

	private String saslJaasConfig;

	@Activate
	public void activate(KafkaEndpoint kafkaEndpoint) {
	kafkaBootstrapServers = requireNonNull(kafkaEndpoint.kafkaBootstrapServers());
	requestTimeout = kafkaEndpoint.kafkaRequestTimeout();
	defaultApiTimeout = kafkaEndpoint.kafkaDefaultApiTimeout();
	securityProtocol = kafkaEndpoint.securityProtocol();
	saslMechanism = kafkaEndpoint.saslMechanism();
	saslJaasConfig = kafkaEndpoint.saslJaasConfig();
	}

	@Deactivate
	public void close() {
	closeQuietly(rawProducer);
	closeQuietly(jsonProducer);
	}

	@Override
	public <T extends GeneratedMessage> MessageSender<T> createSender() {
	return new KafkaMessageSender<>(buildKafkaProducer());
	}

	@Override
	public <T> Closeable createPoller(String topicName, Reset reset, HandlerAdapter<?>... adapters) {
	return createPoller(topicName, reset, null, adapters);
	}

	@Override
	public Closeable createPoller(String topicName, Reset reset, @Nullable String assign, HandlerAdapter<?>... adapters) {
	KafkaConsumer<String, byte[]> consumer = createConsumer(ByteArrayDeserializer.class, reset);
	TopicPartition topicPartition = new TopicPartition(topicName, PARTITION);
	Collection<TopicPartition> topicPartitions = singleton(topicPartition);
	consumer.assign(topicPartitions);
	if (assign != null) {
	consumer.seek(topicPartition, offset(assign));
	} else if (reset == Reset.earliest) {
	consumer.seekToBeginning(topicPartitions);
	} else {
	consumer.seekToEnd(topicPartitions);
	}
	Closeable poller = new KafkaMessagePoller(consumer, adapters);
	LOG.info("Created poller for reset {}, topicName {}, assign {}", reset, topicName, assign);
	return poller;
	}

	@Override
	public <T> JsonMessageSender<T> createJsonSender() {
	return new KafkaJsonMessageSender<>(buildJsonKafkaProducer());
	}

	@Override
	public <T> Closeable createJsonPoller(String topicName, Reset reset, MessageHandler<T> handler, Class<T> type) {
	KafkaConsumer<String, String> consumer = createConsumer(StringDeserializer.class, reset);
	TopicPartition topicPartition = new TopicPartition(topicName, PARTITION);
	Collection<TopicPartition> topicPartitions = singleton(topicPartition);
	consumer.assign(topicPartitions);
	if (reset == Reset.earliest) {
	consumer.seekToBeginning(topicPartitions);
	} else {
	consumer.seekToEnd(topicPartitions);
	}
	return new KafkaJsonMessagePoller<>(consumer, handler, type);
	}

	@Override
	public void assertTopic(String topic) throws MessagingException {
	Map<String, List<PartitionInfo>> topics;
	try (KafkaConsumer<String, String> consumer = createConsumer(StringDeserializer.class, Reset.latest)) {
	topics = consumer.listTopics();
	} catch (Exception e) {
	throw new MessagingException(format("Unable to load topic stats for %s", topic), e);
	}
	if (! topics.containsKey(topic)) {
	throw new MessagingException(format("Topic %s does not exist", topic));
	}
	}

	@Override
	public long retrieveOffset(String topicName, Reset reset) {
	try (KafkaConsumer<String, String> consumer = createConsumer(StringDeserializer.class, reset)) {;
	TopicPartition topicPartition = new TopicPartition(topicName, PARTITION);
	Map<TopicPartition, Long> offsets = getOffsets(reset, consumer, topicPartition);
	Long offset = offsets.get(topicPartition);
	return offset;
	}
	}

	private Map<TopicPartition, Long> getOffsets(Reset reset, KafkaConsumer<String, String> consumer,
	TopicPartition topicPartition) {
	Collection<TopicPartition> topicPartitions = singleton(topicPartition);
	return (reset == Reset.earliest) //
	? consumer.beginningOffsets(topicPartitions)
	: consumer.endOffsets(topicPartitions);
	}

	@Override
	public String assignTo(long offset) {
	return format("%s:%s", PARTITION, offset);
	}

	protected <T> KafkaConsumer<String, T> createConsumer(Class<? extends Deserializer<?>> deserializer, Reset reset) {
	String groupId = UUID.randomUUID().toString();
	ClassLoader oldClassloader = Thread.currentThread().getContextClassLoader();
	Thread.currentThread().setContextClassLoader(KafkaConsumer.class.getClassLoader());
	try {
	return new KafkaConsumer<>(consumerConfig(deserializer, groupId, reset));
	} finally {
	Thread.currentThread().setContextClassLoader(oldClassloader);
	}
	}

	private void closeQuietly(final Closeable closeable) {
	try {
	if (closeable != null) {
	closeable.close();
	}
	} catch (final IOException ioe) {
	// ignore
	}
	}

	@Nonnull
	private synchronized KafkaProducer<String, byte[]> buildKafkaProducer() {
	if (rawProducer == null) {
	ClassLoader oldClassloader = Thread.currentThread().getContextClassLoader();
	Thread.currentThread().setContextClassLoader(KafkaProducer.class.getClassLoader());
	try {
	rawProducer = new KafkaProducer<>(producerConfig(ByteArraySerializer.class));
	} finally {
	Thread.currentThread().setContextClassLoader(oldClassloader);
	}

	rawProducer = new KafkaProducer<>(producerConfig(ByteArraySerializer.class));
	}
	return rawProducer;
	}

	@Nonnull
	private synchronized KafkaProducer<String, String> buildJsonKafkaProducer() {
	if (jsonProducer == null) {
	jsonProducer = new KafkaProducer<>(producerConfig(StringSerializer.class));
	}
	return jsonProducer;
	}

	private Map<String, Object> consumerConfig(Object deserializer, String consumerGroupId, Reset reset) {
	Map<String, Object> config = commonConfig();
	config.put(GROUP_ID_CONFIG, consumerGroupId);
	config.put(ENABLE_AUTO_COMMIT_CONFIG, false);
	config.put(DEFAULT_API_TIMEOUT_MS_CONFIG, defaultApiTimeout);
	config.put(KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
	config.put(VALUE_DESERIALIZER_CLASS_CONFIG, deserializer);
	config.put(AUTO_OFFSET_RESET_CONFIG, reset.name());
	return config;
	}

	private Map<String, Object> producerConfig(Object serializer) {
	Map<String, Object> config = commonConfig();
	config.put(REQUEST_TIMEOUT_MS_CONFIG, requestTimeout);
	config.put(KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class);
	config.put(VALUE_SERIALIZER_CLASS_CONFIG, serializer);
	config.put(ACKS_CONFIG, "all");
	return unmodifiableMap(config);
	}

	private Map<String, Object> commonConfig() {
	Map<String, Object> config = new HashMap<>();
	config.put(BOOTSTRAP_SERVERS_CONFIG, kafkaBootstrapServers);
	config.put(SASL_MECHANISM, saslMechanism);
	config.put(SECURITY_PROTOCOL_CONFIG, securityProtocol);
	config.put(SASL_JAAS_CONFIG, saslJaasConfig);
	return config;
	}

	private long offset(String assign) {
	String[] chunks = assign.split(":");
	if (chunks.length != 2) {
	throw new IllegalArgumentException(format("Illegal assign %s", assign));
	}
	return Long.parseLong(chunks[1]);
	}


	}