io.github.bucket4j.distributed.proxy.DefaultAsyncBucketProxy Maven / Gradle / Ivy
The newest version!
/*-
* ========================LICENSE_START=================================
* Bucket4j
* %%
* Copyright (C) 2015 - 2020 Vladimir Bukhtoyarov
* %%
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
* =========================LICENSE_END==================================
*/
package io.github.bucket4j.distributed.proxy;
import io.github.bucket4j.*;
import io.github.bucket4j.distributed.AsyncBucketProxy;
import io.github.bucket4j.distributed.AsyncOptimizationController;
import io.github.bucket4j.distributed.AsyncVerboseBucket;
import io.github.bucket4j.distributed.remote.CommandResult;
import io.github.bucket4j.distributed.remote.RemoteCommand;
import io.github.bucket4j.distributed.remote.RemoteVerboseResult;
import io.github.bucket4j.distributed.remote.commands.*;
import java.time.Duration;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Supplier;
import static io.github.bucket4j.LimitChecker.*;
public class DefaultAsyncBucketProxy implements AsyncBucketProxy, AsyncOptimizationController {
private final AsyncCommandExecutor commandExecutor;
private final RecoveryStrategy recoveryStrategy;
private final Supplier> configurationSupplier;
private final BucketListener listener;
private final ImplicitConfigurationReplacement implicitConfigurationReplacement;
private final AtomicBoolean wasInitialized;
@Override
public AsyncVerboseBucket asVerbose() {
return asyncVerboseView;
}
@Override
public AsyncBucketProxy toListenable(BucketListener listener) {
return new DefaultAsyncBucketProxy(commandExecutor, recoveryStrategy, configurationSupplier, implicitConfigurationReplacement, wasInitialized, listener);
}
@Override
public SchedulingBucket asScheduler() {
return schedulingBucketView;
}
public DefaultAsyncBucketProxy(AsyncCommandExecutor commandExecutor, RecoveryStrategy recoveryStrategy, Supplier> configurationSupplier,
ImplicitConfigurationReplacement implicitConfigurationReplacement, BucketListener listener) {
this(commandExecutor, recoveryStrategy, configurationSupplier, implicitConfigurationReplacement, new AtomicBoolean(false), listener);
}
private DefaultAsyncBucketProxy(AsyncCommandExecutor commandExecutor, RecoveryStrategy recoveryStrategy, Supplier> configurationSupplier, ImplicitConfigurationReplacement implicitConfigurationReplacement, AtomicBoolean wasInitialized, BucketListener listener) {
this.commandExecutor = Objects.requireNonNull(commandExecutor);
this.recoveryStrategy = recoveryStrategy;
this.configurationSupplier = configurationSupplier;
this.implicitConfigurationReplacement = implicitConfigurationReplacement;
this.wasInitialized = wasInitialized;
if (listener == null) {
throw BucketExceptions.nullListener();
}
this.listener = listener;
}
private final AsyncVerboseBucket asyncVerboseView = new AsyncVerboseBucket() {
@Override
public CompletableFuture> tryConsume(long tokensToConsume) {
checkTokensToConsume(tokensToConsume);
VerboseCommand command = TryConsumeCommand.create(tokensToConsume).asVerbose();
return execute(command).thenApply(consumed -> {
if (consumed.getValue()) {
listener.onConsumed(tokensToConsume);
} else {
listener.onRejected(tokensToConsume);
}
return consumed.asLocal();
});
}
@Override
public CompletableFuture> consumeIgnoringRateLimits(long tokensToConsume) {
checkTokensToConsume(tokensToConsume);
VerboseCommand command = VerboseCommand.from(new ConsumeIgnoringRateLimitsCommand(tokensToConsume));
return execute(command).thenApply(penaltyNanos -> {
if (penaltyNanos.getValue() == INFINITY_DURATION) {
throw BucketExceptions.reservationOverflow();
}
listener.onConsumed(tokensToConsume);
return penaltyNanos.asLocal();
});
}
@Override
public CompletableFuture> tryConsumeAndReturnRemaining(long tokensToConsume) {
checkTokensToConsume(tokensToConsume);
VerboseCommand command = VerboseCommand.from(new TryConsumeAndReturnRemainingTokensCommand(tokensToConsume));
return execute(command).thenApply(probe -> {
if (probe.getValue().isConsumed()) {
listener.onConsumed(tokensToConsume);
} else {
listener.onRejected(tokensToConsume);
}
return probe.asLocal();
});
}
@Override
public CompletableFuture> estimateAbilityToConsume(long numTokens) {
checkTokensToConsume(numTokens);
return execute(VerboseCommand.from(new EstimateAbilityToConsumeCommand(numTokens)))
.thenApply(RemoteVerboseResult::asLocal);
}
@Override
public CompletableFuture> tryConsumeAsMuchAsPossible() {
VerboseCommand command = VerboseCommand.from(new ConsumeAsMuchAsPossibleCommand (UNLIMITED_AMOUNT));
return execute(command).thenApply(consumedTokens -> {
long actuallyConsumedTokens = consumedTokens.getValue();
if (actuallyConsumedTokens > 0) {
listener.onConsumed(actuallyConsumedTokens);
}
return consumedTokens.asLocal();
});
}
@Override
public CompletableFuture> tryConsumeAsMuchAsPossible(long limit) {
checkTokensToConsume(limit);
VerboseCommand verboseCommand = VerboseCommand.from(new ConsumeAsMuchAsPossibleCommand(limit));
return execute(verboseCommand).thenApply(consumedTokens -> {
long actuallyConsumedTokens = consumedTokens.getValue();
if (actuallyConsumedTokens > 0) {
listener.onConsumed(actuallyConsumedTokens);
}
return consumedTokens.asLocal();
});
}
@Override
public CompletableFuture> addTokens(long tokensToAdd) {
checkTokensToAdd(tokensToAdd);
VerboseCommand verboseCommand = VerboseCommand.from(new AddTokensCommand(tokensToAdd));
return execute(verboseCommand).thenApply(RemoteVerboseResult::asLocal);
}
@Override
public CompletableFuture> forceAddTokens(long tokensToAdd) {
checkTokensToAdd(tokensToAdd);
VerboseCommand verboseCommand = VerboseCommand.from(new ForceAddTokensCommand(tokensToAdd));
return execute(verboseCommand).thenApply(RemoteVerboseResult::asLocal);
}
@Override
public CompletableFuture> reset() {
VerboseCommand verboseCommand = VerboseCommand.from(new ResetCommand());
return execute(verboseCommand).thenApply(RemoteVerboseResult::asLocal);
}
@Override
public CompletableFuture> replaceConfiguration(BucketConfiguration newConfiguration, TokensInheritanceStrategy tokensInheritanceStrategy) {
checkConfiguration(newConfiguration);
checkMigrationMode(tokensInheritanceStrategy);
VerboseCommand command = VerboseCommand.from(new ReplaceConfigurationCommand(newConfiguration, tokensInheritanceStrategy));
return execute(command).thenApply(RemoteVerboseResult::asLocal);
}
@Override
public CompletableFuture> getAvailableTokens() {
VerboseCommand command = VerboseCommand.from(new GetAvailableTokensCommand());
return execute(command).thenApply(RemoteVerboseResult::asLocal);
}
};
private final SchedulingBucket schedulingBucketView = new SchedulingBucket() {
@Override
public CompletableFuture tryConsume(long tokensToConsume, long maxWaitTimeNanos, ScheduledExecutorService scheduler) {
checkMaxWaitTime(maxWaitTimeNanos);
checkTokensToConsume(tokensToConsume);
checkScheduler(scheduler);
CompletableFuture resultFuture = new CompletableFuture<>();
ReserveAndCalculateTimeToSleepCommand consumeCommand = new ReserveAndCalculateTimeToSleepCommand(tokensToConsume, maxWaitTimeNanos);
CompletableFuture reservationFuture = execute(consumeCommand);
reservationFuture.whenComplete((nanosToSleep, exception) -> {
if (exception != null) {
resultFuture.completeExceptionally(exception);
return;
}
if (nanosToSleep == INFINITY_DURATION) {
resultFuture.complete(false);
listener.onRejected(tokensToConsume);
return;
}
if (nanosToSleep == 0L) {
resultFuture.complete(true);
listener.onConsumed(tokensToConsume);
return;
}
try {
listener.onConsumed(tokensToConsume);
listener.onDelayed(nanosToSleep);
Runnable delayedCompletion = () -> resultFuture.complete(true);
scheduler.schedule(delayedCompletion, nanosToSleep, TimeUnit.NANOSECONDS);
} catch (Throwable t) {
resultFuture.completeExceptionally(t);
}
});
return resultFuture;
}
@Override
public CompletableFuture consume(long tokensToConsume, ScheduledExecutorService scheduler) {
checkTokensToConsume(tokensToConsume);
checkScheduler(scheduler);
CompletableFuture resultFuture = new CompletableFuture<>();
ReserveAndCalculateTimeToSleepCommand consumeCommand = new ReserveAndCalculateTimeToSleepCommand(tokensToConsume, INFINITY_DURATION);
CompletableFuture reservationFuture = execute(consumeCommand);
reservationFuture.whenComplete((nanosToSleep, exception) -> {
if (exception != null) {
resultFuture.completeExceptionally(exception);
return;
}
if (nanosToSleep == INFINITY_DURATION) {
resultFuture.completeExceptionally(BucketExceptions.reservationOverflow());
return;
}
if (nanosToSleep == 0L) {
resultFuture.complete(null);
listener.onConsumed(tokensToConsume);
return;
}
try {
listener.onConsumed(tokensToConsume);
listener.onDelayed(nanosToSleep);
Runnable delayedCompletion = () -> resultFuture.complete(null);
scheduler.schedule(delayedCompletion, nanosToSleep, TimeUnit.NANOSECONDS);
} catch (Throwable t) {
resultFuture.completeExceptionally(t);
}
});
return resultFuture;
}
@Override
public VerboseSchedulingBucket asVerbose() {
return verboseSchedulingView;
}
};
private final VerboseSchedulingBucket verboseSchedulingView = new VerboseSchedulingBucket() {
@Override
public CompletableFuture> tryConsume(long tokensToConsume, long maxWaitTimeNanos, ScheduledExecutorService scheduler) {
checkMaxWaitTime(maxWaitTimeNanos);
checkTokensToConsume(tokensToConsume);
checkScheduler(scheduler);
CompletableFuture> resultFuture = new CompletableFuture<>();
ReserveAndCalculateTimeToSleepCommand consumeCommand = new ReserveAndCalculateTimeToSleepCommand(tokensToConsume, maxWaitTimeNanos);
CompletableFuture> reservationFuture = execute(consumeCommand.asVerbose());
reservationFuture.whenComplete((RemoteVerboseResult nanosToSleepVerbose, Throwable exception) -> {
if (exception != null) {
resultFuture.completeExceptionally(exception);
return;
}
long nanosToSleep = nanosToSleepVerbose.getValue();
if (nanosToSleep == INFINITY_DURATION) {
resultFuture.complete(nanosToSleepVerbose.withValue(false).asLocal());
listener.onRejected(tokensToConsume);
return;
}
if (nanosToSleep == 0L) {
resultFuture.complete(nanosToSleepVerbose.withValue(true).asLocal());
listener.onConsumed(tokensToConsume);
return;
}
try {
listener.onConsumed(tokensToConsume);
listener.onDelayed(nanosToSleep);
Runnable delayedCompletion = () -> resultFuture.complete(nanosToSleepVerbose.withValue(true).asLocal());
scheduler.schedule(delayedCompletion, nanosToSleep, TimeUnit.NANOSECONDS);
} catch (Throwable t) {
resultFuture.completeExceptionally(t);
}
});
return resultFuture;
}
@Override
public CompletableFuture> consume(long tokensToConsume, ScheduledExecutorService scheduler) {
checkTokensToConsume(tokensToConsume);
checkScheduler(scheduler);
CompletableFuture> resultFuture = new CompletableFuture<>();
ReserveAndCalculateTimeToSleepCommand consumeCommand = new ReserveAndCalculateTimeToSleepCommand(tokensToConsume, INFINITY_DURATION);
CompletableFuture> reservationFuture = execute(consumeCommand.asVerbose());
reservationFuture.whenComplete((RemoteVerboseResult nanosToSleepVerbose, Throwable exception) -> {
if (exception != null) {
resultFuture.completeExceptionally(exception);
return;
}
long nanosToSleep = nanosToSleepVerbose.getValue();
if (nanosToSleep == INFINITY_DURATION) {
resultFuture.completeExceptionally(BucketExceptions.reservationOverflow());
return;
}
if (nanosToSleep == 0L) {
resultFuture.complete(nanosToSleepVerbose.withValue((Void) null).asLocal());
listener.onConsumed(tokensToConsume);
return;
}
try {
listener.onConsumed(tokensToConsume);
listener.onDelayed(nanosToSleep);
Runnable delayedCompletion = () -> resultFuture.complete(nanosToSleepVerbose.withValue((Void) null).asLocal());
scheduler.schedule(delayedCompletion, nanosToSleep, TimeUnit.NANOSECONDS);
} catch (Throwable t) {
resultFuture.completeExceptionally(t);
}
});
return resultFuture;
}
};
@Override
public CompletableFuture consumeIgnoringRateLimits(long tokensToConsume) {
checkTokensToConsume(tokensToConsume);
return execute(new ConsumeIgnoringRateLimitsCommand(tokensToConsume)).thenApply(penaltyNanos -> {
if (penaltyNanos == INFINITY_DURATION) {
throw BucketExceptions.reservationOverflow();
}
listener.onConsumed(tokensToConsume);
return penaltyNanos;
});
}
@Override
public CompletableFuture tryConsume(long tokensToConsume) {
checkTokensToConsume(tokensToConsume);
return execute(TryConsumeCommand.create(tokensToConsume)).thenApply(consumed -> {
if (consumed) {
listener.onConsumed(tokensToConsume);
} else {
listener.onRejected(tokensToConsume);
}
return consumed;
});
}
@Override
public CompletableFuture tryConsumeAndReturnRemaining(long tokensToConsume) {
checkTokensToConsume(tokensToConsume);
return execute(new TryConsumeAndReturnRemainingTokensCommand(tokensToConsume)).thenApply(probe -> {
if (probe.isConsumed()) {
listener.onConsumed(tokensToConsume);
} else {
listener.onRejected(tokensToConsume);
}
return probe;
});
}
@Override
public CompletableFuture estimateAbilityToConsume(long numTokens) {
checkTokensToConsume(numTokens);
return execute(new EstimateAbilityToConsumeCommand(numTokens));
}
@Override
public CompletableFuture tryConsumeAsMuchAsPossible() {
return execute(new ConsumeAsMuchAsPossibleCommand(UNLIMITED_AMOUNT)).thenApply(consumedTokens -> {
if (consumedTokens > 0) {
listener.onConsumed(consumedTokens);
}
return consumedTokens;
});
}
@Override
public CompletableFuture tryConsumeAsMuchAsPossible(long limit) {
checkTokensToConsume(limit);
return execute(new ConsumeAsMuchAsPossibleCommand(limit)).thenApply(consumedTokens -> {
if (consumedTokens > 0) {
listener.onConsumed(consumedTokens);
}
return consumedTokens;
});
}
@Override
public CompletableFuture replaceConfiguration(BucketConfiguration newConfiguration, TokensInheritanceStrategy tokensInheritanceStrategy) {
checkConfiguration(newConfiguration);
checkMigrationMode(tokensInheritanceStrategy);
ReplaceConfigurationCommand replaceConfigCommand = new ReplaceConfigurationCommand(newConfiguration, tokensInheritanceStrategy);
CompletableFuture result = execute(replaceConfigCommand);
return result.thenApply(nothing -> null);
}
@Override
public CompletableFuture addTokens(long tokensToAdd) {
checkTokensToAdd(tokensToAdd);
CompletableFuture future = execute(new AddTokensCommand(tokensToAdd));
return future.thenApply(nothing -> null);
}
@Override
public CompletableFuture forceAddTokens(long tokensToAdd) {
checkTokensToAdd(tokensToAdd);
CompletableFuture future = execute(new ForceAddTokensCommand(tokensToAdd));
return future.thenApply(nothing -> null);
}
@Override
public CompletableFuture reset() {
CompletableFuture future = execute(new ResetCommand());
return future.thenApply(nothing -> null);
}
@Override
public CompletableFuture getAvailableTokens() {
return execute(new GetAvailableTokensCommand());
}
@Override
public AsyncOptimizationController getOptimizationController() {
return this;
}
@Override
public CompletableFuture syncByCondition(long unsynchronizedTokens, Duration timeSinceLastSync) {
return execute(new SyncCommand(unsynchronizedTokens, timeSinceLastSync.toNanos())).thenApply(nothing -> null);
}
private CompletableFuture execute(RemoteCommand command) {
RemoteCommand commandToExecute = implicitConfigurationReplacement == null? command :
new CheckConfigurationVersionAndExecuteCommand<>(command, implicitConfigurationReplacement.getDesiredConfigurationVersion());
boolean wasInitializedBeforeExecution = wasInitialized.get();
CompletableFuture> futureResult = commandExecutor.executeAsync(commandToExecute);
return futureResult.thenCompose(cmdResult -> {
if (!cmdResult.isBucketNotFound() && !cmdResult.isConfigurationNeedToBeReplaced()) {
T resultDate = cmdResult.getData();
return CompletableFuture.completedFuture(resultDate);
}
// the bucket was removed or lost, or not initialized yet, or needs to upgrade configuration
if (cmdResult.isBucketNotFound() && recoveryStrategy == RecoveryStrategy.THROW_BUCKET_NOT_FOUND_EXCEPTION && wasInitializedBeforeExecution) {
CompletableFuture failedFuture = new CompletableFuture<>();
failedFuture.completeExceptionally(new BucketNotFoundException());
return failedFuture;
}
// fetch actual configuration
CompletableFuture configurationFuture;
try {
configurationFuture = configurationSupplier.get();
} catch (Throwable t) {
CompletableFuture failedFuture = new CompletableFuture<>();
failedFuture.completeExceptionally(t);
return failedFuture;
}
if (configurationFuture == null) {
CompletableFuture failedFuture = new CompletableFuture<>();
failedFuture.completeExceptionally(BucketExceptions.nullConfigurationFuture());
return failedFuture;
}
// retry command execution
return configurationFuture.thenCompose(configuration -> {
if (configuration == null) {
CompletableFuture failedFuture = new CompletableFuture<>();
failedFuture.completeExceptionally(BucketExceptions.nullConfiguration());
return failedFuture;
}
RemoteCommand initAndExecuteCommand = implicitConfigurationReplacement == null?
new CreateInitialStateAndExecuteCommand<>(configuration, command) :
new CreateInitialStateWithVersionOrReplaceConfigurationAndExecuteCommand<>(configuration, command, implicitConfigurationReplacement.getDesiredConfigurationVersion(), implicitConfigurationReplacement.getTokensInheritanceStrategy());
return commandExecutor.executeAsync(initAndExecuteCommand).thenApply(initAndExecuteCmdResult -> {
wasInitialized.set(true);
return initAndExecuteCmdResult.getData();
});
});
});
}
}