diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java index 1cc07789c..9f9c33475 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java @@ -2188,6 +2188,10 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware, options.getComment().ifPresent(aggregateIterable::comment); + if(options.hasExecutionTimeLimit()) { + aggregateIterable = aggregateIterable.maxTime(options.getMaxTime().toMillis(), TimeUnit.MILLISECONDS); + } + MongoIterable iterable = aggregateIterable.map(val -> { rawResult.add(val); diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java index de3883f04..08e0410f3 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java @@ -1073,6 +1073,10 @@ public class ReactiveMongoTemplate implements ReactiveMongoOperations, Applicati .map(Collation::toMongoCollation) // .ifPresent(cursor::collation); + if(options.hasExecutionTimeLimit()) { + cursor = cursor.maxTime(options.getMaxTime().toMillis(), TimeUnit.MILLISECONDS); + } + return Flux.from(cursor).map(readCallback::doWith); } diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOptions.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOptions.java index 3a890ceb7..1be08a387 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOptions.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOptions.java @@ -15,6 +15,7 @@ */ package org.springframework.data.mongodb.core.aggregation; +import java.time.Duration; import java.util.Optional; import org.bson.Document; @@ -45,12 +46,14 @@ public class AggregationOptions { private static final String ALLOW_DISK_USE = "allowDiskUse"; private static final String COLLATION = "collation"; private static final String COMMENT = "comment"; + private static final String MAX_TIME = "maxTimeMS"; private final boolean allowDiskUse; private final boolean explain; private final Optional cursor; private final Optional collation; private final Optional comment; + private Duration maxTime = Duration.ZERO; /** * Creates a new {@link AggregationOptions}. @@ -129,7 +132,11 @@ public class AggregationOptions { : null; String comment = document.getString(COMMENT); - return new AggregationOptions(allowDiskUse, explain, cursor, collation, comment); + AggregationOptions options = new AggregationOptions(allowDiskUse, explain, cursor, collation, comment); + if (document.containsKey(MAX_TIME)) { + options.maxTime = Duration.ofMillis(document.getLong(MAX_TIME)); + } + return options; } /** @@ -206,6 +213,14 @@ public class AggregationOptions { return comment; } + /** + * @return the time limit for processing. {@link Duration#ZERO} is used for the default unbounded behavior. + * @since 2.3 + */ + public Duration getMaxTime() { + return maxTime; + } + /** * Returns a new potentially adjusted copy for the given {@code aggregationCommandObject} with the configuration * applied. @@ -233,6 +248,10 @@ public class AggregationOptions { collation.map(Collation::toDocument).ifPresent(val -> result.append(COLLATION, val)); } + if (hasExecutionTimeLimit() && !result.containsKey(MAX_TIME)) { + result.append(MAX_TIME, maxTime.toMillis()); + } + return result; } @@ -251,9 +270,17 @@ public class AggregationOptions { collation.ifPresent(val -> document.append(COLLATION, val.toDocument())); comment.ifPresent(val -> document.append(COMMENT, val)); + if (hasExecutionTimeLimit()) { + document.append(MAX_TIME, maxTime.toMillis()); + } + return document; } + public boolean hasExecutionTimeLimit() { + return !maxTime.isZero() && !maxTime.isNegative(); + } + /* (non-Javadoc) * @see java.lang.Object#toString() */ @@ -279,6 +306,7 @@ public class AggregationOptions { private @Nullable Document cursor; private @Nullable Collation collation; private @Nullable String comment; + private @Nullable Duration maxTime; /** * Defines whether to off-load intensive sort-operations to disk. @@ -355,13 +383,33 @@ public class AggregationOptions { return this; } + /** + * Set the time limit for processing. + * + * @param maxTime {@link Duration#ZERO} is used for the default unbounded behavior. {@link Duration#isNegative() + * Negative} values will be ignored. + * @return this. + * @sinve 2.3 + */ + public Builder maxTime(@Nullable Duration maxTime) { + + this.maxTime = maxTime; + return this; + } + /** * Returns a new {@link AggregationOptions} instance with the given configuration. * * @return */ public AggregationOptions build() { - return new AggregationOptions(allowDiskUse, explain, cursor, collation, comment); + + AggregationOptions options = new AggregationOptions(allowDiskUse, explain, cursor, collation, comment); + if (maxTime != null) { + options.maxTime = maxTime; + } + + return options; } } } diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java index 6b8813734..c19ac80cc 100644 --- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java +++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java @@ -22,6 +22,7 @@ import static org.springframework.data.mongodb.test.util.Assertions.*; import lombok.Data; import java.math.BigInteger; +import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -29,6 +30,7 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.TimeUnit; import java.util.regex.Pattern; import org.assertj.core.api.Assertions; @@ -177,6 +179,7 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests { when(aggregateIterable.allowDiskUse(any())).thenReturn(aggregateIterable); when(aggregateIterable.batchSize(anyInt())).thenReturn(aggregateIterable); when(aggregateIterable.map(any())).thenReturn(aggregateIterable); + when(aggregateIterable.maxTime(anyLong(), any())).thenReturn(aggregateIterable); when(aggregateIterable.into(any())).thenReturn(Collections.emptyList()); when(distinctIterable.collation(any())).thenReturn(distinctIterable); when(distinctIterable.map(any())).thenReturn(distinctIterable); @@ -1261,6 +1264,27 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests { verify(aggregateIterable).collation(eq(com.mongodb.client.model.Collation.builder().locale("fr").build())); } + @Test // DATAMONGO-2390 + public void aggregateShouldNoApplyZeroOrNegativeMaxTime() { + + template.aggregate( + newAggregation(Sith.class, project("id")).withOptions(newAggregationOptions().maxTime(Duration.ZERO).build()), + AutogenerateableId.class, Document.class); + template.aggregate(newAggregation(Sith.class, project("id")).withOptions( + newAggregationOptions().maxTime(Duration.ofSeconds(-1)).build()), AutogenerateableId.class, Document.class); + + verify(aggregateIterable, never()).maxTime(anyLong(), any()); + } + + @Test // DATAMONGO-2390 + public void aggregateShouldApplyMaxTimeIfSet() { + + template.aggregate(newAggregation(Sith.class, project("id")).withOptions( + newAggregationOptions().maxTime(Duration.ofSeconds(10)).build()), AutogenerateableId.class, Document.class); + + verify(aggregateIterable).maxTime(eq(10000L), eq(TimeUnit.MILLISECONDS)); + } + @Test // DATAMONGO-1854 public void findAndReplaceShouldUseCollationWhenPresent() { diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveMongoTemplateUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveMongoTemplateUnitTests.java index 6e5becebd..a3489ca42 100644 --- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveMongoTemplateUnitTests.java +++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveMongoTemplateUnitTests.java @@ -23,11 +23,13 @@ import lombok.Data; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; +import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.TimeUnit; import org.assertj.core.api.Assertions; import org.bson.Document; @@ -147,6 +149,7 @@ public class ReactiveMongoTemplateUnitTests { when(findPublisher.first()).thenReturn(findPublisher); when(aggregatePublisher.allowDiskUse(anyBoolean())).thenReturn(aggregatePublisher); when(aggregatePublisher.collation(any())).thenReturn(aggregatePublisher); + when(aggregatePublisher.maxTime(anyLong(), any())).thenReturn(aggregatePublisher); when(aggregatePublisher.first()).thenReturn(findPublisher); this.mappingContext = new MongoMappingContext(); @@ -582,6 +585,36 @@ public class ReactiveMongoTemplateUnitTests { verify(aggregatePublisher).comment("expensive"); } + @Test // DATAMONGO-2390 + public void aggregateShouldNoApplyZeroOrNegativeMaxTime() { + + template + .aggregate(newAggregation(MongoTemplateUnitTests.Sith.class, project("id")).withOptions( + newAggregationOptions().maxTime(Duration.ZERO).build()), AutogenerateableId.class, Document.class) + .subscribe(); + template + .aggregate( + newAggregation(MongoTemplateUnitTests.Sith.class, project("id")) + .withOptions(newAggregationOptions().maxTime(Duration.ofSeconds(-1)).build()), + AutogenerateableId.class, Document.class) + .subscribe(); + + verify(aggregatePublisher, never()).maxTime(anyLong(), any()); + } + + @Test // DATAMONGO-2390 + public void aggregateShouldApplyMaxTimeIfSet() { + + template + .aggregate( + newAggregation(MongoTemplateUnitTests.Sith.class, project("id")) + .withOptions(newAggregationOptions().maxTime(Duration.ofSeconds(10)).build()), + AutogenerateableId.class, Document.class) + .subscribe(); + + verify(aggregatePublisher).maxTime(eq(10000L), eq(TimeUnit.MILLISECONDS)); + } + @Test // DATAMONGO-18545 public void findAndReplaceShouldUseCollationWhenPresent() {