Allow watching contests from before it begins

Cargo.lock

@@ -2369,6 +2369,7 @@ dependencies = [
  "codeforces",
  "dashmap",
  "lazy_static",
+ "log",
  "regex",
  "reqwest",
  "serde",

youmubot-cf/Cargo.toml

@@ -16,6 +16,7 @@ regex = "1.5.6"
 lazy_static = "1.4.0"
 chrono = { version = "0.4.19", features = ["serde"] }
 dashmap = "5.3.4"
+log = "0.4"
 
 youmubot-prelude = { path = "../youmubot-prelude" }
 youmubot-db = { path = "../youmubot-db" }

youmubot-cf/src/hook.rs

@@ -7,6 +7,7 @@ use serenity::{
     builder::CreateEmbed, framework::standard::CommandError, model::channel::Message,
     utils::MessageBuilder,
 };
+use std::collections::HashMap as StdHashMap;
 use std::time::Instant;
 use youmubot_prelude::*;
 
@@ -31,7 +32,7 @@ enum ContestOrProblem {
 /// Caches the contest list.
 pub struct ContestCache {
     contests: HashMap<u64, (Contest, Option<Vec<Problem>>)>,
-    all_list: RwLock<(Vec<Contest>, Instant)>,
+    all_list: RwLock<(StdHashMap<u64, Contest>, Instant)>,
     http: Client,
 }
 
@@ -42,7 +43,7 @@ impl TypeMapKey for ContestCache {
 impl ContestCache {
     /// Creates a new, empty cache.
     pub(crate) async fn new(http: Client) -> Result<Self> {
-        let contests_list = Contest::list(&*http, true).await?;
+        let contests_list = Self::fetch_contest_list(http.clone()).await?;
         Ok(Self {
             contests: HashMap::new(),
             all_list: RwLock::new((contests_list, Instant::now())),
@@ -50,6 +51,19 @@ impl ContestCache {
         })
     }
 
+    async fn fetch_contest_list(http: Client) -> Result<StdHashMap<u64, Contest>> {
+        log::info!("Fetching contest list, this might take a few seconds to complete...");
+        let gyms = Contest::list(&*http, true).await?;
+        let contests = Contest::list(&*http, false).await?;
+        let r: StdHashMap<u64, Contest> = gyms
+            .into_iter()
+            .chain(contests.into_iter())
+            .map(|v| (v.id, v))
+            .collect();
+        log::info!("Contest list fetched, {} contests indexed", r.len());
+        Ok(r)
+    }
+
     /// Gets a contest from the cache, fetching from upstream if possible.
     pub async fn get(&self, contest_id: u64) -> Result<(Contest, Option<Vec<Problem>>)> {
         if let Some(v) = self.contests.get(&contest_id) {
@@ -81,14 +95,16 @@ impl ContestCache {
         if Instant::now() - last_updated > std::time::Duration::from_secs(60 * 60) {
             // We update at most once an hour.
             let mut v = self.all_list.write().await;
-            *v = (Contest::list(&*self.http, true).await?, Instant::now());
+            *v = (
+                Self::fetch_contest_list(self.http.clone()).await?,
+                Instant::now(),
+            );
         }
         self.all_list
             .read()
             .await
             .0
-            .iter()
+            .get(&contest_id)
-            .find(|v| v.id == contest_id)
             .cloned()
             .ok_or_else(|| Error::msg("Contest not found"))
     }

youmubot-cf/src/lib.rs

@@ -36,6 +36,7 @@ pub async fn setup(path: &std::path::Path, data: &mut TypeMap, announcers: &mut
     let client = Arc::new(codeforces::Client::new());
     data.insert::<hook::ContestCache>(hook::ContestCache::new(client.clone()).await.unwrap());
     data.insert::<CFClient>(client);
+    data.insert::<live::WatchData>(live::WatchData::new());
     announcers.add("codeforces", announcer::Announcer);
 }
 

youmubot-cf/src/live.rs

@@ -1,4 +1,5 @@
-use crate::{db::CfSavedUsers, CFClient};
+use crate::{db::CfSavedUsers, hook::ContestCache, CFClient};
+use chrono::TimeZone;
 use codeforces::{Contest, ContestPhase, Problem, ProblemResult, ProblemResultType, RanklistRow};
 use serenity::{
     model::{
@@ -7,7 +8,8 @@ use serenity::{
     },
     utils::MessageBuilder,
 };
-use std::collections::HashMap;
+use std::collections::HashSet as Set;
+use std::{collections::HashMap, sync::Arc, sync::Mutex as SyncMutex};
 use youmubot_prelude::*;
 
 struct MemberResult {
@@ -16,6 +18,43 @@ struct MemberResult {
     row: Option<RanklistRow>,
 }
 
+/// The structure storing watch-specific stored data.
+#[derive(Debug, Clone)]
+pub(crate) struct WatchData {
+    /// Lists of contests being watched.
+    watching_list: Arc<SyncMutex<Set<u64>>>,
+}
+
+impl TypeMapKey for WatchData {
+    type Value = WatchData;
+}
+
+struct CreateContest(u64, Arc<SyncMutex<Set<u64>>>);
+
+impl Drop for CreateContest {
+    fn drop(&mut self) {
+        self.1.lock().unwrap().remove(&self.0);
+    }
+}
+
+impl WatchData {
+    pub fn new() -> Self {
+        Self {
+            watching_list: Arc::new(SyncMutex::new(Set::new())),
+        }
+    }
+
+    fn watch(&self, contest_id: u64) -> Option<CreateContest> {
+        let mut s = self.watching_list.lock().unwrap();
+        if s.contains(&contest_id) {
+            None
+        } else {
+            s.insert(contest_id);
+            Some(CreateContest(contest_id, self.watching_list.clone()))
+        }
+    }
+}
+
 /// Watch and commentate a contest.
 ///
 /// Does the thing on a channel, block until the contest ends.
@@ -28,9 +67,67 @@ pub async fn watch_contest(
     let data = ctx.data.read().await;
     let db = CfSavedUsers::open(&*data).borrow()?.clone();
     let member_cache = data.get::<member_cache::MemberCache>().unwrap().clone();
+
+    let watch_data = data.get::<WatchData>().unwrap().clone();
+    let _lock = match watch_data.watch(contest_id) {
+        Some(t) => t,
+        None => {
+            channel
+                .send_message(ctx, |f| {
+                    f.content(format!("Contest is already being watched: {}", contest_id))
+                })
+                .await?;
+            return Ok(());
+        }
+    };
+
+    let contest = match data.get::<ContestCache>().unwrap().get(contest_id).await {
+        Ok((p, _)) => p,
+        Err(e) => {
+            channel
+                .send_message(ctx, |f| {
+                    f.content(format!("Cannot get info about contest: {}", e))
+                })
+                .await?;
+            return Ok(());
+        }
+    };
+
+    if contest.phase == ContestPhase::Before {
+        let start_time = match contest.start_time_seconds {
+            Some(s) => chrono::Utc.timestamp(s as i64, 0),
+            None => {
+                channel
+                    .send_message(ctx, |f| {
+                        f.content(format!(
+                            "Contest **{}** found, but we don't know when it will begin!",
+                            contest.name
+                        ))
+                    })
+                    .await?;
+                return Ok(());
+            }
+        };
+        channel
+                .send_message(ctx, |f| {
+                    f.content(format!("Contest **{}** found, but has not started yet. Youmu will start watching as soon as it begins! (which is in about {})", contest.name, start_time.format("<t:%s:R>")))
+                })
+                .await?;
+        tokio::time::sleep(
+            (start_time - chrono::Utc::now() + chrono::Duration::seconds(30))
+                .to_std()
+                .unwrap(),
+        )
+        .await
+    }
+
     let mut msg = channel
         .send_message(&ctx, |e| e.content("Youmu is building the member list..."))
         .await?;
+
+    let http = data.get::<CFClient>().unwrap();
+    let (mut contest, problems, _) =
+        Contest::standings(&*http, contest_id, |f| f.limit(1, 1)).await?;
     // Collect an initial member list.
     // This never changes during the scan.
     let mut member_results: HashMap<UserId, MemberResult> = db
@@ -55,10 +152,6 @@ pub async fn watch_contest(
         .collect()
         .await;
 
-    let http = data.get::<CFClient>().unwrap();
-    let (mut contest, problems, _) =
-        Contest::standings(&*http, contest_id, |f| f.limit(1, 1)).await?;
-
     msg.edit(&ctx, |e| {
         e.content(format!(
             "Youmu is watching contest **{}**, with the following members: {}",