[Good first issue] Fix the flaky test AsyncNetworkBandwidthLimiterTest

[superhx]

Gradle Test Run :s3stream:test > Gradle Test Executor 7 > AsyncNetworkBandwidthLimiterTest > testThrottleConsume4() FAILED
    org.opentest4j.AssertionFailedError: expected: <95> but was: <85>
        at app//org.junit.jupiter.api.AssertionFailureBuilder.build(AssertionFailureBuilder.java:151)
        at app//org.junit.jupiter.api.AssertionFailureBuilder.buildAndThrow(AssertionFailureBuilder.java:132)
        at app//org.junit.jupiter.api.AssertEquals.failNotEqual(AssertEquals.java:197)
        at app//org.junit.jupiter.api.AssertEquals.assertEquals(AssertEquals.java:166)
        at app//org.junit.jupiter.api.AssertEquals.assertEquals(AssertEquals.java:161)
        at app//org.junit.jupiter.api.Assertions.assertEquals(Assertions.java:628)
        at app//com.automq.stream.s3.AsyncNetworkBandwidthLimiterTest.lambda$testThrottleConsume4$3(AsyncNetworkBandwidthLimiterTest.java:95)
        at java.base@17.0.14/java.util.concurrent.CompletableFuture.uniWhenComplete(CompletableFuture.java:863)
        at java.base@17.0.14/java.util.concurrent.CompletableFuture$UniWhenComplete.tryFire(CompletableFuture.java:841)
        at java.base@17.0.14/java.util.concurrent.CompletableFuture.postComplete(CompletableFuture.java:510)
        at java.base@17.0.14/java.util.concurrent.CompletableFuture.waitingGet(CompletableFuture.java:1909)
        at java.base@17.0.14/java.util.concurrent.CompletableFuture.join(CompletableFuture.java:2117)
        at app//com.automq.stream.s3.AsyncNetworkBandwidthLimiterTest.testThrottleConsume4(AsyncNetworkBandwidthLimiterTest.java:97)

[Vist233]

Hi @superhx ,

I've been looking into the flaky behavior of the testThrottleConsume4 test in AsyncNetworkBandwidthLimiterTest, which sometimes failed with org.opentest4j.AssertionFailedError: expected: <95> but was: <85>.

Analyzing the original code structure:

CompletableFuture<Void> cf = bucket.consume(ThrottleStrategy.CATCH_UP, 5);
bucket.consume(ThrottleStrategy.CATCH_UP, 10);
CompletableFuture<Void> result = cf.whenComplete((v, e) -> {
    Assertions.assertEquals(95, bucket.getAvailableTokens()); // Assertion here
});
cf.join();
result.join();

