spring-boot-demo/rate-limit/src/main/java/com/ankurm/ratelimit/RedisRateLimitConfig.java

package com.ankurm.ratelimit;

import java.time.Duration;

import io.github.bucket4j.TimeMeter;
import io.github.bucket4j.distributed.ExpirationAfterWriteStrategy;
import io.github.bucket4j.distributed.proxy.ProxyManager;
import io.github.bucket4j.redis.lettuce.Bucket4jLettuce;
import io.lettuce.core.RedisClient;
import io.lettuce.core.api.StatefulRedisConnection;
import io.lettuce.core.codec.ByteArrayCodec;
import io.lettuce.core.codec.RedisCodec;
import io.lettuce.core.codec.StringCodec;

import org.springframework.beans.factory.ObjectProvider;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.connection.lettuce.LettuceConnectionFactory;

/**
 * Wires Bucket4j to the Lettuce client that Spring Boot already configured, so
 * {@code spring.data.redis.host/port/password/ssl} apply to the rate limiter too.
 */
@Configuration(proxyBeanMethods = false)
class RedisRateLimitConfig {

    /**
     * A second connection on Boot's RedisClient, with String keys and byte[] values. Bucket4j
     * stores each bucket as a binary blob, so the value codec must not touch the bytes.
     */
    @Bean(destroyMethod = "close")
    StatefulRedisConnection<String, byte[]> rateLimitConnection(LettuceConnectionFactory factory) {
        RedisClient client = (RedisClient) factory.getRequiredNativeClient();
        StatefulRedisConnection<String, byte[]> connection =
                client.connect(RedisCodec.of(StringCodec.UTF8, ByteArrayCodec.INSTANCE));
        connection.setTimeout(Duration.ofMillis(500));
        return connection;
    }

    /**
     * The proxy manager is the only Redis-aware object. Everything else talks to
     * {@link ProxyManager}, so the same code runs on Hazelcast, JDBC or Redisson by swapping this
     * bean.
     *
     * <p>The clock is a bean so a test can move time without sleeping. Production gets the system
     * clock.
     */
    @Bean
    ProxyManager<String> rateLimitProxyManager(StatefulRedisConnection<String, byte[]> connection,
                                               ObjectProvider<TimeMeter> clock) {
        return Bucket4jLettuce.casBasedBuilder(connection)
                // Delete a bucket's key once it has been full for 10 seconds: idle callers cost no memory.
                .expirationAfterWrite(
                        ExpirationAfterWriteStrategy.basedOnTimeForRefillingBucketUpToMax(Duration.ofSeconds(10)))
                .clientClock(clock.getIfAvailable(() -> TimeMeter.SYSTEM_MILLISECONDS))
                .build();
    }
}