Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance. Project price only 1 $
You can buy this project and download/modify it how often you want.
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
package com.azure.cosmos.implementation.throughputControl;
import com.azure.cosmos.BridgeInternal;
import com.azure.cosmos.ConnectionMode;
import com.azure.cosmos.CosmosException;
import com.azure.cosmos.implementation.ResourceType;
import com.azure.cosmos.implementation.RxDocumentServiceRequest;
import com.azure.cosmos.implementation.Utils;
import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair;
import com.azure.cosmos.implementation.caches.AsyncCache;
import com.azure.cosmos.implementation.caches.RxClientCollectionCache;
import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache;
import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal;
import com.azure.cosmos.implementation.throughputControl.controller.IThroughputController;
import com.azure.cosmos.implementation.throughputControl.controller.container.EmptyThroughputContainerController;
import com.azure.cosmos.implementation.throughputControl.controller.container.IThroughputContainerController;
import com.azure.cosmos.implementation.throughputControl.controller.container.ThroughputContainerController;
import com.azure.cosmos.implementation.throughputControl.exceptions.ThroughputControlInitializationException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.Exceptions;
import reactor.core.publisher.Mono;
import java.util.concurrent.ConcurrentHashMap;
import static com.azure.cosmos.implementation.Exceptions.isNameCacheStale;
import static com.azure.cosmos.implementation.Exceptions.isPartitionKeyMismatchException;
import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument;
import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
/**
* This is the entrance class for the whole throughput control work flow pipeline.
* The pipeline will consist of controllers which is implementation of {@link IThroughputController} and {@link ThroughputRequestThrottler}.
*
* Following is a high-level diagram of the pipeline:
*
* +-------------------+
* | Client |
* +-------------------+
* |
* |
* |
* +---------------------------------------------------------+
* | ThroughputControlStore |
* +---------------------------------------------------------+
* / \
* / \
* / \
* +---------------------------------------+ +---------------------------------------+
* | Container A controller | | Container B controller |
* +---------------------------------------+ ... +---------------------------------------+
* / \
* / \
* / \
* +--------------------+ +---------------------+
* | Group 1 controller | | Group 2 controller |
* +--------------------+ ... +---------------------+
* | \
* | \-------------\
* | \
* +---------------------------+ +-----------------------------+
* |Global Request controller | OR | PkRanges Request controller |
* +--------------------------+ +------------------------------+
* | / \
* | / \
* | / \
* +------------------+ +------------------+ +------------------+
* |Request throttler | |Request throttler | |Request throttler |
* +------------------+ +------------------+ ... +------------------+
*
*
*/
public class ThroughputControlStore {
private static final Logger logger = LoggerFactory.getLogger(ThroughputControlStore.class);
private final RxClientCollectionCache collectionCache;
private final ConnectionMode connectionMode;
private final AsyncCache containerControllerCache;
private final ConcurrentHashMap containerMap;
private final RxPartitionKeyRangeCache partitionKeyRangeCache;
private final LinkedCancellationTokenSource cancellationTokenSource;
private final ConcurrentHashMap cancellationTokenMap;
public ThroughputControlStore(
RxClientCollectionCache collectionCache,
ConnectionMode connectionMode,
RxPartitionKeyRangeCache partitionKeyRangeCache) {
checkNotNull(collectionCache,"RxClientCollectionCache can not be null");
checkNotNull(partitionKeyRangeCache, "PartitionKeyRangeCache can not be null");
this.collectionCache = collectionCache;
this.connectionMode = connectionMode;
this.containerControllerCache = new AsyncCache<>();
this.containerMap = new ConcurrentHashMap<>();
this.partitionKeyRangeCache = partitionKeyRangeCache;
this.cancellationTokenSource = new LinkedCancellationTokenSource();
this.cancellationTokenMap = new ConcurrentHashMap<>();
}
public void enableThroughputControlGroup(ThroughputControlGroupInternal group, Mono throughputQueryMono) {
checkNotNull(group, "Throughput control group cannot be null");
String containerNameLink = Utils.trimBeginningAndEndingSlashes(BridgeInternal.extractContainerSelfLink(group.getTargetContainer()));
this.containerMap.compute(containerNameLink, (key, throughputControlContainerProperties) -> {
if (throughputControlContainerProperties == null) {
throughputControlContainerProperties = new ContainerThroughputControlGroupProperties(containerNameLink);
}
int groupSizeBefore = throughputControlContainerProperties.getThroughputControlGroupSet().size();
Pair stateAfterEnabling =
throughputControlContainerProperties.enableThroughputControlGroup(group, throughputQueryMono);
int groupSizeAfter = stateAfterEnabling.getLeft();
boolean wasGroupConfigUpdated = stateAfterEnabling.getRight();
if ((groupSizeAfter > groupSizeBefore && groupSizeAfter == 1) || wasGroupConfigUpdated) {
// This is the first enabled group for the target container or an existing group was modified
// Clean the current cache in case we have built EmptyThroughputContainerController or an existing
// group was modified
this.containerControllerCache.remove(containerNameLink);
}
return throughputControlContainerProperties;
});
}
public Mono processRequest(RxDocumentServiceRequest request, Mono originalRequestMono) {
checkNotNull(request, "Request can not be null");
checkNotNull(originalRequestMono, "originalRequestMono can not be null");
// Currently, we will only target two resource types.
// If in the future we find other useful scenarios for throughput control, add more resource type here.
if (request.getResourceType() != ResourceType.Document && request.getResourceType() != ResourceType.StoredProcedure) {
return originalRequestMono;
}
String collectionNameLink = Utils.getCollectionName(request.getResourceAddress());
return this.resolveContainerController(collectionNameLink)
.flatMap(containerController -> {
if (containerController.canHandleRequest(request)) {
return containerController.processRequest(request, originalRequestMono)
.doOnError(throwable -> this.handleException(collectionNameLink, request, throwable));
}
// Unable to find container controller to handle the request,
// It is caused by control store out of sync or the request has staled info.
// We will handle the first scenario by creating a new container controller,
// while fall back to original request Mono for the second scenario.
return this.updateControllerAndRetry(collectionNameLink, request, originalRequestMono);
})
.onErrorResume(throwable -> {
Exception unwrappedException = Utils.as(Exceptions.unwrap(throwable), Exception.class);
if (unwrappedException instanceof ThroughputControlInitializationException) {
if (this.shouldContinueRequestOnInitError(request, collectionNameLink, unwrappedException)) {
return originalRequestMono;
}
return Mono.error(unwrappedException.getCause());
}
return Mono.error(throwable);
});
}
private boolean shouldContinueRequestOnInitError(RxDocumentServiceRequest request, String collectionNameLink, Throwable throwable) {
if (throwable instanceof ThroughputControlInitializationException) {
ContainerThroughputControlGroupProperties throughputControlContainerProperties = this.containerMap.get(collectionNameLink);
checkNotNull(
throughputControlContainerProperties,
"Throughput control container properties should not be null");
checkArgument(
throughputControlContainerProperties.getThroughputControlGroupSet().size() > 0,
"There should be more than one throughput control group");
return throughputControlContainerProperties.allowRequestToContinueOnInitError(request);
}
return false;
}
private Mono updateControllerAndRetry(
String containerNameLink,
RxDocumentServiceRequest request,
Mono originalRequestMono) {
return this.shouldRefreshContainerController(containerNameLink, request)
.flatMap(shouldRefresh -> {
if (shouldRefresh) {
this.cancellationTokenMap.compute(containerNameLink, (key, cancellationToken) -> {
if (cancellationToken != null) {
cancellationToken.cancel();
}
return null;
});
this.containerControllerCache.refresh(containerNameLink, () -> this.createAndInitContainerController(containerNameLink));
return this.resolveContainerController(containerNameLink)
.flatMap(updatedContainerController -> {
if (updatedContainerController.canHandleRequest(request)) {
return updatedContainerController.processRequest(request, originalRequestMono)
.doOnError(throwable -> this.handleException(containerNameLink, request, throwable));
} else {
// still can not handle the request
logger.warn(
"Can not find container controller to process request {} with collectionRid {} ",
request.getActivityId(),
request.requestContext.resolvedCollectionRid);
return originalRequestMono;
}
});
}
return originalRequestMono;
});
}
private Mono resolveContainerController(String containerNameLink) {
checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container name link can not be null or empty");
return this.containerControllerCache.getAsync(
containerNameLink,
null,
() -> this.createAndInitContainerController(containerNameLink))
.onErrorResume(throwable -> Mono.error(new ThroughputControlInitializationException(throwable)));
}
private Mono createAndInitContainerController(String containerNameLink) {
checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container link should not be null or empty");
if (this.containerMap.containsKey(containerNameLink)) {
return Mono.just(this.containerMap.get(containerNameLink))
.flatMap(throughputControlContainerProperties -> {
LinkedCancellationToken parentToken =
this.cancellationTokenMap.compute(
containerNameLink,
(key, cancellationToken) -> this.cancellationTokenSource.getToken());
ThroughputContainerController containerController =
new ThroughputContainerController(
this.collectionCache,
this.connectionMode,
throughputControlContainerProperties.getThroughputControlGroupSet(),
this.partitionKeyRangeCache,
parentToken,
throughputControlContainerProperties.getThroughputQueryMono());
return containerController.init();
});
} else {
return Mono.just(new EmptyThroughputContainerController())
.flatMap(EmptyThroughputContainerController::init);
}
}
private Mono shouldRefreshContainerController(String containerLink, RxDocumentServiceRequest request) {
// TODO: populate diagnostics
return this.collectionCache.resolveByNameAsync(null, containerLink, null)
.flatMap(documentCollection ->
Mono.just(StringUtils.equals(documentCollection.getResourceId(), request.requestContext.resolvedCollectionRid)));
}
private void handleException(String containerNameLink, RxDocumentServiceRequest request, Throwable throwable) {
checkArgument(StringUtils.isNotEmpty(containerNameLink), "Container name link can not be null nor empty");
checkNotNull(request, "Request can not be null");
checkNotNull(throwable, "Exception can not be null");
CosmosException cosmosException = Utils.as(Exceptions.unwrap(throwable), CosmosException.class);
if (cosmosException != null &&
(isNameCacheStale(cosmosException) || isPartitionKeyMismatchException(cosmosException))) {
this.cancellationTokenMap.compute(containerNameLink,(key, cancellationToken) -> {
if (cancellationToken != null) {
cancellationToken.cancel();
}
return null;
});
String containerLink = Utils.getCollectionName(request.getResourceAddress());
this.collectionCache.refresh(null, containerLink, null);
this.containerControllerCache.refresh(
containerLink,
() -> createAndInitContainerController(containerLink)
);
}
}
public void close() {
this.cancellationTokenSource.close();
}
}
