Merge pull request #207 from storytold/brandon-ai-streamer-support

Rate limiter support for AI streamers

crates/service/web/storyteller_web/config/storyteller-web.development.env

@@ -5,6 +5,9 @@ COOKIE_SECURE=false
 COOKIE_DOMAIN=.jungle.horse
 WEBSITE_HOMEPAGE_REDIRECT=http://jungle.horse
 
+# Test the rate limiter bypass (obviously these are dummy values)
+AI_STREAMER_USERNAMES="PewDiePie,xQc"
+
 # Stripe configs
 # All secrets are development and okay to check in
 STRIPE_CHECKOUT_SUCCESS_URL_PATH=/checkout_success

crates/service/web/storyteller_web/src/configs/app_startup/mod.rs

@@ -0,0 +1,2 @@
+pub mod redis_rate_limiters;
+pub mod username_set;

crates/service/web/storyteller_web/src/configs/app_startup/redis_rate_limiters.rs

@@ -0,0 +1,102 @@
+use std::time::Duration;
+
+use log::info;
+
+use config::common_env::CommonEnv;
+use errors::AnyhowResult;
+use limitation::Limiter;
+use crate::configs::app_startup::username_set::UsernameSet;
+
+use crate::http_server::web_utils::redis_rate_limiter::RedisRateLimiter;
+use crate::server_state::RedisRateLimiters;
+
+
+/// Build the various rate limiters
+pub fn configure_redis_rate_limiters(common_env: &CommonEnv) -> AnyhowResult<RedisRateLimiters> {
+  // TODO(bt,2023-12-13): Kill CommonEnv.
+
+  info!("Setting up Redis rate limiters...");
+
+  let logged_out_redis_rate_limiter = {
+    let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_LOGGED_OUT_ENABLED", true);
+    let limiter_max_requests = easyenv::get_env_num("LIMITER_LOGGED_OUT_MAX_REQUESTS", 3)?;
+    let limiter_window_seconds = easyenv::get_env_num("LIMITER_LOGGED_OUT_WINDOW_SECONDS", 10)?;
+
+    let limiter = Limiter::build(&common_env.redis_0_connection_string)
+        .limit(limiter_max_requests)
+        .period(Duration::from_secs(limiter_window_seconds))
+        .finish()?;
+
+    RedisRateLimiter::new(limiter, "logged_out", limiter_enabled)
+  };
+
+  let logged_in_redis_rate_limiter = {
+    let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_LOGGED_IN_ENABLED", true);
+    let limiter_max_requests = easyenv::get_env_num("LIMITER_LOGGED_IN_MAX_REQUESTS", 3)?;
+    let limiter_window_seconds = easyenv::get_env_num("LIMITER_LOGGED_IN_WINDOW_SECONDS", 10)?;
+
+    let limiter = Limiter::build(&common_env.redis_0_connection_string)
+        .limit(limiter_max_requests)
+        .period(Duration::from_secs(limiter_window_seconds))
+        .finish()?;
+
+    RedisRateLimiter::new(limiter, "logged_in", limiter_enabled)
+  };
+
+  let api_high_priority_redis_rate_limiter = {
+    let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_API_HIGH_PRIORITY_ENABLED", true);
+    let limiter_max_requests = easyenv::get_env_num("LIMITER_API_HIGH_PRIORITY_MAX_REQUESTS", 30)?;
+    let limiter_window_seconds = easyenv::get_env_num("LIMITER_API_HIGH_PRIORITY_WINDOW_SECONDS", 30)?;
+
+    let limiter = Limiter::build(&common_env.redis_0_connection_string)
+        .limit(limiter_max_requests)
+        .period(Duration::from_secs(limiter_window_seconds))
+        .finish()?;
+
+    RedisRateLimiter::new(limiter, "api_high_priority", limiter_enabled)
+  };
+
+  // This is for the following users:
+  // https://www.notion.so/storytellerai/dd88609558a24196b4ddeeef6079da98
+  let api_ai_streamers_redis_rate_limiter = {
+    let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_API_AI_STREAMERS_ENABLED", true);
+    let limiter_max_requests = easyenv::get_env_num("LIMITER_API_AI_STREAMERS_MAX_REQUESTS", 30)?;
+    let limiter_window_seconds = easyenv::get_env_num("LIMITER_API_AI_STREAMERS_WINDOW_SECONDS", 30)?;
+
+    let limiter = Limiter::build(&common_env.redis_0_connection_string)
+        .limit(limiter_max_requests)
+        .period(Duration::from_secs(limiter_window_seconds))
+        .finish()?;
+
+    RedisRateLimiter::new(limiter, "api_ai_streamers", limiter_enabled)
+  };
+
+  let ai_streamer_usernames =
+      UsernameSet::from_comma_separated(&easyenv::get_env_string_or_default("AI_STREAMER_USERNAMES", ""));
+
+  info!("AI Streamers that can bypass the normal rate limiter ({}): {:?}",
+    ai_streamer_usernames.len(),
+    ai_streamer_usernames.list_names());
+
+  let model_upload_rate_limiter = {
+    let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_MODEL_UPLOAD_ENABLED", true);
+    let limiter_max_requests = easyenv::get_env_num("LIMITER_MODEL_UPLOAD_MAX_REQUESTS", 3)?;
+    let limiter_window_seconds = easyenv::get_env_num("LIMITER_MODEL_UPLOAD_WINDOW_SECONDS", 10)?;
+
+    let limiter = Limiter::build(&common_env.redis_0_connection_string)
+        .limit(limiter_max_requests)
+        .period(Duration::from_secs(limiter_window_seconds))
+        .finish()?;
+
+    RedisRateLimiter::new(limiter, "model_upload", limiter_enabled)
+  };
+
+  Ok(RedisRateLimiters {
+    logged_out: logged_out_redis_rate_limiter,
+    logged_in: logged_in_redis_rate_limiter,
+    api_high_priority: api_high_priority_redis_rate_limiter,
+    api_ai_streamers: api_ai_streamers_redis_rate_limiter,
+    api_ai_streamer_username_set: ai_streamer_usernames,
+    model_upload: model_upload_rate_limiter,
+  })
+}

