GH-1101: Overide transactionIdPrefix

Resolves #1101

Author: garyrussell

spring-kafka/src/main/java/org/springframework/kafka/core/DefaultKafkaProducerFactory.java

@@ -24,6 +24,7 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Future;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
@@ -99,7 +100,7 @@ public class DefaultKafkaProducerFactory<K, V> implements ProducerFactory<K, V>,
 
 	private final AtomicInteger transactionIdSuffix = new AtomicInteger();
 
-	private final BlockingQueue<CloseSafeProducer<K, V>> cache = new LinkedBlockingQueue<>();
+	private final Map<String, BlockingQueue<CloseSafeProducer<K, V>>> cache = new ConcurrentHashMap<>();
 
 	private final Map<String, CloseSafeProducer<K, V>> consumerProducers = new HashMap<>();
 
@@ -264,16 +265,19 @@ public void destroy() {
 		if (producerToClose != null) {
 			producerToClose.getDelegate().close(this.physicalCloseTimeout);
 		}
-		producerToClose = this.cache.poll();
-		while (producerToClose != null) {
-			try {
-				producerToClose.getDelegate().close(this.physicalCloseTimeout);
-			}
-			catch (Exception e) {
-				LOGGER.error(e, "Exception while closing producer");
+		this.cache.values().forEach(queue -> {
+			CloseSafeProducer<K, V> next = queue.poll();
+			while (next != null) {
+				try {
+					next.getDelegate().close(this.physicalCloseTimeout);
+				}
+				catch (Exception e) {
+					LOGGER.error(e, "Exception while closing producer");
+				}
+				next = queue.poll();
 			}
-			producerToClose = this.cache.poll();
-		}
+		});
+		this.cache.clear();
 		synchronized (this.consumerProducers) {
 			this.consumerProducers.forEach(
 					(k, v) -> v.getDelegate().close(this.physicalCloseTimeout));
@@ -314,12 +318,18 @@ public boolean isRunning() {
 
 	@Override
 	public Producer<K, V> createProducer() {
-		if (this.transactionIdPrefix != null) {
+		return createProducer(this.transactionIdPrefix);
+	}
+
+	@Override
+	public Producer<K, V> createProducer(@Nullable String txIdPrefixArg) {
+		String txIdPrefix = txIdPrefixArg == null ? this.transactionIdPrefix : txIdPrefixArg;
+		if (txIdPrefix != null) {
 			if (this.producerPerConsumerPartition) {
-				return createTransactionalProducerForPartition();
+				return createTransactionalProducerForPartition(txIdPrefix);
 			}
 			else {
-				return createTransactionalProducer();
+				return createTransactionalProducer(txIdPrefix);
 			}
 		}
 		if (this.producerPerThread) {
@@ -349,15 +359,20 @@ protected Producer<K, V> createKafkaProducer() {
 		return new KafkaProducer<>(this.configs, this.keySerializer, this.valueSerializer);
 	}
 
-	Producer<K, V> createTransactionalProducerForPartition() {
+	protected Producer<K, V> createTransactionalProducerForPartition() {
+		return createTransactionalProducerForPartition(this.transactionIdPrefix);
+	}
+
+	protected Producer<K, V> createTransactionalProducerForPartition(String txIdPrefix) {
 		String suffix = TransactionSupport.getTransactionIdSuffix();
 		if (suffix == null) {
-			return createTransactionalProducer();
+			return createTransactionalProducer(txIdPrefix);
 		}
 		else {
 			synchronized (this.consumerProducers) {
 				if (!this.consumerProducers.containsKey(suffix)) {
-					CloseSafeProducer<K, V> newProducer = doCreateTxProducer(suffix, this::removeConsumerProducer);
+					CloseSafeProducer<K, V> newProducer = doCreateTxProducer(txIdPrefix, suffix,
+							this::removeConsumerProducer);
 					this.consumerProducers.put(suffix, newProducer);
 					return newProducer;
 				}
@@ -387,29 +402,43 @@ private void removeConsumerProducer(CloseSafeProducer<K, V> producerToRemove) {
 	 * @since 1.3
 	 */
 	protected Producer<K, V> createTransactionalProducer() {
-		Producer<K, V> cachedProducer = this.cache.poll();
+		return createTransactionalProducer(this.transactionIdPrefix);
+	}
+
+	protected Producer<K, V> createTransactionalProducer(String txIdPrefix) {
+		BlockingQueue<CloseSafeProducer<K, V>> queue = getCache(txIdPrefix);
+		Producer<K, V> cachedProducer = queue.poll();
 		if (cachedProducer == null) {
-			return doCreateTxProducer("" + this.transactionIdSuffix.getAndIncrement(), null);
+			return doCreateTxProducer(txIdPrefix, "" + this.transactionIdSuffix.getAndIncrement(), null);
 		}
 		else {
 			return cachedProducer;
 		}
 	}
 
-	private CloseSafeProducer<K, V> doCreateTxProducer(String suffix,
+	private CloseSafeProducer<K, V> doCreateTxProducer(String prefix, String suffix,
 			@Nullable Consumer<CloseSafeProducer<K, V>> remover) {
 
 		Producer<K, V> newProducer;
 		Map<String, Object> newProducerConfigs = new HashMap<>(this.configs);
-		newProducerConfigs.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, this.transactionIdPrefix + suffix);
+		newProducerConfigs.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, prefix + suffix);
 		newProducer = new KafkaProducer<>(newProducerConfigs, this.keySerializer, this.valueSerializer);
 		newProducer.initTransactions();
-		return new CloseSafeProducer<>(newProducer, this.cache, remover,
+		return new CloseSafeProducer<>(newProducer, this.cache.get(prefix), remover,
 				(String) newProducerConfigs.get(ProducerConfig.TRANSACTIONAL_ID_CONFIG));
 	}
 
+	@Nullable
 	protected BlockingQueue<CloseSafeProducer<K, V>> getCache() {
-		return this.cache;
+		return getCache(this.transactionIdPrefix);
+	}
+
+	@Nullable
+	protected BlockingQueue<CloseSafeProducer<K, V>> getCache(String txIdPrefix) {
+		if (txIdPrefix == null) {
+			return null;
+		}
+		return this.cache.computeIfAbsent(txIdPrefix, txId -> new LinkedBlockingQueue<>());
 	}
 
 	@Override

spring-kafka/src/main/java/org/springframework/kafka/core/KafkaTemplate.java

@@ -78,6 +78,8 @@ public class KafkaTemplate<K, V> implements KafkaOperations<K, V> {
 
 	private volatile ProducerListener<K, V> producerListener = new LoggingProducerListener<K, V>();
 
+	private String transactionIdPrefix;
+
 
 	/**
 	 * Create an instance using the supplied producer factory and autoFlush false.
@@ -154,6 +156,19 @@ public boolean isTransactional() {
 		return this.transactional;
 	}
 
+	public String getTransactionIdPrefix() {
+		return this.transactionIdPrefix;
+	}
+
+	/**
+	 * Set a transaction id prefix to override the prefix in the producer factory.
+	 * @param transactionIdPrefix the prefix.
+	 * @since 2.3
+	 */
+	public void setTransactionIdPrefix(String transactionIdPrefix) {
+		this.transactionIdPrefix = transactionIdPrefix;
+	}
+
 	/**
 	 * Return the producer factory used by this template.
 	 * @return the factory.
@@ -277,7 +292,7 @@ public <T> T executeInTransaction(OperationsCallback<K, V, T> callback) {
 			transactionIdSuffix = null;
 		}
 
-		producer = this.producerFactory.createProducer();
+		producer = this.producerFactory.createProducer(this.transactionIdPrefix);
 
 		try {
 			producer.beginTransaction();
@@ -427,11 +442,11 @@ private Producer<K, V> getTheProducer() {
 				return producer;
 			}
 			KafkaResourceHolder<K, V> holder = ProducerFactoryUtils
-					.getTransactionalResourceHolder(this.producerFactory);
+					.getTransactionalResourceHolder(this.producerFactory, this.transactionIdPrefix);
 			return holder.getProducer();
 		}
 		else {
-			return this.producerFactory.createProducer();
+			return this.producerFactory.createProducer(this.transactionIdPrefix);
 		}
 	}
 

spring-kafka/src/main/java/org/springframework/kafka/core/ProducerFactory.java

@@ -34,6 +34,16 @@ public interface ProducerFactory<K, V> {
 	 */
 	Producer<K, V> createProducer();
 
+	/**
+	 * Create a producer with an overridden transaction id prefix.
+	 * @param txIdPrefix the transaction id prefix.
+	 * @return the producer.
+	 * @since 2.3
+	 */
+	default Producer<K, V> createProducer(@SuppressWarnings("unused") String txIdPrefix) {
+		throw new UnsupportedOperationException("This factory does not support this method");
+	}
+
 	/**
 	 * Return true if the factory supports transactions.
 	 * @return true if transactional.

spring-kafka/src/main/java/org/springframework/kafka/core/ProducerFactoryUtils.java

@@ -50,13 +50,29 @@ private ProducerFactoryUtils() {
 	public static <K, V> KafkaResourceHolder<K, V> getTransactionalResourceHolder(
 			final ProducerFactory<K, V> producerFactory) {
 
+		return getTransactionalResourceHolder(producerFactory, null);
+	}
+
+	/**
+	 * Obtain a Producer that is synchronized with the current transaction, if any.
+	 * @param producerFactory the ProducerFactory to obtain a Channel for
+	 * @param txIdPrefix the transaction id prefix; if null, the producer factory
+	 * prefix is used.
+	 * @param <K> the key type.
+	 * @param <V> the value type.
+	 * @return the resource holder.
+	 * @since 2.3
+	 */
+	public static <K, V> KafkaResourceHolder<K, V> getTransactionalResourceHolder(
+			final ProducerFactory<K, V> producerFactory, @Nullable String txIdPrefix) {
+
 		Assert.notNull(producerFactory, "ProducerFactory must not be null");
 
 		@SuppressWarnings("unchecked")
 		KafkaResourceHolder<K, V> resourceHolder = (KafkaResourceHolder<K, V>) TransactionSynchronizationManager
 				.getResource(producerFactory);
 		if (resourceHolder == null) {
-			Producer<K, V> producer = producerFactory.createProducer();
+			Producer<K, V> producer = producerFactory.createProducer(txIdPrefix);
 
 			try {
 				producer.beginTransaction();

spring-kafka/src/main/java/org/springframework/kafka/transaction/KafkaTransactionManager.java

@@ -70,6 +70,8 @@ public class KafkaTransactionManager<K, V> extends AbstractPlatformTransactionMa
 
 	private final ProducerFactory<K, V> producerFactory;
 
+	private String transactionIdPrefix;
+
 	/**
 	 * Create a new KafkaTransactionManager, given a ProducerFactory.
 	 * Transaction synchronization is turned off by default, as this manager might be used alongside a datastore-based
@@ -84,6 +86,15 @@ public KafkaTransactionManager(ProducerFactory<K, V> producerFactory) {
 		this.producerFactory = producerFactory;
 	}
 
+	/**
+	 * Set a transaction id prefix to override the prefix in the producer factory.
+	 * @param transactionIdPrefix the prefix.
+	 * @since 2.3
+	 */
+	public void setTransactionIdPrefix(String transactionIdPrefix) {
+		this.transactionIdPrefix = transactionIdPrefix;
+	}
+
 	/**
 	 * Get the producer factory.
 	 * @return the producerFactory
@@ -132,7 +143,8 @@ protected void doBegin(Object transaction, TransactionDefinition definition) {
 		KafkaTransactionObject<K, V> txObject = (KafkaTransactionObject<K, V>) transaction;
 		KafkaResourceHolder<K, V> resourceHolder = null;
 		try {
-			resourceHolder = ProducerFactoryUtils.getTransactionalResourceHolder(getProducerFactory());
+			resourceHolder = ProducerFactoryUtils.getTransactionalResourceHolder(getProducerFactory(),
+					this.transactionIdPrefix);
 			if (logger.isDebugEnabled()) {
 				logger.debug("Created Kafka transaction on producer [" + resourceHolder.getProducer() + "]");
 			}

spring-kafka/src/test/java/org/springframework/kafka/core/DefaultKafkaConsumerFactoryTests.java

@@ -20,7 +20,6 @@
 
 import java.util.Collections;
 import java.util.Map;
-import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
@@ -101,10 +100,12 @@ public void testNestedTxProducerIsCached() throws Exception {
 		container.start();
 		try {
 			ListenableFuture<SendResult<Integer, String>> future = template.send("txCache1", "foo");
-			future.get();
-			assertThat(KafkaTestUtils.getPropertyValue(pf, "cache", BlockingQueue.class)).hasSize(0);
+			future.get(10, TimeUnit.SECONDS);
+			pf.getCache();
+			assertThat(KafkaTestUtils.getPropertyValue(pf, "cache", Map.class)).hasSize(0);
 			assertThat(latch.await(30, TimeUnit.SECONDS)).isTrue();
-			assertThat(KafkaTestUtils.getPropertyValue(pfTx, "cache", BlockingQueue.class)).hasSize(1);
+			assertThat(KafkaTestUtils.getPropertyValue(pfTx, "cache", Map.class)).hasSize(1);
+			assertThat(pfTx.getCache()).hasSize(1);
 		}
 		finally {
 			container.stop();
@@ -140,10 +141,10 @@ public void testContainerTxProducerIsNotCached() throws Exception {
 		container.start();
 		try {
 			ListenableFuture<SendResult<Integer, String>> future = template.send("txCache2", "foo");
-			future.get();
-			assertThat(KafkaTestUtils.getPropertyValue(pf, "cache", BlockingQueue.class)).hasSize(0);
+			future.get(10, TimeUnit.SECONDS);
+			assertThat(KafkaTestUtils.getPropertyValue(pf, "cache", Map.class)).hasSize(0);
 			assertThat(latch.await(30, TimeUnit.SECONDS)).isTrue();
-			assertThat(KafkaTestUtils.getPropertyValue(pfTx, "cache", BlockingQueue.class)).hasSize(0);
+			assertThat(KafkaTestUtils.getPropertyValue(pfTx, "cache", Map.class)).hasSize(0);
 		}
 		finally {
 			container.stop();

spring-kafka/src/test/java/org/springframework/kafka/core/DefaultKafkaProducerFactoryTests.java

@@ -25,6 +25,7 @@
 
 import java.time.Duration;
 import java.util.HashMap;
+import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -55,9 +56,9 @@ public void testProducerClosedAfterBadTransition() throws Exception {
 		DefaultKafkaProducerFactory pf = new DefaultKafkaProducerFactory(new HashMap<>()) {
 
 			@Override
-			protected Producer createTransactionalProducer() {
+			protected Producer createTransactionalProducer(String txIdPrefix) {
 				producer.initTransactions();
-				BlockingQueue<Producer> cache = getCache();
+				BlockingQueue<Producer> cache = getCache(txIdPrefix);
 				Producer cached = cache.poll();
 				return cached == null ? new CloseSafeProducer(producer, cache) : cached;
 			}
@@ -80,8 +81,10 @@ protected Producer createTransactionalProducer() {
 			kafkaTemplate.send("foo", "bar");
 			return null;
 		});
-		BlockingQueue cache = KafkaTestUtils.getPropertyValue(pf, "cache", BlockingQueue.class);
+		Map<?, ?> cache = KafkaTestUtils.getPropertyValue(pf, "cache", Map.class);
 		assertThat(cache).hasSize(1);
+		Queue queue = (Queue) cache.get("foo");
+		assertThat(queue).hasSize(1);
 		try {
 			transactionTemplate.execute(s -> {
 				return null;
@@ -90,7 +93,7 @@ protected Producer createTransactionalProducer() {
 		catch (CannotCreateTransactionException e) {
 			assertThat(e.getCause().getMessage()).contains("Invalid transition");
 		}
-		assertThat(cache).hasSize(0);
+		assertThat(queue).hasSize(0);
 
 		InOrder inOrder = inOrder(producer);
 		inOrder.verify(producer).initTransactions();
@@ -119,7 +122,7 @@ protected Producer createKafkaProducer() {
 		assertThat(aProducer).isNotNull();
 		aProducer.close();
 		assertThat(KafkaTestUtils.getPropertyValue(pf, "producer")).isNotNull();
-		Queue cache = KafkaTestUtils.getPropertyValue(pf, "cache", Queue.class);
+		Map<?, ?> cache = KafkaTestUtils.getPropertyValue(pf, "cache", Map.class);
 		assertThat(cache.size()).isEqualTo(0);
 		pf.reset();
 		assertThat(KafkaTestUtils.getPropertyValue(pf, "producer")).isNull();
@@ -134,9 +137,9 @@ public void testResetTx() throws Exception {
 		DefaultKafkaProducerFactory pf = new DefaultKafkaProducerFactory(new HashMap<>()) {
 
 			@Override
-			protected Producer createTransactionalProducer() {
+			protected Producer createTransactionalProducer(String txIdPrefix) {
 				producer.initTransactions();
-				BlockingQueue<Producer> cache = getCache();
+				BlockingQueue<Producer> cache = getCache(txIdPrefix);
 				Producer cached = cache.poll();
 				return cached == null ? new CloseSafeProducer(producer, cache) : cached;
 			}
@@ -148,10 +151,12 @@ protected Producer createTransactionalProducer() {
 		assertThat(aProducer).isNotNull();
 		aProducer.close();
 		assertThat(KafkaTestUtils.getPropertyValue(pf, "producer")).isNull();
-		Queue cache = KafkaTestUtils.getPropertyValue(pf, "cache", Queue.class);
+		Map<?, ?> cache = KafkaTestUtils.getPropertyValue(pf, "cache", Map.class);
 		assertThat(cache.size()).isEqualTo(1);
+		Queue queue = (Queue) cache.get("foo");
+		assertThat(queue.size()).isEqualTo(1);
 		pf.onApplicationEvent(new ContextStoppedEvent(ctx));
-		assertThat(cache.size()).isEqualTo(0);
+		assertThat(queue.size()).isEqualTo(0);
 		verify(producer).close(any(Duration.class));
 	}