1
\$\begingroup\$

It's a simple WebFilter for limiting RPS from one of my projects.

There's only one OOB BucketResolver that has one shared bucket for all requests, but the design hints at the possibility of resolving different buckets for different exchanges (for example, based on API keys).

What do you think?

package com.example.reactivecrptapi.filter;

import com.example.reactivecrptapi.service.bucketResolver.BucketResolver;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.ConsumptionProbe;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;
import org.apache.commons.lang3.time.DurationFormatUtils;

import java.nio.charset.StandardCharsets;
import java.text.MessageFormat;
import java.time.Duration;

/**
 * A {@link WebFilter} that applies a time limiting strategy to incoming requests
 */
@Component
public class RateLimitingWebFilter implements WebFilter {
    private final BucketResolver bucketResolver;

    public RateLimitingWebFilter(BucketResolver bucketResolver) {
        this.bucketResolver = bucketResolver;
    }

    /**
     * Uses the injected {@link BucketResolver} to obtain a {@link Bucket}
     * matching the exchange and then attempts to consume one token from the bucket
     * and pass the exchange down the filter chain. If no token is available, sets
     * the {@code 429 Too many requests} response status and completes the exchange
     */
    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        return bucketResolver.resolveBucket(exchange)
                .map(bucket -> bucket.tryConsumeAndReturnRemaining(1))
                .flatMap(probe -> probe.isConsumed() ?
                        chain.filter(exchange) :
                        write429Response(exchange, probe));
    }

    private Mono<Void> write429Response(ServerWebExchange exchange, ConsumptionProbe probe) {
        ServerHttpResponse response = exchange.getResponse();
        response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
        return writeErrorResponseBody(response, probe);
    }

    private Mono<Void> writeErrorResponseBody(ServerHttpResponse response, ConsumptionProbe probe) {
        Duration timeToRefill = Duration.ofNanos(probe.getNanosToWaitForRefill());
        String humanReadableTimeToRefill = makeReadable(timeToRefill);
        String message = MessageFormat.format(
                "Too many requests. Please wait for at least {0} and make another attempt",
                humanReadableTimeToRefill
        );
        DataBuffer bodyDataBuffer = DefaultDataBufferFactory.sharedInstance
                .wrap(message.getBytes(StandardCharsets.UTF_8));
        return response.writeWith(Mono.just(bodyDataBuffer));
    }

    private String makeReadable(Duration duration) {
        return DurationFormatUtils.formatDurationWords(
                duration.toMillis(), true, true);
    }
}
package com.example.reactivecrptapi.service.bucketResolver;

import io.github.bucket4j.Bucket;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

/**
 * A type that maps a {@link ServerWebExchange} to a {@link Bucket} of tokens
 */
public interface BucketResolver {
    /**
     * Maps a {@code ServerWebExchange} to a {@code Bucket} of tokens
     *
     * @param exchange an exchange that should be mapped to a
     *                 corresponding {@code Bucket}
     * @return a {@code Mono} of a {@code Bucket} matching the exchange
     */
    Mono<Bucket> resolveBucket(ServerWebExchange exchange);
}
package com.example.reactivecrptapi.service.bucketResolver;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.lang.Nullable;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.time.Duration;

/**
 * A simplistic {@link BucketResolver} implementation that provides one shared {@code Bucket}
 * for all exchanges
 */
@Component
public class SimpleBucketResolver implements BucketResolver {
    private final Bucket sharedBucket;

    public SimpleBucketResolver(@Value("${crpt.rate-limiter.capacity}") long capacity,
                                @Value("${crpt.rate-limiter.duration}") Duration duration) {
        Bandwidth globalLimit = Bandwidth.builder()
                .capacity(capacity)
                .refillIntervally(capacity, duration)
                .build();
        this.sharedBucket = Bucket.builder()
                .addLimit(globalLimit)
                .build();
    }

