mononoke/time_window_counter: use trait objects to represent time_window_counter

Summary: Making a trait out of TimeWindowCounter will help with providing different implementations of load limiting for OSS and FB.

Reviewed By: krallin

Differential Revision: D21329265

fbshipit-source-id: 7f317f8e9118493f3dcbadb0519eaff565cbd882

eden/mononoke/Cargo.toml

@@ -107,6 +107,7 @@ members = [
     "sshrelay",
     "tests/fixtures",
     "tests/utils",
+    "time_window_counter",
     "tunables",
     "tunables/tunables-derive",
 ]

eden/mononoke/repo_client/unbundle/src/rate_limits.rs

@@ -6,19 +6,20 @@
  */
 
 use crate::{BundleResolverError, PostResolveAction, PostResolvePush, PostResolvePushRebase};
-use anyhow::format_err;
+use anyhow::{format_err, Result};
 use cloned::cloned;
 use context::CoreContext;
+use futures::future::{FutureExt as NewFutureExt, TryFutureExt};
 use futures_ext::{BoxFuture, FutureExt};
 use futures_old::{future::join_all, Future, IntoFuture};
 use limits::types::{RateLimit, RateLimitStatus};
 use mononoke_types::BonsaiChangeset;
-use ratelim::time_window_counter::TimeWindowCounter;
 use scuba_ext::ScubaSampleBuilderExt;
 use sha2::{Digest, Sha256};
 use slog::debug;
 use std::collections::HashMap;
 use std::time::Duration;
+use time_window_counter::{BoxGlobalTimeWindowCounter, GlobalTimeWindowCounterBuilder};
 use tokio::util::FutureExt as TokioFutureExt;
 
 const TIME_WINDOW_MIN: u32 = 10;
@@ -108,7 +109,7 @@ fn build_counters(
     category: &str,
     limit: &RateLimit,
     groups: HashMap<&str, u64>,
-) -> Vec<(TimeWindowCounter, String, u64)> {
+) -> Vec<(BoxGlobalTimeWindowCounter, String, u64)> {
     groups
         .into_iter()
         .map(|(author, count)| {
@@ -118,8 +119,13 @@ fn build_counters(
                 "Associating key {:?} with author {:?}", key, author
             );
 
-            let counter =
-                TimeWindowCounter::new(ctx.fb, category, key, TIME_WINDOW_MIN, TIME_WINDOW_MAX);
+            let counter = GlobalTimeWindowCounterBuilder::build(
+                ctx.fb,
+                category,
+                key,
+                TIME_WINDOW_MIN,
+                TIME_WINDOW_MAX,
+            );
             (counter, author.to_owned(), count)
         })
         .collect()
@@ -128,7 +134,7 @@ fn build_counters(
 fn dispatch_counter_checks_and_bumps(
     ctx: CoreContext,
     limit: &RateLimit,
-    counters: Vec<(TimeWindowCounter, String, u64)>,
+    counters: Vec<(BoxGlobalTimeWindowCounter, String, u64)>,
     enforced: bool,
 ) -> Vec<BoxFuture<(), (String, f64)>> {
     let max_value = limit.max_value as f64;
@@ -138,16 +144,17 @@ fn dispatch_counter_checks_and_bumps(
         .into_iter()
         .map(move |(counter, author, bump)| {
             cloned!(ctx);
-            counter
-                .get(interval)
-                .then(move |res| {
+            async move { Ok((counter.get(interval).await?, counter)) }
+                .boxed()
+                .compat()
+                .then(move |res: Result<_>| {
                     // NOTE: We only bump after we've allowed a response. This is reasonable for
                     // this kind of limit.
                     let mut scuba = ctx.scuba().clone();
                     scuba.add("author", author.clone());
 
                     match res {
-                        Ok(count) => {
+                        Ok((count, counter)) => {
                             scuba.add("rate_limit_status", count);
 
                             if count <= max_value {

eden/mononoke/time_window_counter/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "time_window_counter"
+edition = "2018"
+version = "0.1.0"
+authors = ['Facebook']
+license = "GPLv2+"
+include = ["src/**/*.rs"]
+
+[dependencies]
+fbinit = { git = "https://github.com/facebookexperimental/rust-shed.git", branch = "master" }
+anyhow = "1.0"
+async-trait = "0.1.29"
+futures = { version = "0.3", features = ["async-await", "compat"] }

eden/mononoke/time_window_counter/src/lib.rs

@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This software may be used and distributed according to the terms of the
+ * GNU General Public License version 2.
+ */
+
+#[cfg(fbcode_build)]
+mod facebook;
+
+use anyhow::Result;
+use async_trait::async_trait;
+use std::sync::Arc;
+
+pub type ArcGlobalTimeWindowCounter = Arc<dyn GlobalTimeWindowCounter + Send + Sync + 'static>;
+pub type BoxGlobalTimeWindowCounter = Box<dyn GlobalTimeWindowCounter + Send + Sync + 'static>;
+
+#[async_trait]
+pub trait GlobalTimeWindowCounter {
+    async fn get(&self, time_window: u32) -> Result<f64>;
+
+    fn bump(&self, value: f64);
+}
+
+pub struct GlobalTimeWindowCounterBuilder {}
+
+#[cfg(not(fbcode_build))]
+mod r#impl {
+    use super::*;
+
+    use fbinit::FacebookInit;
+
+    struct AlwaysZeroCounter {}
+
+    #[async_trait]
+    impl GlobalTimeWindowCounter for AlwaysZeroCounter {
+        async fn get(&self, _time_window: u32) -> Result<f64> {
+            Ok(0.0)
+        }
+
+        fn bump(&self, _value: f64) {}
+    }
+
+    impl GlobalTimeWindowCounterBuilder {
+        pub fn build(
+            _fb: FacebookInit,
+            _category: impl AsRef<str>,
+            _key: impl AsRef<str>,
+            _min_time_window: u32,
+            _max_time_window: u32,
+        ) -> BoxGlobalTimeWindowCounter {
+            Box::new(AlwaysZeroCounter {})
+        }
+    }
+}