From 3c78b1fdf34392ba277f34a51cc0be049824ce0c Mon Sep 17 00:00:00 2001
From: "ivan.druzhinin" <ivan.druzhinin@exactprosystems.com>
Date: Wed, 13 Jul 2022 17:05:18 +0400
Subject: [PATCH] Different limiter implementation

---
 .../com/exactpro/th2/conn/RateLimiter.kt      | 52 ++++++-------------
 1 file changed, 16 insertions(+), 36 deletions(-)

diff --git a/src/main/kotlin/com/exactpro/th2/conn/RateLimiter.kt b/src/main/kotlin/com/exactpro/th2/conn/RateLimiter.kt
index 802f954..a8ad8f6 100644
--- a/src/main/kotlin/com/exactpro/th2/conn/RateLimiter.kt
+++ b/src/main/kotlin/com/exactpro/th2/conn/RateLimiter.kt
@@ -16,49 +16,29 @@
 
 package com.exactpro.th2.conn
 
-import java.lang.Double.min
 import java.util.concurrent.TimeUnit.SECONDS
-import java.util.concurrent.locks.LockSupport
+import java.util.concurrent.locks.LockSupport.parkNanos
 import javax.annotation.concurrent.ThreadSafe
 
 @ThreadSafe
-class RateLimiter(rate: Int) {
-    init {
-        require(rate > 0) { "rate must be positive" }
-    }
-
-    private val maxPermits = rate.toDouble()
-    private val permitDuration = SECONDS.toNanos(1) / maxPermits
-    private var freePermits = 0.0
-    private var syncTime = 0L
-
-    fun acquire() = acquire(1)
+class RateLimiter(rate: Int, refillsPerSecond: Int) {
+    constructor(rate: Int) : this(rate, 10)
 
-    fun acquire(permits: Int) {
-        var currentTime = System.nanoTime()
-        val waitUntilTime = getWaitUntilTime(permits, currentTime)
+    private var tokens = 0L
+    private val maxTokens = rate / refillsPerSecond
+    private var nextRefillTime = 0L
+    private val refillInterval = SECONDS.toNanos(1) / refillsPerSecond
 
-        while (waitUntilTime > currentTime) {
-            LockSupport.parkNanos(waitUntilTime - currentTime)
-            currentTime = System.nanoTime()
-        }
+    fun acquire() {
+        while (!tryAcquire()) parkNanos(1_000)
     }
 
-    private fun getWaitUntilTime(permits: Int, currentTime: Long): Long = synchronized(this) {
-        if (currentTime > syncTime) {
-            val newPermits = (currentTime - syncTime) / permitDuration
-            freePermits = min(maxPermits, freePermits + newPermits)
-            syncTime = currentTime
-        }
-
-        val waitUntilTime = syncTime
-        val stalePermits = min(permits.toDouble(), freePermits)
-        val freshPermits = permits - stalePermits
-        val syncTimeOffset = (freshPermits * permitDuration).toLong()
-
-        syncTime += syncTimeOffset
-        freePermits -= stalePermits
-
-        return waitUntilTime
+    fun tryAcquire(): Boolean = synchronized(this) {
+        if (tokens++ < maxTokens) return true
+        val currentTime = System.nanoTime()
+        if (currentTime < nextRefillTime) return false
+        nextRefillTime = currentTime + refillInterval
+        tokens = 1
+        return true
     }
 }
\ No newline at end of file