diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java index 8910984401c2f..a856fcb27db9d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java @@ -27,7 +27,7 @@ import java.net.URI; import java.net.URISyntaxException; import java.util.List; -import java.util.concurrent.ConcurrentMap; +import java.util.Map; /** * Provides a client-side logical representation of the Azure Cosmos DB @@ -1570,7 +1570,7 @@ Flux> readAllDocuments( CosmosQueryRequestOptions options ); - ConcurrentMap getQueryPlanCache(); + Map getQueryPlanCache(); /** * Gets the collection cache. diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index eb7f152eb17b0..ccf4e48ef5403 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -235,7 +235,7 @@ public static String getEnvironmentName() { public static boolean isQueryPlanCachingEnabled() { // Queryplan caching will be disabled by default - return getJVMConfigAsBoolean(QUERYPLAN_CACHING_ENABLED, false); + return getJVMConfigAsBoolean(QUERYPLAN_CACHING_ENABLED, true); } public static int getAddressRefreshResponseTimeoutInSeconds() { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java index d5b4e90265bac..fc98770bbec98 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java @@ -245,4 +245,6 @@ public static final class PartitionedQueryExecutionInfo { public static final class QueryExecutionContext { public static final String INCREMENTAL_FEED_HEADER_VALUE = "Incremental feed"; } + + public static final int QUERYPLAN_CACHE_SIZE = 1000; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index a794eb66202de..ee7fd8ca91e5a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -16,6 +16,7 @@ import com.azure.cosmos.implementation.batch.PartitionKeyRangeServerBatchRequest; import com.azure.cosmos.implementation.batch.ServerBatchRequest; import com.azure.cosmos.implementation.batch.SinglePartitionKeyServerBatchRequest; +import com.azure.cosmos.implementation.caches.SizeLimitingLRUCache; import com.azure.cosmos.implementation.caches.RxClientCollectionCache; import com.azure.cosmos.implementation.caches.RxCollectionCache; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; @@ -142,7 +143,7 @@ public class RxDocumentClientImpl implements AsyncDocumentClient, IAuthorization private RxPartitionKeyRangeCache partitionKeyRangeCache; private Map> resourceTokensMap; private final boolean contentResponseOnWriteEnabled; - private ConcurrentMap queryPlanCache; + private Map queryPlanCache; private final AtomicBoolean closed = new AtomicBoolean(false); private final int clientId; @@ -362,7 +363,7 @@ private RxDocumentClientImpl(URI serviceEndpoint, this.retryPolicy = new RetryPolicy(this, this.globalEndpointManager, this.connectionPolicy); this.resetSessionTokenRetryPolicy = retryPolicy; CpuMemoryMonitor.register(this); - this.queryPlanCache = new ConcurrentHashMap<>(); + this.queryPlanCache = Collections.synchronizedMap(new SizeLimitingLRUCache(Constants.QUERYPLAN_CACHE_SIZE)); } catch (RuntimeException e) { logger.error("unexpected failure in initializing client.", e); close(); @@ -2462,7 +2463,7 @@ public Flux> readAllDocuments( } @Override - public ConcurrentMap getQueryPlanCache() { + public Map getQueryPlanCache() { return queryPlanCache; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/SizeLimitingLRUCache.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/SizeLimitingLRUCache.java new file mode 100644 index 0000000000000..5e73def189e2e --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/SizeLimitingLRUCache.java @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation.caches; + +import com.azure.cosmos.implementation.query.PartitionedQueryExecutionInfo; + +import java.util.LinkedHashMap; +import java.util.Map; + +/** + * LRU Cache using LinkedHashMap that limits the number of entries + */ +public class SizeLimitingLRUCache extends LinkedHashMap { + + private static final long serialVersionUID = 1L; + private final int maxEntries; + + public SizeLimitingLRUCache(int maxEntries) { + this.maxEntries = maxEntries; + } + + public SizeLimitingLRUCache(int initialCapacity, float loadFactor, int maxEntries) { + super(initialCapacity, loadFactor); + this.maxEntries = maxEntries; + } + + public SizeLimitingLRUCache( + Map m, int maxEntries) { + super(m); + this.maxEntries = maxEntries; + } + + public SizeLimitingLRUCache(int initialCapacity, float loadFactor, boolean accessOrder, int maxEntries) { + super(initialCapacity, loadFactor, accessOrder); + this.maxEntries = maxEntries; + } + + public SizeLimitingLRUCache(int initialCapacity, int maxEntries) { + super(initialCapacity); + this.maxEntries = maxEntries; + } + + @Override + protected boolean removeEldestEntry( + Map.Entry eldest) { + return size() > maxEntries; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java index 62e6751f8090f..8d9c3c6df2215 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java @@ -34,7 +34,6 @@ import java.util.List; import java.util.Map; import java.util.UUID; -import java.util.concurrent.ConcurrentMap; import java.util.stream.Collectors; /** @@ -71,7 +70,7 @@ private static Mono>,QueryInfo>> ge String resourceLink, DocumentCollection collection, DefaultDocumentQueryExecutionContext queryExecutionContext, boolean queryPlanCachingEnabled, - ConcurrentMap queryPlanCache) { + Map queryPlanCache) { // The partitionKeyRangeIdInternal is no more a public API on // FeedOptions, but have the below condition @@ -99,6 +98,7 @@ private static Mono>,QueryInfo>> ge Instant endTime = Instant.now(); // endTime for query plan diagnostics PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = queryPlanCache.get(query.getQueryText()); if (partitionedQueryExecutionInfo != null) { + logger.info("Skipping query plan round trip by using the cached plan"); return getTargetRangesFromQueryPlan(cosmosQueryRequestOptions, collection, queryExecutionContext, partitionedQueryExecutionInfo, startTime, endTime); } @@ -117,7 +117,7 @@ private static Mono>,QueryInfo>> ge Instant endTime = Instant.now(); - if (queryPlanCachingEnabled) { + if (queryPlanCachingEnabled && isScopedToSinglePartition(cosmosQueryRequestOptions)) { tryCacheQueryPlan(query, partitionedQueryExecutionInfo, queryPlanCache); } @@ -169,14 +169,9 @@ private static Mono>, QueryInfo>> g private static void tryCacheQueryPlan( SqlQuerySpec query, PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, - ConcurrentMap queryPlanCache) { + Map queryPlanCache) { QueryInfo queryInfo = partitionedQueryExecutionInfo.getQueryInfo(); if (canCacheQuery(queryInfo) && !queryPlanCache.containsKey(query.getQueryText())) { - if (queryPlanCache.size() > MAX_CACHE_SIZE) { - // Clearing query plan cache if size is above max size. This can be optimized in future by using - // a threadsafe LRU cache - queryPlanCache.clear(); - } queryPlanCache.put(query.getQueryText(), partitionedQueryExecutionInfo); } } @@ -188,7 +183,9 @@ private static boolean canCacheQuery(QueryInfo queryInfo) { && !queryInfo.hasGroupBy() && !queryInfo.hasLimit() && !queryInfo.hasTop() - && !queryInfo.hasOffset(); + && !queryInfo.hasOffset() + && !queryInfo.hasDCount() + && !queryInfo.hasOrderBy(); } private static boolean isScopedToSinglePartition(CosmosQueryRequestOptions cosmosQueryRequestOptions) { @@ -208,7 +205,7 @@ public static Flux queryPlanCache) { + Map queryPlanCache) { // return proxy Flux> collectionObs = Flux.just(new Utils.ValueHolder<>(null)); diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/QueryValidationTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/QueryValidationTests.java index 9655d25223356..cb81f7fc94031 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/QueryValidationTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/QueryValidationTests.java @@ -239,7 +239,6 @@ public void queryOptionNullValidation() { private Object[][] query() { return new Object[][]{ new Object[] { "Select * from c "}, - new Object[] { "select * from c order by c.prop ASC"}, }; } @@ -299,7 +298,6 @@ public void queryPlanCacheSinglePartitionParameterizedQueriesCorrectness() { createdContainer); documentsInserted.addAll(pk2Docs); - CosmosQueryRequestOptions options = new CosmosQueryRequestOptions(); options.setPartitionKey(new PartitionKey(pk2)); @@ -325,6 +323,20 @@ public void queryPlanCacheSinglePartitionParameterizedQueriesCorrectness() { // Top query should not be cached assertThat(contextClient.getQueryPlanCache().containsKey(sqlQuerySpec.getQueryText())).isFalse(); + // group by should not be cached + sqlQuerySpec.setQueryText("select max(c.id) from c order by c.name group by c.name"); + values1 = queryAndGetResults(sqlQuerySpec, options, TestObject.class); + assertThat(contextClient.getQueryPlanCache().containsKey(sqlQuerySpec.getQueryText())).isFalse(); + + // distinct queries should not be cached + sqlQuerySpec.setQueryText("SELECT distinct c.name from c"); + values1 = queryAndGetResults(sqlQuerySpec, options, TestObject.class); + assertThat(contextClient.getQueryPlanCache().containsKey(sqlQuerySpec.getQueryText())).isFalse(); + + //order by query should not be cached + sqlQuerySpec.setQueryText("select * from c order by c.name"); + values1 = queryAndGetResults(sqlQuerySpec, options, TestObject.class); + assertThat(contextClient.getQueryPlanCache().containsKey(sqlQuerySpec.getQueryText())).isFalse(); } @Test(groups = {"simple"}, timeOut = TIMEOUT * 40) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplatePartitionIT.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplatePartitionIT.java index 26a1a849b9c7e..2331cfdad611f 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplatePartitionIT.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplatePartitionIT.java @@ -42,8 +42,8 @@ import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.UUID; -import java.util.concurrent.ConcurrentMap; import static com.azure.spring.data.cosmos.common.TestConstants.ADDRESSES; import static com.azure.spring.data.cosmos.common.TestConstants.FIRST_NAME; @@ -144,7 +144,7 @@ public void testFindWithPartitionWithQueryPlanCachingEnabled() { CosmosAsyncClient cosmosAsyncClient = cosmosFactory.getCosmosAsyncClient(); AsyncDocumentClient asyncDocumentClient = CosmosBridgeInternal.getAsyncDocumentClient(cosmosAsyncClient); - ConcurrentMap initialCache = asyncDocumentClient.getQueryPlanCache(); + Map initialCache = asyncDocumentClient.getQueryPlanCache(); assertThat(initialCache.containsKey(sqlQuerySpec.getQueryText())).isTrue(); int initialSize = initialCache.size(); @@ -158,7 +158,7 @@ public void testFindWithPartitionWithQueryPlanCachingEnabled() { result = TestUtils.toList(cosmosTemplate.find(query, PartitionPerson.class, PartitionPerson.class.getSimpleName())); - ConcurrentMap postQueryCallCache = asyncDocumentClient.getQueryPlanCache(); + Map postQueryCallCache = asyncDocumentClient.getQueryPlanCache(); assertThat(postQueryCallCache.containsKey(sqlQuerySpec.getQueryText())).isTrue(); assertThat(postQueryCallCache.size()).isEqualTo(initialSize); assertThat(result.size()).isEqualTo(1); diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplatePartitionIT.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplatePartitionIT.java index 3c980e9318ac2..86d906023ad74 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplatePartitionIT.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplatePartitionIT.java @@ -39,8 +39,8 @@ import reactor.test.StepVerifier; import java.util.Collections; +import java.util.Map; import java.util.UUID; -import java.util.concurrent.ConcurrentMap; import static org.assertj.core.api.Assertions.assertThat; import static org.hamcrest.Matchers.equalTo; @@ -126,7 +126,7 @@ public void testFindWithPartitionWithQueryPlanCachingEnabled() { CosmosAsyncClient cosmosAsyncClient = cosmosFactory.getCosmosAsyncClient(); AsyncDocumentClient asyncDocumentClient = CosmosBridgeInternal.getAsyncDocumentClient(cosmosAsyncClient); - ConcurrentMap initialCache = asyncDocumentClient.getQueryPlanCache(); + Map initialCache = asyncDocumentClient.getQueryPlanCache(); assertThat(initialCache.containsKey(sqlQuerySpec.getQueryText())).isTrue(); int initialSize = initialCache.size(); @@ -145,7 +145,7 @@ public void testFindWithPartitionWithQueryPlanCachingEnabled() { Assert.assertThat(actual.getZipCode(), is(equalTo(TEST_PERSON_2.getZipCode()))); }).verifyComplete(); - ConcurrentMap postQueryCallCache = asyncDocumentClient.getQueryPlanCache(); + Map postQueryCallCache = asyncDocumentClient.getQueryPlanCache(); assertThat(postQueryCallCache.containsKey(sqlQuerySpec.getQueryText())).isTrue(); assertThat(postQueryCallCache.size()).isEqualTo(initialSize); }