Add weighting to TTS Services

Author: GnomedDev

src/main.rs

@@ -6,6 +6,7 @@ use std::{
 
 use anyhow::Ok;
 use parking_lot::Mutex;
+use small_fixed_array::FixedArray;
 
 use poise::serenity_prelude as serenity;
 use serenity::small_fixed_array::FixedString;
@@ -56,7 +57,7 @@ async fn main_(start_time: std::time::SystemTime) -> Result<()> {
     let http = Arc::new(http_builder.build());
 
     println!("Performing big startup join");
-    let tts_service = || config.main.tts_services[0].clone();
+    let tts_service = || config.tts_services[0].url.clone();
     let (
         ws_connections,
         webhooks,
@@ -73,7 +74,7 @@ async fn main_(start_time: std::time::SystemTime) -> Result<()> {
         shard_count,
         premium_user,
     ) = tokio::try_join!(
-        setup_ws_stream(&config.main),
+        setup_ws_stream(&config.tts_services),
         get_webhooks(&http, config.webhooks),
         create_db_handler!(pool.clone(), "guilds", "guild_id"),
         create_db_handler!(pool.clone(), "userinfo", "user_id"),
@@ -111,43 +112,65 @@ async fn main_(start_time: std::time::SystemTime) -> Result<()> {
     tokio::spawn(analytics.clone().start());
 
     let data = Arc::new(Data {
-        pool,
+        analytics,
+        guilds_db,
+        userinfo_db,
+        nickname_db,
+        user_voice_db,
+        guild_voice_db,
+
+        entitlement_cache: mini_moka::sync::Cache::builder()
+            .time_to_live(Duration::from_hours(1))
+            .build(),
+        startup_message,
+        premium_avatar_url: FixedString::from_string_trunc(premium_user.face()),
         system_info: Mutex::new(sysinfo::System::new()),
-        bot_list_tokens: Mutex::new(config.bot_list_tokens),
+        start_time,
+        reqwest,
+        regex_cache: RegexCache::new()?,
+        webhooks,
+        pool,
 
+        service_weight_lookups: FixedArray::try_from(
+            config
+                .tts_services
+                .iter()
+                .enumerate()
+                .flat_map(|(index, service)| {
+                    (1..=service.weight.get()).map(move |_| index.try_into().unwrap())
+                })
+                .collect::<Box<[_]>>(),
+        )
+        .unwrap(),
+        tts_services: FixedArray::try_from(
+            config
+                .tts_services
+                .into_iter()
+                .map(|service| service.url)
+                .collect::<Box<[_]>>(),
+        )
+        .unwrap(),
         ws_connections,
+        voice_connections: Mutex::default(),
+
+        config: config.main,
+        premium_config: config.premium,
         runners: OnceLock::new(), // Filled in later
+
+        website_info: Mutex::new(config.website_info),
+        bot_list_tokens: Mutex::new(config.bot_list_tokens),
         fully_started: AtomicBool::new(false),
-        voice_connections: parking_lot::Mutex::default(),
         update_startup_lock: tokio::sync::Mutex::new(()),
-        entitlement_cache: mini_moka::sync::Cache::builder()
-            .time_to_live(Duration::from_hours(1))
-            .build(),
 
-        gtts_voices,
         espeak_voices,
-        translation_languages,
-        gcloud_voices: prepare_gcloud_voices(gcloud_voices),
+        gtts_voices,
         polly_voices: polly_voices
             .into_iter()
             .map(|v| (v.id.clone(), v))
             .collect::<BTreeMap<_, _>>(),
+        gcloud_voices: prepare_gcloud_voices(gcloud_voices),
 
-        config: config.main,
-        premium_config: config.premium,
-        website_info: Mutex::new(config.website_info),
-        reqwest,
-        premium_avatar_url: FixedString::from_string_trunc(premium_user.face()),
-        analytics,
-        webhooks,
-        start_time,
-        startup_message,
-        regex_cache: RegexCache::new()?,
-        guilds_db,
-        userinfo_db,
-        nickname_db,
-        user_voice_db,
-        guild_voice_db,
+        translation_languages,
     });
 
     let framework_options = poise::FrameworkOptions {

src/startup.rs

@@ -12,7 +12,7 @@ use tokio_tungstenite::tungstenite::Message;
 use tts_core::{
     opt_ext::OptionTryUnwrap as _,
     structs::{
-        Data, GoogleGender, GoogleVoice, MainConfig, Result, TTSMode, WebhookConfig,
+        Data, GoogleGender, GoogleVoice, Result, TTSMode, TTSServiceConfig, WebhookConfig,
         WebhookConfigRaw,
     },
     voice,
@@ -121,14 +121,18 @@ async fn connect_ws_stream(mut url: reqwest::Url) -> Result<voice::RawWSStream>
     Ok(tokio_tungstenite::connect_async(url).await?.0)
 }
 
-pub async fn setup_ws_stream(config: &MainConfig) -> Result<FixedArray<voice::LockedWSStream, u8>> {
-    let tasks = config.tts_services.iter().map(async |url| {
-        let stream = connect_ws_stream(url.clone()).await?;
-        anyhow::Ok(voice::LockedWSStream::new(stream))
-    });
+pub async fn setup_ws_stream(
+    tts_services: &[TTSServiceConfig],
+) -> Result<FixedArray<voice::LockedWSStream, u8>> {
+    let tasks = tts_services
+        .iter()
+        .map(async |TTSServiceConfig { url, .. }| {
+            let stream = connect_ws_stream(url.clone()).await?;
+            anyhow::Ok(voice::LockedWSStream::new(stream))
+        });
 
     let streams = futures::future::try_join_all(tasks).await?;
-    println!("Connected to {} tts-services", config.tts_services.len());
+    println!("Connected to {} tts-services", tts_services.len());
     Ok(streams.trunc_into())
 }
 
@@ -176,13 +180,13 @@ pub fn start_ws_health_checks(data: &Arc<Data>) {
             if check_ws_healthy(&mut rng, &mut ws_tx).await {
                 tracing::debug!("Health check passed for tts-service-{index}");
             } else {
-                let url = data.config.tts_services[index].clone();
+                let url = data.tts_services[index].clone();
                 *ws_tx = reconnect_ws_stream(&url, index).await;
             }
         }
     };
 
-    for index in 0..data.ws_connections.len() {
+    for index in 0..data.tts_services.len() {
         tokio::spawn(health_check(data.clone(), index));
     }
 }

tts_commands/src/other.rs

@@ -9,7 +9,7 @@ use poise::{
 
 use aformat::ToArrayString;
 use tts_core::{
-    common::{build_invite_components, fetch_audio, prepare_url, select_tts_index},
+    common::{build_invite_components, fetch_audio, prepare_url},
     constants::OPTION_SEPERATORS,
     opt_ext::OptionTryUnwrap,
     require_guild,
@@ -123,12 +123,12 @@ async fn tts_(ctx: Context<'_>, author: &serenity::User, message: &str) -> Comma
         let speaking_rate = data.speaking_rate(author.id, mode).await?;
 
         let tts_service_index = match ctx.guild_id() {
-            Some(guild_id) => select_tts_index(guild_id, data.config.tts_services.len()),
+            Some(guild_id) => data.select_tts_index(guild_id),
             None => 0,
         };
 
         let url = prepare_url(
-            data.config.tts_services[tts_service_index].clone(),
+            data.tts_services[tts_service_index].clone(),
             message,
             &voice,
             mode,

tts_commands/src/owner.rs

@@ -156,13 +156,15 @@ pub async fn info_(ctx: Context<'_>) -> CommandResult {
         .await?;
 
     let voice_debug = voice::debug_info(&data, guild_id);
+    let tts_service_url = &data.tts_services[data.select_tts_index(guild_id)];
 
     let embed = CreateEmbed::default()
         .title("TTS Bot Debug Info")
         .description(format!(
             "
 Shard ID: `{shard_id}`
 Voice Connection: `{voice_debug:?}`
+TTS Service URL: `{tts_service_url}`
 
 Server Data: `{guild_row:?}`
 User Data: `{user_row:?}`

tts_core/src/common.rs

@@ -11,13 +11,6 @@ pub(crate) fn timestamp_in_future(ts: serenity::Timestamp) -> bool {
     *ts > chrono::Utc::now()
 }
 
-#[must_use]
-pub fn select_tts_index(guild_id: serenity::GuildId, tts_services: u8) -> u8 {
-    (guild_id.get() % u64::from(tts_services))
-        .try_into()
-        .unwrap()
-}
-
 /// Builds components for invite command and invite link in DMs
 ///
 /// This has to allocate internally due to lifetime issues with trying CPS.

tts_core/src/structs.rs

@@ -48,6 +48,8 @@ macro_rules! into_static_display {
 pub struct Config {
     #[serde(rename = "Main")]
     pub main: MainConfig,
+    #[serde(rename = "TTS-Services")]
+    pub tts_services: FixedArray<TTSServiceConfig, u8>,
     #[serde(rename = "Webhook-Info")]
     pub webhooks: WebhookConfigRaw,
     #[serde(rename = "Website-Info")]
@@ -60,7 +62,6 @@ pub struct Config {
 
 #[derive(serde::Deserialize)]
 pub struct MainConfig {
-    pub tts_services: FixedArray<reqwest::Url, u8>,
     pub website_url: Option<reqwest::Url>,
     pub announcements_channel: ChannelId,
     pub main_server_invite: FixedString,
@@ -74,6 +75,12 @@ pub struct MainConfig {
     pub gtts_disabled: AtomicBool,
 }
 
+#[derive(serde::Deserialize)]
+pub struct TTSServiceConfig {
+    pub url: reqwest::Url,
+    pub weight: NonZeroU8,
+}
+
 #[derive(serde::Deserialize)]
 pub struct PostgresConfig {
     pub host: String,
@@ -204,11 +211,14 @@ pub struct Data {
     pub webhooks: WebhookConfig,
     pub pool: sqlx::PgPool,
 
+    pub service_weight_lookups: FixedArray<u8, u8>, // Maps weighted index to non-weighted.
+    pub tts_services: FixedArray<reqwest::Url, u8>,
     pub ws_connections: FixedArray<voice::LockedWSStream, u8>,
     pub voice_connections: Mutex<HashMap<serenity::GuildId, voice::ConnectionEntry>>,
-    pub runners: OnceLock<Arc<DashMap<serenity::ShardId, serenity::ShardRunnerMetadata>>>,
+
     pub config: MainConfig,
     pub premium_config: Option<PremiumConfig>,
+    pub runners: OnceLock<Arc<DashMap<serenity::ShardId, serenity::ShardRunnerMetadata>>>,
 
     // Startup information
     pub website_info: Mutex<Option<WebsiteInfo>>,
@@ -231,6 +241,13 @@ impl std::fmt::Debug for Data {
 }
 
 impl Data {
+    #[must_use]
+    pub fn select_tts_index(&self, guild_id: serenity::GuildId) -> u8 {
+        self.service_weight_lookups[(guild_id.get() % u64::from(self.service_weight_lookups.len()))
+            .try_into()
+            .unwrap()]
+    }
+
     pub async fn speaking_rate(&self, user_id: UserId, mode: TTSMode) -> Result<f32> {
         let row = self.user_voice_db.get((user_id.into(), mode)).await?;
 

tts_core/src/voice/mod.rs

@@ -21,7 +21,6 @@ use tokio::{net::TcpStream, sync::Mutex as TMutex};
 use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, tungstenite::Message as RawWSMessage};
 
 use crate::{
-    common::select_tts_index,
     structs::Data,
     voice::models::{WSConnectionInfo, WSMessageFrame},
 };
@@ -112,7 +111,7 @@ pub async fn start_connection(data: &Data, ctx: VCContext) -> StartConnectionRes
         //
         // It is important that `rx` is dropped AFTER the leave notifier is triggered, the `rx` drop
         // will any pending leave notifiers and therefore trigger them.
-        let ws_tx = &data.ws_connections[select_tts_index(guild_id, data.ws_connections.len())];
+        let ws_tx = &data.ws_connections[data.select_tts_index(guild_id)];
         let leave_notifier = ws_task(ctx, ws_tx, &mut rx, connect_tx).await;
 
         data.voice_connections.lock().remove(&guild_id);

tts_migrations/src/lib.rs

@@ -240,7 +240,7 @@ pub async fn load_db_and_conf() -> Result<(sqlx::PgPool, Config)> {
     run(&mut config_toml, &pool).await?;
 
     let config: Config = config_toml.try_into()?;
-    if config.main.tts_services.is_empty() {
+    if config.tts_services.is_empty() {
         return Err(anyhow::anyhow!("No TTS services are configured"));
     }