How to use Bucket4j with postgresql

1.1k views Asked by At

A use case arose in which I needed to rate-limit requests for specific endpoints per user in a spring boot application that does not have an API gateway and has no plans to add one,the use case is as follows:

(1) I have a user name obtained through a JWT token.
(2) I limit each user to 60 requests per day (value is stored in db and can be changed).
-- I know I have to utilize a HandlerInterceptor at this point. 
(3) Save the user's status to a postgresql database (Can be retrieved for additional evaluation per new requests) 
(4) Save previous day's status information for archival purposes(Create a new Status each ne wday)

so I began searching. My first guess was to use resilience4j, but I later discovered that it does not work server side, then I discovered Repose Rate limit, but it did not have the applicable stuff for my use case, and after some digging, I discovered Bucket4j.

I scoured the internet for tutorials and even read the bucket4j documentation, but I didn't find one that explained it (most tutorials, I discovered, pligarise from each other), nor did the documentation provide any assistance; it just threw some functions in my face and said, hey, you can use these, but no other explanation is provided.

Here is one of my attempts to figure stuff out:

@Service
@RequiredArgsConstructor

public class RateLimitingService {
    private final DataSource dsService;

    private final Map<UUID, Bucket> bucketCache = new ConcurrentHashMap<UUID, Bucket>();
    private final UserPlanMappingRepository userPlanMappingRepository;

    public Bucket resolveBucket(final UUID userId) {
        Bucket t = bucketCache.computeIfAbsent(userId, this::newBucket);
        return t;
    }

    public void deleteIfExists(final UUID userId) {
        bucketCache.remove(userId);
    }

    private Bucket newBucket(UUID userId) {

        final var plan = userPlanMappingRepository.findByUserIdAndIsActive(userId, true).get().getPlan();
        final Integer limitPerHour = plan.getLimitPerHour();

        Long key = 1L;
        PostgreSQLadvisoryLockBasedProxyManager proxyManager = new PostgreSQLadvisoryLockBasedProxyManager(new SQLProxyConfiguration(dsService));

        BucketConfiguration bucketConfiguration = BucketConfiguration.builder()
                .addLimit(Bandwidth.classic(limitPerHour, Refill.intervally(limitPerHour, Duration.ofHours(1))))
                .build();

        return proxyManager.builder().build(key, bucketConfiguration);

    }
}

The Bean class for the DataSource:

@Configuration
@AllArgsConstructor
public class DataSourceConfig {

    Environment env;

    @Bean(name = "dsService")
    @Primary
    public DataSource createDataSourceService() {

        DriverManagerDataSource dataSource = new DriverManagerDataSource();
        dataSource.setDriverClassName(env.getProperty("spring.jpa.database-platform"));
        dataSource.setUrl(env.getProperty("spring.datasource.url"));
        dataSource.setUsername(env.getProperty("spring.datasource.username"));
        dataSource.setPassword(env.getProperty("spring.datasource.password"));

        return dataSource;
    }
    
}

And as per the documentation, I created the Sql for the store:

CREATE TABLE IF NOT EXISTS buckets (
  id BIGINT PRIMARY KEY,
  state BYTEA
);

My main points are that

  1. In the state, what am I supposed to store, I know that the Token based Bucket Algorithm usually stores a hash that includes the "total amount of remainig tokens", "Instant of the time that last transaction happened"
  2. how to identify the user if the table only takes a Long value and a state, can I add additional columns like a user_id column, and how to make this.
  3. Am I overengineering by using Bucket4j, should I build the rate limiter myself, the 2nd option feels like I am recreating the wheel.
1

There are 1 answers

0
Vladimir Bukhtoyarov On

Answers:

  1. The whole state of bucket like available-tokens, last-refill time, configuration and statistics is stored in the "state" column.

  2. User id shold be stored inside "id" column. It is obviosly that currently it is impossible to do, because your id is UUID which has 16 bytes length, while BIGINT in postgresql is 8 bytes, but you can ask maintainers via discussions https://github.com/bucket4j/bucket4j/discussions for implementing varioative ID strategy, instead of hard-coding BIGINT.

  3. The answer is too long to be accomodated inside comment, it will take article.

The general recommendation about code style:

  1. You should store reference to proxyManager somewhere, it is not recommended to create it each time when you need to create a bucket.
  2. You should not cache buckets, it is useless and leads to risk to come in OutOfMemoryError. Just remove bucketCache and create bucket when you need to check limit.
  3. Because your limiting rules are stored inside database, you should this variant of buid https://github.com/bucket4j/bucket4j/blob/master/bucket4j-core/src/main/java/io/github/bucket4j/distributed/proxy/RemoteBucketBuilder.java#L96 to successfully follow previos recomendation.