From 790d899953670f72bcbe38c67d16600ec4435d1b Mon Sep 17 00:00:00 2001
From: Marius Bogoevici <marius@vetinari.local>
Date: Fri, 27 Mar 2015 16:12:07 -0400
Subject: [PATCH] INTEXT-159 DefaultConnectionFactory thread safety

- protect creation of connections with the lock as well, ensuring that it doesn't clash with refreshMetadata();
---
 .../kafka/core/DefaultConnectionFactory.java  | 45 +++++++++++--------
 1 file changed, 27 insertions(+), 18 deletions(-)

diff --git a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/core/DefaultConnectionFactory.java b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/core/DefaultConnectionFactory.java
index 37daf9b499..b50fdcc895 100644
--- a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/core/DefaultConnectionFactory.java
+++ b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/core/DefaultConnectionFactory.java
@@ -58,8 +58,6 @@ public class DefaultConnectionFactory implements InitializingBean, ConnectionFac
 
 	private final GetBrokersByPartitionFunction getBrokersByPartitionFunction = new GetBrokersByPartitionFunction();
 
-	private final ConnectionInstantiationFunction connectionInstantiationFunction = new ConnectionInstantiationFunction();
-
 	private final Configuration configuration;
 
 	private final AtomicReference<MetadataCache> metadataCacheHolder =
@@ -137,7 +135,33 @@ public class DefaultConnectionFactory implements InitializingBean, ConnectionFac
 	 */
 	@Override
 	public Connection connect(BrokerAddress brokerAddress) {
-		return this.kafkaBrokersCache.getIfAbsentPutWithKey(brokerAddress, connectionInstantiationFunction);
+		Connection connection = null;
+		try {
+			this.lock.readLock().lock();
+			connection = this.kafkaBrokersCache.get(brokerAddress);
+		}
+		finally {
+			this.lock.readLock().unlock();
+		}
+		if (connection == null) {
+			try {
+				this.lock.writeLock().lock();
+				connection = this.kafkaBrokersCache.get(brokerAddress);
+				if (connection == null) {
+					connection = new DefaultConnection(brokerAddress,
+							DefaultConnectionFactory.this.configuration.getClientId(),
+							DefaultConnectionFactory.this.configuration.getBufferSize(),
+							DefaultConnectionFactory.this.configuration.getSocketTimeout(),
+							DefaultConnectionFactory.this.configuration.getMinBytes(),
+							DefaultConnectionFactory.this.configuration.getMaxWait());
+					kafkaBrokersCache.put(brokerAddress, connection);
+				}
+			}
+			finally {
+				this.lock.writeLock().unlock();
+			}
+		}
+		return connection;
 	}
 
 	/**
@@ -223,21 +247,6 @@ public class DefaultConnectionFactory implements InitializingBean, ConnectionFac
 		}
 	}
 
-	@SuppressWarnings("serial")
-	private class ConnectionInstantiationFunction implements Function<BrokerAddress, Connection> {
-
-		@Override
-		public Connection valueOf(BrokerAddress brokerAddress) {
-			return new DefaultConnection(brokerAddress,
-					DefaultConnectionFactory.this.configuration.getClientId(),
-					DefaultConnectionFactory.this.configuration.getBufferSize(),
-					DefaultConnectionFactory.this.configuration.getSocketTimeout(),
-					DefaultConnectionFactory.this.configuration.getMinBytes(),
-					DefaultConnectionFactory.this.configuration.getMaxWait());
-		}
-
-	}
-
 	@SuppressWarnings("serial")
 	private class GetBrokersByPartitionFunction implements Function<Partition, BrokerAddress> {