All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.github.bucket4j.redis.redisson.cas.RedissonBasedProxyManager Maven / Gradle / Ivy

Go to download

bucket4j - is a java implementation of token bucket algorithm for rate limiting

The newest version!
/*-
 * ========================LICENSE_START=================================
 * Bucket4j
 * %%
 * Copyright (C) 2015 - 2021 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.redisson.cas;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

import org.redisson.api.RFuture;
import org.redisson.client.RedisException;
import org.redisson.client.codec.ByteArrayCodec;
import org.redisson.client.protocol.RedisCommand;
import org.redisson.client.protocol.RedisCommands;
import org.redisson.client.protocol.convertor.BooleanNotNullReplayConvertor;
import org.redisson.command.CommandAsyncExecutor;

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.distributed.serialization.Mapper;
import io.github.bucket4j.redis.AbstractRedisProxyManagerBuilder;
import io.github.bucket4j.redis.consts.LuaScripts;
import io.github.bucket4j.redis.redisson.Bucket4jRedisson;
import io.netty.buffer.ByteBuf;

public class RedissonBasedProxyManager extends AbstractCompareAndSwapBasedProxyManager {

    public static final RedisCommand SET = new RedisCommand<>("SET", new BooleanNotNullReplayConvertor());

    private final CommandAsyncExecutor commandExecutor;
    private final ExpirationAfterWriteStrategy expirationStrategy;

    private final Mapper keyMapper;

    /**
     * @deprecated use {@link Bucket4jRedisson#casBasedBuilder(CommandAsyncExecutor)}
     */
    @Deprecated
    public static RedissonBasedProxyManagerBuilder builderFor(CommandAsyncExecutor commandExecutor) {
        return new RedissonBasedProxyManagerBuilder<>(Mapper.STRING, commandExecutor);
    }

    public static class RedissonBasedProxyManagerBuilder extends AbstractRedisProxyManagerBuilder> {

        private final CommandAsyncExecutor commandExecutor;
        private Mapper keyMapper;

        private RedissonBasedProxyManagerBuilder(Mapper keyMapper, CommandAsyncExecutor commandExecutor) {
            this.keyMapper = Objects.requireNonNull(keyMapper);
            this.commandExecutor = Objects.requireNonNull(commandExecutor);
        }

        public  RedissonBasedProxyManagerBuilder withKeyMapper(Mapper keyMapper) {
            this.keyMapper = (Mapper) Objects.requireNonNull(keyMapper);
            return (RedissonBasedProxyManagerBuilder) this;
        }

        public RedissonBasedProxyManager build() {
            return new RedissonBasedProxyManager<>(this);
        }

    }

    public RedissonBasedProxyManager(Bucket4jRedisson.RedissonBasedProxyManagerBuilder builder) {
        super(builder.getClientSideConfig());
        this.commandExecutor = builder.getCommandExecutor();
        this.expirationStrategy = builder.getExpirationAfterWrite().orElse(ExpirationAfterWriteStrategy.none());
        this.keyMapper = builder.getKeyMapper();
    }

    private RedissonBasedProxyManager(RedissonBasedProxyManagerBuilder builder) {
        super(builder.getClientSideConfig());
        this.commandExecutor = builder.commandExecutor;
        this.expirationStrategy = builder.getNotNullExpirationStrategy();
        this.keyMapper = builder.keyMapper;
    }

    @Override
    public boolean isExpireAfterWriteSupported() {
        return true;
    }

    @Override
    protected CompareAndSwapOperation beginCompareAndSwapOperation(K key) {
        String stringKey = keyMapper.toString(key);
        List keys = Collections.singletonList(stringKey);
        return new CompareAndSwapOperation() {
            @Override
            public Optional getStateData(Optional timeoutNanos) {
                RFuture persistedState = commandExecutor.readAsync(stringKey, ByteArrayCodec.INSTANCE, RedisCommands.GET, stringKey);
                return Optional.ofNullable(getWithTimeout(persistedState, timeoutNanos));
            }

            @Override
            public boolean compareAndSwap(byte[] originalData, byte[] newData, RemoteBucketState newState, Optional timeoutNanos) {
                long ttlMillis = expirationStrategy.calculateTimeToLiveMillis(newState, currentTimeNanos());
                if (ttlMillis > 0) {
                    if (originalData == null) {
                        // Redisson prohibits the usage null as values, so "replace" must not be used in such cases
                        RFuture redissonFuture = commandExecutor.writeAsync(stringKey, ByteArrayCodec.INSTANCE, SET, stringKey, encodeByteArray(newData), "PX", ttlMillis, "NX");
                        return getWithTimeout(redissonFuture, timeoutNanos);
                    } else {
                        Object[] params = new Object[] {originalData, newData, ttlMillis};
                        RFuture redissonFuture = commandExecutor.evalWriteAsync(stringKey, ByteArrayCodec.INSTANCE, RedisCommands.EVAL_BOOLEAN, LuaScripts.SCRIPT_COMPARE_AND_SWAP_PX, keys, params);
                        return getWithTimeout(redissonFuture, timeoutNanos);
                    }
                } else {
                    if (originalData == null) {
                        // Redisson prohibits the usage null as values, so "replace" must not be used in such cases
                        RFuture redissonFuture = commandExecutor.writeAsync(stringKey, ByteArrayCodec.INSTANCE, SET, stringKey, encodeByteArray(newData), "NX");
                        return getWithTimeout(redissonFuture, timeoutNanos);
                    } else {
                        Object[] params = new Object[] {originalData, newData};
                        RFuture redissonFuture = commandExecutor.evalWriteAsync(stringKey, ByteArrayCodec.INSTANCE, RedisCommands.EVAL_BOOLEAN, LuaScripts.SCRIPT_COMPARE_AND_SWAP, keys, params);
                        return getWithTimeout(redissonFuture, timeoutNanos);
                    }
                }
            }
        };
    }



    @Override
    protected AsyncCompareAndSwapOperation beginAsyncCompareAndSwapOperation(K key) {
        String stringKey = keyMapper.toString(key);
        List keys = Collections.singletonList(stringKey);
        return new AsyncCompareAndSwapOperation() {
            @Override
            public CompletableFuture> getStateData(Optional timeoutNanos) {
                RFuture redissonFuture = commandExecutor.readAsync(stringKey, ByteArrayCodec.INSTANCE, RedisCommands.GET, stringKey);
                if (timeoutNanos.isEmpty()) {
                    return convertFuture(redissonFuture, timeoutNanos)
                        .thenApply((byte[] resultBytes) -> Optional.ofNullable(resultBytes));
                } else {
                    return convertFuture(redissonFuture, timeoutNanos)
                        .thenApply((byte[] resultBytes) -> Optional.ofNullable(resultBytes));
                }
            }
            @Override
            public CompletableFuture compareAndSwap(byte[] originalData, byte[] newData, RemoteBucketState newState, Optional timeoutNanos) {
                long ttlMillis = expirationStrategy.calculateTimeToLiveMillis(newState, currentTimeNanos());
                if (ttlMillis > 0) {
                    if (originalData == null) {
                        RFuture redissonFuture = commandExecutor.writeAsync(stringKey, ByteArrayCodec.INSTANCE, SET, stringKey, encodeByteArray(newData), "PX", ttlMillis, "NX");
                        return convertFuture(redissonFuture, timeoutNanos);
                    } else {
                        Object[] params = new Object[] {encodeByteArray(originalData), encodeByteArray(newData), ttlMillis};
                        RFuture redissonFuture = commandExecutor.evalWriteAsync(stringKey, ByteArrayCodec.INSTANCE,
                                RedisCommands.EVAL_BOOLEAN, LuaScripts.SCRIPT_COMPARE_AND_SWAP_PX, keys, params);
                        return convertFuture(redissonFuture, timeoutNanos);
                    }
                } else {
                    if (originalData == null) {
                        RFuture redissonFuture = commandExecutor.writeAsync(stringKey, ByteArrayCodec.INSTANCE, SET, stringKey, encodeByteArray(newData), "NX");
                        return convertFuture(redissonFuture, timeoutNanos);
                    } else {
                        Object[] params = new Object[] {encodeByteArray(originalData), encodeByteArray(newData)};
                        RFuture redissonFuture = commandExecutor.evalWriteAsync(stringKey, ByteArrayCodec.INSTANCE,
                                RedisCommands.EVAL_BOOLEAN, LuaScripts.SCRIPT_COMPARE_AND_SWAP, keys, params);
                        return convertFuture(redissonFuture, timeoutNanos);
                    }
                }
            }
        };
    }

    @Override
    public void removeProxy(K key) {
        RFuture future = commandExecutor.writeAsync(keyMapper.toString(key), RedisCommands.DEL_VOID, key);
        commandExecutor.get(future);
    }

    @Override
    protected CompletableFuture removeAsync(K key) {
        RFuture redissonFuture = commandExecutor.writeAsync(keyMapper.toString(key), RedisCommands.DEL_VOID, key);
        return convertFuture(redissonFuture, Optional.empty()).thenApply(bytes -> null);
    }

    @Override
    public boolean isAsyncModeSupported() {
        return true;
    }

    private  CompletableFuture convertFuture(RFuture redissonFuture, Optional timeoutNanos) {
        if (timeoutNanos.isEmpty()) {
            return redissonFuture.toCompletableFuture();
        } else {
            return redissonFuture.toCompletableFuture().orTimeout(timeoutNanos.get(), TimeUnit.NANOSECONDS);
        }
    }

    private  T getWithTimeout(RFuture redissonFuture, Optional timeoutNanos) {
        if (timeoutNanos.isEmpty()) {
            return commandExecutor.get(redissonFuture);
        } else {
            try {
                return redissonFuture.get(timeoutNanos.get(), TimeUnit.NANOSECONDS);
            } catch (InterruptedException e) {
                redissonFuture.cancel(true);
                Thread.currentThread().interrupt();
                throw new RedisException(e);
            } catch (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) {
                    throw (RedisException) e.getCause();
                }
                throw new RedisException(e);
            }
        }
    }

    public ByteBuf encodeByteArray(byte[] value) {
        try {
            return ByteArrayCodec.INSTANCE.getValueEncoder().encode(value);
        } catch (IOException e) {
            throw new IllegalArgumentException(e);
        }
    }

}
    

    

    
            
    
            







    
© 2015 - 2024 Weber Informatics LLC | Privacy Policy