DATAMONGO-2001 - Polishing.

Extract count aggregation pipeline setup to AggregationUtil. Fix count extraction if aggregation returns no results. Fix nullability of Query argument in ReactiveMongoTemplate.count(…). Improve synchronization of multi-threaded aggregation count test to prevent commit before all threads have issued a count query and to await thread completion.

Upgrade to MongoDB 4.0.0-rc4.

Original pull request: #568.

.travis.yml

@@ -13,12 +13,12 @@ before_install:
   - |-
     downloads/mongodb-linux-x86_64-ubuntu1604-${MONGO_VERSION}/bin/mongo --eval "rs.initiate({_id: 'rs0', members:[{_id: 0, host: '127.0.0.1:27017'}]});"
     sleep 15
-    
+
 env:
   matrix:
     - PROFILE=ci
   global:
-    - MONGO_VERSION=4.0.0-rc3
+    - MONGO_VERSION=4.0.0-rc4
 
 addons:
   apt:

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/AggregationUtil.java

@@ -17,6 +17,8 @@ package org.springframework.data.mongodb.core;
 
 import lombok.AllArgsConstructor;
 
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
@@ -24,13 +26,19 @@ import java.util.stream.Collectors;
 import org.bson.Document;
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mongodb.core.aggregation.Aggregation;
+import org.springframework.data.mongodb.core.aggregation.AggregationOperation;
 import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext;
+import org.springframework.data.mongodb.core.aggregation.AggregationOptions;
+import org.springframework.data.mongodb.core.aggregation.CountOperation;
 import org.springframework.data.mongodb.core.aggregation.TypeBasedAggregationOperationContext;
 import org.springframework.data.mongodb.core.aggregation.TypedAggregation;
 import org.springframework.data.mongodb.core.convert.QueryMapper;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
