From 39ae2fb944ebeea7542a5b0340c50b3bdd3e0c79 Mon Sep 17 00:00:00 2001 From: Mia Date: Thu, 31 Jul 2025 20:37:37 +0000 Subject: [PATCH] feat: caching --- Cargo.lock | 3 + consumer/src/backfill/mod.rs | 19 ++- consumer/src/backfill/repo.rs | 9 +- consumer/src/indexer/mod.rs | 72 +++++++-- dataloader-rs/src/async_cached.rs | 200 +++++++++++++++++++++++++ dataloader-rs/src/lib.rs | 1 + parakeet-db/Cargo.toml | 1 + parakeet-db/src/models.rs | 31 ++-- parakeet/Cargo.toml | 2 + parakeet/src/cache.rs | 150 +++++++++++++++++++ parakeet/src/config.rs | 1 + parakeet/src/hydration/feedgen.rs | 9 +- parakeet/src/hydration/labeler.rs | 37 +++-- parakeet/src/hydration/posts.rs | 19 ++- parakeet/src/hydration/profile.rs | 38 +++-- parakeet/src/loaders.rs | 239 ++++++++++++++++-------------- parakeet/src/main.rs | 8 + 17 files changed, 651 insertions(+), 188 deletions(-) create mode 100644 dataloader-rs/src/async_cached.rs create mode 100644 parakeet/src/cache.rs diff --git a/Cargo.lock b/Cargo.lock index 2515b514..c25506ea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2770,8 +2770,10 @@ dependencies = [ "multibase", "parakeet-db", "parakeet-index", + "redis", "reqwest", "serde", + "serde_ipld_dagcbor", "serde_json", "tokio", "tower-http", @@ -2786,6 +2788,7 @@ dependencies = [ "chrono", "diesel", "postgres-types", + "serde", "serde_json", ] diff --git a/consumer/src/backfill/mod.rs b/consumer/src/backfill/mod.rs index b52a33c8..fe73ebb1 100644 --- a/consumer/src/backfill/mod.rs +++ b/consumer/src/backfill/mod.rs @@ -97,12 +97,13 @@ impl BackfillManager { let mut inner = self.inner.clone(); let mut conn = self.pool.get().await?; + let mut rc = self.redis.clone(); tracker.spawn(async move { let _p = p; tracing::trace!("backfilling {job}"); - if let Err(e) = backfill_actor(&mut conn, &mut inner, &job).await { + if let Err(e) = backfill_actor(&mut conn, &mut rc, &mut inner, &job).await { tracing::error!(did = &job, "backfill failed: {e}"); counter!("backfill_failure").increment(1); @@ -132,6 +133,7 @@ impl BackfillManager { #[instrument(skip(conn, inner))] async fn backfill_actor( conn: &mut Object, + rc: &mut MultiplexedConnection, inner: &mut BackfillManagerInner, did: &str, ) -> eyre::Result<()> { @@ -140,7 +142,7 @@ async fn backfill_actor( tracing::trace!("loading repo"); - let (commit, mut deltas, copies) = repo::insert_repo(&mut t, &inner.tmp_dir, did).await?; + let (commit, mut deltas, copies) = repo::insert_repo(&mut t, rc, &inner.tmp_dir, did).await?; db::actor_set_repo_state(&mut t, did, &commit.rev, commit.data).await?; @@ -152,7 +154,7 @@ async fn backfill_actor( ) .await?; - handle_backfill_rows(&mut t, &mut deltas, did, &commit.rev).await?; + handle_backfill_rows(&mut t, rc, &mut deltas, did, &commit.rev).await?; tracing::trace!("insertion finished"); @@ -192,20 +194,21 @@ async fn backfill_actor( async fn handle_backfill_rows( conn: &mut Transaction<'_>, + rc: &mut MultiplexedConnection, deltas: &mut impl AggregateDeltaStore, repo: &str, rev: &str, -) -> Result<(), tokio_postgres::Error> { +) -> eyre::Result<()> { // `pull_backfill_rows` filters out anything before the last commit we pulled let backfill_rows = db::backfill_rows_get(conn, repo, rev).await?; for row in backfill_rows { // blindly unwrap-ing this CID as we've already parsed it and re-serialized it - let repo_cid = Cid::from_str(&row.cid).unwrap(); + let repo_cid = Cid::from_str(&row.cid)?; db::actor_set_repo_state(conn, repo, &row.repo_ver, repo_cid).await?; // again, we've serialized this. - let items: Vec = serde_json::from_value(row.data).unwrap(); + let items: Vec = serde_json::from_value(row.data)?; for item in items { let Some((_, rkey)) = item.at_uri.rsplit_once("/") else { @@ -218,11 +221,13 @@ async fn handle_backfill_rows( continue; }; - indexer::index_op(conn, deltas, repo, cid, record, &item.at_uri, rkey).await? + indexer::index_op(conn, rc, deltas, repo, cid, record, &item.at_uri, rkey) + .await? } BackfillItemInner::Delete => { indexer::index_op_delete( conn, + rc, deltas, repo, item.collection, diff --git a/consumer/src/backfill/repo.rs b/consumer/src/backfill/repo.rs index 60557dd2..95d916ed 100644 --- a/consumer/src/backfill/repo.rs +++ b/consumer/src/backfill/repo.rs @@ -10,6 +10,7 @@ use ipld_core::cid::Cid; use iroh_car::CarReader; use metrics::counter; use parakeet_index::AggregateType; +use redis::aio::MultiplexedConnection; use std::collections::HashMap; use std::path::Path; use tokio::io::BufReader; @@ -18,6 +19,7 @@ type BackfillDeltaStore = HashMap<(String, i32), i32>; pub async fn insert_repo( t: &mut Transaction<'_>, + rc: &mut MultiplexedConnection, tmp_dir: &Path, repo: &str, ) -> eyre::Result<(CarCommitEntry, BackfillDeltaStore, CopyStore)> { @@ -54,7 +56,7 @@ pub async fn insert_repo( } CarEntry::Record(record) => { if let Some(path) = mst_nodes.remove(&cid) { - record_index(t, &mut copies, &mut deltas, repo, &path, cid, record).await?; + record_index(t, rc, &mut copies, &mut deltas, repo, &path, cid, record).await?; } else { records.insert(cid, record); } @@ -84,7 +86,7 @@ pub async fn insert_repo( for (cid, record) in records { if let Some(path) = mst_nodes.remove(&cid) { - record_index(t, &mut copies, &mut deltas, repo, &path, cid, record).await?; + record_index(t, rc, &mut copies, &mut deltas, repo, &path, cid, record).await?; } else { tracing::warn!("couldn't find MST node for record {cid}") } @@ -97,6 +99,7 @@ pub async fn insert_repo( async fn record_index( t: &mut Transaction<'_>, + rc: &mut MultiplexedConnection, copies: &mut CopyStore, deltas: &mut BackfillDeltaStore, did: &str, @@ -190,7 +193,7 @@ async fn record_index( copies.push_record(&at_uri, cid); copies.verifications.push((at_uri, cid, rec)); } - _ => indexer::index_op(t, deltas, did, cid, record, &at_uri, rkey).await?, + _ => indexer::index_op(t, rc, deltas, did, cid, record, &at_uri, rkey).await?, } Ok(()) diff --git a/consumer/src/indexer/mod.rs b/consumer/src/indexer/mod.rs index d20aa2aa..b207bfdc 100644 --- a/consumer/src/indexer/mod.rs +++ b/consumer/src/indexer/mod.rs @@ -394,7 +394,7 @@ async fn index_commit( db::actor_set_repo_state(&mut t, &commit.repo, &commit.rev, commit.commit).await?; for op in &commit.ops { - process_op(&mut t, &mut state.idxc_tx, &commit.repo, op, &blocks).await?; + process_op(&mut t, rc, &mut state.idxc_tx, &commit.repo, op, &blocks).await?; } t.commit().await?; @@ -461,11 +461,12 @@ fn process_op_bf( #[inline(always)] async fn process_op( conn: &mut Transaction<'_>, + rc: &mut MultiplexedConnection, deltas: &mut impl AggregateDeltaStore, repo: &str, op: &CommitOp, blocks: &HashMap>, -) -> Result<(), tokio_postgres::Error> { +) -> eyre::Result<()> { let Some((collection_raw, rkey)) = op.path.split_once("/") else { tracing::warn!("op contained invalid path {}", op.path); return Ok(()); @@ -490,9 +491,9 @@ async fn process_op( return Ok(()); }; - index_op(conn, deltas, repo, cid, decoded, &full_path, rkey).await?; + index_op(conn, rc, deltas, repo, cid, decoded, &full_path, rkey).await?; } else if op.action == "delete" { - index_op_delete(conn, deltas, repo, collection, &full_path, rkey).await?; + index_op_delete(conn, rc, deltas, repo, collection, &full_path, rkey).await?; } else { tracing::warn!("op contained invalid action {}", op.action); } @@ -517,13 +518,14 @@ fn decode_op(op: &CommitOp, blocks: &HashMap>) -> Option<(Cid, Reco pub async fn index_op( conn: &mut Transaction<'_>, + rc: &mut MultiplexedConnection, deltas: &mut impl AggregateDeltaStore, repo: &str, cid: Cid, record: RecordTypes, at_uri: &str, rkey: &str, -) -> Result<(), tokio_postgres::Error> { +) -> eyre::Result<()> { match record { RecordTypes::AppBskyActorProfile(record) => { if rkey == "self" { @@ -533,11 +535,14 @@ pub async fn index_op( if let Some(labels) = labels { db::maintain_self_labels(conn, repo, Some(cid), at_uri, labels).await?; } + + redis::AsyncTypedCommands::del(rc, format!("profile#{repo}")).await?; } } RecordTypes::AppBskyActorStatus(record) => { if rkey == "self" { db::status_upsert(conn, repo, record).await?; + redis::AsyncTypedCommands::del(rc, format!("profile#{repo}")).await?; } } RecordTypes::AppBskyFeedGenerator(record) => { @@ -550,6 +555,8 @@ pub async fn index_op( if did_insert { deltas.incr(repo, AggregateType::ProfileFeed).await; + } else { + redis::AsyncTypedCommands::del(rc, format!("feedgen#{at_uri}")).await?; } } RecordTypes::AppBskyFeedLike(record) => { @@ -618,6 +625,9 @@ pub async fn index_op( disable_effective, ) .await?; + + // TODO: should we purge embed#{at_uri} for everything in detached_embeding_uris? + // maybe postgate_maintain_detaches should return a list of uris? } RecordTypes::AppBskyFeedRepost(record) => { deltas @@ -633,6 +643,7 @@ pub async fn index_op( } db::threadgate_upsert(conn, at_uri, cid, record).await?; + redis::AsyncTypedCommands::del(rc, format!("post#{at_uri}")).await?; } RecordTypes::AppBskyGraphBlock(record) => { db::block_insert(conn, rkey, repo, record).await?; @@ -658,6 +669,8 @@ pub async fn index_op( if did_insert { deltas.incr(repo, AggregateType::ProfileList).await; + } else { + redis::AsyncTypedCommands::del(rc, format!("list#{at_uri}")).await?; } } RecordTypes::AppBskyGraphListBlock(record) => { @@ -671,6 +684,7 @@ pub async fn index_op( return Ok(()); } + redis::AsyncTypedCommands::del(rc, format!("list#{}", &record.list)).await?; db::list_item_insert(conn, at_uri, record).await?; } RecordTypes::AppBskyGraphStarterPack(record) => { @@ -678,6 +692,8 @@ pub async fn index_op( if did_insert { deltas.incr(repo, AggregateType::ProfileStarterpack).await; + } else { + redis::AsyncTypedCommands::del(rc, format!("starterpacks#{at_uri}")).await?; } } RecordTypes::AppBskyGraphVerification(record) => { @@ -691,16 +707,20 @@ pub async fn index_op( if let Some(labels) = labels { db::maintain_self_labels(conn, repo, Some(cid), at_uri, labels).await?; } + + redis::AsyncTypedCommands::del(rc, format!("labeler#{repo}")).await?; } } RecordTypes::AppBskyNotificationDeclaration(record) => { if rkey == "self" { db::notif_decl_upsert(conn, repo, record).await?; + redis::AsyncTypedCommands::del(rc, format!("profile#{repo}")).await?; } } RecordTypes::ChatBskyActorDeclaration(record) => { if rkey == "self" { db::chat_decl_upsert(conn, repo, record).await?; + redis::AsyncTypedCommands::del(rc, format!("profile#{repo}")).await?; } } } @@ -712,17 +732,25 @@ pub async fn index_op( pub async fn index_op_delete( conn: &mut Transaction<'_>, + rc: &mut MultiplexedConnection, deltas: &mut impl AggregateDeltaStore, repo: &str, collection: CollectionType, at_uri: &str, rkey: &str, -) -> Result<(), tokio_postgres::Error> { +) -> eyre::Result<()> { match collection { - CollectionType::BskyProfile => db::profile_delete(conn, repo).await?, - CollectionType::BskyStatus => db::status_delete(conn, repo).await?, + CollectionType::BskyProfile => { + redis::AsyncTypedCommands::del(rc, format!("profile#{repo}")).await?; + db::profile_delete(conn, repo).await? + } + CollectionType::BskyStatus => { + redis::AsyncTypedCommands::del(rc, format!("profile#{repo}")).await?; + db::status_delete(conn, repo).await? + } CollectionType::BskyBlock => db::block_delete(conn, rkey, repo).await?, CollectionType::BskyFeedGen => { + redis::AsyncTypedCommands::del(rc, format!("feedgen#{at_uri}")).await?; let count = db::feedgen_delete(conn, at_uri).await?; deltas .add_delta(repo, AggregateType::ProfileFeed, -(count as i32)) @@ -739,6 +767,7 @@ pub async fn index_op_delete( let post_info = db::post_get_info_for_delete(conn, at_uri).await?; db::post_delete(conn, at_uri).await?; + redis::AsyncTypedCommands::del(rc, format!("post#{at_uri}")).await?; if let Some((reply_to, embed)) = post_info { deltas.decr(repo, AggregateType::ProfilePost).await; @@ -759,7 +788,10 @@ pub async fn index_op_delete( } 0 } - CollectionType::BskyFeedThreadgate => db::threadgate_delete(conn, at_uri).await?, + CollectionType::BskyFeedThreadgate => { + redis::AsyncTypedCommands::del(rc, format!("post#{at_uri}")).await?; + db::threadgate_delete(conn, at_uri).await? + } CollectionType::BskyFollow => { if let Some(followee) = db::follow_delete(conn, rkey, repo).await? { deltas.decr(&followee, AggregateType::Follower).await; @@ -768,6 +800,7 @@ pub async fn index_op_delete( 0 } CollectionType::BskyList => { + redis::AsyncTypedCommands::del(rc, format!("list#{at_uri}")).await?; let count = db::list_delete(conn, at_uri).await?; deltas .add_delta(repo, AggregateType::ProfileList, -(count as i32)) @@ -775,8 +808,12 @@ pub async fn index_op_delete( count } CollectionType::BskyListBlock => db::list_block_delete(conn, at_uri).await?, - CollectionType::BskyListItem => db::list_item_delete(conn, at_uri).await?, + CollectionType::BskyListItem => { + redis::AsyncTypedCommands::del(rc, format!("list#{at_uri}")).await?; + db::list_item_delete(conn, at_uri).await? + } CollectionType::BskyStarterPack => { + redis::AsyncTypedCommands::del(rc, format!("starterpacks#{at_uri}")).await?; let count = db::starter_pack_delete(conn, at_uri).await?; deltas .add_delta(repo, AggregateType::ProfileStarterpack, -(count as i32)) @@ -784,9 +821,18 @@ pub async fn index_op_delete( count } CollectionType::BskyVerification => db::verification_delete(conn, at_uri).await?, - CollectionType::BskyLabelerService => db::labeler_delete(conn, at_uri).await?, - CollectionType::BskyNotificationDeclaration => db::notif_decl_delete(conn, repo).await?, - CollectionType::ChatActorDecl => db::chat_decl_delete(conn, repo).await?, + CollectionType::BskyLabelerService => { + redis::AsyncTypedCommands::del(rc, format!("labeler#{repo}")).await?; + db::labeler_delete(conn, at_uri).await? + } + CollectionType::BskyNotificationDeclaration => { + redis::AsyncTypedCommands::del(rc, format!("profile#{repo}")).await?; + db::notif_decl_delete(conn, repo).await? + } + CollectionType::ChatActorDecl => { + redis::AsyncTypedCommands::del(rc, format!("profile#{repo}")).await?; + db::chat_decl_delete(conn, repo).await? + } _ => unreachable!(), }; diff --git a/dataloader-rs/src/async_cached.rs b/dataloader-rs/src/async_cached.rs new file mode 100644 index 00000000..619425aa --- /dev/null +++ b/dataloader-rs/src/async_cached.rs @@ -0,0 +1,200 @@ +use crate::runtime::{Arc, Mutex}; +use crate::{yield_fn, BatchFn, WaitForWorkFn}; +use std::collections::{HashMap, HashSet}; +use std::hash::Hash; +use std::iter::IntoIterator; + +pub trait AsyncCache { + type Key; + type Val; + async fn get(&mut self, key: &Self::Key) -> Option; + async fn insert(&mut self, key: Self::Key, val: Self::Val); + async fn remove(&mut self, key: &Self::Key) -> Option; + async fn clear(&mut self); +} + +struct State +where + C: AsyncCache, +{ + completed: C, + pending: HashSet, +} + +impl State +where + C: AsyncCache, +{ + fn with_cache(cache: C) -> Self { + State { + completed: cache, + pending: HashSet::new(), + } + } +} + +#[derive(Clone)] +pub struct Loader +where + K: Eq + Hash + Clone, + V: Clone, + F: BatchFn, + C: AsyncCache, +{ + state: Arc>>, + load_fn: Arc>, + wait_for_work_fn: Arc, + max_batch_size: usize, +} + +impl Loader +where + K: Eq + Hash + Clone, + V: Clone, + F: BatchFn, + C: AsyncCache, +{ + pub fn new(load_fn: F, cache: C) -> Self { + Loader { + state: Arc::new(Mutex::new(State::with_cache(cache))), + load_fn: Arc::new(Mutex::new(load_fn)), + max_batch_size: 200, + wait_for_work_fn: Arc::new(yield_fn(10)), + } + } + + pub fn with_max_batch_size(mut self, max_batch_size: usize) -> Self { + self.max_batch_size = max_batch_size; + self + } + + pub fn with_yield_count(mut self, yield_count: usize) -> Self { + self.wait_for_work_fn = Arc::new(yield_fn(yield_count)); + self + } + + /// Replaces the yielding for work behavior with an arbitrary future. Rather than yielding + /// the runtime repeatedly this will generate and `.await` a future of your choice. + /// ***This is incompatible with*** [`Self::with_yield_count()`]. + pub fn with_custom_wait_for_work(mut self, wait_for_work_fn: impl WaitForWorkFn) -> Self { + self.wait_for_work_fn = Arc::new(wait_for_work_fn); + self + } + + pub fn max_batch_size(&self) -> usize { + self.max_batch_size + } + + pub async fn load(&self, key: K) -> Option { + let mut state = self.state.lock().await; + if let Some(v) = state.completed.get(&key).await { + return Some(v.clone()); + } + + if !state.pending.contains(&key) { + state.pending.insert(key.clone()); + if state.pending.len() >= self.max_batch_size { + let keys = state.pending.drain().collect::>(); + let mut load_fn = self.load_fn.lock().await; + let load_ret = load_fn.load(keys.as_ref()).await; + drop(load_fn); + for (k, v) in load_ret.into_iter() { + state.completed.insert(k, v).await; + } + return state.completed.get(&key).await.clone(); + } + } + drop(state); + + (self.wait_for_work_fn)().await; + + let mut state = self.state.lock().await; + if let Some(v) = state.completed.get(&key).await { + return Some(v.clone()); + } + + if !state.pending.is_empty() { + let keys = state.pending.drain().collect::>(); + let mut load_fn = self.load_fn.lock().await; + let load_ret = load_fn.load(keys.as_ref()).await; + drop(load_fn); + for (k, v) in load_ret.into_iter() { + state.completed.insert(k, v).await; + } + } + + state.completed.get(&key).await.clone() + } + + pub async fn load_many(&self, keys: Vec) -> HashMap { + let mut state = self.state.lock().await; + let mut ret = HashMap::new(); + let mut rest = Vec::new(); + for key in keys.into_iter() { + if let Some(v) = state.completed.get(&key).await.clone() { + ret.insert(key, v); + continue; + } + if !state.pending.contains(&key) { + state.pending.insert(key.clone()); + + if state.pending.len() >= self.max_batch_size { + let keys = state.pending.drain().collect::>(); + let mut load_fn = self.load_fn.lock().await; + let load_ret = load_fn.load(keys.as_ref()).await; + drop(load_fn); + for (k, v) in load_ret.into_iter() { + state.completed.insert(k, v).await; + } + } + } + rest.push(key); + } + drop(state); + + (self.wait_for_work_fn)().await; + + if !rest.is_empty() { + let mut state = self.state.lock().await; + if !state.pending.is_empty() { + let keys = state.pending.drain().collect::>(); + let mut load_fn = self.load_fn.lock().await; + let load_ret = load_fn.load(keys.as_ref()).await; + drop(load_fn); + for (k, v) in load_ret.into_iter() { + state.completed.insert(k, v).await; + } + } + + for key in rest.into_iter() { + if let Some(v) = state.completed.get(&key).await.clone() { + ret.insert(key, v); + } + } + } + + ret + } + + pub async fn prime(&self, key: K, val: V) { + let mut state = self.state.lock().await; + state.completed.insert(key, val).await; + } + + pub async fn prime_many(&self, values: impl IntoIterator) { + let mut state = self.state.lock().await; + for (k, v) in values.into_iter() { + state.completed.insert(k, v).await; + } + } + + pub async fn clear(&self, key: K) { + let mut state = self.state.lock().await; + state.completed.remove(&key).await; + } + + pub async fn clear_all(&self) { + let mut state = self.state.lock().await; + state.completed.clear().await + } +} diff --git a/dataloader-rs/src/lib.rs b/dataloader-rs/src/lib.rs index b06de808..36c7fc73 100644 --- a/dataloader-rs/src/lib.rs +++ b/dataloader-rs/src/lib.rs @@ -1,5 +1,6 @@ #![allow(async_fn_in_trait)] +pub mod async_cached; mod batch_fn; pub mod cached; pub mod non_cached; diff --git a/parakeet-db/Cargo.toml b/parakeet-db/Cargo.toml index c95ab4e7..3227c7ed 100644 --- a/parakeet-db/Cargo.toml +++ b/parakeet-db/Cargo.toml @@ -7,6 +7,7 @@ edition = "2021" chrono = { version = "0.4.39", features = ["serde"] } diesel = { version = "2.2.6", features = ["chrono", "serde_json"], optional = true } postgres-types = { version = "0.2.9", optional = true } +serde = { version = "1.0.217", features = ["derive"] } serde_json = "1.0.134" [features] diff --git a/parakeet-db/src/models.rs b/parakeet-db/src/models.rs index f47282b2..1689e5bd 100644 --- a/parakeet-db/src/models.rs +++ b/parakeet-db/src/models.rs @@ -1,6 +1,7 @@ use crate::types::*; use chrono::prelude::*; use diesel::prelude::*; +use serde::{Deserialize, Serialize}; #[derive(Debug, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::actors)] @@ -16,7 +17,7 @@ pub struct Actor { pub last_indexed: Option, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::profiles)] #[diesel(primary_key(did))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -40,7 +41,7 @@ pub struct Profile { pub indexed_at: NaiveDateTime, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::lists)] #[diesel(primary_key(at_uri))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -85,7 +86,7 @@ pub struct ListItem { pub indexed_at: NaiveDateTime, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::feedgens)] #[diesel(primary_key(at_uri))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -121,7 +122,7 @@ pub struct BackfillRow { pub indexed_at: NaiveDateTime, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::posts)] #[diesel(primary_key(at_uri))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -148,7 +149,7 @@ pub struct Post { pub indexed_at: NaiveDateTime, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::post_embed_images)] #[diesel(primary_key(post_uri, seq))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -164,7 +165,7 @@ pub struct PostEmbedImage { pub height: Option, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::post_embed_video)] #[diesel(primary_key(post_uri))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -179,7 +180,7 @@ pub struct PostEmbedVideo { pub height: Option, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::post_embed_video_captions)] #[diesel(primary_key(post_uri, language))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -191,7 +192,7 @@ pub struct PostEmbedVideoCaption { pub cid: String, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::post_embed_ext)] #[diesel(primary_key(post_uri))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -206,7 +207,7 @@ pub struct PostEmbedExt { pub thumb_cid: Option, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::post_embed_record)] #[diesel(primary_key(post_uri))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -236,7 +237,7 @@ pub struct Postgate { pub indexed_at: NaiveDateTime, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::threadgates)] #[diesel(primary_key(post_uri))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -255,7 +256,7 @@ pub struct Threadgate { pub indexed_at: NaiveDateTime, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::starterpacks)] #[diesel(primary_key(at_uri))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -275,7 +276,7 @@ pub struct StaterPack { pub indexed_at: NaiveDateTime, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::labelers)] #[diesel(primary_key(did))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -291,7 +292,7 @@ pub struct LabelerService { pub indexed_at: NaiveDateTime, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable, Associations)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable, Associations)] #[diesel(table_name = crate::schema::labeler_defs)] #[diesel(belongs_to(LabelerService, foreign_key = labeler))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -329,7 +330,7 @@ pub struct Label { pub indexed_at: NaiveDateTime, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::verification)] #[diesel(primary_key(at_uri))] #[diesel(check_for_backend(diesel::pg::Pg))] @@ -346,7 +347,7 @@ pub struct VerificationEntry { pub indexed_at: NaiveDateTime, } -#[derive(Clone, Debug, Queryable, Selectable, Identifiable)] +#[derive(Clone, Debug, Serialize, Deserialize, Queryable, Selectable, Identifiable)] #[diesel(table_name = crate::schema::statuses)] #[diesel(primary_key(did))] #[diesel(check_for_backend(diesel::pg::Pg))] diff --git a/parakeet/Cargo.toml b/parakeet/Cargo.toml index 3d1a8608..043a3a13 100644 --- a/parakeet/Cargo.toml +++ b/parakeet/Cargo.toml @@ -23,8 +23,10 @@ lexica = { path = "../lexica" } multibase = "0.9.1" parakeet-db = { path = "../parakeet-db" } parakeet-index = { path = "../parakeet-index" } +redis = { version = "0.32", features = ["tokio-native-tls-comp"] } reqwest = { version = "0.12", features = ["json"] } serde = { version = "1.0.217", features = ["derive"] } +serde_ipld_dagcbor = "0.6.1" serde_json = "1.0.134" tokio = { version = "1.42.0", features = ["full"] } tower-http = { version = "0.6.2", features = ["cors", "trace"] } diff --git a/parakeet/src/cache.rs b/parakeet/src/cache.rs new file mode 100644 index 00000000..3d6ed56d --- /dev/null +++ b/parakeet/src/cache.rs @@ -0,0 +1,150 @@ +use dataloader::async_cached::AsyncCache; +use redis::aio::MultiplexedConnection; +use redis::AsyncTypedCommands; +use serde::{Deserialize, Serialize}; +use std::marker::PhantomData; + +/// General Loader Cache +pub struct LoaderCache { + conn: MultiplexedConnection, + exp: Option, + _phantom: PhantomData, +} + +impl LoaderCache { + pub fn new(conn: &MultiplexedConnection, exp: Option) -> Self { + LoaderCache:: { + conn: conn.clone(), + exp, + _phantom: PhantomData, + } + } +} + +impl AsyncCache for LoaderCache +where + V: for<'a> Deserialize<'a> + Serialize, +{ + type Key = String; + type Val = V; + + async fn get(&mut self, key: &Self::Key) -> Option { + let res: Option> = redis::AsyncCommands::get(&mut self.conn, &key).await.ok()?; + + match serde_ipld_dagcbor::from_slice(&res?) { + Ok(v) => Some(v), + Err(err) => { + tracing::error!(key, "failed to decode cache value: {err}"); + None + } + } + } + + async fn insert(&mut self, key: Self::Key, val: Self::Val) { + let data = match serde_ipld_dagcbor::to_vec(&val) { + Ok(data) => data, + Err(err) => { + tracing::error!(key, "failed to encode cache value: {err}"); + return; + } + }; + + if let Some(exp) = self.exp { + self.conn.set_ex(key, data, exp).await.unwrap(); + } else { + self.conn.set(key, data).await.unwrap(); + } + } + + async fn remove(&mut self, key: &Self::Key) -> Option { + let res: Option> = redis::AsyncCommands::get_del(&mut self.conn, &key) + .await + .ok()?; + + match serde_ipld_dagcbor::from_slice(&res?) { + Ok(v) => Some(v), + Err(err) => { + tracing::error!(key, "failed to decode cache value: {err}"); + None + } + } + } + + async fn clear(&mut self) {} +} + +/// A Loader Cache in with a key prefix +pub struct PrefixedLoaderCache { + conn: MultiplexedConnection, + prefix: String, + exp: Option, + _phantom: PhantomData, +} + +impl PrefixedLoaderCache { + pub fn new(conn: &MultiplexedConnection, prefix: String, exp: Option) -> Self { + PrefixedLoaderCache { + conn: conn.clone(), + prefix, + exp, + _phantom: PhantomData, + } + } +} + +impl AsyncCache for PrefixedLoaderCache +where + V: for<'a> Deserialize<'a> + Serialize, +{ + type Key = String; + type Val = V; + + async fn get(&mut self, key: &Self::Key) -> Option { + let key = format!("{}#{}", self.prefix, key); + + let res: Option> = redis::AsyncCommands::get(&mut self.conn, &key).await.ok()?; + + match serde_ipld_dagcbor::from_slice(&res?) { + Ok(v) => Some(v), + Err(err) => { + tracing::error!(key, "failed to decode cache value: {err}"); + None + } + } + } + + async fn insert(&mut self, key: Self::Key, val: Self::Val) { + let key = format!("{}#{}", self.prefix, key); + let data = match serde_ipld_dagcbor::to_vec(&val) { + Ok(data) => data, + Err(err) => { + tracing::error!(key = &key, "failed to encode cache value: {err}"); + return; + } + }; + + if let Some(exp) = self.exp { + self.conn.set_ex(key, data, exp).await.unwrap(); + } else { + self.conn.set(key, data).await.unwrap(); + } + } + + async fn remove(&mut self, key: &Self::Key) -> Option { + let key = format!("{}#{}", self.prefix, key); + + let res: Option> = redis::AsyncCommands::get_del(&mut self.conn, &key) + .await + .ok()?; + + match serde_ipld_dagcbor::from_slice(&res?) { + Ok(v) => Some(v), + Err(err) => { + tracing::error!(key, "failed to decode cache value: {err}"); + None + } + } + } + + async fn clear(&mut self) {} +} diff --git a/parakeet/src/config.rs b/parakeet/src/config.rs index d0704882..c07e66c1 100644 --- a/parakeet/src/config.rs +++ b/parakeet/src/config.rs @@ -15,6 +15,7 @@ pub(crate) fn load_config() -> eyre::Result { pub struct Config { pub index_uri: String, pub database_url: String, + pub cache_uri: String, #[serde(default)] pub server: ConfigServer, pub service: ConfigService, diff --git a/parakeet/src/hydration/feedgen.rs b/parakeet/src/hydration/feedgen.rs index 265df021..dde52b5b 100644 --- a/parakeet/src/hydration/feedgen.rs +++ b/parakeet/src/hydration/feedgen.rs @@ -43,7 +43,8 @@ fn build_feedgen( impl super::StatefulHydrator<'_> { pub async fn hydrate_feedgen(&self, feedgen: String) -> Option { let labels = self.get_label(&feedgen).await; - let (feedgen, likes) = self.loaders.feedgen.load(feedgen).await?; + let likes = self.loaders.like.load(feedgen.clone()).await; + let feedgen = self.loaders.feedgen.load(feedgen).await?; let profile = self.hydrate_profile(feedgen.owner.clone()).await?; Some(build_feedgen(feedgen, profile, labels, likes, &self.cdn)) @@ -51,20 +52,22 @@ impl super::StatefulHydrator<'_> { pub async fn hydrate_feedgens(&self, feedgens: Vec) -> HashMap { let labels = self.get_label_many(&feedgens).await; + let mut likes = self.loaders.like.load_many(feedgens.clone()).await; let feedgens = self.loaders.feedgen.load_many(feedgens).await; let creators = feedgens .values() - .map(|(feedgen, _)| feedgen.owner.clone()) + .map(|feedgen| feedgen.owner.clone()) .collect(); let creators = self.hydrate_profiles(creators).await; feedgens .into_iter() - .filter_map(|(uri, (feedgen, likes))| { + .filter_map(|(uri, feedgen)| { let creator = creators.get(&feedgen.owner).cloned()?; let labels = labels.get(&uri).cloned().unwrap_or_default(); + let likes = likes.remove(&uri); Some(( uri, diff --git a/parakeet/src/hydration/labeler.rs b/parakeet/src/hydration/labeler.rs index e80d1f62..51768aa6 100644 --- a/parakeet/src/hydration/labeler.rs +++ b/parakeet/src/hydration/labeler.rs @@ -92,7 +92,8 @@ fn build_view_detailed( impl StatefulHydrator<'_> { pub async fn hydrate_labeler(&self, labeler: String) -> Option { let labels = self.get_label(&labeler).await; - let (labeler, _, likes) = self.loaders.labeler.load(labeler).await?; + let likes = self.loaders.like.load(make_labeler_uri(&labeler)).await; + let (labeler, _) = self.loaders.labeler.load(labeler).await?; let creator = self.hydrate_profile(labeler.did.clone()).await?; Some(build_view(labeler, creator, labels, likes)) @@ -102,17 +103,19 @@ impl StatefulHydrator<'_> { let labels = self.get_label_many(&labelers).await; let labelers = self.loaders.labeler.load_many(labelers).await; - let creators = labelers + let (creators, uris) = labelers .values() - .map(|(labeler, _, _)| labeler.did.clone()) - .collect(); + .map(|(labeler, _)| (labeler.did.clone(), make_labeler_uri(&labeler.did))) + .unzip::<_, _, Vec<_>, Vec<_>>(); let creators = self.hydrate_profiles(creators).await; + let mut likes = self.loaders.like.load_many(uris.clone()).await; labelers .into_iter() - .filter_map(|(k, (labeler, _, likes))| { + .filter_map(|(k, (labeler, _))| { let creator = creators.get(&labeler.did).cloned()?; let labels = labels.get(&k).cloned().unwrap_or_default(); + let likes = likes.remove(&make_labeler_uri(&labeler.did)); Some((k, build_view(labeler, creator, labels, likes))) }) @@ -121,7 +124,8 @@ impl StatefulHydrator<'_> { pub async fn hydrate_labeler_detailed(&self, labeler: String) -> Option { let labels = self.get_label(&labeler).await; - let (labeler, defs, likes) = self.loaders.labeler.load(labeler).await?; + let likes = self.loaders.like.load(make_labeler_uri(&labeler)).await; + let (labeler, defs) = self.loaders.labeler.load(labeler).await?; let creator = self.hydrate_profile(labeler.did.clone()).await?; Some(build_view_detailed(labeler, defs, creator, labels, likes)) @@ -134,23 +138,28 @@ impl StatefulHydrator<'_> { let labels = self.get_label_many(&labelers).await; let labelers = self.loaders.labeler.load_many(labelers).await; - let creators = labelers + let (creators, uris) = labelers .values() - .map(|(labeler, _, _)| labeler.did.clone()) - .collect(); + .map(|(labeler, _)| (labeler.did.clone(), make_labeler_uri(&labeler.did))) + .unzip::<_, _, Vec<_>, Vec<_>>(); let creators = self.hydrate_profiles(creators).await; + let mut likes = self.loaders.like.load_many(uris.clone()).await; labelers .into_iter() - .filter_map(|(k, (labeler, defs, likes))| { + .filter_map(|(k, (labeler, defs))| { let creator = creators.get(&labeler.did).cloned()?; let labels = labels.get(&k).cloned().unwrap_or_default(); + let likes = likes.remove(&make_labeler_uri(&labeler.did)); - Some(( - k, - build_view_detailed(labeler, defs, creator, labels, likes), - )) + let view = build_view_detailed(labeler, defs, creator, labels, likes); + + Some((k, view)) }) .collect() } } + +fn make_labeler_uri(did: &str) -> String { + format!("at://{did}/app.bsky.labeler.service/self") +} diff --git a/parakeet/src/hydration/posts.rs b/parakeet/src/hydration/posts.rs index fd430f8e..cd1ae87f 100644 --- a/parakeet/src/hydration/posts.rs +++ b/parakeet/src/hydration/posts.rs @@ -99,7 +99,8 @@ impl StatefulHydrator<'_> { } pub async fn hydrate_post(&self, post: String) -> Option { - let (post, threadgate, stats) = self.loaders.posts.load(post).await?; + let stats = self.loaders.post_stats.load(post.clone()).await; + let (post, threadgate) = self.loaders.posts.load(post).await?; let embed = self.hydrate_embed(post.at_uri.clone()).await; let author = self.hydrate_profile_basic(post.did.clone()).await?; let threadgate = self.hydrate_threadgate(threadgate).await; @@ -111,11 +112,12 @@ impl StatefulHydrator<'_> { } pub async fn hydrate_posts(&self, posts: Vec) -> HashMap { + let stats = self.loaders.post_stats.load_many(posts.clone()).await; let posts = self.loaders.posts.load_many(posts).await; let (authors, post_uris) = posts .values() - .map(|(post, _, _)| (post.did.clone(), post.at_uri.clone())) + .map(|(post, _)| (post.did.clone(), post.at_uri.clone())) .unzip::<_, _, Vec<_>, Vec<_>>(); let authors = self.hydrate_profiles_basic(authors).await; @@ -123,7 +125,7 @@ impl StatefulHydrator<'_> { let threadgates = posts .values() - .filter_map(|(_, threadgate, _)| threadgate.clone()) + .filter_map(|(_, threadgate)| threadgate.clone()) .collect(); let threadgates = self.hydrate_threadgates(threadgates).await; @@ -131,11 +133,12 @@ impl StatefulHydrator<'_> { posts .into_iter() - .filter_map(|(uri, (post, threadgate, stats))| { + .filter_map(|(uri, (post, threadgate))| { let author = authors.get(&post.did)?; let embed = embeds.get(&uri).cloned(); let threadgate = threadgate.and_then(|tg| threadgates.get(&tg.at_uri).cloned()); let labels = post_labels.get(&uri).cloned().unwrap_or_default(); + let stats = stats.get(&uri).cloned(); Some(( uri, @@ -146,11 +149,12 @@ impl StatefulHydrator<'_> { } pub async fn hydrate_feed_posts(&self, posts: Vec) -> HashMap { + let stats = self.loaders.post_stats.load_many(posts.clone()).await; let posts = self.loaders.posts.load_many(posts).await; let (authors, post_uris) = posts .values() - .map(|(post, _, _)| (post.did.clone(), post.at_uri.clone())) + .map(|(post, _)| (post.did.clone(), post.at_uri.clone())) .unzip::<_, _, Vec<_>, Vec<_>>(); let authors = self.hydrate_profiles_basic(authors).await; @@ -160,7 +164,7 @@ impl StatefulHydrator<'_> { let reply_refs = posts .values() - .flat_map(|(post, _, _)| [post.parent_uri.clone(), post.root_uri.clone()]) + .flat_map(|(post, _)| [post.parent_uri.clone(), post.root_uri.clone()]) .flatten() .collect::>(); @@ -168,7 +172,7 @@ impl StatefulHydrator<'_> { posts .into_iter() - .filter_map(|(post_uri, (post, _, stats))| { + .filter_map(|(post_uri, (post, _))| { let author = authors.get(&post.did)?; let root = post.root_uri.as_ref().and_then(|uri| reply_posts.get(uri)); @@ -199,6 +203,7 @@ impl StatefulHydrator<'_> { let embed = embeds.get(&post_uri).cloned(); let labels = post_labels.get(&post_uri).cloned().unwrap_or_default(); + let stats = stats.get(&post_uri).cloned(); let post = build_postview(post, author.to_owned(), labels, embed, None, stats); Some(( diff --git a/parakeet/src/hydration/profile.rs b/parakeet/src/hydration/profile.rs index 7777f645..965d9f9f 100644 --- a/parakeet/src/hydration/profile.rs +++ b/parakeet/src/hydration/profile.rs @@ -152,7 +152,8 @@ fn build_status(status: models::Status, cdn: &BskyCdn) -> Option { } fn build_basic( - (handle, profile, chat_decl, is_labeler, stats, status, notif_decl): ProfileLoaderRet, + (handle, profile, chat_decl, is_labeler, status, notif_decl): ProfileLoaderRet, + stats: Option, labels: Vec, verifications: Option>, cdn: &BskyCdn, @@ -176,7 +177,8 @@ fn build_basic( } fn build_profile( - (handle, profile, chat_decl, is_labeler, stats, status, notif_decl): ProfileLoaderRet, + (handle, profile, chat_decl, is_labeler, status, notif_decl): ProfileLoaderRet, + stats: Option, labels: Vec, verifications: Option>, cdn: &BskyCdn, @@ -202,7 +204,8 @@ fn build_profile( } fn build_detailed( - (handle, profile, chat_decl, is_labeler, stats, status, notif_decl): ProfileLoaderRet, + (handle, profile, chat_decl, is_labeler, status, notif_decl): ProfileLoaderRet, + stats: Option, labels: Vec, verifications: Option>, cdn: &BskyCdn, @@ -235,9 +238,10 @@ impl super::StatefulHydrator<'_> { pub async fn hydrate_profile_basic(&self, did: String) -> Option { let labels = self.get_profile_label(&did).await; let verif = self.loaders.verification.load(did.clone()).await; + let stats = self.loaders.profile_stats.load(did.clone()).await; let profile_info = self.loaders.profile.load(did).await?; - Some(build_basic(profile_info, labels, verif, &self.cdn)) + Some(build_basic(profile_info, stats, labels, verif, &self.cdn)) } pub async fn hydrate_profiles_basic( @@ -246,6 +250,7 @@ impl super::StatefulHydrator<'_> { ) -> HashMap { let labels = self.get_profile_label_many(&dids).await; let verif = self.loaders.verification.load_many(dids.clone()).await; + let stats = self.loaders.profile_stats.load_many(dids.clone()).await; let profiles = self.loaders.profile.load_many(dids).await; profiles @@ -253,8 +258,9 @@ impl super::StatefulHydrator<'_> { .map(|(k, profile_info)| { let labels = labels.get(&k).cloned().unwrap_or_default(); let verif = verif.get(&k).cloned(); + let stats = stats.get(&k).cloned(); - let v = build_basic(profile_info, labels, verif, &self.cdn); + let v = build_basic(profile_info, stats, labels, verif, &self.cdn); (k, v) }) .collect() @@ -262,16 +268,17 @@ impl super::StatefulHydrator<'_> { pub async fn hydrate_profile(&self, did: String) -> Option { let labels = self.get_profile_label(&did).await; - let verif = self.loaders.verification.load(did.clone()).await; + let stats = self.loaders.profile_stats.load(did.clone()).await; let profile_info = self.loaders.profile.load(did).await?; - Some(build_profile(profile_info, labels, verif, &self.cdn)) + Some(build_profile(profile_info, stats, labels, verif, &self.cdn)) } pub async fn hydrate_profiles(&self, dids: Vec) -> HashMap { let labels = self.get_profile_label_many(&dids).await; let verif = self.loaders.verification.load_many(dids.clone()).await; + let stats = self.loaders.profile_stats.load_many(dids.clone()).await; let profiles = self.loaders.profile.load_many(dids).await; profiles @@ -279,8 +286,9 @@ impl super::StatefulHydrator<'_> { .map(|(k, profile_info)| { let labels = labels.get(&k).cloned().unwrap_or_default(); let verif = verif.get(&k).cloned(); + let stats = stats.get(&k).cloned(); - let v = build_profile(profile_info, labels, verif, &self.cdn); + let v = build_profile(profile_info, stats, labels, verif, &self.cdn); (k, v) }) .collect() @@ -288,11 +296,17 @@ impl super::StatefulHydrator<'_> { pub async fn hydrate_profile_detailed(&self, did: String) -> Option { let labels = self.get_profile_label(&did).await; - let verif = self.loaders.verification.load(did.clone()).await; + let stats = self.loaders.profile_stats.load(did.clone()).await; let profile_info = self.loaders.profile.load(did).await?; - Some(build_detailed(profile_info, labels, verif, &self.cdn)) + Some(build_detailed( + profile_info, + stats, + labels, + verif, + &self.cdn, + )) } pub async fn hydrate_profiles_detailed( @@ -301,6 +315,7 @@ impl super::StatefulHydrator<'_> { ) -> HashMap { let labels = self.get_profile_label_many(&dids).await; let verif = self.loaders.verification.load_many(dids.clone()).await; + let stats = self.loaders.profile_stats.load_many(dids.clone()).await; let profiles = self.loaders.profile.load_many(dids).await; profiles @@ -308,8 +323,9 @@ impl super::StatefulHydrator<'_> { .map(|(k, profile_info)| { let labels = labels.get(&k).cloned().unwrap_or_default(); let verif = verif.get(&k).cloned(); + let stats = stats.get(&k).cloned(); - let v = build_detailed(profile_info, labels, verif, &self.cdn); + let v = build_detailed(profile_info, stats, labels, verif, &self.cdn); (k, v) }) .collect() diff --git a/parakeet/src/loaders.rs b/parakeet/src/loaders.rs index 1f0a3d17..d5e98a99 100644 --- a/parakeet/src/loaders.rs +++ b/parakeet/src/loaders.rs @@ -1,5 +1,7 @@ +use crate::cache::PrefixedLoaderCache; use crate::xrpc::extract::LabelConfigItem; -use dataloader::cached::Loader; +use dataloader::async_cached::Loader; +use dataloader::non_cached::Loader as NonCachedLoader; use dataloader::BatchFn; use diesel::prelude::*; use diesel_async::pooled_connection::deadpool::Pool; @@ -7,37 +9,88 @@ use diesel_async::{AsyncPgConnection, RunQueryDsl}; use itertools::Itertools; use lexica::app_bsky::actor::{ChatAllowIncoming, ProfileAllowSubscriptions}; use parakeet_db::{models, schema}; +use redis::aio::MultiplexedConnection; +use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::str::FromStr; +type CachingLoader = Loader>; + +fn new_plc_loader( + load_fn: F, + conn: &MultiplexedConnection, + prefix: &str, + exp: u64, +) -> Loader> +where + V: Clone + Serialize + for<'a> Deserialize<'a>, + F: BatchFn, +{ + Loader::new( + load_fn, + PrefixedLoaderCache::new(conn, prefix.to_string(), Some(exp)), + ) +} + pub struct Dataloaders { - pub embed: Loader, - pub feedgen: Loader, - pub handle: Loader, + pub embed: CachingLoader, + pub feedgen: CachingLoader, + pub handle: CachingLoader, pub label: LabelLoader, - pub labeler: Loader, - pub list: Loader, - pub posts: Loader, - pub profile: Loader, - pub starterpacks: Loader, - pub verification: Loader, VerificationLoader>, + pub labeler: CachingLoader, + pub list: CachingLoader, + pub like: NonCachedLoader, + pub posts: CachingLoader, + pub post_stats: NonCachedLoader, + pub profile: CachingLoader, + pub profile_stats: NonCachedLoader, + pub starterpacks: CachingLoader, + pub verification: CachingLoader, VerificationLoader>, } impl Dataloaders { - // for the moment, we set up memory cached loaders - // we should build a redis/valkey backend at some point in the future. - pub fn new(pool: Pool, idxc: parakeet_index::Client) -> Dataloaders { + #[rustfmt::skip] + pub fn new( + pool: Pool, + rc: MultiplexedConnection, + idxc: parakeet_index::Client, + ) -> Dataloaders { Dataloaders { - embed: Loader::new(EmbedLoader(pool.clone())), - feedgen: Loader::new(FeedGenLoader(pool.clone(), idxc.clone())), - handle: Loader::new(HandleLoader(pool.clone())), + embed: new_plc_loader(EmbedLoader(pool.clone()), &rc, "embed", 3600), + feedgen: new_plc_loader(FeedGenLoader(pool.clone(), idxc.clone()), &rc, "feedgen", 600), + handle: new_plc_loader(HandleLoader(pool.clone()), &rc, "handle", 60), label: LabelLoader(pool.clone()), // CARE: never cache this. - labeler: Loader::new(LabelServiceLoader(pool.clone(), idxc.clone())), - list: Loader::new(ListLoader(pool.clone())), - posts: Loader::new(PostLoader(pool.clone(), idxc.clone())), - profile: Loader::new(ProfileLoader(pool.clone(), idxc.clone())), - starterpacks: Loader::new(StarterPackLoader(pool.clone())), - verification: Loader::new(VerificationLoader(pool.clone())), + labeler: new_plc_loader(LabelServiceLoader(pool.clone(), idxc.clone()), &rc, "labeler", 600), + like: NonCachedLoader::new(LikeLoader(idxc.clone())), + list: new_plc_loader(ListLoader(pool.clone()), &rc, "list", 600), + posts: new_plc_loader(PostLoader(pool.clone()), &rc, "post", 3600), + post_stats: NonCachedLoader::new(PostStatsLoader(idxc.clone())), + profile: new_plc_loader(ProfileLoader(pool.clone()), &rc, "profile", 3600), + profile_stats: NonCachedLoader::new(ProfileStatsLoader(idxc.clone())), + starterpacks: new_plc_loader(StarterPackLoader(pool.clone()), &rc, "starterpacks", 600), + verification: new_plc_loader(VerificationLoader(pool.clone()), &rc, "verification", 60), + } + } +} + +pub struct LikeLoader(parakeet_index::Client); +impl BatchFn for LikeLoader { + async fn load(&mut self, keys: &[String]) -> HashMap { + let res = self + .0 + .get_like_count_many(parakeet_index::GetStatsManyReq { + uris: keys.to_vec(), + }) + .await + .map(|v| v.into_inner()); + + match res { + Ok(data) => data + .entries + .into_iter() + .map(|(k, v)| (k, v.likes)) + .collect(), + Err(_) => HashMap::new(), } } } @@ -66,13 +119,12 @@ impl BatchFn for HandleLoader { } } -pub struct ProfileLoader(Pool, parakeet_index::Client); +pub struct ProfileLoader(Pool); pub type ProfileLoaderRet = ( Option, models::Profile, Option, bool, - Option, Option, Option, ); @@ -115,17 +167,6 @@ impl BatchFn for ProfileLoader { )>(&mut conn) .await; - let stats_req = parakeet_index::GetStatsManyReq { - uris: keys.to_vec(), - }; - let mut stats = self - .1 - .get_profile_stats_many(stats_req) - .await - .unwrap() - .into_inner() - .entries; - match res { Ok(res) => HashMap::from_iter(res.into_iter().map( |(did, handle, profile, chat_decl, labeler_cid, status, notif_decl)| { @@ -133,17 +174,8 @@ impl BatchFn for ProfileLoader { let notif_decl = notif_decl.and_then(|v| ProfileAllowSubscriptions::from_str(&v).ok()); let is_labeler = labeler_cid.is_some(); - let maybe_stats = stats.remove(&did); - - let val = ( - handle, - profile, - chat_decl, - is_labeler, - maybe_stats, - status, - notif_decl, - ); + + let val = (handle, profile, chat_decl, is_labeler, status, notif_decl); (did, val) }, @@ -156,6 +188,22 @@ impl BatchFn for ProfileLoader { } } +pub struct ProfileStatsLoader(parakeet_index::Client); +impl BatchFn for ProfileStatsLoader { + async fn load(&mut self, keys: &[String]) -> HashMap { + let stats_req = parakeet_index::GetStatsManyReq { + uris: keys.to_vec(), + }; + + self.0 + .get_profile_stats_many(stats_req) + .await + .unwrap() + .into_inner() + .entries + } +} + pub struct ListLoader(Pool); type ListLoaderRet = (models::List, i64); impl BatchFn for ListLoader { @@ -189,9 +237,8 @@ impl BatchFn for ListLoader { } pub struct FeedGenLoader(Pool, parakeet_index::Client); -type FeedGenLoaderRet = (models::FeedGen, Option); -impl BatchFn for FeedGenLoader { - async fn load(&mut self, keys: &[String]) -> HashMap { +impl BatchFn for FeedGenLoader { + async fn load(&mut self, keys: &[String]) -> HashMap { let mut conn = self.0.get().await.unwrap(); let res = schema::feedgens::table @@ -200,23 +247,11 @@ impl BatchFn for FeedGenLoader { .load(&mut conn) .await; - let stats_req = parakeet_index::GetStatsManyReq { - uris: keys.to_vec(), - }; - let mut stats = self - .1 - .get_like_count_many(stats_req) - .await - .unwrap() - .into_inner() - .entries; - match res { - Ok(res) => HashMap::from_iter(res.into_iter().map(|feedgen| { - let likes = stats.remove(&feedgen.at_uri).map(|v| v.likes); - - (feedgen.at_uri.clone(), (feedgen, likes)) - })), + Ok(res) => HashMap::from_iter( + res.into_iter() + .map(|feedgen| (feedgen.at_uri.clone(), feedgen)), + ), Err(e) => { tracing::error!("feedgen load failed: {e}"); HashMap::new() @@ -225,12 +260,8 @@ impl BatchFn for FeedGenLoader { } } -pub struct PostLoader(Pool, parakeet_index::Client); -type PostLoaderRet = ( - models::Post, - Option, - Option, -); +pub struct PostLoader(Pool); +type PostLoaderRet = (models::Post, Option); impl BatchFn for PostLoader { async fn load(&mut self, keys: &[String]) -> HashMap { let mut conn = self.0.get().await.unwrap(); @@ -245,33 +276,37 @@ impl BatchFn for PostLoader { .load(&mut conn) .await; + match res { + Ok(res) => HashMap::from_iter( + res.into_iter() + .map(|(post, threadgate)| (post.at_uri.clone(), (post, threadgate))), + ), + Err(e) => { + tracing::error!("post load failed: {e}"); + HashMap::new() + } + } + } +} + +pub struct PostStatsLoader(parakeet_index::Client); +impl BatchFn for PostStatsLoader { + async fn load(&mut self, keys: &[String]) -> HashMap { let stats_req = parakeet_index::GetStatsManyReq { uris: keys.to_vec(), }; - let mut stats = self - .1 + + self.0 .get_post_stats_many(stats_req) .await .unwrap() .into_inner() - .entries; - - match res { - Ok(res) => HashMap::from_iter(res.into_iter().map(|(post, threadgate)| { - let maybe_stats = stats.remove(&post.at_uri); - - (post.at_uri.clone(), (post, threadgate, maybe_stats)) - })), - Err(e) => { - tracing::error!("post load failed: {e}"); - HashMap::new() - } - } + .entries } } pub struct EmbedLoader(Pool); -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub enum EmbedLoaderRet { Images(Vec), Video(models::PostEmbedVideo), @@ -385,11 +420,7 @@ impl BatchFn for StarterPackLoader { } pub struct LabelServiceLoader(Pool, parakeet_index::Client); -type LabelServiceLoaderRet = ( - models::LabelerService, - Vec, - Option, -); +type LabelServiceLoaderRet = (models::LabelerService, Vec); impl BatchFn for LabelServiceLoader { async fn load(&mut self, keys: &[String]) -> HashMap { let mut conn = self.0.get().await.unwrap(); @@ -408,32 +439,10 @@ impl BatchFn for LabelServiceLoader { let defs = defs.grouped_by(&labelers); - let uris = keys - .iter() - .map(|v| format!("at://{v}/app.bsky.labeler.service/self")) - .collect(); - let stats_req = parakeet_index::GetStatsManyReq { uris }; - let mut stats = self - .1 - .get_like_count_many(stats_req) - .await - .unwrap() - .into_inner() - .entries; - labelers .into_iter() .zip(defs) - .map(|(labeler, defs)| { - let likes = stats - .remove(&format!( - "at://{}/app.bsky.labeler.service/self", - &labeler.did - )) - .map(|v| v.likes); - - (labeler.did.clone(), (labeler, defs, likes)) - }) + .map(|(labeler, defs)| (labeler.did.clone(), (labeler, defs))) .collect() } } diff --git a/parakeet/src/main.rs b/parakeet/src/main.rs index db02f1ba..7b386766 100644 --- a/parakeet/src/main.rs +++ b/parakeet/src/main.rs @@ -3,12 +3,14 @@ use diesel_async::pooled_connection::deadpool::Pool; use diesel_async::pooled_connection::AsyncDieselConnectionManager; use diesel_async::AsyncPgConnection; use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness}; +use redis::aio::MultiplexedConnection; use std::sync::Arc; use tower_http::cors::{AllowHeaders, AllowOrigin, CorsLayer}; use tower_http::trace::TraceLayer; const MIGRATIONS: EmbeddedMigrations = embed_migrations!(); +mod cache; mod config; mod db; mod hydration; @@ -18,6 +20,7 @@ mod xrpc; #[derive(Clone)] pub struct GlobalState { pub pool: Pool, + pub redis_mp: MultiplexedConnection, pub dataloaders: Arc, pub resolver: Arc, pub index_client: parakeet_index::Client, @@ -46,10 +49,14 @@ async fn main() -> eyre::Result<()> { tracing::info!("database migrations complete"); } + let redis_client = redis::Client::open(conf.cache_uri)?; + let redis_mp = redis_client.get_multiplexed_tokio_connection().await?; + let index_client = parakeet_index::Client::connect(conf.index_uri).await?; let dataloaders = Arc::new(loaders::Dataloaders::new( pool.clone(), + redis_mp.clone(), index_client.clone(), )); let resolver = Arc::new(did_resolver::Resolver::new(did_resolver::ResolverOpts { @@ -82,6 +89,7 @@ async fn main() -> eyre::Result<()> { .layer(cors) .with_state(GlobalState { pool, + redis_mp, dataloaders, resolver, index_client, -- 2.51.2