Adding benchmarks for the async rate limiters.

Signed-off-by: Hiram Chirino <[email protected]>

limitador/Cargo.toml

@@ -53,7 +53,7 @@ base64 = { version = "0.22", optional = true }
 
 [dev-dependencies]
 serial_test = "3.0"
-criterion = { version = "0.5.1", features = ["html_reports"] }
+criterion = { version = "0.5.1", features = ["html_reports", "async_tokio"] }
 redis-test = { version = "0.4.0", features = ["aio"] }
 paste = "1"
 rand = "0.8"

limitador/benches/bench.rs

@@ -5,11 +5,12 @@ use limitador::limit::Limit;
 #[cfg(feature = "disk_storage")]
 use limitador::storage::disk::{DiskStorage, OptimizeFor};
 use limitador::storage::in_memory::InMemoryStorage;
-use limitador::storage::CounterStorage;
-use limitador::RateLimiter;
+use limitador::storage::{AsyncCounterStorage, CounterStorage};
+use limitador::{AsyncRateLimiter, RateLimiter};
 use rand::SeedableRng;
 use std::collections::HashMap;
 use std::fmt::{Display, Formatter};
+use limitador::storage::redis::{CachedRedisStorageBuilder};
 
 const SEED: u64 = 42;
 
@@ -138,6 +139,38 @@ fn bench_disk(c: &mut Criterion) {
     group.finish();
 }
 
+
+async fn x() {
+    let storage_builder = CachedRedisStorageBuilder::new("redis://127.0.0.1:6379");
+    let storage = storage_builder.build().await.expect("We need a Redis running locally");
+    _ = storage.clear().await;
+
+    // TODO: write rest of benchmark..???
+}
+
+#[cfg(feature = "redis_storage")]
+fn bench_cached_redis(c: &mut Criterion) {
+    let runtime = tokio::runtime::Builder::new_current_thread()
+        .build()
+        .unwrap();
+
+    let mut group = c.benchmark_group("Cached Redis");
+    for scenario in TEST_SCENARIOS {
+        group.bench_with_input(
+            BenchmarkId::new("is_rate_limited", scenario),
+            scenario,
+            |b: &mut Bencher, test_scenario: &&TestScenario| {
+                let storage = runtime.block_on(async {
+                    let storage_builder = CachedRedisStorageBuilder::new("redis://127.0.0.1:6379");
+                    return storage_builder.build().await.expect("We need a Redis running locally");
+                });
+                async_bench_is_rate_limited(&runtime, b, test_scenario, Box::new(storage));
+            },
+        );
+    }
+    group.finish();
+}
+
 #[cfg(feature = "redis_storage")]
 fn bench_redis(c: &mut Criterion) {
     let mut group = c.benchmark_group("Redis");
@@ -195,6 +228,35 @@ fn bench_is_rate_limited(
     })
 }
 
+
+fn async_bench_is_rate_limited(
+    runtime: &tokio::runtime::Runtime,
+    b: &mut Bencher,
+    test_scenario: &TestScenario,
+    storage: Box<dyn AsyncCounterStorage>,
+) {
+    runtime.block_on(async {
+        storage.clear().await.unwrap();
+    });
+
+    b.to_async(runtime).iter_with_setup(
+        move || {
+            let (rate_limiter, call_params) = generate_async_test_data(test_scenario, storage);
+            let rng = Box::new(rand::rngs::StdRng::seed_from_u64(SEED));
+            return (rate_limiter, call_params, rng)
+        },
+        |(rate_limiter, call_params, mut rng)| async move {
+            let params = call_params.choose(rng.as_mut()).unwrap();
+            rate_limiter
+                .is_rate_limited(
+                    &params.namespace.to_owned().into(),
+                    &params.values,
+                    params.delta,
+                ).await
+                .unwrap()
+        });
+}
+
 fn bench_update_counters(
     b: &mut Bencher,
     test_scenario: &TestScenario,
@@ -255,6 +317,40 @@ fn generate_test_data(
     scenario: &TestScenario,
     storage: Box<dyn CounterStorage>,
 ) -> (RateLimiter, Vec<TestCallParams>) {
+    let rate_limiter = RateLimiter::new_with_storage(storage);
+
+    let (test_limits, call_params) = generate_test_limits(scenario);
+    for limit in test_limits {
+        rate_limiter.add_limit(limit);
+    }
+
+    (rate_limiter, call_params)
+}
+
+
+// Notice that this function creates all the limits with the same conditions and
+// variables. Also, all the conditions have the same format: "cond_x == 1".
+// That's to simplify things, those are not the aspects that should have the
+// greatest impact on performance.
+// The limits generated are big enough to avoid being rate-limited during the
+// benchmark.
+// Note that with this test data each request only increases one counter, we can
+// that as another variable in the future.
+fn generate_async_test_data(
+    scenario: &TestScenario,
+    storage: Box<dyn AsyncCounterStorage>,
+) -> (AsyncRateLimiter, Vec<TestCallParams>) {
+    let rate_limiter = AsyncRateLimiter::new_with_storage(storage);
+
+    let (test_limits, call_params) = generate_test_limits(scenario);
+    for limit in test_limits {
+        rate_limiter.add_limit(limit);
+    }
+
+    (rate_limiter, call_params)
+}
+
+fn generate_test_limits(scenario: &TestScenario) -> (Vec<Limit>, Vec<TestCallParams>) {
     let mut test_values: HashMap<String, String> = HashMap::new();
 
     let mut conditions = vec![];
@@ -293,12 +389,5 @@ fn generate_test_data(
             delta: 1,
         });
     }
-
-    let rate_limiter = RateLimiter::new_with_storage(storage);
-
-    for limit in test_limits {
-        rate_limiter.add_limit(limit);
-    }
-
-    (rate_limiter, call_params)
+    (test_limits, call_params)
 }