Repository: alibaba/Sentinel
Pull request #3106 (2023-04-25 11:33:02)

Add implementation of token bucket for flow-control
<!--  Thanks for submitting a pull request! Here are some tips for you:
1. Please make sure you have read and understood the contributing guidelines: https://github.com/alibaba/Sentinel/blob/master/CONTRIBUTING.md
2. Please make sure the PR has a corresponding issue.
-->

### Describe what this PR does / why we need it

This PR adds a simple implementation of the token bucket algorithm.
The biggest advantage of sentinel's sliding window data structure is that it can handle some extreme high concurrency scenarios, so we pay some performance costs, but in some scenarios, we need more statistical dimensions and do not care about extreme high concurrency scenarios. In these scenarios, the implementation of a lightweight token bucket algorithm may be a better choice.

### Does this pull request fix one issue?

<!--If that, add "Fixes #xxxx" below in the next line. For example, Fixes #15. Otherwise, add "NONE" -->
This PR is a pre-work to fix the following issues:
#263
#2231

### Describe how you did it

In the implementation, the number of tokens in the token bucket is not continuously updated, it will be calculated at the time of acquisition, thereby reducing the overall performance costs.

### Describe how to verify it

Run the test cases.

### Special notes for reviews


CLAassistant: [![CLA assistant check](https://cla-assistant.io/pull/badge/not_signed)](https://cla-assistant.io/alibaba/Sentinel?pullRequest=3106) <br/>Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our [Contributor License Agreement](https://cla-assistant.io/alibaba/Sentinel?pullRequest=3106) before we can accept your contribution.<br/><sub>You have signed the CLA already but the status is still pending? Let us [recheck](https://cla-assistant.io/check/alibaba/Sentinel?pullRequest=3106) it.</sub>

Patch:
diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/AbstractTokenBucket.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/AbstractTokenBucket.java
new file mode 100644
index 00000000..d7bd1fb4
--- /dev/null
+++ b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/AbstractTokenBucket.java
@@ -0,0 +1,117 @@
+/*
+ * Copyright 1999-2023 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow.tokenbucket;
+
+import com.alibaba.csp.sentinel.util.AssertUtil;
+import com.alibaba.csp.sentinel.util.TimeUtil;
+
+/**
+ * @author LearningGp
+ */
+public class AbstractTokenBucket implements TokenBucket{
+    protected final long MAX_UNIT_PRODUCE_NUM = Long.MAX_VALUE;
+
+    /**
+     * Number of tokens left in the bucket
+     */
+    protected volatile long currentTokenNum;
+
+    /**
+     * Time of next production token
+     */
+    protected volatile long nextProduceTime;
+
+    /**
+     * Number of tokens produced per unit of time
+     */
+    protected final long unitProduceNum;
+
+    /**
+     * Maximum number of tokens stored in the bucket
+     */
+    protected final long maxTokenNum;
+
+    protected final long intervalInMs;
+    protected final long startTime;
+
+    public AbstractTokenBucket(long unitProduceNum, long maxTokenNum, boolean fullStart, long intervalInMs) {
+        AssertUtil.isTrue(unitProduceNum > 0 && intervalInMs > 0 && unitProduceNum < MAX_UNIT_PRODUCE_NUM,
+                "Illegal unitProduceNum or intervalInSeconds");
+        AssertUtil.isTrue(maxTokenNum > 0, "Illegal maxTokenNum");
+        this.unitProduceNum = unitProduceNum;
+        this.maxTokenNum = maxTokenNum;
+        this.intervalInMs = intervalInMs;
+        this.startTime = TimeUtil.currentTimeMillis();
+        this.nextProduceTime = startTime;
+        if (fullStart) {
+            this.currentTokenNum = maxTokenNum;
+        } else {
+            //The token will be filled when the first request arrives (including the initial token)
+            this.currentTokenNum = 0;
+        }
+    }
+
+    @Override
+    public boolean tryConsume(long tokenNum) {
+        if (tokenNum <= 0) {
+            return true;
+        }
+        if (tokenNum > maxTokenNum) {
+            return false;
+        }
+        long currentTimestamp = TimeUtil.currentTimeMillis();
+        refreshCurrentTokenNum(currentTimestamp);
+        if (tokenNum <= currentTokenNum) {
+            currentTokenNum -= tokenNum;
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    @Override
+    public void refreshCurrentTokenNum(long currentTimestamp) {
+        if (nextProduceTime > currentTimestamp) {
+            return;
+        }
+        currentTokenNum = Math.min(maxTokenNum, currentTokenNum + calProducedTokenNum(currentTimestamp));
+        updateNextProduceTime(currentTimestamp);
+    }
+
+    protected long calProducedTokenNum(long currentTimestamp) {
+        if (nextProduceTime > currentTimestamp) {
+            return 0;
+        }
+        long nextRefreshUnitCount = (nextProduceTime - startTime) / intervalInMs;
+        long currentUnitCount = (currentTimestamp - startTime) / intervalInMs;
+        long unitCount = currentUnitCount - nextRefreshUnitCount + 1;
+        return unitCount * unitProduceNum;
+    }
+
+    protected void updateNextProduceTime(long currentTimestamp) {
+        nextProduceTime = intervalInMs - ((currentTimestamp - startTime) % intervalInMs) + currentTimestamp;
+    }
+
+    public long refreshTokenAndGetCurrentTokenNum() {
+        refreshCurrentTokenNum(TimeUtil.currentTimeMillis());
+        return currentTokenNum;
+    }
+
+    public long getCurrentTokenNum() {
+        return currentTokenNum;
+    }
+
+}
diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/DefaultTokenBucket.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/DefaultTokenBucket.java
new file mode 100644
index 00000000..6105d213
--- /dev/null
+++ b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/DefaultTokenBucket.java
@@ -0,0 +1,30 @@
+/*
+ * Copyright 1999-2023 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow.tokenbucket;
+
+/**
+ * @author LearningGp
+ */
+public class DefaultTokenBucket extends AbstractTokenBucket{
+
+    public DefaultTokenBucket(long unitProduceNum, long maxTokenNum, long intervalInMs){
+        super(unitProduceNum, maxTokenNum, false, intervalInMs);
+    }
+
+    public DefaultTokenBucket(long unitProduceNum, long maxTokenNum, boolean fullStart, long intervalInMs){
+        super(unitProduceNum, maxTokenNum, fullStart, intervalInMs);
+    }
+}
diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/StrictTokenBucket.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/StrictTokenBucket.java
new file mode 100644
index 00000000..787e881f
--- /dev/null
+++ b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/StrictTokenBucket.java
@@ -0,0 +1,69 @@
+/*
+ * Copyright 1999-2023 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow.tokenbucket;
+
+import com.alibaba.csp.sentinel.util.TimeUtil;
+
+/**
+ * @author LearningGp
+ */
+public class StrictTokenBucket extends AbstractTokenBucket{
+
+    final private Object refreshLock = new Object();
+    final private Object consumeLock = new Object();
+
+    public StrictTokenBucket(long unitProduceNum, long maxTokenNum, long intervalInMs) {
+        super(unitProduceNum, maxTokenNum, false, intervalInMs);
+    }
+
+    public StrictTokenBucket(long unitProduceNum, long maxTokenNum, boolean fullStart, long intervalInMs) {
+        super(unitProduceNum, maxTokenNum, fullStart, intervalInMs);
+    }
+
+    @Override
+    public boolean tryConsume(long tokenNum) {
+        if (tokenNum > maxTokenNum) {
+            return false;
+        }
+        long currentTimestamp = TimeUtil.currentTimeMillis();
+        refreshCurrentTokenNum(currentTimestamp);
+        if (tokenNum <= currentTokenNum) {
+            synchronized (consumeLock) {
+                if (tokenNum <= currentTokenNum) {
+                    currentTokenNum -= tokenNum;
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
+    @Override
+    public void refreshCurrentTokenNum(long currentTimestamp) {
+        if (nextProduceTime > currentTimestamp) {
+            return;
+        }
+        long producedTokenNum = calProducedTokenNum(currentTimestamp);
+        synchronized (refreshLock) {
+            if (nextProduceTime > currentTimestamp) {
+                return;
+            }
+            currentTokenNum = Math.min(maxTokenNum, currentTokenNum + producedTokenNum);
+            updateNextProduceTime(currentTimestamp);
+        }
+    }
+
+}
diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/TokenBucket.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/TokenBucket.java
new file mode 100644
index 00000000..3f35b051
--- /dev/null
+++ b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/TokenBucket.java
@@ -0,0 +1,27 @@
+/*
+ * Copyright 1999-2023 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow.tokenbucket;
+
+/**
+ * @author LearningGp
+ */
+public interface TokenBucket {
+
+    boolean tryConsume(long tokenNum);
+
+    void refreshCurrentTokenNum(long timestamp);
+
+}


Test patch:
diff --git a/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/TokenBucketTest.java b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/TokenBucketTest.java
new file mode 100644
index 00000000..5d07531d
--- /dev/null
+++ b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/tokenbucket/TokenBucketTest.java
@@ -0,0 +1,125 @@
+/*
+ * Copyright 1999-2023 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow.tokenbucket;
+
+import com.alibaba.csp.sentinel.concurrent.NamedThreadFactory;
+import com.alibaba.csp.sentinel.test.AbstractTimeBasedTest;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+
+import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author LearningGp
+ */
+public class TokenBucketTest extends AbstractTimeBasedTest {
+
+    private static ThreadPoolExecutor threadPoolExecutor;
+
+    @BeforeClass
+    public static void beforeClass() throws Exception {
+        threadPoolExecutor = new ThreadPoolExecutor(64, 64, 0,
+                TimeUnit.SECONDS, new LinkedBlockingQueue<Runnable>(),
+                new NamedThreadFactory("sentinel-token-bucket-test", true),
+                new ThreadPoolExecutor.AbortPolicy());
+    }
+
+    @AfterClass
+    public static void afterClass() throws Exception {
+        threadPoolExecutor.shutdownNow();
+    }
+
+    @Test
+    public void testForDefaultTokenBucket() throws InterruptedException {
+        long unitProduceNum = 1;
+        long maxTokenNum = 2;
+        long intervalInMs = 1000;
+        long testStart = System.currentTimeMillis();
+        setCurrentMillis(testStart);
+
+        DefaultTokenBucket defaultTokenBucket = new DefaultTokenBucket(unitProduceNum, maxTokenNum, intervalInMs);
+
+        assertTrue(defaultTokenBucket.tryConsume(1));
+        assertFalse(defaultTokenBucket.tryConsume(1));
+
+        DefaultTokenBucket defaultTokenBucketFullStart = new DefaultTokenBucket(unitProduceNum, maxTokenNum,
+                true, intervalInMs);
+
+        assertTrue(defaultTokenBucketFullStart.tryConsume(2));
+        assertFalse(defaultTokenBucketFullStart.tryConsume(1));
+
+        sleep(1000);
+        assertTrue(defaultTokenBucket.tryConsume(1));
+        assertFalse(defaultTokenBucket.tryConsume(1));
+
+        sleep(1000);
+        assertTrue(defaultTokenBucketFullStart.tryConsume(2));
+        assertFalse(defaultTokenBucketFullStart.tryConsume(1));
+    }
+
+    @Test
+    public void testForStrictTokenBucket() throws InterruptedException {
+        long unitProduceNum = 5;
+        long maxTokenNum = 10;
+        long intervalInMs = 1000;
+        final int n = 64;
+        long testStart = System.currentTimeMillis();
+        setCurrentMillis(testStart);
+
+        final AtomicLong passNum = new AtomicLong();
+        final AtomicLong passNumFullStart = new AtomicLong();
+        final CountDownLatch countDownLatch = new CountDownLatch(n);
+        final CountDownLatch countDownLatchFullStart = new CountDownLatch(n);
+        final StrictTokenBucket strictTokenBucket = new StrictTokenBucket(unitProduceNum, maxTokenNum, intervalInMs);
+        final StrictTokenBucket strictTokenBucketFullStart = new StrictTokenBucket(unitProduceNum, maxTokenNum, true,
+                intervalInMs);
+
+        for (int i = 0; i < n; i++) {
+            threadPoolExecutor.execute(new Runnable() {
+                @Override
+                public void run() {
+                    if (strictTokenBucket.tryConsume(1)) {
+                        passNum.incrementAndGet();
+                    }
+                    countDownLatch.countDown();
+                }
+            });
+            threadPoolExecutor.execute(new Runnable() {
+                @Override
+                public void run() {
+                    if (strictTokenBucketFullStart.tryConsume(1)) {
+                        passNumFullStart.incrementAndGet();
+                    }
+                    countDownLatchFullStart.countDown();
+                }
+            });
+        }
+
+        countDownLatch.await();
+        countDownLatchFullStart.await();
+        assertEquals(5, passNum.longValue());
+        assertEquals(10, passNumFullStart.longValue());
+    }
+
+}