write redis with tts use count (needs connection config)

Cargo.lock

@@ -886,6 +886,16 @@ dependencies = [
  "generic-array",
 ]
 
+[[package]]
+name = "combine"
+version = "4.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cc4369b5e4c0cddf64ad8981c0111e7df4f7078f4d6ba98fb31f2e17c4c57b7e"
+dependencies = [
+ "bytes 1.0.1",
+ "memchr",
+]
+
 [[package]]
 name = "concurrent-queue"
 version = "1.2.2"
@@ -1145,6 +1155,12 @@ version = "0.15.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "77c90badedccf4105eca100756a0b1289e191f6fcbdadd3cee1d2f614f97da8f"
 
+[[package]]
+name = "dtoa"
+version = "0.4.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "56899898ce76aaf4a0f24d914c97ea6ed976d42fec6ad33fcbb0a1103e07b2b0"
+
 [[package]]
 name = "easyenv"
 version = "0.0.5"
@@ -2394,6 +2410,27 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "r2d2"
+version = "0.8.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "545c5bc2b880973c9c10e4067418407a0ccaa3091781d1671d46eb35107cb26f"
+dependencies = [
+ "log",
+ "parking_lot",
+ "scheduled-thread-pool",
+]
+
+[[package]]
+name = "r2d2_redis"
+version = "0.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "182473b876b0b93e353682ec58e207dd1cb4a62278bbe0045fe52b86b74363bb"
+dependencies = [
+ "r2d2",
+ "redis",
+]
+
 [[package]]
 name = "radium"
 version = "0.5.3"
@@ -2518,6 +2555,21 @@ dependencies = [
  "rand_core 0.3.1",
 ]
 
+[[package]]
+name = "redis"
+version = "0.20.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0a32cb439c4e89c1e6415e5b3b23d9d8cc6dc1bf5a8cade19adbd5418de803be"
+dependencies = [
+ "async-trait",
+ "combine",
+ "dtoa",
+ "itoa",
+ "percent-encoding 2.1.0",
+ "sha1",
+ "url 2.2.2",
+]
+
 [[package]]
 name = "redox_syscall"
 version = "0.2.8"
@@ -2711,6 +2763,15 @@ dependencies = [
  "winapi 0.3.9",
 ]
 
+[[package]]
+name = "scheduled-thread-pool"
+version = "0.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc6f74fd1204073fa02d5d5d68bec8021be4c38690b61264b2fdb48083d0e7d7"
+dependencies = [
+ "parking_lot",
+]
+
 [[package]]
 name = "scopeguard"
 version = "1.1.0"
