[server] Fixed rare race condition for TokenBucket tryConsume()

TokenBucket relies on consumption itself to update/refill its tokens. There is a rare
race condition that could happen and cause very few consumption requests to be rejected
even when the token bucket is refilled and has sufficient tokens. The race is caused by
the old implementation of tryConsume()

return noRetryTryConsume(tokensToConsume) || (update() && noRetryTryConsume(tokensToConsume));

This is problematic because inside update() there is a small window where token is just updated
by another thread so the first nextUpdateTime check will be false and hence returning false from
update() and causing the tryConsume() to return false when in fact there are plenty of token
available and was just refilled.

An alternative fix would be to check the bucket within update() but we might as well always try
to consume again.

internal/venice-common/src/main/java/com/linkedin/venice/throttle/TokenBucket.java

@@ -71,10 +71,10 @@ public TokenBucket(long capacity, long refillAmount, long refillInterval, TimeUn
   }
 
   /**
-   *
-   * @return true if tokens may have been added, false if short circuited and no tokens were added
+   * Check and add tokens if conditions are met. Note that token may have been updated by another thread even when the
+   * function is short-circuited. Consumers of the token bucket should always retry.
    */
-  private boolean update() {
+  private void update() {
     if (clock.millis() > nextUpdateTime) {
       synchronized (this) {
         long timeNow = clock.millis();
@@ -99,9 +99,6 @@ private boolean update() {
           nextUpdateTime = timeNow + refillIntervalMs;
         }
       }
-      return true;
-    } else {
-      return false;
     }
   }
 
@@ -115,7 +112,12 @@ public long getStaleTokenCount() {
   }
 
   public boolean tryConsume(long tokensToConsume) {
-    return noRetryTryConsume(tokensToConsume) || (update() && noRetryTryConsume(tokensToConsume));
+    if (noRetryTryConsume(tokensToConsume)) {
+      return true;
+    } else {
+      update();
+      return noRetryTryConsume(tokensToConsume);
+    }
   }
 
   private boolean noRetryTryConsume(long tokensToConsume) {