Kafka native offset implementation

Fixes #79 (https://github.com/spring-projects/spring-integration-kafka/pull/79)

Adds an implementation for offset management based on Kafka's own offset management support

Update copyright and use fill name in author tag.

Refactoring the native offset manager

Remove unused imports

Polishing

spring-integration-kafka/src/main/java/org/springframework/integration/kafka/core/AbstractConfiguration.java

@@ -47,8 +47,9 @@ public abstract class AbstractConfiguration implements InitializingBean, Configu
 
 	private int socketTimeout = KafkaConsumerDefaults.SOCKET_TIMEOUT_INT;
 
-	private int fetchMetadataTimeout = KafkaConsumerDefaults.FETCH_METADATA_TIMEOUT;
+	private int backoff = KafkaConsumerDefaults.BACKOFF_INCREMENT_INT;
 
+	private int fetchMetadataTimeout = KafkaConsumerDefaults.FETCH_METADATA_TIMEOUT;
 
 	/**
 	 * The minimum amount of data that a server fetch operation will wait for before returning,
@@ -126,6 +127,18 @@ public abstract class AbstractConfiguration implements InitializingBean, Configu
 		this.socketTimeout = socketTimeout;
 	}
 
+	@Override
+	public int getBackOff() {
+		return backoff;
+	}
+
+	/**
+	 * The retry backoff time for this client
+	 * @param backoff the retry backoff
+	 */
+	public void setBackoff(int backoff) {
+		this.backoff = backoff;
+	}
 
 	/**
 	 * The timeout on fetching metadata (e.g. partition leaders)

spring-integration-kafka/src/main/java/org/springframework/integration/kafka/core/Configuration.java

@@ -66,6 +66,13 @@ public interface Configuration {
 	 */
 	int getSocketTimeout();
 
+	/**
+	 * The retry backoff for this client
+	 * @return the retry backoff
+	 * @since 1.3
+	 */
+	int getBackOff();
+
 	/**
 	 * The timeout on fetching metadata (e.g. partition leaders)
 	 * @return the fetch metadata timeout
@@ -90,4 +97,5 @@ public interface Configuration {
 	 */
 	String getDefaultTopic();
 
+
 }

spring-integration-kafka/src/main/java/org/springframework/integration/kafka/core/DefaultConnection.java

@@ -317,7 +317,7 @@ public class DefaultConnection implements Connection {
 		@Override
 		public Pair<TopicAndPartition, OffsetAndMetadata> value(Partition partition, Long offset) {
 			return Tuples.pair(new TopicAndPartition(partition.getTopic(), partition.getId()),
-					new OffsetAndMetadata(offset, OffsetAndMetadata.NoMetadata(), ErrorMapping.NoError()));
+					new OffsetAndMetadata(offset, OffsetAndMetadata.NoMetadata(), OffsetAndMetadata.InvalidTime()));
 		}
 
 	}

spring-integration-kafka/src/main/java/org/springframework/integration/kafka/core/KafkaConsumerDefaults.java

