diff options
Diffstat (limited to 'src/main.rs')
| -rw-r--r-- | src/main.rs | 57 |
1 files changed, 48 insertions, 9 deletions
diff --git a/src/main.rs b/src/main.rs index 4a89315..7874c86 100644 --- a/src/main.rs +++ b/src/main.rs | |||
| @@ -1,4 +1,6 @@ | |||
| 1 | use std::collections::HashMap; | 1 | use std::collections::HashMap; |
| 2 | use std::sync::Arc; | ||
| 3 | use std::time::{Duration, Instant}; | ||
| 2 | 4 | ||
| 3 | use poise::serenity_prelude as serenity; | 5 | use poise::serenity_prelude as serenity; |
| 4 | use songbird::SerenityInit; | 6 | use songbird::SerenityInit; |
| @@ -17,6 +19,17 @@ mod voice; | |||
| 17 | 19 | ||
| 18 | use database::DbPool; | 20 | use database::DbPool; |
| 19 | 21 | ||
| 22 | /// How long a search cache entry stays fresh (10 seconds). | ||
| 23 | const AC_SEARCH_CACHE_TTL: Duration = Duration::from_secs(10); | ||
| 24 | |||
| 25 | /// Per-user autocomplete stash: (UserId, filename) → SearchHit. | ||
| 26 | type AcCache = tokio::sync::Mutex<HashMap<(serenity::UserId, String), copyparty::SearchHit>>; | ||
| 27 | |||
| 28 | /// Shared autocomplete search cache: query → (hits, cached_at). | ||
| 29 | type AcSearchCache = Arc< | ||
| 30 | tokio::sync::Mutex<HashMap<String, (Vec<copyparty::SearchHit>, Instant)>>, | ||
| 31 | >; | ||
| 32 | |||
| 20 | pub struct Data { | 33 | pub struct Data { |
| 21 | pub db_pool: DbPool, | 34 | pub db_pool: DbPool, |
| 22 | pub config: config::Config, | 35 | pub config: config::Config, |
| @@ -30,7 +43,12 @@ pub struct Data { | |||
| 30 | /// Mirrors Python's per-cog `_cached_urls[user_id]` map: when two users | 43 | /// Mirrors Python's per-cog `_cached_urls[user_id]` map: when two users |
| 31 | /// search the same filename simultaneously, each invocation must resolve its | 44 | /// search the same filename simultaneously, each invocation must resolve its |
| 32 | /// own stashed `SearchHit`, so the cache is keyed by `(UserId, filename)`. | 45 | /// own stashed `SearchHit`, so the cache is keyed by `(UserId, filename)`. |
| 33 | pub ac_cache: tokio::sync::Mutex<HashMap<(serenity::UserId, String), copyparty::SearchHit>>, | 46 | pub ac_cache: AcCache, |
| 47 | /// Autocomplete search results cache (query → hits). | ||
| 48 | /// | ||
| 49 | /// Populated by background tasks when the first search times out, so the | ||
| 50 | /// next keystroke's autocomplete can serve real results immediately. | ||
| 51 | pub ac_search_cache: AcSearchCache, | ||
| 34 | /// Active music track per guild (only one music source at a time). | 52 | /// Active music track per guild (only one music source at a time). |
| 35 | pub music_tracks: tokio::sync::Mutex<HashMap<u64, songbird::tracks::TrackHandle>>, | 53 | pub music_tracks: tokio::sync::Mutex<HashMap<u64, songbird::tracks::TrackHandle>>, |
| 36 | } | 54 | } |
| @@ -63,13 +81,36 @@ impl Data { | |||
| 63 | .get(&(user_id, key.to_string())) | 81 | .get(&(user_id, key.to_string())) |
| 64 | .cloned() | 82 | .cloned() |
| 65 | } | 83 | } |
| 84 | |||
| 85 | /// Look up search cache for an autocomplete query. Returns `None` if | ||
| 86 | /// uncached or stale. | ||
| 87 | pub async fn ac_search_cache_lookup( | ||
| 88 | &self, | ||
| 89 | query: &str, | ||
| 90 | ) -> Option<Vec<copyparty::SearchHit>> { | ||
| 91 | let cache = self.ac_search_cache.lock().await; | ||
| 92 | let (hits, cached_at) = cache.get(query)?; | ||
| 93 | if cached_at.elapsed() > AC_SEARCH_CACHE_TTL { | ||
| 94 | None | ||
| 95 | } else { | ||
| 96 | Some(hits.clone()) | ||
| 97 | } | ||
| 98 | } | ||
| 99 | |||
| 100 | /// Store search results into the shared cache. | ||
| 101 | pub async fn ac_search_cache_store(&self, query: String, hits: Vec<copyparty::SearchHit>) { | ||
| 102 | self.ac_search_cache | ||
| 103 | .lock() | ||
| 104 | .await | ||
| 105 | .insert(query, (hits, Instant::now())); | ||
| 106 | } | ||
| 66 | } | 107 | } |
| 67 | 108 | ||
| 68 | #[allow(dead_code)] | 109 | #[allow(dead_code)] |
| 69 | type Error = error::Error; | 110 | type Error = error::Error; |
| 70 | 111 | ||
| 71 | #[tokio::main] | 112 | #[tokio::main] |
| 72 | async fn main() { | 113 | async fn main() -> Result<(), Box<dyn std::error::Error>> { |
| 73 | let _ = rustls::crypto::ring::default_provider().install_default(); | 114 | let _ = rustls::crypto::ring::default_provider().install_default(); |
| 74 | 115 | ||
| 75 | tracing_subscriber::fmt() | 116 | tracing_subscriber::fmt() |
| @@ -84,10 +125,7 @@ async fn main() { | |||
| 84 | config::check_env_vars(); | 125 | config::check_env_vars(); |
| 85 | let config = config::Config::from_env(); | 126 | let config = config::Config::from_env(); |
| 86 | 127 | ||
| 87 | let db_pool = database::init_pool(&config.database_path) | 128 | let db_pool = database::init_pool(&config.database_path).await?; |
| 88 | .await | ||
| 89 | .expect("Failed to initialize database"); | ||
| 90 | |||
| 91 | let http_client = reqwest::Client::new(); | 129 | let http_client = reqwest::Client::new(); |
| 92 | 130 | ||
| 93 | let copyparty = copyparty::CopypartyClient::new( | 131 | let copyparty = copyparty::CopypartyClient::new( |
| @@ -132,6 +170,7 @@ async fn main() { | |||
| 132 | volume: std::sync::atomic::AtomicU64::new(1.0f64.to_bits()), | 170 | volume: std::sync::atomic::AtomicU64::new(1.0f64.to_bits()), |
| 133 | triggers: tokio::sync::Mutex::new(triggers), | 171 | triggers: tokio::sync::Mutex::new(triggers), |
| 134 | ac_cache: tokio::sync::Mutex::new(HashMap::new()), | 172 | ac_cache: tokio::sync::Mutex::new(HashMap::new()), |
| 173 | ac_search_cache: Arc::new(tokio::sync::Mutex::new(HashMap::new())), | ||
| 135 | music_tracks: tokio::sync::Mutex::new(HashMap::new()), | 174 | music_tracks: tokio::sync::Mutex::new(HashMap::new()), |
| 136 | }; | 175 | }; |
| 137 | 176 | ||
| @@ -164,10 +203,10 @@ async fn main() { | |||
| 164 | serenity::ClientBuilder::new(&config_for_client.token, serenity::GatewayIntents::all()) | 203 | serenity::ClientBuilder::new(&config_for_client.token, serenity::GatewayIntents::all()) |
| 165 | .framework(framework) | 204 | .framework(framework) |
| 166 | .register_songbird() | 205 | .register_songbird() |
| 167 | .await | 206 | .await?; |
| 168 | .expect("Error creating client"); | ||
| 169 | 207 | ||
| 170 | client.start().await.expect("Error starting client"); | 208 | client.start().await?; |
| 209 | Ok(()) | ||
| 171 | } | 210 | } |
| 172 | 211 | ||
| 173 | async fn handle_event(ctx: &serenity::Context, event: &serenity::FullEvent, user_data: &Data) { | 212 | async fn handle_event(ctx: &serenity::Context, event: &serenity::FullEvent, user_data: &Data) { |
