diff --git a/Cargo.lock b/Cargo.lock index 9a72fa1aa2..236ed8af02 100644 --- a/Cargo.lock +++ b/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", diff --git a/Cargo.toml b/Cargo.toml index 3b45184ee3..7b91fead08 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/http_server/endpoints/tts/enqueue_infer_tts.rs b/src/http_server/endpoints/tts/enqueue_infer_tts.rs index 0ae039be73..86639b3c4a 100644 --- a/src/http_server/endpoints/tts/enqueue_infer_tts.rs +++ b/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(); diff --git a/src/main.rs b/src/main.rs index f1a7a1225a..40ad895615 100644 --- a/src/main.rs +++ b/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, diff --git a/src/server_state.rs b/src/server_state.rs index 1968307d7a..73b6658904 100644 --- a/src/server_state.rs +++ b/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, + pub cookie_manager: CookieManager, pub session_checker: SessionChecker, diff --git a/src/util/mod.rs b/src/util/mod.rs index 6ac0370ea4..46340aa50f 100644 --- a/src/util/mod.rs +++ b/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; \ No newline at end of file diff --git a/src/util/redis_keys.rs b/src/util/redis_keys.rs new file mode 100644 index 0000000000..5cc2ea353c --- /dev/null +++ b/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) + } +}