@@ -42,7 +42,9 @@ public abstract class KafkaConsumerDefaults {
 
 	public static final String FETCH_SIZE = Integer.toString(FETCH_SIZE_INT);
 
-	public static final String BACKOFF_INCREMENT = "1000";
+	public static final int BACKOFF_INCREMENT_INT = 1000;
+
+	public static final String BACKOFF_INCREMENT = Integer.toString(BACKOFF_INCREMENT_INT);
 
 	public static final String QUEUED_CHUNKS_MAX = "100";
 

spring-integration-kafka/src/main/java/org/springframework/integration/kafka/listener/KafkaNativeOffsetManager.java

@@ -0,0 +1,231 @@
+/*
+ * Copyright 2015 the original author or authors.
+ *
+ * Licensed 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.springframework.integration.kafka.listener;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.I0Itec.zkclient.ZkClient;
+
+import org.springframework.beans.factory.InitializingBean;
+import org.springframework.integration.kafka.core.BrokerAddress;
+import org.springframework.integration.kafka.core.Configuration;
+import org.springframework.integration.kafka.core.Connection;
+import org.springframework.integration.kafka.core.ConnectionFactory;
+import org.springframework.integration.kafka.core.ConsumerException;
+import org.springframework.integration.kafka.core.DefaultConnectionFactory;
+import org.springframework.integration.kafka.core.KafkaConsumerDefaults;
+import org.springframework.integration.kafka.core.Partition;
+import org.springframework.integration.kafka.core.PartitionNotFoundException;
+import org.springframework.integration.kafka.core.Result;
+import org.springframework.integration.kafka.support.ZookeeperConnect;
+import org.springframework.retry.RetryCallback;
+import org.springframework.retry.RetryContext;
+import org.springframework.retry.backoff.ExponentialBackOffPolicy;
+import org.springframework.retry.listener.RetryListenerSupport;
+import org.springframework.retry.policy.SimpleRetryPolicy;
+import org.springframework.retry.support.RetryTemplate;
+import org.springframework.util.Assert;
+
+import com.gs.collections.impl.factory.Maps;
+
+import kafka.client.ClientUtils$;
+import kafka.common.ErrorMapping;
+import kafka.common.OffsetAndMetadata;
+import kafka.network.BlockingChannel;
+import kafka.utils.ZKStringSerializer$;
+
+/**
+ * Implementation of an {@link OffsetManager} that uses kafka native 'topic' offset storage.
+ *
+ * @author Chris Lemper
+ * @author Marius Bogoevici
+ * @since 1.3
+ */
+public class KafkaNativeOffsetManager extends AbstractOffsetManager implements InitializingBean {
+
+	private static final String PARTITION_ATTRIBUTE = "partition";
+
+	private final Map<Partition, BrokerAddress> offsetManagerBrokerAddressCache = new ConcurrentHashMap<>();
+
+	private final ZkClient zkClient;
+
+	private RetryTemplate retryTemplate;
+
+	/**
+	 * @param connectionFactory a Kafka connection factory
+	 * @param zookeeperConnect the zookeeper connection information
+	 */
+	public KafkaNativeOffsetManager(ConnectionFactory connectionFactory, ZookeeperConnect zookeeperConnect) {
+		this(connectionFactory, zookeeperConnect, Collections.<Partition, Long>emptyMap());
+	}
+
+	/**
+	 * @param connectionFactory a Kafka connection factory
+	 * @param zookeeperConnect zookeeper connection for retrieving
+	 * @param initialOffsets a map of partitions to initial offsets
+	 */
+	public KafkaNativeOffsetManager(ConnectionFactory connectionFactory, ZookeeperConnect zookeeperConnect,
+									Map<Partition, Long> initialOffsets) {
+		super(connectionFactory, initialOffsets);
+		Assert.notNull(zookeeperConnect, "'zookeeperConnect' must not be null.");
+		this.zkClient = new ZkClient(zookeeperConnect.getZkConnect(),
+				Integer.parseInt(zookeeperConnect.getZkSessionTimeout()),
+				Integer.parseInt(zookeeperConnect.getZkConnectionTimeout()),
+				ZKStringSerializer$.MODULE$);
+	}
+
+	public void setRetryTemplate(RetryTemplate retryTemplate) {
+		this.retryTemplate = retryTemplate;
+	}
+
+	@Override
+	public void afterPropertiesSet() throws Exception {
+		if (this.retryTemplate == null) {
+			this.retryTemplate = new RetryTemplate();
+			this.retryTemplate.registerListener(new ResetOffsetManagerBrokerAddressRetryListener());
+			final SimpleRetryPolicy retryPolicy = new SimpleRetryPolicy();
+			retryPolicy.setMaxAttempts(5);
+			this.retryTemplate.setRetryPolicy(retryPolicy);
+			final ExponentialBackOffPolicy backOffPolicy = new ExponentialBackOffPolicy();
+			backOffPolicy.setInitialInterval(125L);
+			backOffPolicy.setMaxInterval(5000L);
+			backOffPolicy.setMultiplier(2);
+			this.retryTemplate.setBackOffPolicy(backOffPolicy);
+		}
+	}
+
+	@Override
+	protected Long doGetOffset(final Partition partition) {
+		final Long offset = this.retryTemplate.execute(new RetryCallback<Long, RuntimeException>() {
+
+			@Override
+			public Long doWithRetry(RetryContext context) throws RuntimeException {
+				context.setAttribute(PARTITION_ATTRIBUTE, partition);
+				Result<Long> result = getOffsetManagerConnection(partition).fetchStoredOffsetsForConsumer(
+						getConsumerId(), partition);
+				checkResultForErrors(result, partition);
+				return result.getResult(partition);
+			}
+
+			@Override
+			public String toString() {
+				return String.format("fetchStoredOffsetsForConsumer(%s, %s)", getConsumerId(), partition);
+			}
+
+		});
+
+		if (offset != null && offset < 0) {
+			return null;
+		}
+
+		return offset;
+	}
+
+	@Override
+	protected void doUpdateOffset(final Partition partition, final long offset) {
+		this.retryTemplate.execute(new RetryCallback<Void, RuntimeException>() {
+
+			@Override
+			public Void doWithRetry(RetryContext context) throws RuntimeException {
+				context.setAttribute(PARTITION_ATTRIBUTE, partition);
+				Result<Void> result = getOffsetManagerConnection(partition).commitOffsetsForConsumer(
+						getConsumerId(), Maps.immutable.of(partition, offset).castToMap());
+				checkResultForErrors(result, partition);
+				return null;
+			}
+
+			@Override
+			public String toString() {
+				return String.format("commitOffsetsForConsumer(%s, %s, %s)", getConsumerId(), partition, offset);
+			}
+
+		});
+	}
+
+	@Override
+	protected void doRemoveOffset(Partition partition) {
+		doUpdateOffset(partition, OffsetAndMetadata.InvalidOffset());
+	}
+
+	@Override
+	public void close() throws IOException {
+		this.zkClient.close();
+	}
+
+	@Override
+	public void flush() throws IOException {
+		// this function left intentionally blank
+	}
+
+	private BrokerAddress getOffsetManagerBrokerAddress(final Partition partition) {
+		BrokerAddress brokerAddress = this.offsetManagerBrokerAddressCache.get(partition);
+		if (brokerAddress == null) {
+			int socketTimeoutMs = KafkaConsumerDefaults.SOCKET_TIMEOUT_INT;
+			int retryBackOffMs = KafkaConsumerDefaults.BACKOFF_INCREMENT_INT;
+			if (this.connectionFactory instanceof DefaultConnectionFactory) {
+				Configuration configuration = ((DefaultConnectionFactory) connectionFactory).getConfiguration();
+				socketTimeoutMs = configuration.getSocketTimeout();
+				retryBackOffMs = configuration.getBackOff();
+			}
+			final BlockingChannel channel = ClientUtils$.MODULE$.channelToOffsetManager(
+					getConsumerId(), zkClient, socketTimeoutMs, retryBackOffMs);
+			brokerAddress = new BrokerAddress(channel.host(), channel.port());
+			if (log.isDebugEnabled()) {
+				log.debug(String.format("Offset manager for [%s] is at [%s].", partition,
+						brokerAddress));
+			}
+			this.offsetManagerBrokerAddressCache.put(partition, brokerAddress);
+			channel.disconnect();
+		}
+		return brokerAddress;
+	}
+
+	private Connection getOffsetManagerConnection(final Partition partition) {
+		return this.connectionFactory.connect(getOffsetManagerBrokerAddress(partition));
+	}
+
+	private void checkResultForErrors(final Result<?> result, final Partition partition) {
+		if (result.getErrors().containsKey(partition)) {
+			final short errorCode = result.getError(partition);
+			if (errorCode == ErrorMapping.UnknownTopicOrPartitionCode()) {
+				throw new PartitionNotFoundException(partition);
+			}
+			else if (errorCode != ErrorMapping.NoError()) {
+				throw new ConsumerException(ErrorMapping.exceptionFor(errorCode));
+			}
+		}
+	}
+
+	private class ResetOffsetManagerBrokerAddressRetryListener extends RetryListenerSupport {
+
+		@Override
+		@SuppressWarnings("unchecked")
+		public <T, E extends Throwable> void onError(RetryContext context, RetryCallback<T, E> callback, Throwable t) {
+			if (log.isWarnEnabled()) {
+				log.warn("Retrying kafka operation [" + callback + "] due to [" + t + "]", t);
+			}
+			Partition partition = (Partition) context.getAttribute(PARTITION_ATTRIBUTE);
+			if (partition != null) {
+				offsetManagerBrokerAddressCache.remove(partition);
+			}
+		}
+	}
+
+}

spring-integration-kafka/src/main/java/org/springframework/integration/kafka/outbound/KafkaProducerMessageHandler.java

@@ -19,12 +19,10 @@ package org.springframework.integration.kafka.outbound;
 import org.springframework.expression.EvaluationContext;
 import org.springframework.expression.Expression;
 import org.springframework.integration.expression.ExpressionUtils;
-import org.springframework.integration.expression.IntegrationEvaluationContextAware;
 import org.springframework.integration.handler.AbstractMessageHandler;
 import org.springframework.integration.kafka.support.KafkaHeaders;
 import org.springframework.integration.kafka.support.KafkaProducerContext;
 import org.springframework.messaging.Message;
-import org.springframework.util.Assert;
 
 /**
  * @author Soby Chacko
@@ -58,6 +56,7 @@ public class KafkaProducerMessageHandler extends AbstractMessageHandler {
 	}
 
 	/**
+	 * @param partitionExpression an expression that returns a partition id
 	 * @deprecated as of 1.3, {@link #setPartitionIdExpression(Expression)} should be used instead
 	 */
 	@Deprecated

spring-integration-kafka/src/test/java/org/springframework/integration/kafka/offset/KafkaNativeOffsetManagerTests.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2015 the original author or authors.
+ *
+ * Licensed 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.springframework.integration.kafka.offset;
+
+import java.util.Map;
+
+import org.springframework.integration.kafka.core.DefaultConnectionFactory;
+import org.springframework.integration.kafka.core.Partition;
+import org.springframework.integration.kafka.core.ZookeeperConfiguration;
+import org.springframework.integration.kafka.listener.KafkaNativeOffsetManager;
+import org.springframework.integration.kafka.listener.OffsetManager;
+import org.springframework.integration.kafka.support.ZookeeperConnect;
+
+/**
+ * @author Chris Lemper
+ * @since 1.3
+ */
+public class KafkaNativeOffsetManagerTests extends AbstractOffsetManagerTests {
+
+	@Override
+	protected OffsetManager createOffsetManager(long referenceTimestamp, String consumerId,
+			Map<Partition, Long> initialOffsets) throws Exception {
+		ZookeeperConnect zookeeperConnect = new ZookeeperConnect(kafkaRule.getZookeeperConnectionString());
+		KafkaNativeOffsetManager kafkaNativeOffsetManager =
+				new KafkaNativeOffsetManager(new DefaultConnectionFactory(new ZookeeperConfiguration(zookeeperConnect)),
+						zookeeperConnect,
+				initialOffsets);
+		kafkaNativeOffsetManager.setConsumerId(consumerId);
+		kafkaNativeOffsetManager.afterPropertiesSet();
+		kafkaNativeOffsetManager.setReferenceTimestamp(referenceTimestamp);
+		return kafkaNativeOffsetManager;
+	}
+
+}

spring-integration-kafka/src/test/java/org/springframework/integration/kafka/offset/KafkaTopicOffsetManagerTests.java

@@ -26,7 +26,7 @@ import org.springframework.integration.kafka.support.ZookeeperConnect;
 /**
  * @author Marius Bogoevici
  */
-public class KafkaOffsetManagerTests extends AbstractOffsetManagerTests {
+public class KafkaTopicOffsetManagerTests extends AbstractOffsetManagerTests {
 
 	@Override
 	protected OffsetManager createOffsetManager(long referenceTimestamp, String consumerId,

spring-integration-kafka/src/test/java/org/springframework/integration/kafka/rule/KafkaEmbedded.java

@@ -93,6 +93,7 @@ public class KafkaEmbedded extends ExternalResource implements KafkaRule {
 			Properties brokerConfigProperties = TestUtils.createBrokerConfig(i, kafkaPorts.get(i),controlledShutdown);
 			brokerConfigProperties.setProperty("replica.socket.timeout.ms","1000");
 			brokerConfigProperties.setProperty("controller.socket.timeout.ms","1000");
+			brokerConfigProperties.setProperty("offsets.topic.replication.factor","1");
 			KafkaServer server = TestUtils.createServer(new KafkaConfig(brokerConfigProperties), SystemTime$.MODULE$);
 			kafkaServers.add(server);
 		}