From f1289c46e6ad52f785ff8d8a5ed5effa6302ffee Mon Sep 17 00:00:00 2001
From: Oliver Gierke <ogierke@vmware.com>
Date: Thu, 6 Sep 2012 16:22:57 +0200
Subject: [PATCH] DATAMONGO-535 - Fixed broken MongoDBUtils resource
 synchronization.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

We now make sure we really re-use the database instance bound to the thread avoiding duplicate lookups of Mongo.getDb(…). This doesn't seem to gain a big performance benefit anymore as the Mongo instance caches the DB instances internally anyway.
---
 .../data/mongodb/core/MongoDbUtils.java       | 23 +++++++--------
 .../mongodb/core/MongoDbUtilsUnitTests.java   | 29 +++++++++++++++----
 2 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoDbUtils.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoDbUtils.java
index 182015244..5a39a0202 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoDbUtils.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoDbUtils.java
@@ -79,21 +79,18 @@ public abstract class MongoDbUtils {
 
 		DbHolder dbHolder = (DbHolder) TransactionSynchronizationManager.getResource(mongo);
 
-		if (dbHolder != null && !dbHolder.isEmpty()) {
+		// Do we have a populated holder and TX sync active?
+		if (dbHolder != null && !dbHolder.isEmpty() && TransactionSynchronizationManager.isSynchronizationActive()) {
 
-			DB db = null;
+			DB db = dbHolder.getDB(databaseName);
 
-			if (TransactionSynchronizationManager.isSynchronizationActive() && dbHolder.doesNotHoldNonDefaultDB()) {
+			// DB found but not yet synchronized
+			if (db != null && !dbHolder.isSynchronizedWithTransaction()) {
 
-				db = dbHolder.getDB(databaseName);
+				LOGGER.debug("Registering Spring transaction synchronization for existing MongoDB {}.", databaseName);
 
-				if (db != null && !dbHolder.isSynchronizedWithTransaction()) {
-
-					LOGGER.debug("Registering Spring transaction synchronization for existing MongoDB {}.", databaseName);
-
-					TransactionSynchronizationManager.registerSynchronization(new MongoSynchronization(dbHolder, mongo));
-					dbHolder.setSynchronizedWithTransaction(true);
-				}
+				TransactionSynchronizationManager.registerSynchronization(new MongoSynchronization(dbHolder, mongo));
+				dbHolder.setSynchronizedWithTransaction(true);
 			}
 
 			if (db != null) {
@@ -101,6 +98,7 @@ public abstract class MongoDbUtils {
 			}
 		}
 
+		// Lookup fresh database instance
 		LOGGER.debug("Getting Mongo Database name=[{}]", databaseName);
 
 		DB db = mongo.getDB(databaseName);
@@ -117,8 +115,7 @@ public abstract class MongoDbUtils {
 			}
 		}
 
-		// Use same Session for further Mongo actions within the transaction.
-		// Thread object will get removed by synchronization at transaction completion.
+		// TX sync active, bind new database to thread
 		if (TransactionSynchronizationManager.isSynchronizationActive()) {
 
 			LOGGER.debug("Registering Spring transaction synchronization for MongoDB instance {}.", databaseName);
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoDbUtilsUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoDbUtilsUnitTests.java
index 571e82b92..507680c26 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoDbUtilsUnitTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoDbUtilsUnitTests.java
@@ -17,10 +17,17 @@ package org.springframework.data.mongodb.core;
 
 import static org.hamcrest.CoreMatchers.*;
 import static org.junit.Assert.*;
+import static org.mockito.Matchers.*;
+import static org.mockito.Mockito.*;
 
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.runners.MockitoJUnitRunner;
+import org.mockito.stubbing.Answer;
 import org.springframework.transaction.support.TransactionSynchronizationManager;
 
 import com.mongodb.DB;
@@ -31,13 +38,21 @@ import com.mongodb.Mongo;
  * 
  * @author Oliver Gierke
  */
+@RunWith(MockitoJUnitRunner.class)
 public class MongoDbUtilsUnitTests {
 
+	@Mock
 	Mongo mongo;
 
 	@Before
 	public void setUp() throws Exception {
-		this.mongo = new Mongo();
+
+		when(mongo.getDB(anyString())).then(new Answer<DB>() {
+			public DB answer(InvocationOnMock invocation) throws Throwable {
+				return mock(DB.class);
+			}
+		});
+
 		TransactionSynchronizationManager.initSynchronization();
 	}
 
@@ -55,11 +70,15 @@ public class MongoDbUtilsUnitTests {
 	public void returnsNewInstanceForDifferentDatabaseName() {
 
 		DB first = MongoDbUtils.getDB(mongo, "first");
-		assertThat(first, is(notNullValue()));
-		assertThat(MongoDbUtils.getDB(mongo, "first"), is(first));
-
 		DB second = MongoDbUtils.getDB(mongo, "second");
 		assertThat(second, is(not(first)));
-		assertThat(MongoDbUtils.getDB(mongo, "second"), is(second));
+	}
+
+	@Test
+	public void returnsSameInstanceForSameDatabaseName() {
+
+		DB first = MongoDbUtils.getDB(mongo, "first");
+		assertThat(first, is(notNullValue()));
+		assertThat(MongoDbUtils.getDB(mongo, "first"), is(sameInstance(first)));
 	}
 }