pikadick/commands/
quizizz.rs

1use crate::{
2    checks::ENABLED_CHECK,
3    util::LoadingReaction,
4    ClientDataKey,
5};
6use anyhow::Context as _;
7use serenity::{
8    framework::standard::{
9        macros::command,
10        Args,
11        CommandResult,
12    },
13    model::prelude::*,
14    prelude::*,
15};
16use std::{
17    collections::BinaryHeap,
18    sync::Arc,
19    time::{
20        Duration,
21        Instant,
22    },
23};
24use tokio::sync::{
25    watch::{
26        Receiver as WatchReceiver,
27        Sender as WatchSender,
28    },
29    Notify,
30};
31use tracing::{
32    error,
33    info,
34};
35
36pub type SearchResult = Result<Option<String>, Arc<anyhow::Error>>;
37
38const MAX_TRIES: usize = 1_000;
39const MAX_CODE: u32 = 999_999;
40
41const LIMIT_REACHED_MSG: &str = "Reached limit while searching for quizizz code, quitting...";
42
43#[derive(Clone, Debug)]
44pub struct QuizizzClient {
45    finder_task_wakeup: Arc<Notify>,
46    // finder_task_interest: Arc<AtomicU64>,
47    finder_task_rx: WatchReceiver<SearchResult>,
48}
49
50impl QuizizzClient {
51    /// Make a new [`QuizizzClient`].
52    pub fn new() -> Self {
53        let finder_task_wakeup = Arc::new(Notify::new());
54        let wakeup = finder_task_wakeup.clone();
55        let (watch_tx, finder_task_rx) = tokio::sync::watch::channel(Ok(None));
56
57        tokio::spawn(finder_task(watch_tx, wakeup));
58
59        Self {
60            finder_task_wakeup,
61            finder_task_rx,
62        }
63    }
64
65    /// Get the next searched code.
66    ///
67    /// `None` signifies that the task ran out of tries.
68    pub async fn search_for_code(&self) -> SearchResult {
69        let mut finder_task_rx = self.finder_task_rx.clone();
70
71        // Mark current value as seen
72        finder_task_rx.borrow_and_update();
73
74        // Wake up task
75        //
76        // You might be wondering why not `notify_waiters`,
77        // as the current code will potentially make the task do an extra unecessary lookup
78        // if two requests come in at once.
79        // This is because there is an edge case where the task may have sent its response already,
80        // but still be processing and caching the other requests when it is woken up.
81        // Since it is not waiting, this will make the request task hang until another request wakes up the finder task.
82        // It is therefore better to use `notify_one` to ensure the task is always woken up when it needs to be, even spuriously.
83        // The finder task is also equipped with a caching mechanism,
84        // so spurious wakeups will likely quickly pull a value from there instead of making web requests.
85        self.finder_task_wakeup.notify_one();
86
87        // Wait for new value
88        finder_task_rx
89            .changed()
90            .await
91            .context("failed to get response from finder task")?;
92
93        // Return new value
94        let ret = finder_task_rx.borrow_and_update().clone();
95        ret
96    }
97}
98
99impl Default for QuizizzClient {
100    fn default() -> Self {
101        Self::new()
102    }
103}
104
105/// A Cache for quizzizz codes
106#[derive(Debug)]
107pub struct CodeCache {
108    cache: BinaryHeap<(std::cmp::Reverse<Instant>, String)>,
109}
110
111impl CodeCache {
112    /// Make a new cache
113    pub fn new() -> Self {
114        // Worst case caches `MAX_TRIES - 1` entries, since we gather MAX_TRIES entries and return one on success.
115        Self {
116            cache: BinaryHeap::with_capacity(MAX_TRIES),
117        }
118    }
119
120    /// Get the # of entries
121    pub fn len(&self) -> usize {
122        self.cache.len()
123    }
124
125    /// Returns true if it is empty
126    pub fn is_empty(&self) -> bool {
127        self.cache.is_empty()
128    }
129
130    /// Trim the cache
131    pub fn trim(&mut self) {
132        while let Some((time, _)) = self.cache.peek() {
133            if time.0.elapsed() > Duration::from_secs(10 * 60) {
134                self.cache.pop();
135            } else {
136                // The newest value has not expired.
137                // Exit the peek loop.
138                break;
139            }
140        }
141    }
142
143    /// Trim the cache and pop a code if it exists
144    pub fn trim_pop(&mut self) -> Option<String> {
145        self.trim();
146        Some(self.cache.pop()?.1)
147    }
148
149    /// Add a code to the cache
150    pub fn push(&mut self, code_str: String) {
151        self.cache
152            .push((std::cmp::Reverse(Instant::now()), code_str));
153    }
154}
155
156impl Default for CodeCache {
157    fn default() -> Self {
158        Self::new()
159    }
160}
161
162async fn finder_task(watch_tx: WatchSender<SearchResult>, wakeup: Arc<Notify>) {
163    let client = quizizz::Client::new();
164    let mut cache = CodeCache::new();
165
166    while tokio::select! {
167        _ = wakeup.notified() => true,
168        _ = watch_tx.closed() => false,
169    } {
170        // Try cache first
171        if let Some(code_str) = cache.trim_pop() {
172            let _ = watch_tx.send(Ok(Some(code_str))).is_ok();
173            continue;
174        }
175
176        // Generate start code
177        let mut code: u32 = rand::random::<u32>() % MAX_CODE;
178        info!(start_code = code);
179
180        // Spawn parallel guesses
181        let (tx, mut rx) = tokio::sync::mpsc::channel(MAX_TRIES);
182        for _ in 0..MAX_TRIES {
183            let code_str = format!("{:06}", code);
184            code = code.wrapping_add(1);
185
186            let client = client.clone();
187            let tx = tx.clone();
188            tokio::spawn(async move {
189                let check_room_result = client
190                    .check_room(&code_str)
191                    .await
192                    .and_then(|r| r.error_for_response())
193                    .map(|res| res.room);
194
195                let _ = tx.send((code_str, check_room_result)).await.is_ok();
196            });
197        }
198        drop(tx);
199
200        // Process parallel guess responses
201        let mut sent_response = false;
202        while let Some((code_str, check_room_result)) = rx.recv().await {
203            match check_room_result {
204                Ok(Some(room)) if room.is_running() => {
205                    if !sent_response {
206                        let _ = watch_tx.send(Ok(Some(code_str))).is_ok();
207                        sent_response = true;
208                    } else {
209                        // Cache extra results
210                        cache.push(code_str);
211                    }
212                }
213                Ok(None | Some(_)) => {
214                    // Pass
215                    // room data is missing / room is not running
216                }
217                Err(quizizz::Error::InvalidGenericResponse(e))
218                    if e.is_room_not_found() || e.is_player_login_required() =>
219                {
220                    // Pass
221                    // the room was not found / the player needs to be logged in to access this game
222                }
223                Err(e) => {
224                    let e = Err(e)
225                        .with_context(|| {
226                            format!("failed to search for quizizz code '{}'", code_str)
227                        })
228                        .map_err(Arc::new);
229                    error!("{:?}", e);
230                    if !sent_response {
231                        let _ = watch_tx.send(e).is_ok();
232                        sent_response = true;
233                    }
234                }
235            }
236        }
237        if !sent_response {
238            let _ = watch_tx.send(Ok(None)).is_ok();
239        }
240
241        info!("quizizz has {} cached entries", cache.len());
242    }
243}
244
245#[command]
246#[description("Locate a quizizz code")]
247#[bucket("quizizz")]
248#[checks(Enabled)]
249async fn quizizz(ctx: &Context, msg: &Message, _args: Args) -> CommandResult {
250    let data_lock = ctx.data.read().await;
251    let client_data = data_lock
252        .get::<ClientDataKey>()
253        .expect("failed to get client data");
254    let client = client_data.quizizz_client.clone();
255    drop(data_lock);
256
257    let mut loading = LoadingReaction::new(ctx.http.clone(), msg);
258
259    match client.search_for_code().await {
260        Ok(Some(code_str)) => {
261            info!("located quizizz code '{}'", code_str);
262            loading.send_ok();
263            msg.channel_id
264                .say(&ctx.http, format!("Located quizizz code: {}", code_str))
265                .await?;
266        }
267        Ok(None) => {
268            info!("quizziz finder reached limit");
269            msg.channel_id.say(&ctx.http, LIMIT_REACHED_MSG).await?;
270        }
271        Err(e) => {
272            error!("{:?}", e);
273            msg.channel_id.say(&ctx.http, format!("{:?}", e)).await?;
274        }
275    }
276
277    Ok(())
278}
pikadick/commands/quizizz.rs

pikadick/commands/
quizizz.rs