@@ -3128,6 +3189,7 @@ dependencies = [
  "magic-crypt",
  "md-5",
  "pulldown-cmark",
+ "r2d2_redis",
  "rand 0.8.3",
  "regex",
  "ring",

Cargo.toml

@@ -55,7 +55,7 @@ uuid = { version = "0.8.2", features = ["v4"] }
 #tokio = { version = "1.6.0", features = ["macros"] }
 tokio = { version = "0.2.25", features = ["macros"] }
 
-# Database
+# MySql Database
 # "offline" supports typesafe builds without a live database
 # "rustls" doesn't appear to work with IP addresses, per:
 # https://github.com/launchbadge/sqlx/issues/1095#issuecomment-798050828
@@ -64,6 +64,9 @@ tokio = { version = "0.2.25", features = ["macros"] }
 #sqlx = { version = "0.5", features = [ "mysql", "runtime-actix-native-tls", "offline", "chrono" ] }
 sqlx = { version = "0.4", features = [ "mysql", "runtime-tokio-native-tls", "offline", "chrono" ] }
 
+# Redis Database
+r2d2_redis = "0.14.0"
+
 # Files
 data-encoding = "2.3.2"
 ring = "0.17.0-alpha.10"

src/http_server/endpoints/tts/enqueue_infer_tts.rs

@@ -7,20 +7,22 @@ use actix_web::http::StatusCode;
 use actix_web::{Responder, web, HttpResponse, error, HttpRequest};
 use crate::database_helpers::enums::{DownloadUrlType, CreatorSetVisibility, W2lTemplateType};
 use crate::http_server::web_utils::ip_address::get_request_ip;
+use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
 use crate::server_state::ServerState;
 use crate::util::random_prefix_crockford_token::random_prefix_crockford_token;
+use crate::util::redis_keys::RedisKeys;
+use crate::validations::check_for_slurs::contains_slurs;
 use crate::validations::model_uploads::validate_model_title;
 use crate::validations::passwords::validate_passwords;
 use crate::validations::username::validate_username;
 use derive_more::{Display, Error};
 use log::{info, warn, log};
+use r2d2_redis::redis::Commands;
 use regex::Regex;
 use sqlx::error::DatabaseError;
 use sqlx::error::Error::Database;
 use sqlx::mysql::MySqlDatabaseError;
 use std::sync::Arc;
-use crate::validations::check_for_slurs::contains_slurs;
-use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
 
 #[derive(Deserialize)]
 pub struct InferTtsRequest {
@@ -91,6 +93,21 @@ pub async fn infer_tts_handler(
   // TODO(bt): CHECK DATABASE!
   let model_token = request.tts_model_token.to_string();
 
+  let mut redis = server_state.redis_pool
+      .get()
+      .map_err(|e| {
+        warn!("redis error: {:?}", e);
+        InferTtsError::ServerError
+      })?;
+
+  let redis_count_key = RedisKeys::tts_model_usage_count(&model_token);
+
+  redis.incr(&redis_count_key, 1)
+      .map_err(|e| {
+        warn!("redis error: {:?}", e);
+        InferTtsError::ServerError
+      })?;
+
   let ip_address = get_request_ip(&http_request);
   let creator_set_visibility = "public".to_string();
 

src/main.rs

@@ -81,6 +81,9 @@ use crate::server_state::{ServerState, EnvConfig};
 use crate::shared_constants::{DEFAULT_RUST_LOG, DEFAULT_MYSQL_PASSWORD};
 use crate::util::encrypted_sort_id::SortKeyCrypto;
 use log::{info};
+use r2d2_redis::RedisConnectionManager;
+use r2d2_redis::r2d2;
+use r2d2_redis::redis::Commands;
 use sqlx::MySqlPool;
 use sqlx::mysql::MySqlPoolOptions;
 use std::sync::Arc;
@@ -149,6 +152,11 @@ async fn main() -> AnyhowResult<()> {
     mysql_pool: pool.clone(), // NB: Pool is clone/sync/send-safe
   };
 
+  let redis_manager = RedisConnectionManager::new("redis://localhost")?;
+
+  let redis_pool = r2d2::Pool::builder()
+      .build(redis_manager)?;
+
   info!("Reading env vars and setting up utils...");
 
   let bind_address = easyenv::get_env_string_or_default("BIND_ADDRESS", DEFAULT_BIND_ADDRESS);
@@ -204,6 +212,7 @@ async fn main() -> AnyhowResult<()> {
     },
     hostname: server_hostname,
     mysql_pool: pool,
+    redis_pool,
     firehose_publisher,
     cookie_manager,
     session_checker,

src/server_state.rs

@@ -3,6 +3,7 @@ use crate::common_queries::firehose_publisher::FirehosePublisher;
 use crate::http_server::web_utils::cookie_manager::CookieManager;
 use crate::http_server::web_utils::session_checker::SessionChecker;
 use crate::util::encrypted_sort_id::SortKeyCrypto;
+use r2d2_redis::{r2d2, RedisConnectionManager};
 use sqlx::MySqlPool;
 
 /// State that is injected into every endpoint.
@@ -15,6 +16,8 @@ pub struct ServerState {
 
   pub mysql_pool: MySqlPool,
 
+  pub redis_pool: r2d2::Pool<RedisConnectionManager>,
+
   pub cookie_manager: CookieManager,
 
   pub session_checker: SessionChecker,

src/util/mod.rs

@@ -7,4 +7,5 @@ pub mod markdown_to_html;
 pub mod random_crockford_token;
 pub mod random_prefix_crockford_token;
 pub mod random_uuid;
+pub mod redis_keys;
 pub mod semi_persistent_cache_dir;

src/util/redis_keys.rs

@@ -0,0 +1,12 @@
+/// Centralize all of the Redis key naming and dereferencing.
+pub struct RedisKeys {}
+
+impl RedisKeys {
+
+  /// This is a counter incremented with TTS inference.
+  /// We can use a job to read MySql and fix counts if they get out of sync.
+  pub fn tts_model_usage_count(model_token: &str) -> String {
+    // "model use count"
+    format!("ttsUseCount:{}", model_token)
+  }
+}