+import org.springframework.data.mongodb.core.query.CriteriaDefinition;
+import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
 import org.springframework.util.ObjectUtils;
 
 /**
@@ -107,6 +115,53 @@ class AggregationUtil {
 		return command;
 	}
 
+	/**
+	 * Create a {@code $count} aggregation for {@link Query} and optionally a {@link Class entity class}.
+	 *
+	 * @param query must not be {@literal null}.
+	 * @param entityClass can be {@literal null} if the {@link Query} object is empty.
+	 * @return the {@link Aggregation} pipeline definition to run a {@code $count} aggregation.
+	 */
+	Aggregation createCountAggregation(Query query, @Nullable Class<?> entityClass) {
+
+		List<AggregationOperation> pipeline = computeCountAggregationPipeline(query, entityClass);
+
+		Aggregation aggregation = entityClass != null ? Aggregation.newAggregation(entityClass, pipeline)
+				: Aggregation.newAggregation(pipeline);
+		aggregation.withOptions(AggregationOptions.builder().collation(query.getCollation().orElse(null)).build());
+
+		return aggregation;
+	}
+
+	private List<AggregationOperation> computeCountAggregationPipeline(Query query, @Nullable Class<?> entityType) {
+
+		CountOperation count = Aggregation.count().as("totalEntityCount");
+		if (query.getQueryObject().isEmpty()) {
+			return Collections.singletonList(count);
+		}
+
+		Assert.notNull(entityType, "Entity type must not be null!");
+
+		Document mappedQuery = queryMapper.getMappedObject(query.getQueryObject(),
+				mappingContext.getPersistentEntity(entityType));
+
+		CriteriaDefinition criteria = new CriteriaDefinition() {
+
+			@Override
+			public Document getCriteriaObject() {
+				return mappedQuery;
+			}
+
+			@Nullable
+			@Override
+			public String getKey() {
+				return null;
+			}
+		};
+
+		return Arrays.asList(Aggregation.match(criteria), count);
+	}
+
 	private List<Document> mapAggregationPipeline(List<Document> pipeline) {
 
 		return pipeline.stream().map(val -> queryMapper.getMappedObject(val, Optional.empty()))

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java

@@ -67,11 +67,9 @@ import org.springframework.data.mongodb.SessionSynchronization;
 import org.springframework.data.mongodb.core.BulkOperations.BulkMode;
 import org.springframework.data.mongodb.core.DefaultBulkOperations.BulkOperationContext;
 import org.springframework.data.mongodb.core.aggregation.Aggregation;
-import org.springframework.data.mongodb.core.aggregation.AggregationOperation;
 import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext;
 import org.springframework.data.mongodb.core.aggregation.AggregationOptions;
 import org.springframework.data.mongodb.core.aggregation.AggregationResults;
-import org.springframework.data.mongodb.core.aggregation.CountOperation;
 import org.springframework.data.mongodb.core.aggregation.Fields;
 import org.springframework.data.mongodb.core.aggregation.TypeBasedAggregationOperationContext;
 import org.springframework.data.mongodb.core.aggregation.TypedAggregation;
@@ -107,7 +105,6 @@ import org.springframework.data.mongodb.core.mapreduce.MapReduceOptions;
 import org.springframework.data.mongodb.core.mapreduce.MapReduceResults;
 import org.springframework.data.mongodb.core.query.Collation;
 import org.springframework.data.mongodb.core.query.Criteria;
-import org.springframework.data.mongodb.core.query.CriteriaDefinition;
 import org.springframework.data.mongodb.core.query.Meta;
 import org.springframework.data.mongodb.core.query.NearQuery;
 import org.springframework.data.mongodb.core.query.Query;
@@ -3542,49 +3539,25 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
 		 * @see org.springframework.data.mongodb.core.MongoTemplate#count(org.springframework.data.mongodb.core.query.Query, java.lang.Class, java.lang.String)
 		 */
 		@Override
-		public long count(Query query, @Nullable Class<?> entityClass, String collection) {
+		@SuppressWarnings("unchecked")
+		public long count(Query query, @Nullable Class<?> entityClass, String collectionName) {
 
 			if (!session.hasActiveTransaction()) {
-				return super.count(query, entityClass, collection);
+				return super.count(query, entityClass, collectionName);
 			}
 
-			List<AggregationOperation> pipeline = computeCountAggregationPipeline(query, entityClass);
+			AggregationUtil aggregationUtil = new AggregationUtil(delegate.queryMapper, delegate.mappingContext);
+			Aggregation aggregation = aggregationUtil.createCountAggregation(query, entityClass);
+			AggregationResults<Document> aggregationResults = aggregate(aggregation, collectionName, Document.class);
 
-			Aggregation aggregation = entityClass != null ? Aggregation.newAggregation(entityClass, pipeline)
-					: Aggregation.newAggregation(pipeline);
-			aggregation.withOptions(AggregationOptions.builder().collation(query.getCollation().orElse(null)).build());
+			List<Document> result = (List<Document>) aggregationResults.getRawResults().getOrDefault("results",
+					Collections.emptyList());
 
-			AggregationResults<Document> aggregationResults = aggregate(aggregation, collection, Document.class);
-			return ((List<Document>) aggregationResults.getRawResults().getOrDefault("results",
-					Collections.singletonList(new Document("totalEntityCount", 0)))).get(0).get("totalEntityCount", Number.class)
-							.longValue();
-		}
-
-		private List<AggregationOperation> computeCountAggregationPipeline(Query query, @Nullable Class<?> entityType) {
-
-			CountOperation count = Aggregation.count().as("totalEntityCount");
-			if (query.getQueryObject().isEmpty()) {
-				return Arrays.asList(count);
+			if (result.isEmpty()) {
+				return 0;
 			}
 
-			Document mappedQuery = delegate.queryMapper.getMappedObject(query.getQueryObject(),
-					delegate.getPersistentEntity(entityType));
-
-			CriteriaDefinition criteria = new CriteriaDefinition() {
-
-				@Override
-				public Document getCriteriaObject() {
-					return mappedQuery;
-				}
-
-				@Nullable
-				@Override
-				public String getKey() {
-					return null;
-				}
-			};
-
-			return Arrays.asList(Aggregation.match(criteria), count);
+			return result.get(0).get("totalEntityCount", Number.class).longValue();
 		}
 	}
 }

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java

@@ -1149,8 +1149,9 @@ public class ReactiveMongoTemplate implements ReactiveMongoOperations, Applicati
 	 * (non-Javadoc)
 	 * @see org.springframework.data.mongodb.core.ReactiveMongoOperations#count(org.springframework.data.mongodb.core.query.Query, java.lang.Class, java.lang.String)
 	 */
-	public Mono<Long> count(@Nullable Query query, @Nullable Class<?> entityClass, String collectionName) {
+	public Mono<Long> count(Query query, @Nullable Class<?> entityClass, String collectionName) {
 
+		Assert.notNull(query, "Query must not be null!");
 		Assert.hasText(collectionName, "Collection name must not be null or empty!");
 
 		return createMono(collectionName, collection -> {
@@ -3275,22 +3276,19 @@ public class ReactiveMongoTemplate implements ReactiveMongoOperations, Applicati
 		 * @see org.springframework.data.mongodb.core.ReactiveMongoTemplate#count(org.springframework.data.mongodb.core.query.Query, java.lang.Class, java.lang.String)
 		 */
 		@Override
-		public Mono<Long> count(@Nullable Query query, @Nullable Class<?> entityClass, String collectionName) {
+		public Mono<Long> count(Query query, @Nullable Class<?> entityClass, String collectionName) {
 
 			if (!session.hasActiveTransaction()) {
 				return super.count(query, entityClass, collectionName);
 			}
 
-			List<AggregationOperation> pipeline = computeCountAggregationPipeline(query, entityClass);
-
-			Aggregation aggregation = entityClass != null ? Aggregation.newAggregation(entityClass, pipeline)
-					: Aggregation.newAggregation(pipeline);
-			aggregation.withOptions(AggregationOptions.builder().collation(query.getCollation().orElse(null)).build());
+			AggregationUtil aggregationUtil = new AggregationUtil(delegate.queryMapper, delegate.mappingContext);
+			Aggregation aggregation = aggregationUtil.createCountAggregation(query, entityClass);
 
 			return aggregate(aggregation, collectionName, Document.class) //
-					.defaultIfEmpty(new Document("totalEntityCount", 0)) //
 					.next() //
-					.map(it -> it.get("totalEntityCount", Number.class).longValue());
+					.map(it -> it.get("totalEntityCount", Number.class).longValue()) //
+					.defaultIfEmpty(0L);
 		}
 
 		private List<AggregationOperation> computeCountAggregationPipeline(@Nullable Query query,

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveClientSessionTests.java

@@ -16,6 +16,8 @@
 package org.springframework.data.mongodb.core;
 
 import static org.assertj.core.api.Assertions.*;
+import static org.springframework.data.mongodb.core.query.Criteria.*;
+import static org.springframework.data.mongodb.core.query.Query.*;
 
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
@@ -138,7 +140,7 @@ public class ReactiveClientSessionTests {
 	}
 
 	@Test // DATAMONGO-2001
-	public void countShouldOnlyReturnCorrectly() {
+	public void countInTransactionShouldReturnCount() {
 
 		ClientSession session = Mono
 				.from(client.startSession(ClientSessionOptions.builder().causallyConsistent(true).build())).block();
@@ -148,11 +150,23 @@ public class ReactiveClientSessionTests {
 			session.startTransaction();
 
 			return action.insert(new Document("_id", "id-2").append("value", "in transaction"), COLLECTION_NAME) //
-					.then(action.count(new Query(), Document.class, COLLECTION_NAME)) //
+					.then(action.count(query(where("value").is("in transaction")), Document.class, COLLECTION_NAME)) //
 					.flatMap(it -> Mono.from(session.commitTransaction()).then(Mono.just(it)));
 
 		}).as(StepVerifier::create) //
-				.expectNext(2L) //
+				.expectNext(1L) //
+				.verifyComplete();
+
+		template.withSession(() -> session).execute(action -> {
+
+			session.startTransaction();
+
+			return action.insert(new Document("value", "in transaction"), COLLECTION_NAME) //
+					.then(action.count(query(where("value").is("foo")), Document.class, COLLECTION_NAME)) //
+					.flatMap(it -> Mono.from(session.commitTransaction()).then(Mono.just(it)));
+
+		}).as(StepVerifier::create) //
+				.expectNext(0L) //
 				.verifyComplete();
 	}
 

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/SessionBoundMongoTemplateTests.java

@@ -19,6 +19,8 @@ import static org.assertj.core.api.Assertions.*;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.*;
+import static org.springframework.data.mongodb.core.query.Criteria.*;
+import static org.springframework.data.mongodb.core.query.Query.*;
 
 import lombok.Data;
 
@@ -30,6 +32,8 @@ import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 import org.aopalliance.aop.Advice;
 import org.bson.Document;
@@ -268,7 +272,7 @@ public class SessionBoundMongoTemplateTests {
 	}
 
 	@Test // DATAMONGO-2001
-	public void countShouldOnlyReturnCorrectly() throws InterruptedException {
+	public void countShouldWorkInTransactions() {
 
 		if (!template.collectionExists(Person.class)) {
 			template.createCollection(Person.class);
@@ -276,44 +280,77 @@ public class SessionBoundMongoTemplateTests {
 			template.remove(Person.class).all();
 		}
 
-		List<Object> resultList = new CopyOnWriteArrayList<>();
+		ClientSession session = client.startSession();
+		session.startTransaction();
 
-		int nrThreads = 2;
-		CountDownLatch countDownLatch = new CountDownLatch(nrThreads);
+		MongoTemplate sessionBound = template.withSession(session);
 
-		for (int i = 0; i < nrThreads; i++) {
+		sessionBound.save(new Person("Kylar Stern"));
 
-			new Thread(() -> {
+		assertThat(sessionBound.query(Person.class).matching(query(where("firstName").is("foobar"))).count()).isZero();
+		assertThat(sessionBound.query(Person.class).matching(query(where("firstName").is("Kylar Stern"))).count()).isOne();
+		assertThat(sessionBound.query(Person.class).count()).isOne();
 
-				ClientSession session = client.startSession();
-				session.startTransaction();
+		session.commitTransaction();
+		session.close();
+	}
 
-				try {
+	@Test // DATAMONGO-2001
+	public void countShouldReturnIsolatedCount() throws InterruptedException {
 
-					MongoTemplate sessionBound = template.withSession(session);
-
-					try {
-						sessionBound.save(new Person("Kylar Stern"));
-					} finally {
-						countDownLatch.countDown();
-					}
-
-					countDownLatch.await(1, TimeUnit.SECONDS);
-
-					resultList.add(Long.valueOf(sessionBound.query(Person.class).count()));
-				} catch (Exception e) {
-					resultList.add(e);
-				}
-
-				session.commitTransaction();
-				session.close();
-			}).start();
+		if (!template.collectionExists(Person.class)) {
+			template.createCollection(Person.class);
+		} else {
+			template.remove(Person.class).all();
 		}
 
-		countDownLatch.await();
+		int nrThreads = 2;
+		CountDownLatch savedInTransaction = new CountDownLatch(nrThreads);
+		CountDownLatch beforeCommit = new CountDownLatch(nrThreads);
+		List<Object> resultList = new CopyOnWriteArrayList<>();
+
+		Runnable runnable = () -> {
+
+			ClientSession session = client.startSession();
+			session.startTransaction();
+
+			try {
+				MongoTemplate sessionBound = template.withSession(session);
+
+				try {
+					sessionBound.save(new Person("Kylar Stern"));
+				} finally {
+					savedInTransaction.countDown();
+				}
+
+				savedInTransaction.await(1, TimeUnit.SECONDS);
+
+				try {
+					resultList.add(sessionBound.query(Person.class).count());
+				} finally {
+					beforeCommit.countDown();
+				}
+
+				beforeCommit.await(1, TimeUnit.SECONDS);
+			} catch (Exception e) {
+				resultList.add(e);
+			}
+
+			session.commitTransaction();
+			session.close();
+		};
+
+		List<Thread> threads = IntStream.range(0, nrThreads) //
+				.mapToObj(i -> new Thread(runnable)) //
+				.peek(Thread::start) //
+				.collect(Collectors.toList());
+
+		for (Thread thread : threads) {
+			thread.join();
+		}
 
 		assertThat(template.query(Person.class).count()).isEqualTo(2L);
-		assertThat(resultList).allMatch(it -> it.equals(1L));
+		assertThat(resultList).hasSize(nrThreads).allMatch(it -> it.equals(1L));
 	}
 
 	@Data

src/main/asciidoc/reference/client-session-transactions.adoc

@@ -215,7 +215,7 @@ public class StateService {
 NOTE: `@Transactional(readOnly = true)` advises `MongoTransactionManager` to also start a transaction that adds the
  `ClientSession` to outgoing requests.
 
-== Reactive transactions
+== Reactive Transactions
 
 Same as with the reactive `ClientSession` support, the `ReactiveMongoTemplate` offers dedicated methods for operating
 within a transaction without having to worry about the commit/abort actions depending on the operations outcome.
@@ -299,14 +299,14 @@ error occurs here.
 
 == Special behavior inside transactions
 
-Inside transactions MongoDB server has a slightly different behavior.
+Inside transactions, MongoDB server has a slightly different behavior.
 
 *Connection Settings*
 
-The MongoDB drivers offer a dedicated replica set name configuration option turing the driver into an auto detection
+The MongoDB drivers offer a dedicated replica set name configuration option turing the driver into auto detection
 mode. This option helps identifying replica set master nodes and command routing during a transaction.
 
-INFO: Make sure to add `replicaSet` to the MongoDB Uri. Please refer to https://docs.mongodb.com/manual/reference/connection-string/#connections-connection-options[connection string options] for further details.
+NOTE: Make sure to add `replicaSet` to the MongoDB URI. Please refer to https://docs.mongodb.com/manual/reference/connection-string/#connections-connection-options[connection string options] for further details.
 
 *Collection Operations*
 
@@ -321,9 +321,9 @@ The server responds with _error 50851_ when issuing a `count` command inside of
 Once `MongoTemplate` detects an active transaction, all exposed `count()` methods are converted and delegated to the
 aggregation framework using `$match` and `$count` operators, preserving `Query` settings, such as `collation`.
 
-====
-The following snippet of `count` inside the session bound closure
+The following snippet shows `count` usage inside the session-bound closure:
 
+====
 [source,javascript]
 ----
 session.startTransaction();
@@ -333,23 +333,27 @@ template.withSession(session)
         action.count(query(where("state").is("active")), Step.class)
         ...
 ----
+====
 
-runs:
+The snippet above materializes in the following command:
 
+====
 [source,javascript]
 ----
 db.collection.aggregate(
    [
       { $match: { state: "active" } },
-      { $group: { _id: null, count: { $sum: 1 } } }
+      { $count: "totalEntityCount" }
    ]
 )
 ----
+====
 
 instead of:
 
+====
 [source,javascript]
 ----
 db.collection.find( { state: "active" } ).count()
 ----
-====
+====