io.github.bucket4j.redis.lettuce.cas.LettuceBasedProxyManager Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of bucket4j_jdk17-lettuce Show documentation
Show all versions of bucket4j_jdk17-lettuce Show documentation
bucket4j - is a java implementation of token bucket algorithm for rate limiting
The newest version!
/*-
* ========================LICENSE_START=================================
* Bucket4j
* %%
* Copyright (C) 2015 - 2022 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.redis.lettuce.cas;
import java.nio.charset.StandardCharsets;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import io.github.bucket4j.distributed.ExpirationAfterWriteStrategy;
import io.github.bucket4j.distributed.proxy.generic.compare_and_swap.AbstractCompareAndSwapBasedProxyManager;
import io.github.bucket4j.distributed.proxy.generic.compare_and_swap.AsyncCompareAndSwapOperation;
import io.github.bucket4j.distributed.proxy.generic.compare_and_swap.CompareAndSwapOperation;
import io.github.bucket4j.distributed.remote.RemoteBucketState;
import io.github.bucket4j.redis.AbstractRedisProxyManagerBuilder;
import io.github.bucket4j.redis.consts.LuaScripts;
import io.github.bucket4j.redis.lettuce.Bucket4jLettuce;
import io.github.bucket4j.redis.lettuce.RedisApi;
import io.lettuce.core.RedisClient;
import io.lettuce.core.RedisException;
import io.lettuce.core.RedisFuture;
import io.lettuce.core.ScriptOutputType;
import io.lettuce.core.api.StatefulRedisConnection;
import io.lettuce.core.api.async.RedisAsyncCommands;
import io.lettuce.core.cluster.RedisClusterClient;
import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
import io.lettuce.core.cluster.api.async.RedisAdvancedClusterAsyncCommands;
import io.lettuce.core.codec.ByteArrayCodec;
public class LettuceBasedProxyManager extends AbstractCompareAndSwapBasedProxyManager {
private final RedisApi redisApi;
private final ExpirationAfterWriteStrategy expirationStrategy;
/**
* @deprecated use {@link Bucket4jLettuce#casBasedBuilder(RedisAsyncCommands)}
*/
@Deprecated
public static LettuceBasedProxyManagerBuilder builderFor(RedisAsyncCommands redisAsyncCommands) {
Objects.requireNonNull(redisAsyncCommands);
RedisApi redisApi = new RedisApi<>() {
@Override
public RedisFuture eval(String script, ScriptOutputType scriptOutputType, K[] keys, byte[][] params) {
return redisAsyncCommands.eval(script, scriptOutputType, keys, params);
}
@Override
public RedisFuture get(K key) {
return redisAsyncCommands.get(key);
}
@Override
public RedisFuture delete(K key) {
return redisAsyncCommands.del(key);
}
};
return new LettuceBasedProxyManagerBuilder<>(redisApi);
}
/**
* @deprecated use {@link Bucket4jLettuce#casBasedBuilder(StatefulRedisConnection)}
*/
@Deprecated
public static LettuceBasedProxyManagerBuilder builderFor(StatefulRedisConnection statefulRedisConnection) {
return builderFor(statefulRedisConnection.async());
}
/**
* @deprecated use {@link Bucket4jLettuce#casBasedBuilder(RedisClient)}
*/
@Deprecated
public static LettuceBasedProxyManagerBuilder builderFor(RedisClient redisClient) {
return builderFor(redisClient.connect(ByteArrayCodec.INSTANCE));
}
/**
* @deprecated use {@link Bucket4jLettuce#casBasedBuilder(RedisClusterClient)}
*/
@Deprecated
public static LettuceBasedProxyManagerBuilder builderFor(RedisClusterClient redisClient) {
return builderFor(redisClient.connect(ByteArrayCodec.INSTANCE));
}
/**
* @deprecated use {@link Bucket4jLettuce#casBasedBuilder(StatefulRedisClusterConnection)}
*/
@Deprecated
public static LettuceBasedProxyManagerBuilder builderFor(StatefulRedisClusterConnection connection) {
return builderFor(connection.async());
}
/**
* @deprecated use {@link Bucket4jLettuce#casBasedBuilder(RedisAdvancedClusterAsyncCommands)}
*/
@Deprecated
public static LettuceBasedProxyManagerBuilder builderFor(RedisAdvancedClusterAsyncCommands redisAsyncCommands) {
Objects.requireNonNull(redisAsyncCommands);
RedisApi redisApi = new RedisApi<>() {
@Override
public RedisFuture eval(String script, ScriptOutputType scriptOutputType, K[] keys, byte[][] params) {
return redisAsyncCommands.eval(script, scriptOutputType, keys, params);
}
@Override
public RedisFuture get(K key) {
return redisAsyncCommands.get(key);
}
@Override
public RedisFuture delete(K key) {
return redisAsyncCommands.del(key);
}
};
return new LettuceBasedProxyManagerBuilder<>(redisApi);
}
@Override
public boolean isExpireAfterWriteSupported() {
return true;
}
public static class LettuceBasedProxyManagerBuilder extends AbstractRedisProxyManagerBuilder> {
private final RedisApi redisApi;
private LettuceBasedProxyManagerBuilder(RedisApi redisApi) {
this.redisApi = Objects.requireNonNull(redisApi);
}
public LettuceBasedProxyManager build() {
return new LettuceBasedProxyManager<>(this);
}
}
public LettuceBasedProxyManager(LettuceBasedProxyManagerBuilder builder) {
super(builder.getClientSideConfig());
this.expirationStrategy = builder.getNotNullExpirationStrategy();
this.redisApi = builder.redisApi;
}
public LettuceBasedProxyManager(Bucket4jLettuce.LettuceBasedProxyManagerBuilder builder) {
super(builder.getClientSideConfig());
this.expirationStrategy = builder.getExpirationAfterWrite().orElse(ExpirationAfterWriteStrategy.none());
this.redisApi = builder.getRedisApi();
}
@Override
protected CompareAndSwapOperation beginCompareAndSwapOperation(K key) {
@SuppressWarnings("unchecked")
K[] keys = (K[]) new Object[]{key};
return new CompareAndSwapOperation() {
@Override
public Optional getStateData(Optional timeoutNanos) {
RedisFuture stateFuture = redisApi.get(key);
return Optional.ofNullable(getFutureValue(stateFuture, timeoutNanos));
}
@Override
public boolean compareAndSwap(byte[] originalData, byte[] newData, RemoteBucketState newState, Optional timeoutNanos) {
return getFutureValue(compareAndSwapFuture(keys, originalData, newData, newState), timeoutNanos);
}
};
}
@Override
protected AsyncCompareAndSwapOperation beginAsyncCompareAndSwapOperation(K key) {
@SuppressWarnings("unchecked")
K[] keys = (K[]) new Object[]{key};
return new AsyncCompareAndSwapOperation() {
@Override
public CompletableFuture> getStateData(Optional timeoutNanos) {
RedisFuture stateFuture = redisApi.get(key);
return convertToCompletableFuture(stateFuture, timeoutNanos)
.thenApply(Optional::ofNullable);
}
@Override
public CompletableFuture compareAndSwap(byte[] originalData, byte[] newData, RemoteBucketState newState, Optional timeoutNanos) {
return convertToCompletableFuture(compareAndSwapFuture(keys, originalData, newData, newState), timeoutNanos);
}
};
}
@Override
public void removeProxy(K key) {
RedisFuture future = redisApi.delete(key);
getFutureValue(future, Optional.empty());
}
@Override
protected CompletableFuture removeAsync(K key) {
RedisFuture future = redisApi.delete(key);
return convertToCompletableFuture(future, Optional.empty()).thenApply(bytes -> null);
}
@Override
public boolean isAsyncModeSupported() {
return true;
}
private RedisFuture compareAndSwapFuture(K[] keys, byte[] originalData, byte[] newData, RemoteBucketState newState) {
long ttlMillis = expirationStrategy.calculateTimeToLiveMillis(newState, currentTimeNanos());
if (ttlMillis > 0) {
if (originalData == null) {
// nulls are prohibited as values, so "replace" must not be used in such cases
byte[][] params = {newData, encodeLong(ttlMillis)};
return redisApi.eval(LuaScripts.SCRIPT_SET_NX_PX, ScriptOutputType.BOOLEAN, keys, params);
} else {
byte[][] params = {originalData, newData, encodeLong(ttlMillis)};
return redisApi.eval(LuaScripts.SCRIPT_COMPARE_AND_SWAP_PX, ScriptOutputType.BOOLEAN, keys, params);
}
} else {
if (originalData == null) {
// nulls are prohibited as values, so "replace" must not be used in such cases
byte[][] params = {newData};
return redisApi.eval(LuaScripts.SCRIPT_SET_NX, ScriptOutputType.BOOLEAN, keys, params);
} else {
byte[][] params = {originalData, newData};
return redisApi.eval(LuaScripts.SCRIPT_COMPARE_AND_SWAP, ScriptOutputType.BOOLEAN, keys, params);
}
}
}
private CompletableFuture convertToCompletableFuture(RedisFuture redisFuture, Optional timeoutNanos) {
if (timeoutNanos.isEmpty()) {
return redisFuture.toCompletableFuture();
} else {
return redisFuture.toCompletableFuture().orTimeout(timeoutNanos.get(), TimeUnit.NANOSECONDS);
}
}
private V getFutureValue(RedisFuture redisFuture, Optional timeoutNanos) {
try {
if (timeoutNanos.isEmpty()) {
return redisFuture.get();
} else {
return redisFuture.get(timeoutNanos.get(), TimeUnit.NANOSECONDS);
}
} catch (InterruptedException e) {
redisFuture.cancel(true);
Thread.currentThread().interrupt();
throw new RedisException(e);
} catch (java.util.concurrent.TimeoutException e) {
String message = "Violated timeout while waiting for redis future for " + timeoutNanos.get() + "ns";
throw new io.github.bucket4j.TimeoutException(message, timeoutNanos.get(), timeoutNanos.get());
} catch (ExecutionException e) {
if (e.getCause() instanceof RedisException re) {
throw re;
}
throw new RedisException("Unexpected exception while processing command", e.getCause());
}
}
private byte[] encodeLong(Long value) {
return ("" + value).getBytes(StandardCharsets.UTF_8);
}
}