crates/service/web/storyteller_web/src/configs/app_startup/username_set.rs

@@ -0,0 +1,102 @@
+use std::collections::{BTreeSet, HashSet};
+
+#[derive(Clone)]
+pub struct UsernameSet {
+  /// "usernames" in our service are always lowercase.
+  /// "display names" are case sensitive.
+  usernames: HashSet<String>,
+}
+
+impl UsernameSet {
+  pub fn from_comma_separated(comma_separated_list: &str) -> Self {
+    let usernames = comma_separated_list.split(",")
+        .into_iter()
+        .map(|item| item.trim())
+        .map(|item| item.to_lowercase())
+        .filter(|item| !item.is_empty())
+        .collect::<HashSet<_>>();
+    Self  {
+      usernames
+    }
+  }
+
+  pub fn len(&self) -> usize {
+    self.usernames.len()
+  }
+
+  pub fn list_names(&self) -> BTreeSet<&str> {
+    // NB: BTreeSet for sorting, which is good for logging at app startup
+    self.usernames.iter().map(|item| item.as_str()).collect::<_>()
+  }
+
+  pub fn username_is_in_set(&self, username: &str) -> bool {
+    // We lowercase the input username just in case a display name is passed
+    let username = username.to_lowercase();
+    self.usernames.contains(&username)
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use crate::configs::app_startup::username_set::UsernameSet;
+  fn build_set() -> UsernameSet {
+    UsernameSet::from_comma_separated("bob,john,ALICE, BoB,jOhn,  ,,,  , aLiCe, joHn, bob,Alice")
+  }
+  #[test]
+  fn test_deduplication() {
+    let set = build_set();
+    assert_eq!(set.len(), 3);
+  }
+
+  #[test]
+  fn test_in_set() {
+    let set = build_set();
+    assert!(set.username_is_in_set("bob"));
+    assert!(set.username_is_in_set("alice"));
+    assert!(set.username_is_in_set("john"));
+  }
+
+  #[test]
+  fn test_in_set_case_insensitive() {
+    let set = build_set();
+
+    assert!(set.username_is_in_set("Bob"));
+    assert!(set.username_is_in_set("BOB"));
+
+    assert!(set.username_is_in_set("AlIcE"));
+    assert!(set.username_is_in_set("ALICE"));
+
+    assert!(set.username_is_in_set("jOHN"));
+    assert!(set.username_is_in_set("JOHN"));
+  }
+
+  #[test]
+  fn test_not_in_set() {
+    let set = build_set();
+
+    // NB: The build_set() function might have attempted to introduce these,
+    // but the struct constructor should filter blank spaces out
+    assert!(!set.username_is_in_set(""));
+    assert!(!set.username_is_in_set("    "));
+
+    // Obviously not in set
+    assert!(!set.username_is_in_set("mary"));
+    assert!(!set.username_is_in_set("bill"));
+
+    // These usernames are in the set, but not with padding
+    assert!(!set.username_is_in_set("    alice   "));
+    assert!(!set.username_is_in_set("bob  "));
+    assert!(!set.username_is_in_set("  john"));
+  }
+
+  #[test]
+  fn test_username_return() {
+    let set = build_set();
+    let usernames = set.list_names();
+    assert_eq!(usernames.len(), 3);
+    assert!(usernames.contains(&"alice"));
+    assert!(usernames.contains(&"bob"));
+    assert!(usernames.contains(&"john"));
+    assert!(!usernames.contains(&"mary"));
+  }
+}

crates/service/web/storyteller_web/src/configs/mod.rs

@@ -1,2 +1,3 @@
+pub mod app_startup;
 pub mod plans;
 pub mod static_api_tokens;

crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs

@@ -242,7 +242,12 @@ pub async fn enqueue_infer_tts_handler(
 
   // ==================== RATE LIMIT ==================== //
 
+  // TODO(bt,2023-12-13): Unit test this!
   if !disable_rate_limiter {
+    // For AI live streamers
+    // https://www.notion.so/storytellerai/dd88609558a24196b4ddeeef6079da98
+    let mut is_ai_streamer = false;
+
     let mut rate_limiter = match maybe_user_session {
       None => &server_state.redis_rate_limiters.logged_out,
       Some(ref user) => {
@@ -250,6 +255,15 @@ pub async fn enqueue_infer_tts_handler(
           warn!("User is not authorized to use TTS model because they are banned.");
           return Err(InferTtsError::NotAuthorized);
         }
+
+        if server_state
+            .redis_rate_limiters
+            .api_ai_streamer_username_set
+            .username_is_in_set(&user.user.username)
+        {
+          is_ai_streamer = true;
+        }
+
         &server_state.redis_rate_limiters.logged_in
       },
     };
@@ -264,6 +278,11 @@ pub async fn enqueue_infer_tts_handler(
       rate_limiter = &server_state.redis_rate_limiters.api_high_priority;
     }
 
+    if is_ai_streamer {
+      info!("Using AI streamer rate limiter");
+      rate_limiter = &server_state.redis_rate_limiters.api_ai_streamers;
+    }
+
     if let Err(_err) = rate_limiter.rate_limit_request(&http_request) {
       return Err(InferTtsError::RateLimited);
     }

crates/service/web/storyteller_web/src/main.rs

@@ -75,6 +75,7 @@ use crate::billing::internal_product_to_stripe_lookup_impl::InternalProductToStr
 use crate::billing::internal_session_cache_purge_impl::InternalSessionCachePurgeImpl;
 use crate::billing::stripe_internal_subscription_product_lookup_impl::StripeInternalSubscriptionProductLookupImpl;
 use crate::billing::stripe_internal_user_lookup_impl::StripeInternalUserLookupImpl;
+use crate::configs::app_startup::redis_rate_limiters::configure_redis_rate_limiters;
 use crate::configs::static_api_tokens::{StaticApiTokenConfig, StaticApiTokens, StaticApiTokenSet};
 use crate::http_server::middleware::pushback_filter_middleware::PushbackFilter;
 use crate::http_server::web_utils::redis_rate_limiter::RedisRateLimiter;
@@ -173,65 +174,6 @@ async fn main() -> AnyhowResult<()> {
     easyenv::get_env_num("REDIS_CACHE_TTL_SECONDS", 60)?,
   );
 
-  info!("Setting up Redis rate limiters...");
-
-  // Old env vars:
-  //
-  // "LIMITER_ENABLED"
-  // "LIMITER_MAX_REQUESTS"
-  // "LIMITER_WINDOW_SECONDS"
-  //
-  let logged_out_redis_rate_limiter = {
-    let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_LOGGED_OUT_ENABLED", true);
-    let limiter_max_requests = easyenv::get_env_num("LIMITER_LOGGED_OUT_MAX_REQUESTS", 3)?;
-    let limiter_window_seconds = easyenv::get_env_num("LIMITER_LOGGED_OUT_WINDOW_SECONDS", 10)?;
-
-    let limiter = Limiter::build(&common_env.redis_0_connection_string)
-        .limit(limiter_max_requests)
-        .period(Duration::from_secs(limiter_window_seconds))
-        .finish()?;
-
-    RedisRateLimiter::new(limiter, "logged_out", limiter_enabled)
-  };
-
-  let logged_in_redis_rate_limiter = {
-    let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_LOGGED_IN_ENABLED", true);
-    let limiter_max_requests = easyenv::get_env_num("LIMITER_LOGGED_IN_MAX_REQUESTS", 3)?;
-    let limiter_window_seconds = easyenv::get_env_num("LIMITER_LOGGED_IN_WINDOW_SECONDS", 10)?;
-
-    let limiter = Limiter::build(&common_env.redis_0_connection_string)
-        .limit(limiter_max_requests)
-        .period(Duration::from_secs(limiter_window_seconds))
-        .finish()?;
-
-    RedisRateLimiter::new(limiter, "logged_in", limiter_enabled)
-  };
-
-  let api_high_priority_redis_rate_limiter = {
-    let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_API_HIGH_PRIORITY_ENABLED", true);
-    let limiter_max_requests = easyenv::get_env_num("LIMITER_API_HIGH_PRIORITY_MAX_REQUESTS", 30)?;
-    let limiter_window_seconds = easyenv::get_env_num("LIMITER_API_HIGH_PRIORITY_WINDOW_SECONDS", 30)?;
-
-    let limiter = Limiter::build(&common_env.redis_0_connection_string)
-        .limit(limiter_max_requests)
-        .period(Duration::from_secs(limiter_window_seconds))
-        .finish()?;
-
-    RedisRateLimiter::new(limiter, "api_high_priority", limiter_enabled)
-  };
-
-  let model_upload_rate_limiter = {
-    let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_MODEL_UPLOAD_ENABLED", true);
-    let limiter_max_requests = easyenv::get_env_num("LIMITER_MODEL_UPLOAD_MAX_REQUESTS", 3)?;
-    let limiter_window_seconds = easyenv::get_env_num("LIMITER_MODEL_UPLOAD_WINDOW_SECONDS", 10)?;
-
-    let limiter = Limiter::build(&common_env.redis_0_connection_string)
-        .limit(limiter_max_requests)
-        .period(Duration::from_secs(limiter_window_seconds))
-        .finish()?;
-
-    RedisRateLimiter::new(limiter, "model_upload", limiter_enabled)
-  };
 
   info!("Connecting to elasticsearch...");
 
@@ -466,12 +408,7 @@ async fn main() -> AnyhowResult<()> {
     elasticsearch,
     redis_pool,
     redis_ttl_cache,
-    redis_rate_limiters: RedisRateLimiters {
-      logged_out: logged_out_redis_rate_limiter,
-      logged_in: logged_in_redis_rate_limiter,
-      api_high_priority: api_high_priority_redis_rate_limiter,
-      model_upload: model_upload_rate_limiter,
-    },
+    redis_rate_limiters: configure_redis_rate_limiters(&common_env)?,
     firehose_publisher,
     badge_granter,
     avt_cookie_manager,

crates/service/web/storyteller_web/src/server_state.rs

@@ -20,6 +20,7 @@ use url_config::third_party_url_redirector::ThirdPartyUrlRedirector;
 use users_component::cookies::anonymous_visitor_tracking::avt_cookie_manager::AvtCookieManager;
 use users_component::cookies::session::session_cookie_manager::SessionCookieManager;
 use users_component::utils::session_checker::SessionChecker;
+use crate::configs::app_startup::username_set::UsernameSet;
 
 use crate::http_server::endpoints::categories::tts::list_fully_computed_assigned_tts_categories::list_fully_computed_assigned_tts_categories::ModelTokensByCategoryToken;
 use crate::http_server::endpoints::leaderboard::get_leaderboard::LeaderboardInfo;
@@ -140,6 +141,12 @@ pub struct RedisRateLimiters {
   /// constructions)
   pub api_high_priority: RedisRateLimiter,
 
+  /// API rate limit for AI streamers
+  pub api_ai_streamers: RedisRateLimiter,
+
+  /// Usernames of AI streamers
+  pub api_ai_streamer_username_set: UsernameSet,
+
   /// A rate limiter for TTS and W2L uploads
   pub model_upload: RedisRateLimiter,
 }