-3

I have below code

public class TokenBucketRateLimiter implements RateLimiter {
    // Maximum number of tokens that can be accumulated in the bucket.
    private final long threshold;
    // Time unit for the token refill (in milliseconds).
    private final long windowUnit = 1000;
    // Current number of available tokens.
    private final AtomicLong tokens = new AtomicLong(0);
    // Timestamp of the last refill operation.
    private volatile long lastRefillTimestamp = System.currentTimeMillis();

    /**
     * Constructs a TokenBucketRateLimiter with the specified threshold.
     *
     * @param threshold the maximum number of tokens that can be accumulated.
     */
    public TokenBucketRateLimiter(long threshold) {
        this.threshold = threshold;
        // Initialize the bucket to full capacity.
        this.tokens.set(threshold);
    }

    /**
     * Attempts to acquire a token from the bucket.
     *
     * @return true if a token was successfully acquired; false otherwise.
     */
    @Override
    public boolean tryAcquire() {
        long currentTime = System.currentTimeMillis();
        // Calculate the number of tokens to add based on elapsed time since the last refill.
        long tokensToAdd = ((currentTime - lastRefillTimestamp) / windowUnit) * threshold;

        // Refill the bucket with the calculated number of tokens.
        if (tokensToAdd > 0) {
            // Ensure the bucket does not exceed its capacity.
            tokens.set(Math.min(threshold, tokens.addAndGet(tokensToAdd)));
            // Update the refill timestamp.
            lastRefillTimestamp = currentTime;
        }

        // Attempt to acquire a token.
        if (tokens.get() > 0) {
            // Decrement the token count and grant access.
            tokens.decrementAndGet();
            return true;
        }

        // Token bucket is empty; deny access.
        return false;
    }
}

This is simple one node token bucket implementation. Isn't it true that in the below code when I check

        // Attempt to acquire a token.
        if (tokens.get() > 0) {
            // Decrement the token count and grant access.
            tokens.decrementAndGet();
            return true;
        }

If the number of tokens was 1, tokens.get() will return 1. Two threads can still pass this check and can decrement it one after the other making tokens negative and allowing the second call to go through though it should not be allowed?

Is using synchronized the only way here? Or say Reentrant lock api?

Improve this question
3
  • 1
    Please specify your question and provide a more focused code example. Commented Sep 14, 2024 at 18:46
  • The javadoc for AtomicLong methods shows various calls for updating the state of token without locks or synchronization. Commented Sep 14, 2024 at 19:50
  • 1
    @curiousengineer DuncG was 100% right. Yes, there's a race condition, but the Javadoc for AtomicLong shows you how to deal with it correctly (not how you're doing it now). Commented Sep 15, 2024 at 4:26

2 Answers 2

Reset to default
4

You are correct that there's danger of a race condition, and that synchronized or ReentrantLock could fix it. However, you could instead do something like

long prevTokens = tokens.getAndUpdate(t -> Math.max(t-1, 0));
if (prevTokens > 0) { return true; }
Improve this answer
Sign up to request clarification or add additional context in comments.

8 Comments

Thanks, one follow up, I am wondering if this will at all then require synchronized. And if not, can it sometimes falsely say the api needs to be throttled? Just thinking aloud. I think this might work, and I can replace -1 with the tokens I need to grab if it is more than one I think?
Also, that is t in your lambda implementation? Won't that have a race condition?
No, it doesn't require synchronized. No, it can't be wrong. No, it won't have a race condition. If you need more than one, you need to be cleverer about the lambda body, but the same general idea works.
In the code snippet I have, what is the value of t here? Can you please tell me that? Then I can try to connect the dots. thanks. Because I think t is the current token value taht is read and when other thread reads it, by then it will have changed already in some cases. Isn't it?
t is the current value, but getAndUpdate guarantees that it works atomically, that other threads cannot do anything that conflicts. Please read the documentation.
|
0

You right about this code being unsafe.

Correct version is:

    long old = tokens.get();
    while (old > 0) {
        if (tokens.compareAndSet(old, old - 1)) {
            return true;
        }
        old = tokens.get();
    }
Improve this answer

11 Comments

Thanks, there is a minor flaw. The compareAndSet will fail if the old is modified. And let us say even after being modified, it was greater than Zero, we will still throttle the call. One solution would be to introduce retry, but then that can in theory cause a problem
I edited my answer. I forgot to update old after failed attempt to swap.
Thank you and this is where the theoretical notion of infinite loop comes. I think the way to get around it is, exponential backoff and do a retry and abort after some attempts. Correct? Secondly, is there any API in Java doc that can take care of all of this as others are hammering me with that suggestion in here? I doubt
@curiousengineer when more than one thread make an attempt at the same time, at least one of them will make progress. For all practical purposes, where the threads are doing more than executing these seven lines of code, there is enough chance for each thread to eventually proceed. Unless you have so many threads that some of them never get the CPU anyway. But that’s a problem on its own, not a problem of this piece of code.
@curiousengineer all atomic operations work that way. Even those provided by the CPU itself work by trying and re-attempting when contented. It's all a matter of likelihood given by the ratio of time between attempts and the time needed for an attempt. Not even synchronized provides fairness and also ReentrantLock does not by default, which means they all bear the possibility that one thread gets stuck and the other gets the lock again and again when acquiring in short intervals. If fairness is a concern, ReentrantLock constructed as fair will do.
|

Your Answer

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

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.