The placement of the assertion assertEquals(95, ...) inside the whenComplete callback attached only to the first consume(5) future (cf) strongly suggests that the original intent might have been to specifically test the state immediately after the first consume(5) operation completed, at which point the token count would theoretically be 95 (assuming the second consume(10) hadn't executed yet).

However, due to the asynchronous nature and potential race conditions between the whenComplete callback execution and the internal processing of the second consume(10) request, this structure proved to be unreliable, leading to the flaky failures where the assertion sometimes ran after the second consume completed (resulting in 85 tokens).

I've implemented a stable fix (using thenCompose to chain the operations and asserting the final state of 85 after both complete), which ensures the test is reliable:

CompletableFuture<Void> combinedCf = bucket.consume(ThrottleStrategy.CATCH_UP, 5)
    .thenCompose(v -> bucket.consume(ThrottleStrategy.CATCH_UP, 10));
combinedCf.join();
Assertions.assertEquals(85, bucket.getAvailableTokens()); // Assert final state

My question is: Was the original intention truly to capture and assert that specific intermediate state (95 tokens), even with the inherent risk of flakiness due to timing? Or is testing the final, stable state (85 tokens) after both operations complete an acceptable and preferred outcome for this test case?

Understanding the original goal will help ensure the fix aligns with the intended test coverage. If testing the intermediate state is crucial, we might need to explore more sophisticated (and potentially complex) ways to synchronize or observe that specific point in time, though it might be challenging to do reliably.

Thanks for your clarification!

[SCNieh]

@Vist233 Hi, the intention for this test if to check if the consumption from the token bucket of the same priority can be satisfied in the same order as they are requested, i.e. the request of trying to consume 5 tokens will be completed prior to the one try to consume 10 tokens.

The test will be stable for most of the cases because right after the consumption of 1000 tokens, there are no available tokens in the bucket thus makes the following request fall into the callback queue, and will be completed later in the refill thread, and since the refilling of the bucket is single-threaded, so when the first cf completes, the remaining available tokens will always be 95.

However in some extreme cases, when cpu time is scheduled to refilling thread instead of main thread before the consumption of the first 5 tokens (e.g. in a single-core machine), the available token in the bucket can be refilled to 100 thus makes the consumption of 5 and 10 tokens suceed immediately in the main thread, and the assertion will be failed with a 85 available tokens.

A simple approach to resolve the issue is to simply check the completion order of the two request instead of using the remainning token.

[Vist233]

Hi ,

I've fixed the flaky test issue in AsyncNetworkBandwidthLimiterTest.testThrottleConsume4() as well as several "this escape" issues in the codebase.

Flaky test fix

As you suggested in your comment, I've modified the test to check completion order rather than token counts:

Before:

@Test
public void testThrottleConsume4() {
    AsyncNetworkBandwidthLimiter bucket = new AsyncNetworkBandwidthLimiter(
            AsyncNetworkBandwidthLimiter.Type.INBOUND, 100, 1000);
    bucket.consume(ThrottleStrategy.BYPASS, 1000);
    Assertions.assertEquals(-100, bucket.getAvailableTokens());
    CompletableFuture<Void> cf = bucket.consume(ThrottleStrategy.CATCH_UP, 5);
    bucket.consume(ThrottleStrategy.CATCH_UP, 10);
    CompletableFuture<Void> result = cf.whenComplete((v, e) -> {
        Assertions.assertNull(e);
        Assertions.assertEquals(95, bucket.getAvailableTokens()); // This was unstable
    });
    cf.join();
    result.join();
}

After:

@Test
public void testThrottleConsume4() {
    AsyncNetworkBandwidthLimiter bucket = new AsyncNetworkBandwidthLimiter(
            AsyncNetworkBandwidthLimiter.Type.INBOUND, 100, 1000);
    bucket.consume(ThrottleStrategy.BYPASS, 1000);
    Assertions.assertEquals(-100, bucket.getAvailableTokens());
    CompletableFuture<Boolean> firstCompleted = new CompletableFuture<>();
    CompletableFuture<Void> cf1 = bucket.consume(ThrottleStrategy.CATCH_UP, 5);
    cf1 = cf1.thenApply(v -> {
        firstCompleted.complete(true);
        return null;
    });
    CompletableFuture<Void> cf2 = bucket.consume(ThrottleStrategy.CATCH_UP, 10);
    CompletableFuture<Void> result = cf2.thenAccept(v -> {
        Assertions.assertTrue(firstCompleted.isDone(),
                "First request should complete before second request");
    });
    result.join();
}

"this escape" fixes

I've also addressed several "this escape" issues in the codebase: Vist233@dfe5f09

Note that my initial changes to the test in that commit were incorrect. The correct test implementation is the one shown above, which properly focuses on checking completion order rather than token counts.

Test Results

The test now passes consistently in my local environment. Here's the output from my test run:

PS C:\Repo\automq> ./gradlew :s3stream:test --tests com.automq.stream.s3.AsyncNetworkBandwidthLimiterTest.testThrottleConsume4

> Configure project :
Starting build with version 3.9.0-SNAPSHOT (commit id 1d6a36a1) using Gradle 8.8, Java 22 and Scala 2.13.14
Build properties: maxParallelForks=16, maxScalacThreads=8, maxTestRetries=0

> Task :s3stream:test

Gradle Test Run :s3stream:test > Gradle Test Executor 3 > AsyncNetworkBandwidthLimiterTest > testThrottleConsume4() PASSED

BUILD SUCCESSFUL in 1m 3s
8 actionable tasks: 8 executed

Please let me know if these changes address the issue properly or if you'd like me to make any adjustments.

Thank you!

[SCNieh]

@Vist233 We don't usually do the code review in the discussion, pls submit a PR for your fix specifically to this issue. As for other fixes that're not related to this issue, pls submit another issue for further disussion

[Vist233]

pick up