    /**
     * Returns a shared {@code Bucket} of tokens configured according to
     * application properties
     *
     * @param exchange a nullable exchange that will be ignored
     */
    @Override
    public Mono<Bucket> resolveBucket(@Nullable ServerWebExchange exchange) {
        return Mono.just(sharedBucket);
    }
}
# application.yml
# sample rate limit config
crpt:
  rate-limiter:
    capacity: 1
    duration: 1s
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>
    <parent>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-parent</artifactId>
        <version>3.2.3</version>
        <relativePath/> <!-- lookup parent from repository -->
    </parent>
    <groupId>com.example</groupId>
    <artifactId>reactive-crpt-api</artifactId>
    <version>0.0.1-SNAPSHOT</version>
    <name>reactive-crpt-api</name>
    <description>reactive-crpt-api</description>
    <properties>
        <java.version>17</java.version>
        <mapstruct.version>1.5.5.Final</mapstruct.version>
    </properties>
    <dependencies>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-webflux</artifactId>
        </dependency>
        <dependency>
            <groupId>com.bucket4j</groupId>
            <artifactId>bucket4j-core</artifactId>
            <version>8.9.0</version>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-configuration-processor</artifactId>
            <optional>true</optional>
        </dependency>
        <dependency>
            <groupId>org.apache.commons</groupId>
            <artifactId>commons-lang3</artifactId>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-test</artifactId>
            <scope>test</scope>
        </dependency>
    </dependencies>

    <build>
        <plugins>
        <!-- this section used to have Lombok and Mapstruct-related configuration, but I removed it since it's irrelevant here -->
            <plugin>
                <groupId>org.springframework.boot</groupId>
                <artifactId>spring-boot-maven-plugin</artifactId>
            </plugin>
            <plugin>
                <groupId>org.apache.maven.plugins</groupId>
                <artifactId>maven-compiler-plugin</artifactId>
                <version>3.8.1</version>
                <configuration>
                    <source>${java.version}</source>
                    <target>${java.version}</target>
                </configuration>
            </plugin>
        </plugins>
    </build>

</project>

Tests

package com.example.reactivecrptapi.service.bucketResolver;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.BucketState;
import io.github.bucket4j.ConsumptionProbe;
import org.junit.jupiter.api.Test;
import org.springframework.test.util.ReflectionTestUtils;
import reactor.test.StepVerifier;

import java.time.Duration;
import java.util.concurrent.atomic.AtomicReference;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assumptions.assumeThat;

class SimpleBucketResolverTest {

    @Test
    void resolveBucket_returnsExpectedBucket() {
        long capacity = 1;
        Duration duration = Duration.ofHours(1);
        SimpleBucketResolver bucketResolver = new SimpleBucketResolver(capacity, duration);
        Bandwidth expectedLimit = Bandwidth.builder()
                .capacity(capacity)
                .refillIntervally(capacity, duration)
                .build();
        StepVerifier.create(bucketResolver.resolveBucket(null))
                .assertNext(bucket -> assertBucketHasExpectedLimit(bucket, expectedLimit))
                .verifyComplete();
    }

    @SuppressWarnings({"unchecked", "DataFlowIssue"})
    void assertBucketHasExpectedLimit(Bucket bucket, Bandwidth expectedLimit) {
        ConsumptionProbe probe = bucket.tryConsumeAndReturnRemaining(1);
        assertThat(probe.isConsumed()).isTrue();
        assertThat(probe.getRemainingTokens()).isEqualTo(expectedLimit.getCapacity() - 1);
        AtomicReference<BucketState> atomicState =
                ((AtomicReference<BucketState>) ReflectionTestUtils.getField(bucket, "stateRef"));
        assumeThat(atomicState).isNotNull();
        Bandwidth[] limits = atomicState.get().getConfiguration().getBandwidths();
        assertThat(limits).hasSize(1);
        Bandwidth limit = limits[0];
        assertThat(limit.getRefillPeriodNanos()).isEqualTo(expectedLimit.getRefillPeriodNanos());
    }
}
package com.example.reactivecrptapi.filter;

