diff --git a/system/Throttle/Throttler.php b/system/Throttle/Throttler.php index b549a70467ed..db0edf7f8fdb 100644 --- a/system/Throttle/Throttler.php +++ b/system/Throttle/Throttler.php @@ -138,12 +138,12 @@ public function check(string $key, int $capacity, int $seconds, int $cost = 1): $tokenName = $this->prefix . $key; // Check to see if the bucket has even been created yet. - if (($tokens = $this->cache->get($tokenName)) === false) + if (($tokens = $this->cache->get($tokenName)) === null) { // If it hasn't been created, then we'll set it to the maximum // capacity - 1, and save it to the cache. $this->cache->save($tokenName, $capacity - $cost, $seconds); - $this->cache->save($tokenName . 'Time', time()); + $this->cache->save($tokenName . 'Time', time(), $seconds); return true; } @@ -152,12 +152,15 @@ public function check(string $key, int $capacity, int $seconds, int $cost = 1): // based on how long it's been since the last update. $throttleTime = $this->cache->get($tokenName . 'Time'); $elapsed = $this->time() - $throttleTime; + // Number of tokens to add back per second $rate = $capacity / $seconds; - // We must have a minimum wait of 1 second for a new token + // How many seconds till a new token is available. + // We must have a minimum wait of 1 second for a new token. // Primarily stored to allow devs to report back to users. - $this->tokenTime = max(1, $rate); + $newTokenAvailable = (1 / $rate) - $elapsed; + $this->tokenTime = max(1, $newTokenAvailable); // Add tokens based up on number per second that // should be refilled, then checked against capacity @@ -165,19 +168,17 @@ public function check(string $key, int $capacity, int $seconds, int $cost = 1): $tokens += $rate * $elapsed; $tokens = $tokens > $capacity ? $capacity : $tokens; - // If $tokens > 0, then we are save to perform the action, but + // If $tokens > 0, then we are safe to perform the action, but // we need to decrement the number of available tokens. - $response = false; - if ($tokens > 0) { - $response = true; + $this->cache->save($tokenName, $tokens - $cost, $seconds); + $this->cache->save($tokenName . 'Time', time(), $seconds); - $this->cache->save($tokenName, $tokens - $cost, $elapsed); - $this->cache->save($tokenName . 'Time', time()); + return true; } - return $response; + return false; } //-------------------------------------------------------------------- diff --git a/tests/_support/Cache/Handlers/MockHandler.php b/tests/_support/Cache/Handlers/MockHandler.php index 6e6fc3615a7d..5295178faef7 100644 --- a/tests/_support/Cache/Handlers/MockHandler.php +++ b/tests/_support/Cache/Handlers/MockHandler.php @@ -43,7 +43,7 @@ public function get(string $key) return array_key_exists($key, $this->cache) ? $this->cache[$key] - : false; + : null; } //-------------------------------------------------------------------- diff --git a/tests/system/Throttle/ThrottleTest.php b/tests/system/Throttle/ThrottleTest.php index 229303a0bdbd..b4bf8ea4f4b0 100644 --- a/tests/system/Throttle/ThrottleTest.php +++ b/tests/system/Throttle/ThrottleTest.php @@ -16,20 +16,19 @@ public function testTokenTime() { $throttler = new Throttler($this->cache); - // token time should be zero to start + // tokenTime should be 0 to start $this->assertEquals(0, $throttler->getTokenTime()); - // as soon as we try a rate check, token time affected - $rate = 1; // allow 1 per minute - $cost = 1; + // set $rate + $rate = 1; // allow 1 request per minute - // after using one slot, still good - $throttler->check('127.0.0.1', $rate, MINUTE, $cost); + // first check just creates a bucket, so tokenTime should be 0 + $throttler->check('127.0.0.1', $rate, MINUTE); $this->assertEquals(0, $throttler->getTokenTime()); - // after consuming a second, we have to wait - $throttler->check('127.0.0.1', $rate, MINUTE, $cost); - $this->assertEquals(1, $throttler->getTokenTime()); + // additional check affects tokenTime, so tokenTime should be 1 or greater + $throttler->check('127.0.0.1', $rate, MINUTE); + $this->assertGreaterThanOrEqual(1, $throttler->getTokenTime()); } public function testIPSavesBucket()