Only reset token bucket if needed (#821)

Fixes #801

Author: hadley

NEWS.md

@@ -1,5 +1,6 @@
 # httr2 (development version)
 
+* Multiple calls to `req_throttle()` no longer reset the token bucket (#801).
 * New `resps_ok()` returns a logical vector indicating which requests were successful (#807).
 * httr2 will now emit OpenTelemetry traces for all requests when tracing is  enabled. Requires the `otelsdk` package (@atheriel, #729).
 * `req_perform_connection()` no longer errors with `no applicable method for 'close' applied to an object of class "c('httr2_failure', 'httr2_error', 'rlang_error', 'error', 'condition')` (#817).

R/req-throttle.R

@@ -48,13 +48,15 @@ req_throttle <- function(req, rate, capacity, fill_time_s = 60, realm = NULL) {
     capacity <- rate * fill_time_s
   } else {
     check_number_whole(capacity, min = 0)
+    rate <- capacity / fill_time_s
   }
   check_number_decimal(fill_time_s, min = 0)
   check_string(realm, allow_null = TRUE)
 
   realm <- realm %||% url_parse(req$url)$hostname
-  the$throttle[[realm]] <- TokenBucket$new(capacity, fill_time_s)
-
+  if (!throttle_exists(realm, capacity, rate)) {
+    the$throttle[[realm]] <- TokenBucket$new(capacity, rate)
+  }
   req_policies(req, throttle_realm = realm)
 }
 
@@ -92,6 +94,15 @@ throttle_reset <- function(realm = NULL) {
   invisible()
 }
 
+throttle_exists <- function(realm, capacity, fill_rate) {
+  if (!env_has(the$throttle, realm)) {
+    return(FALSE)
+  }
+
+  cur_throttle <- the$throttle[[realm]]
+  cur_throttle$capacity == capacity && cur_throttle$fill_rate == fill_rate
+}
+
 throttle_delay <- function(req) {
   if (!req_policy_exists(req, "throttle_realm")) {
     0
@@ -115,10 +126,11 @@ TokenBucket <- R6::R6Class(
     last_fill = NULL,
     tokens = NULL,
 
-    initialize = function(capacity, fill_time_s) {
+    initialize = function(capacity, fill_rate) {
       self$capacity <- capacity
       self$tokens <- capacity
-      self$fill_rate <- capacity / fill_time_s
+      self$fill_rate <- fill_rate
+
       self$last_fill <- unix_time()
     },
 

tests/testthat/test-req-throttle.R

@@ -27,6 +27,32 @@ test_that("first request isn't throttled", {
   expect_equal(throttle_delay(req), 0.5)
 })
 
+test_that("throttle only reset if parameters change", {
+  on.exit(throttle_reset())
+
+  mock_time <- 0
+  local_mocked_bindings(unix_time = function() mock_time)
+
+  req <- request_test() |> req_throttle(capacity = 1, fill_time_s = 1)
+  expect_equal(throttle_delay(req), 0)
+
+  mock_time <- 0.5
+  expect_equal(throttle_delay(req), 0.5)
+  mock_time <- 1
+
+  req <- request_test() |> req_throttle(capacity = 1, fill_time_s = 1)
+  expect_equal(throttle_delay(req), 1)
+
+  # reset: capacity changed
+  req <- request_test() |> req_throttle(capacity = 2, fill_time_s = 1)
+  expect_equal(throttle_delay(req), 0)
+
+  # reset: fill time changed
+  req <- request_test() |> req_throttle(capacity = 2, fill_time_s = 2)
+  expect_equal(throttle_delay(req), 0)
+})
+
+
 test_that("realm defaults to hostname but can be overridden", {
   on.exit(throttle_reset())
 
@@ -46,7 +72,7 @@ test_that("token bucket respects capacity limits", {
   mock_time <- 0
   local_mocked_bindings(unix_time = function() mock_time)
 
-  bucket <- TokenBucket$new(capacity = 2, fill_time_s = 1)
+  bucket <- TokenBucket$new(capacity = 2, fill_rate = 2)
   expect_equal(bucket$take_token(), 0)
   expect_equal(bucket$tokens, 1)
   expect_equal(bucket$take_token(), 0)
@@ -61,7 +87,7 @@ test_that("token bucket handles fractions correctly", {
   mock_time <- 0
   local_mocked_bindings(unix_time = function() mock_time)
 
-  bucket <- TokenBucket$new(capacity = 2, fill_time_s = 1)
+  bucket <- TokenBucket$new(capacity = 2, fill_rate = 2)
   bucket$tokens <- 0
   expect_equal(bucket$take_token(), 0.5)
   expect_equal(bucket$tokens, -1)