import com.example.reactivecrptapi.service.bucketResolver.SimpleBucketResolver;
import org.junit.jupiter.api.Test;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.autoconfigure.web.reactive.WebFluxTest;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Import;
import org.springframework.http.HttpStatus;
import org.springframework.test.annotation.DirtiesContext;
import org.springframework.test.context.TestPropertySource;
import org.springframework.test.web.reactive.server.WebTestClient;
import org.springframework.web.reactive.function.server.RouterFunction;
import org.springframework.web.reactive.function.server.RouterFunctions;
import org.springframework.web.reactive.function.server.ServerResponse;

import static com.example.reactivecrptapi.filter.RateLimitingComponentTest.RateLimitingComponentTestConfig;
import static org.assertj.core.api.Assertions.assertThat;

@WebFluxTest
@Import(RateLimitingComponentTestConfig.class)
@TestPropertySource(properties = """
        crpt.rate-limiter.capacity=1
        crpt.rate-limiter.duration=30m
        """)
@DirtiesContext(classMode = DirtiesContext.ClassMode.AFTER_EACH_TEST_METHOD)
class RateLimitingComponentTest {
    @Autowired
    WebTestClient client;

    @Test
    void withinLimit_returns2xx() {
        client
                .get()
                .uri("/test")
                .exchange()
                .expectStatus().isOk();
    }

    @Test
    void beyondLimit_returns429() {
        client
                .get()
                .uri("/test")
                .exchange()
                .expectStatus().isOk();

        client
                .get()
                .uri("/test")
                .exchange()
                .expectStatus().isEqualTo(HttpStatus.TOO_MANY_REQUESTS)
                .expectBody(String.class).value(b -> assertThat(b).isNotBlank());
    }

    @Configuration
    @Import({
            RateLimitingWebFilter.class,
            SimpleBucketResolver.class
    })
    static class RateLimitingComponentTestConfig {
        @Bean
        RouterFunction<ServerResponse> testRouterFunction() {
            return RouterFunctions.route()
                    .GET("/test", r -> ServerResponse.ok().build())
                    .build();
        }
    }
}
\$\endgroup\$
1
  • \$\begingroup\$ Don't ask what crpt means. I'm not sure myself \$\endgroup\$
    – Sergey
    Apr 1 at 17:37

1 Answer 1

2
\$\begingroup\$
  1. Instead of using one maven-property for source and value please allow to differnciate using two properties:
    <maven.compiler.target>
    <maven.compiler.source>

This gives the benefit that you do not need the configuration-xml-node of the compiler plugin.

  1. You import the maven configuration from https://repo1.maven.org/maven2/org/springframework/boot/spring-boot-starter-parent/3.2.3/spring-boot-starter-parent-3.2.3.pom. The import transitivly import the plugin

         <plugin>
             <groupId>org.springframework.boot</groupId>
             <artifactId>spring-boot-maven-plugin</artifactId>
         </plugin>
    

You do not need to mention it again in your pom.xml.

\$\endgroup\$
4
  • \$\begingroup\$ Thank you. Yeah, I know. The <plugin> element used to have configuration for Mapstruct and Maven. I removed those since they are not used in the filter and kept the explicit plugin declaration itself \$\endgroup\$
    – Sergey
    Apr 4 at 13:30
  • \$\begingroup\$ Also, didn't you mention justOrEmpty() as my notification suggests? I don't see it in your answer and don't see any edits. Mystery! \$\endgroup\$
    – Sergey
    Apr 4 at 13:34
  • \$\begingroup\$ Yes, forget about it. You wrote Mono.just(sharedBucket); but I reat Mono.just(exchange); and exchange could be null. Would you mind removing the unneccessary parameter form Mono<Bucket> resolveBucket(ServerWebExchange exchange);? \$\endgroup\$
    – Grim
    Apr 4 at 15:31
  • 1
    \$\begingroup\$ It's part of the interface it implements, can't do! Other implementations may actually use it (retrieve API keys etc.) \$\endgroup\$
    – Sergey
    Apr 4 at 18:29

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.