RateLimiter initial implementation (#25253)

GitOrigin-RevId: 852b26f0900216c159c0f463e6ded847cb6a28cc

Author: gautamg795

Cargo.lock

@@ -3,6 +3,7 @@
 
 use std::{
     future::Future,
+    hash::Hash,
     num::TryFromIntError,
     ops::{
         Add,
@@ -27,6 +28,7 @@ pub use governor::nanos::Nanos;
 use governor::{
     middleware::NoOpMiddleware,
     state::{
+        keyed::DefaultKeyedStateStore,
         InMemoryState,
         NotKeyed,
     },
@@ -312,10 +314,24 @@ pub type RateLimiter<RT> = governor::RateLimiter<
     NoOpMiddleware<<RuntimeClock<RT> as governor::clock::Clock>::Instant>,
 >;
 
+pub type KeyedRateLimiter<K, RT> = governor::RateLimiter<
+    K,
+    DefaultKeyedStateStore<K>,
+    RuntimeClock<RT>,
+    NoOpMiddleware<<RuntimeClock<RT> as governor::clock::Clock>::Instant>,
+>;
+
 pub fn new_rate_limiter<RT: Runtime>(runtime: RT, quota: Quota) -> RateLimiter<RT> {
     RateLimiter::direct_with_clock(quota, &RuntimeClock { runtime })
 }
 
+pub fn new_keyed_rate_limiter<RT: Runtime, K: Hash + Eq + Clone>(
+    runtime: RT,
+    quota: Quota,
+) -> KeyedRateLimiter<K, RT> {
+    KeyedRateLimiter::dashmap_with_clock(quota, &RuntimeClock { runtime })
+}
+
 impl<RT: Runtime> governor::clock::Clock for RuntimeClock<RT> {
     type Instant = Nanos;
 

crates/common/src/runtime/mod.rs

@@ -43,3 +43,6 @@ pub mod storage {
 pub mod usage {
     include!(concat!(env!("OUT_DIR"), "/usage.rs"));
 }
+
+pub const FILE_DESCRIPTOR_BYTES: &[u8] =
+    include_bytes!(concat!(env!("OUT_DIR"), "/descriptors.bin"));

crates/pb/src/lib.rs

@@ -2,7 +2,10 @@ use std::{
     ffi::OsStr,
     fs,
     io::Result,
-    path::Path,
+    path::{
+        Path,
+        PathBuf,
+    },
 };
 
 cfg_if::cfg_if! {
@@ -32,6 +35,7 @@ pub fn pb_build() -> Result<()> {
     println!("cargo:rerun-if-changed=protos");
     let mut paths = vec![];
     let mut packages = vec![];
+    let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap());
     for dent in std::fs::read_dir("protos")? {
         let dent = dent?;
         let path = dent.path();
@@ -44,7 +48,9 @@ pub fn pb_build() -> Result<()> {
             ));
         }
     }
-    tonic_build::configure().compile(&paths, &["protos/"])?;
+    tonic_build::configure()
+        .file_descriptor_set_path(out_dir.join("descriptors.bin"))
+        .compile(&paths, &["protos/"])?;
 
     // We sort the package names just so we're generating the lib.rs
     // deterministically to avoid NOOP commits.
@@ -76,6 +82,11 @@ pub fn pb_build() -> Result<()> {
         ));
     }
 
+    lib_file_contents.push_str(&format!(
+        "\npub const FILE_DESCRIPTOR_BYTES: &[u8] =\n    \
+         include_bytes!(concat!(env!(\"OUT_DIR\"), \"/descriptors.bin\"));\n"
+    ));
+
     let out_file = Path::new("src/lib.rs");
     if fs::read_to_string(out_file)? != lib_file_contents {
         fs::write(out_file, lib_file_contents)?;