Something went wrong. Try again.
Source code for my personal quote bot project.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545use crate::storage::{ queue::{DequeueQuote, EnqueueQuotes}, source::SourceQuotes,};
use crate::data::Quote;use crate::data::{Handle, Message};
mod rng { use rand::SeedableRng;
pub struct PrngState { rng: rand::rngs::SmallRng, }
impl PrngState { pub fn from_thread_rng() -> Self { Self { rng: rand::rngs::SmallRng::from_rng(&mut rand::rng()), } }
pub fn from_seed(seed: u64) -> Self { Self { rng: rand::rngs::SmallRng::seed_from_u64(seed), } }
pub fn shuffle_slice<T>(&mut self, slice: &mut [T]) { use rand::seq::SliceRandom; slice.shuffle(&mut self.rng); } }
mod test { #[tokio::test] async fn shuffle_slice() { use super::*;
let mut data = vec![1, 2, 3, 4]; let mut rng = PrngState::from_thread_rng();
rng.shuffle_slice(&mut data); println!("{:?}", data); } }}
pub mod source { use super::*;
// TODO: Should the quote source filters be // generic over the exact manager implementation being used? /// Message to request that a SourceManager source its quotes once again. pub struct SourceQuotes; pub type SourceReply = Result<Vec<Quote>, ()>; impl Message for SourceQuotes { type Reply = SourceReply; }
/// Subtrait of Actor which specifically /// denotes actors that can handle all relevant source messages. pub trait SourceManager: Handle<SourceQuotes> {}
impl<T> SourceManager for T where T: Handle<SourceQuotes> {}
/// Implementation of [SourceManager] which sources quotes from a Vec /// that it holds in memory, without accessing external services. /// Its main purpose is to be used for testing. pub struct MemorySourceManager { quotes: Vec<Quote>, }
impl MemorySourceManager { pub fn new(quotes: impl IntoIterator<Item = impl Into<Quote>>) -> Self { Self { quotes: quotes.into_iter().map(Into::into).collect(), } } }
#[async_trait::async_trait(?Send)] impl Handle<SourceQuotes> for MemorySourceManager { async fn handle(&mut self, _msg: SourceQuotes) -> Result<Vec<Quote>, ()> { // We just clone the quotes we've been holding onto since startup Ok(self.quotes.clone()) } }
/// Uses a [QuoteFilter] to source quotes from the local filesystem /// at the beginning of each cycle. pub struct FsFilterSourceManager { filter: QuoteFilter, }
impl FsFilterSourceManager { pub fn new(filter: QuoteFilter) -> Self { Self { filter } } }
impl Default for FsFilterSourceManager { fn default() -> Self { Self { filter: QuoteFilter { content: r".*".to_string(), // TODO: Maybe make this a compile-time constant for debugging path: "test/**/*.txt".to_string(), _dates: vec![], }, } } }
#[async_trait::async_trait(?Send)] impl Handle<SourceQuotes> for FsFilterSourceManager { async fn handle(&mut self, _msg: SourceQuotes) -> Result<Vec<Quote>, ()> { self.filter.read_files() } }
#[derive(Clone, Debug)] pub struct QuoteFilter { /// "glob-like" pattern of files containing quotes to be read. path: String,
/// Regular expression to further filter the files matched by [path] /// based on their *contents*. content: String,
/// Currently unused. Intended for future functionality. _dates: Vec<String>, }
impl QuoteFilter { /// Creates a new [QuoteFilter] from a "glob"-like pattern provided as input. /// No filtering is applied based on matching files' contents. pub fn new_from_glob(glob: String) -> Self { Self { path: glob, // The following regex matches *any* string: content: ".*".to_string(), _dates: vec![], } }
// TODO: actually leverage async I/O fn read_files(&self) -> Result<Vec<Quote>, ()> { use glob::glob; use grep::{regex, searcher::sinks};
let matcher = regex::RegexMatcher::new(&self.content).map_err(|_| ())?; let mut searcher = grep::searcher::Searcher::new(); let mut results = Vec::new();
for file in glob(&self.path).map_err(|_| ())? { let file = match file { Ok(file) => file, Err(_) => continue, };
let mut matched = false; let sink = sinks::Lossy(|_lnum, _line| { matched = true; Ok(false) });
let search_result = searcher.search_path(&matcher, &file, sink); if !matched || search_result.is_err() { continue; }
let contents = std::fs::read_to_string(file).map_err(|_| ())?; results.push(contents.trim().into()); }
Ok(results) } }}
pub mod queue { use redis::AsyncCommands; use std::collections::VecDeque;
use super::*;
// Messages to interact with the quote queue #[derive(Debug)] pub struct DequeueQuote; pub type DequeueReply = Result<Option<Quote>, ()>; impl Message for DequeueQuote { type Reply = DequeueReply; }
#[derive(Debug)] pub struct EnqueueQuotes(pub Vec<Quote>);
/// In case of a successful enqueue, no further data is required; /// however, if an error occurs partway through, we must return /// the numbers of quotes to have already been pushed up to that point. /// A [usize] is used as that is the maximum amount of items a [Vec] /// of quotes within [EnqueueQuotes] can contain. pub type EnqueueReply = Result<(), usize>; impl Message for EnqueueQuotes { type Reply = EnqueueReply; }
/// Subtrait of Actor which specifically /// denotes actors that can handle all relevant queue messages. pub trait QueueManager: Handle<EnqueueQuotes> + Handle<DequeueQuote> {}
impl<T> QueueManager for T where T: Handle<EnqueueQuotes> + Handle<DequeueQuote> {}
/// A basic implementation of an in-memory queue of quotes. /// Its contents are *not* persisted across application restarts, so it /// is only suited for testing purposes. #[derive(Default)] pub struct MemoryQueueStorage { quotes: VecDeque<Quote>, }
impl MemoryQueueStorage { pub fn new() -> Self { Self { quotes: VecDeque::new(), } } }
#[async_trait::async_trait(?Send)] impl Handle<EnqueueQuotes> for MemoryQueueStorage { async fn handle(&mut self, msg: EnqueueQuotes) -> Result<(), usize> { // Note: as `.push_back` into a `VecDeque` is an infallible operation, // we never need to worry about failure cases in this implementation.
for q in msg.0 { self.quotes.push_back(q); }
Ok(()) } }
#[async_trait::async_trait(?Send)] impl Handle<DequeueQuote> for MemoryQueueStorage { async fn handle(&mut self, _msg: DequeueQuote) -> Result<Option<Quote>, ()> { // Note: this can never fail, since the quotes are stored in memory Ok(self.quotes.pop_front()) } }
/// Implementation of persistent quote queue storage through a Redis server. pub struct RedisQueueStorage { /// Key that points to the Redis queue to be used for fetching and storing quotes. queue_key: String,
/// Underlying client from which to instantiate Redis connections. client: redis::Client, }
impl RedisQueueStorage { /// The default queue key we use to fetch and store Redis quote. const DEFAULT_KEY: &str = "queue:default";
pub fn new(client: redis::Client, queue_key: String) -> Self { Self { queue_key, client } }
pub fn new_from_env() -> Result<Self, ()> { let client = redis::Client::open( std::env::var("REDIS_URL").unwrap_or("redis://localhost".to_string()), ) .map_err(|_| ())?;
let queue_key = std::env::var("REDIS_KEY").unwrap_or(Self::DEFAULT_KEY.to_string());
Ok(Self { queue_key, client }) }
async fn con(&self) -> Result<redis::aio::MultiplexedConnection, ()> { self.client .get_multiplexed_async_connection() .await .map_err(|e| { tracing::error!(error = e.to_string(), "Could not connect to Redis server"); () }) } }
#[async_trait::async_trait(?Send)] impl Handle<EnqueueQuotes> for RedisQueueStorage { #[tracing::instrument(name = "redis::enqueue_many", skip(self, msg))] async fn handle(&mut self, msg: EnqueueQuotes) -> Result<(), usize> { tracing::debug!( to_enqueue = msg.0.len(), "Enqueuing {} quotes...", msg.0.len() );
// We retrieve the connection object ahead of time to error-out early // in case of connection issues and avoid the overhead of fetching // a connection every iteration of the loop. let mut con = self.con().await.map_err(|_| 0usize)?;
for (idx, q) in msg.0.iter().enumerate() { let _: () = con .rpush(&self.queue_key, q.get()) .await // If something goes wrong, we return the number of quotes // **already successfully queued** as part of our error message. .map_err(|e| { tracing::error!( error = e.to_string(), queued_so_far = idx, total_to_queue = msg.0.len(), "Could not enqueue all quotes to redis" ); idx })?; }
Ok(()) } }
#[async_trait::async_trait(?Send)] impl Handle<DequeueQuote> for RedisQueueStorage { #[tracing::instrument(name = "redis::dequeue_one", skip(self, _msg))] async fn handle(&mut self, _msg: DequeueQuote) -> Result<Option<Quote>, ()> { tracing::debug!("Dequeuing quote..."); let quote: Result<Option<String>, ()> = self .con() .await? .lpop(&self.queue_key, None) .await .map_err(|e| { tracing::error!(error = e.to_string(), "Could not dequeue quote from Redis"); () });
match quote { Ok(Some(quote)) => Ok(Some(Quote::from(quote))),
// We need to separate these next two cases, // as the type of quote: `Result<Option<String>>` // does not match the return type: `Result<Option<Quote>>` Ok(None) => Ok(None), Err(()) => Err(()), } } }}
pub struct QuoteCycle<S: source::SourceManager, Q: queue::QueueManager> { rng: rng::PrngState, source_manager: S, queue_manager: Q,}
impl<S: source::SourceManager, Q: queue::QueueManager> QuoteCycle<S, Q> { pub fn new(rng: rng::PrngState, source_manager: S, queue_manager: Q) -> Self { Self { rng, source_manager, queue_manager, } }
pub fn with_thread_rng(source_manager: S, queue_manager: Q) -> Self { Self { rng: rng::PrngState::from_thread_rng(), source_manager, queue_manager, } }}
/// A message to [QuoteCycle] to fetch one more quote from its storage.#[derive(Debug)]pub struct FetchQuote;impl Message for FetchQuote { type Reply = Result<Quote, ()>;}
#[async_trait::async_trait(?Send)]impl<S, Q> Handle<FetchQuote> for QuoteCycle<S, Q>where S: source::SourceManager, Q: queue::QueueManager,{ #[tracing::instrument(name = "quote_cycle::fetch", skip(self))] async fn handle(&mut self, _msg: FetchQuote) -> Result<Quote, ()> { // 1. We query our queue storage for the next quote tracing::debug!("Fetching new quotes..."); if let Some(next_quote) = self .queue_manager .handle(DequeueQuote) .await .map_err(|_| ())? { // if there is a quote, we simply return it and move on tracing::debug!( quote = next_quote.get().to_string(), "Quote dequeued immediately." ); return Ok(next_quote); }
// 2. Otherwise, we must repopulate the queue through our source let mut refreshed_quotes = self.source_manager .handle(SourceQuotes) .await .map_err(|_| { tracing::error!("Could not retrieve new quotes from source"); () })?; tracing::debug!("Retrieved {} quotes.", refreshed_quotes.len());
// 3. We shuffle the newly-sourced quotes tracing::debug!("Shuffling quotes..."); self.rng.shuffle_slice(&mut refreshed_quotes); let refreshed_quotes = refreshed_quotes; // No longer mutable
// TODO: Perhaps we should assert that the new quotes are non-empty? // 4. We enqueue the newly-sourced quotes... tracing::debug!("Enqueuing quotes..."); self.queue_manager .handle(EnqueueQuotes(refreshed_quotes)) .await .map_err(|_| { tracing::error!("Could not enqueue quotes to storage"); () })?;
// 5. and, finally, we return the first among them. match self.queue_manager.handle(DequeueQuote).await { Ok(Some(q)) => Ok(q), Ok(None) => panic!("Newly-enqueued quotes should never be empty"), Err(_) => Err(()), } }}
mod test { #[tokio::test] async fn memory_queue() { use super::Quote; use super::queue::*; use super::*;
let mut queue_manager = MemoryQueueStorage::new();
let sample_quotes = ["Test no.1", "Test no.2", "Test no.3"];
queue_manager .handle(EnqueueQuotes( sample_quotes.iter().cloned().map(Quote::from).collect(), )) .await .expect("In-memory quote queue storage should be valid for insertion");
for text in sample_quotes.iter() { assert_eq!( *text, queue_manager .handle(DequeueQuote) .await .expect("In-memory queue storage should never panic on dequeue") .expect("In-memory queue storage should never be initialized as empty") .get() ); } }
#[tokio::test] async fn memory_source() { use super::source::*; use super::*;
let sample_quotes = ["Minie", "Miney", "Moe", "and", "some", "more"];
let mut source_manager = MemorySourceManager::new(sample_quotes);
let quotes = source_manager .handle(SourceQuotes) .await .expect("In-memory quote queue storage should be valid for insertion");
assert_eq!( sample_quotes.as_slice(), quotes .into_iter() // Since [Quote] doesn't implement any Equality trait, // we map the strings into quotes instead .map(String::from) .collect::<Vec<_>>() .as_slice(), ); }
#[tokio::test] async fn memory_cycle() { use std::{collections::HashMap, ops::AddAssign};
use super::FetchQuote; use super::QuoteCycle; use super::queue::*; use super::source::*; use super::*;
let sample_quotes = ["Minie", "Miney", "Moe"]; let mut cycle = { let source = MemorySourceManager::new(sample_quotes); let queue = MemoryQueueStorage::new();
QuoteCycle::with_thread_rng(source, queue) };
// We loop over `sample_quotes` twice to simulate the queue being exhausted fully, then re-sourced // Since the `cycle` manager will shuffle the quote sequence, we will verify that each // quote appears *exactly* `LOOPS` times throughout these iterations. const LOOPS: usize = 3; let mut quote_counts = HashMap::new(); for _ in 0..(sample_quotes.len() * LOOPS) { let next_quote = cycle.handle(FetchQuote).await.unwrap(); quote_counts .entry(next_quote.get().to_owned()) .or_insert(0) .add_assign(1); }
let quote_counts = quote_counts; // no longer mut assert!( // Note: technically speaking, different quotes could contain equivalent strings, // which would make this test fail; a "more proper" invariant check would ensure // verify that all counts are a multiple of the amount of times `sample_quotes` was chained, // and that the sum of all counts equals the total number of times a `FetchQuote` message was sent. quote_counts.into_values().into_iter().all(|c| c == LOOPS) ); }}