diff --git a/.gitignore b/.gitignore index dbf67885..27e4889e 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ target/ .idea/ .env Config.toml +data/ diff --git a/Cargo.lock b/Cargo.lock index fe319cb5..5aca806d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -401,7 +401,7 @@ version = "0.69.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088" dependencies = [ - "bitflags", + "bitflags 2.8.0", "cexpr", "clang-sys", "itertools 0.12.1", @@ -418,6 +418,12 @@ dependencies = [ "which", ] +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.8.0" @@ -656,6 +662,7 @@ dependencies = [ "metrics", "metrics-exporter-prometheus", "parakeet-db", + "parakeet-index", "reqwest", "serde", "serde_bytes", @@ -663,6 +670,7 @@ dependencies = [ "serde_json", "tokio", "tokio-postgres", + "tokio-stream", "tokio-tungstenite", "tracing", "tracing-subscriber", @@ -712,6 +720,15 @@ dependencies = [ "libc", ] +[[package]] +name = "crc32fast" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a97769d94ddab943e4510d138150169a2758b5ef3eb191a9ee688de3e23ef7b3" +dependencies = [ + "cfg-if", +] + [[package]] name = "crossbeam-epoch" version = "0.9.18" @@ -847,7 +864,7 @@ version = "2.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf1bedf64cdb9643204a36dd15b19a6ce8e7aa7f7b105868e9f1fad5ffa7d12" dependencies = [ - "bitflags", + "bitflags 2.8.0", "byteorder", "chrono", "diesel_derives", @@ -1040,6 +1057,12 @@ dependencies = [ "version_check", ] +[[package]] +name = "fixedbitset" +version = "0.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" + [[package]] name = "flume" version = "0.11.1" @@ -1088,6 +1111,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs2" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "fs_extra" version = "1.3.0" @@ -1196,6 +1229,15 @@ dependencies = [ "slab", ] +[[package]] +name = "fxhash" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c31b6d751ae2c7f11320402d34e41349dd1016f8d5d45e48c4312bc8625af50c" +dependencies = [ + "byteorder", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -1335,7 +1377,7 @@ dependencies = [ "ipconfig", "lru-cache", "once_cell", - "parking_lot", + "parking_lot 0.12.3", "rand", "resolv-conf", "smallvec", @@ -1458,6 +1500,19 @@ dependencies = [ "tower-service", ] +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-tls" version = "0.6.0" @@ -1683,6 +1738,15 @@ version = "0.1.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8fae54786f62fb2918dcfae3d568594e50eb9b5c25bf04371af6fe7516452fb" +[[package]] +name = "instant" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0242819d153cba4b4b05a5a8f2a7e9bbf97b6055b2a002b395c96b5ff3c0222" +dependencies = [ + "cfg-if", +] + [[package]] name = "ipconfig" version = "0.3.2" @@ -1998,6 +2062,12 @@ dependencies = [ "unsigned-varint 0.8.0", ] +[[package]] +name = "multimap" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "defc4c55412d89136f966bbb339008b474350e5e6e78d2714439c386b3137a03" + [[package]] name = "nanorand" version = "0.7.0" @@ -2084,7 +2154,7 @@ version = "0.10.68" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6174bc48f102d208783c2c84bf931bb75927a617866870de8a4ea85597f871f5" dependencies = [ - "bitflags", + "bitflags 2.8.0", "cfg-if", "foreign-types", "libc", @@ -2145,6 +2215,7 @@ dependencies = [ "itertools 0.14.0", "lexica", "parakeet-db", + "parakeet-index", "serde", "serde_json", "tokio", @@ -2162,6 +2233,23 @@ dependencies = [ "serde_json", ] +[[package]] +name = "parakeet-index" +version = "0.1.0" +dependencies = [ + "eyre", + "figment", + "itertools 0.14.0", + "prost", + "serde", + "sled", + "tokio", + "tonic", + "tonic-build", + "tracing", + "tracing-subscriber", +] + [[package]] name = "parakeet-lexgen" version = "0.1.0" @@ -2180,6 +2268,17 @@ version = "2.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" +[[package]] +name = "parking_lot" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d17b78036a60663b797adeaee46f5c9dfebb86948d1255007a1d6be0271ff99" +dependencies = [ + "instant", + "lock_api", + "parking_lot_core 0.8.6", +] + [[package]] name = "parking_lot" version = "0.12.3" @@ -2187,7 +2286,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f1bf18183cf54e8d6059647fc3063646a1801cf30896933ec2311622cc4b9a27" dependencies = [ "lock_api", - "parking_lot_core", + "parking_lot_core 0.9.10", +] + +[[package]] +name = "parking_lot_core" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a2cfe6f0ad2bfc16aefa463b497d5c7a5ecd44a23efa72aa342d90177356dc" +dependencies = [ + "cfg-if", + "instant", + "libc", + "redox_syscall 0.2.16", + "smallvec", + "winapi", ] [[package]] @@ -2198,7 +2311,7 @@ checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8" dependencies = [ "cfg-if", "libc", - "redox_syscall", + "redox_syscall 0.5.8", "smallvec", "windows-targets 0.52.6", ] @@ -2238,6 +2351,16 @@ version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" +[[package]] +name = "petgraph" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772" +dependencies = [ + "fixedbitset", + "indexmap", +] + [[package]] name = "phf" version = "0.11.3" @@ -2256,6 +2379,26 @@ dependencies = [ "siphasher", ] +[[package]] +name = "pin-project" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "677f1add503faace112b9f1373e43e9e054bfdd22ff1a63c1bc485eaec6a6a8a" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e918e4ff8c4549eb882f14b3a4bc8c8bc93de829416eacf579f1207a8fbf861" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "pin-project-lite" version = "0.2.16" @@ -2377,6 +2520,58 @@ dependencies = [ "yansi", ] +[[package]] +name = "prost" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-build" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" +dependencies = [ + "heck", + "itertools 0.14.0", + "log", + "multimap", + "once_cell", + "petgraph", + "prettyplease", + "prost", + "prost-types", + "regex", + "syn", + "tempfile", +] + +[[package]] +name = "prost-derive" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" +dependencies = [ + "anyhow", + "itertools 0.14.0", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "prost-types" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52c2c1bf36ddb1a1c396b3601a3cec27c2462e45f07c386894ec3ccf5332bd16" +dependencies = [ + "prost", +] + [[package]] name = "quanta" version = "0.12.5" @@ -2452,7 +2647,16 @@ version = "11.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c6df7ab838ed27997ba19a4664507e6f82b41fe6e20be42929332156e5e85146" dependencies = [ - "bitflags", + "bitflags 2.8.0", +] + +[[package]] +name = "redox_syscall" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a" +dependencies = [ + "bitflags 1.3.2", ] [[package]] @@ -2461,7 +2665,7 @@ version = "0.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03a862b389f93e68874fbf580b9de08dd02facb9a788ebadaf4a3fd33cf58834" dependencies = [ - "bitflags", + "bitflags 2.8.0", ] [[package]] @@ -2580,7 +2784,7 @@ version = "0.38.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" dependencies = [ - "bitflags", + "bitflags 2.8.0", "errno", "libc", "linux-raw-sys", @@ -2691,7 +2895,7 @@ version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "897b2245f0b511c87893af39b033e5ca9cce68824c4d7e7630b5a1d339658d02" dependencies = [ - "bitflags", + "bitflags 2.8.0", "core-foundation 0.9.4", "core-foundation-sys", "libc", @@ -2704,7 +2908,7 @@ version = "3.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "271720403f46ca04f7ba6f55d438f8bd878d6b8ca0a1046e8228c4145bcbb316" dependencies = [ - "bitflags", + "bitflags 2.8.0", "core-foundation 0.10.0", "core-foundation-sys", "libc", @@ -2885,6 +3089,22 @@ dependencies = [ "autocfg", ] +[[package]] +name = "sled" +version = "0.34.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f96b4737c2ce5987354855aed3797279def4ebf734436c6aa4552cf8e169935" +dependencies = [ + "crc32fast", + "crossbeam-epoch", + "crossbeam-utils", + "fs2", + "fxhash", + "libc", + "log", + "parking_lot 0.11.2", +] + [[package]] name = "smallvec" version = "1.13.2" @@ -2976,7 +3196,7 @@ version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c879d448e9d986b661742763247d3693ed13609438cf3d006f51f5368a5ba6b" dependencies = [ - "bitflags", + "bitflags 2.8.0", "core-foundation 0.9.4", "system-configuration-sys", ] @@ -3090,7 +3310,7 @@ dependencies = [ "bytes", "libc", "mio", - "parking_lot", + "parking_lot 0.12.3", "pin-project-lite", "signal-hook-registry", "socket2", @@ -3132,7 +3352,7 @@ dependencies = [ "futures-channel", "futures-util", "log", - "parking_lot", + "parking_lot 0.12.3", "percent-encoding", "phf", "pin-project-lite", @@ -3155,6 +3375,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eca58d7bba4a75707817a2c44174253f9236b2d5fbd055602e9d5c07c139a047" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "tokio-tungstenite" version = "0.26.1" @@ -3216,6 +3447,49 @@ dependencies = [ "winnow", ] +[[package]] +name = "tonic" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85839f0b32fd242bb3209262371d07feda6d780d16ee9d2bc88581b89da1549b" +dependencies = [ + "async-trait", + "axum", + "base64", + "bytes", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "prost", + "socket2", + "tokio", + "tokio-stream", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tonic-build" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d85f0383fadd15609306383a90e85eaed44169f931a5d2be1b42c76ceff1825e" +dependencies = [ + "prettyplease", + "proc-macro2", + "prost-build", + "prost-types", + "quote", + "syn", +] + [[package]] name = "tower" version = "0.5.2" @@ -3224,9 +3498,12 @@ checksum = "d039ad9159c98b70ecfd540b2573b97f7f52c3e8d9f8ad57a24b916a536975f9" dependencies = [ "futures-core", "futures-util", + "indexmap", "pin-project-lite", + "slab", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", "tracing", @@ -3238,7 +3515,7 @@ version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "403fa3b783d4b626a8ad51d766ab03cb6d2dbfc46b1c5d4448395e6628dc9697" dependencies = [ - "bitflags", + "bitflags 2.8.0", "bytes", "http", "http-body", @@ -3592,7 +3869,7 @@ version = "1.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "372d5b87f58ec45c384ba03563b03544dc5fadc3983e434b286913f5b4a9bb6d" dependencies = [ - "redox_syscall", + "redox_syscall 0.5.8", "wasite", "web-sys", ] diff --git a/Cargo.toml b/Cargo.toml index 4b970c57..debb9357 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,5 +8,6 @@ members = [ "lexica", "parakeet", "parakeet-db", + "parakeet-index", "parakeet-lexgen" ] diff --git a/consumer/Cargo.toml b/consumer/Cargo.toml index 24c6d1f9..21f39641 100644 --- a/consumer/Cargo.toml +++ b/consumer/Cargo.toml @@ -20,6 +20,7 @@ lexica = { path = "../lexica" } metrics = "0.24.1" metrics-exporter-prometheus = "0.16.2" parakeet-db = { path = "../parakeet-db" } +parakeet-index = { path = "../parakeet-index" } reqwest = { version = "0.12.12", features = ["native-tls"] } serde = { version = "1.0.217", features = ["derive"] } serde_bytes = "0.11" @@ -27,6 +28,7 @@ serde_ipld_dagcbor = "0.6.1" serde_json = "1.0.134" tokio = { version = "1.42.0", features = ["full"] } tokio-postgres = { version = "0.7.12", features = ["with-chrono-0_4"] } +tokio-stream = "0.1.17" tokio-tungstenite = { version = "0.26.1", features = ["native-tls"] } tracing = "0.1.40" tracing-subscriber = "0.3.18" diff --git a/consumer/run.sh b/consumer/run.sh new file mode 100755 index 00000000..0f4c3cf6 --- /dev/null +++ b/consumer/run.sh @@ -0,0 +1 @@ +cargo run \ No newline at end of file diff --git a/consumer/src/backfill/mod.rs b/consumer/src/backfill/mod.rs index 0dea750f..4f9556a4 100644 --- a/consumer/src/backfill/mod.rs +++ b/consumer/src/backfill/mod.rs @@ -1,5 +1,5 @@ use crate::config::HistoryMode; -use crate::indexer::types::{BackfillItem, BackfillItemInner, CollectionType, RecordTypes}; +use crate::indexer::types::{AggregateDeltaStore, BackfillItem, BackfillItemInner}; use crate::indexer::{self, db as indexer_db}; use did_resolver::Resolver; use diesel_async::pooled_connection::deadpool::Pool; @@ -9,6 +9,7 @@ use ipld_core::cid::Cid; use metrics::counter; use parakeet_db::types::{ActorStatus, ActorSyncState}; use reqwest::{Client, StatusCode}; +use std::collections::HashMap; use std::str::FromStr; use std::sync::Arc; use tracing::{instrument, Instrument}; @@ -18,12 +19,16 @@ mod repo; mod types; const PDS_SERVICE_ID: &str = "#atproto_pds"; +// There's a 4MiB limit on parakeet-index, so break delta batches up if there's loads. +// this should be plenty low enough to not trigger the size limit. (59k did slightly) +const DELTA_BATCH_SIZE: usize = 32 * 1024; #[derive(Clone)] pub struct BackfillManagerInner { pool: Pool, resolver: Arc, client: Client, + index_client: parakeet_index::Client, } pub struct BackfillManager { @@ -37,6 +42,7 @@ impl BackfillManager { pool: Pool, history_mode: HistoryMode, resolver: Arc, + index_client: parakeet_index::Client, ) -> eyre::Result<(Self, Sender)> { let client = Client::new(); @@ -47,6 +53,7 @@ impl BackfillManager { pool, resolver, client, + index_client, }, rx, do_backfill: history_mode == HistoryMode::BackfillHistory, @@ -61,13 +68,13 @@ impl BackfillManager { if self.do_backfill { for idx in 0..threads { let rx = self.rx.clone(); - let inner = self.inner.clone(); + let mut inner = self.inner.clone(); js.spawn( async move { while let Ok(did) = rx.recv_async().await { tracing::trace!("backfilling {did}"); - if let Err(e) = backfill_actor(&inner, &did).await { + if let Err(e) = backfill_actor(&mut inner, &did).await { tracing::error!(did, "backfill failed: {e}"); counter!("backfill_failure").increment(1); } else { @@ -90,7 +97,7 @@ impl BackfillManager { } #[instrument(skip(inner))] -async fn backfill_actor(inner: &BackfillManagerInner, did: &str) -> eyre::Result<()> { +async fn backfill_actor(inner: &mut BackfillManagerInner, did: &str) -> eyre::Result<()> { let mut conn = inner.pool.get().await?; let (status, sync_state) = db::get_actor_status(&mut conn, did).await?; @@ -161,53 +168,76 @@ async fn backfill_actor(inner: &BackfillManagerInner, did: &str) -> eyre::Result tracing::trace!("repo pulled - inserting"); - conn.transaction::<(), diesel::result::Error, _>(|t| { - Box::pin(async move { - db::defer(t).await?; + let delta_store = conn + .transaction::<_, diesel::result::Error, _>(|t| { + Box::pin(async move { + let mut delta_store = HashMap::new(); - indexer_db::update_repo_version(t, did, &rev, cid).await?; + db::defer(t).await?; - let mut follow_stats = vec![did.to_string()]; + indexer_db::update_repo_version(t, did, &rev, cid).await?; - for (path, (cid, record)) in records { - let Some((collection, rkey)) = path.split_once("/") else { - tracing::warn!("record contained invalid path {}", path); - return Err(diesel::result::Error::RollbackTransaction); - }; + // let mut follow_stats = vec![did.to_string()]; - counter!("backfilled_commits", "collection" => collection.to_string()).increment(1); + for (path, (cid, record)) in records { + let Some((collection, rkey)) = path.split_once("/") else { + tracing::warn!("record contained invalid path {}", path); + return Err(diesel::result::Error::RollbackTransaction); + }; - let full_path = format!("at://{did}/{path}"); + counter!("backfilled_commits", "collection" => collection.to_string()) + .increment(1); - match record { - RecordTypes::AppBskyGraphFollow(record) => { - follow_stats.push(record.subject.clone()); - indexer_db::insert_follow(t, did, &full_path, record).await?; - } - _ => indexer::index_op(t, did, cid, record, &full_path, rkey).await?, + let full_path = format!("at://{did}/{path}"); + + indexer::index_op(t, &mut delta_store, did, cid, record, &full_path, rkey) + .await? } - } - db::update_repo_sync_state(t, did, ActorSyncState::Synced).await?; + db::update_repo_sync_state(t, did, ActorSyncState::Synced).await?; - handle_backfill_rows(t, &mut follow_stats, did, &rev).await?; + handle_backfill_rows(t, &mut delta_store, did, &rev).await?; + tracing::trace!("insertion finished"); - // on second thought, should this be done after the transaction? - // if we're loading a chunky repo, we might be a few seconds+ out of date? - indexer_db::update_follow_stats(t, &follow_stats).await?; + Ok(delta_store) + }) + }) + .await?; - tracing::trace!("insertion finished"); - Ok(()) + // submit the deltas + let delta_store = delta_store + .into_iter() + .map(|((uri, typ), delta)| parakeet_index::AggregateDeltaReq { + typ, + uri: uri.to_string(), + delta, }) - }) - .await?; + .collect::>(); + + let mut read = 0; + + while read < delta_store.len() { + let rem = delta_store.len() - read; + let take = DELTA_BATCH_SIZE.min(rem); + + tracing::debug!("reading & submitting {take} deltas"); + + let deltas = delta_store[read..read + take].to_vec(); + inner + .index_client + .submit_aggregate_delta_batch(parakeet_index::AggregateDeltaBatchReq { deltas }) + .await?; + + read += take; + tracing::debug!("read {read} of {} deltas", delta_store.len()); + } Ok(()) } async fn handle_backfill_rows( conn: &mut AsyncPgConnection, - follow_stats: &mut Vec, + deltas: &mut impl AggregateDeltaStore, repo: &str, rev: &str, ) -> diesel::QueryResult<()> { @@ -233,25 +263,12 @@ async fn handle_backfill_rows( continue; }; - match record { - RecordTypes::AppBskyGraphFollow(follow) => { - follow_stats.push(follow.subject.clone()); - indexer_db::insert_follow(conn, repo, &item.at_uri, follow).await?; - } - _ => indexer::index_op(conn, repo, cid, record, &item.at_uri, rkey).await?, - } + indexer::index_op(conn, deltas, repo, cid, record, &item.at_uri, rkey).await? + } + BackfillItemInner::Delete => { + indexer::index_op_delete(conn, deltas, repo, item.collection, &item.at_uri) + .await? } - BackfillItemInner::Delete => match item.collection { - CollectionType::BskyFollow => { - if let Some(subject) = indexer_db::delete_follow(conn, &item.at_uri).await? - { - follow_stats.push(subject); - } - } - _ => { - indexer::index_op_delete(conn, repo, item.collection, &item.at_uri).await? - } - }, } } } diff --git a/consumer/src/config.rs b/consumer/src/config.rs index 0f128daf..d2d8b638 100644 --- a/consumer/src/config.rs +++ b/consumer/src/config.rs @@ -13,6 +13,7 @@ pub(crate) fn load_config() -> eyre::Result { #[derive(Debug, Deserialize)] pub struct Config { + pub index_uri: String, pub relay_source: String, pub database_url: String, pub plc_directory: Option, diff --git a/consumer/src/indexer/db.rs b/consumer/src/indexer/db.rs index 32bd8d87..f283bfa2 100644 --- a/consumer/src/indexer/db.rs +++ b/consumer/src/indexer/db.rs @@ -522,6 +522,25 @@ pub async fn delete_post(conn: &mut AsyncPgConnection, at_uri: &str) -> QueryRes .await } +pub async fn get_post_info_for_delete( + conn: &mut AsyncPgConnection, + at_uri: &str, +) -> QueryResult, Option)>> { + schema::posts::table + .left_join( + schema::post_embed_record::table + .on(schema::posts::at_uri.eq(schema::post_embed_record::post_uri)), + ) + .select(( + schema::posts::parent_uri, + schema::post_embed_record::uri.nullable(), + )) + .filter(schema::posts::at_uri.eq(at_uri)) + .get_result(conn) + .await + .optional() +} + pub async fn upsert_postgate( conn: &mut AsyncPgConnection, at_uri: &str, @@ -642,11 +661,16 @@ pub async fn insert_like( .await } -pub async fn delete_like(conn: &mut AsyncPgConnection, at_uri: &str) -> QueryResult { +pub async fn delete_like( + conn: &mut AsyncPgConnection, + at_uri: &str, +) -> QueryResult> { diesel::delete(schema::likes::table) .filter(schema::likes::at_uri.eq(at_uri)) - .execute(conn) + .returning(schema::likes::subject) + .get_result(conn) .await + .optional() } pub async fn insert_repost( @@ -669,11 +693,16 @@ pub async fn insert_repost( .await } -pub async fn delete_repost(conn: &mut AsyncPgConnection, at_uri: &str) -> QueryResult { +pub async fn delete_repost( + conn: &mut AsyncPgConnection, + at_uri: &str, +) -> QueryResult> { diesel::delete(schema::reposts::table) .filter(schema::reposts::at_uri.eq(at_uri)) - .execute(conn) + .returning(schema::reposts::post) + .get_result(conn) .await + .optional() } pub async fn upsert_chat_decl( diff --git a/consumer/src/indexer/mod.rs b/consumer/src/indexer/mod.rs index 94176060..970c9254 100644 --- a/consumer/src/indexer/mod.rs +++ b/consumer/src/indexer/mod.rs @@ -1,6 +1,8 @@ use crate::config::HistoryMode; use crate::firehose::{AtpAccountEvent, AtpCommitEvent, AtpIdentityEvent, CommitOp, FirehoseEvent}; -use crate::indexer::types::{BackfillItem, BackfillItemInner, CollectionType, RecordTypes}; +use crate::indexer::types::{ + AggregateDeltaStore, BackfillItem, BackfillItemInner, CollectionType, RecordTypes, +}; use did_resolver::Resolver; use diesel_async::pooled_connection::deadpool::Pool; use diesel_async::{AsyncConnection, AsyncPgConnection}; @@ -9,6 +11,7 @@ use futures::StreamExt; use ipld_core::cid::Cid; use metrics::counter; use parakeet_db::types::{ActorStatus, ActorSyncState}; +use parakeet_index::AggregateType; use std::collections::HashMap; use std::hash::BuildHasher; use std::sync::Arc; @@ -22,6 +25,7 @@ pub mod types; #[derive(Clone)] struct RelayIndexerState { backfill_tx: flume::Sender, + idxc_tx: Sender, resolver: Arc, do_backfill: bool, } @@ -37,6 +41,7 @@ impl RelayIndexer { pub async fn new( pool: Pool, backfill_tx: flume::Sender, + idxc_tx: Sender, resolver: Arc, history_mode: HistoryMode, ) -> eyre::Result<(Self, Sender)> { @@ -48,6 +53,7 @@ impl RelayIndexer { backfill_tx, resolver, do_backfill: history_mode == HistoryMode::BackfillHistory, + idxc_tx, }, rx, hasher: RandomState::default(), @@ -60,7 +66,7 @@ impl RelayIndexer { let (submit, _handles) = (0..threads) .map(|idx| { let pool = self.pool.clone(); - let state = self.state.clone(); + let mut state = self.state.clone(); let (tx, mut rx) = channel(16); let handle = tokio::spawn(async move { @@ -76,7 +82,7 @@ impl RelayIndexer { index_account(&state, &mut conn, account).await } FirehoseEvent::Commit(commit) => { - index_commit(&state, &mut conn, commit).await + index_commit(&mut state, &mut conn, commit).await } FirehoseEvent::Label(_) => unreachable!(), }; @@ -186,7 +192,7 @@ async fn index_account( #[instrument(skip_all, fields(seq = commit.seq, repo = commit.repo, rev = commit.rev))] async fn index_commit( - state: &RelayIndexerState, + state: &mut RelayIndexerState, conn: &mut AsyncPgConnection, commit: AtpCommitEvent, ) -> eyre::Result<()> { @@ -262,7 +268,7 @@ async fn index_commit( } for op in &commit.ops { - process_op(t, &commit.repo, op, &blocks).await?; + process_op(t, &mut state.idxc_tx, &commit.repo, op, &blocks).await?; } } else { let items = commit @@ -333,6 +339,7 @@ fn process_op_bf( #[inline(always)] async fn process_op( conn: &mut AsyncPgConnection, + deltas: &mut impl AggregateDeltaStore, repo: &str, op: &CommitOp, blocks: &HashMap>, @@ -361,9 +368,9 @@ async fn process_op( return Ok(()); }; - index_op(conn, repo, cid, decoded, &full_path, rkey).await?; + index_op(conn, deltas, repo, cid, decoded, &full_path, rkey).await?; } else if op.action == "delete" { - index_op_delete(conn, repo, collection, &full_path).await?; + index_op_delete(conn, deltas, repo, collection, &full_path).await?; } else { tracing::warn!("op contained invalid action {}", op.action); } @@ -388,6 +395,7 @@ fn decode_op(op: &CommitOp, blocks: &HashMap>) -> Option<(Cid, Reco pub async fn index_op( conn: &mut AsyncPgConnection, + deltas: &mut impl AggregateDeltaStore, repo: &str, cid: Cid, record: RecordTypes, @@ -407,14 +415,23 @@ pub async fn index_op( } RecordTypes::AppBskyFeedGenerator(record) => { let labels = record.labels.clone(); - db::upsert_feedgen(conn, repo, cid, at_uri, record).await?; + let count = db::upsert_feedgen(conn, repo, cid, at_uri, record).await?; if let Some(labels) = labels { db::maintain_self_labels(conn, repo, Some(cid), at_uri, labels).await?; } + + deltas + .add_delta(repo, AggregateType::ProfileFeed, count as i32) + .await; } RecordTypes::AppBskyFeedLike(record) => { - db::insert_like(conn, repo, at_uri, record).await?; + let subject = record.subject.uri.clone(); + let count = db::insert_like(conn, repo, at_uri, record).await?; + + deltas + .add_delta(&subject, AggregateType::Like, count as i32) + .await; } RecordTypes::AppBskyFeedPost(record) => { if let Some(records::AppBskyEmbed::RecordWithMedia(embed)) = &record.embed { @@ -423,11 +440,26 @@ pub async fn index_op( } } + let maybe_reply = record.reply.as_ref().map(|v| v.parent.uri.clone()); + let maybe_embed = record.embed.as_ref().and_then(|v| match v { + records::AppBskyEmbed::Record(r) => Some(r.record.uri.clone()), + records::AppBskyEmbed::RecordWithMedia(r) => Some(r.record.record.uri.clone()), + _ => None, + }); + let labels = record.labels.clone(); db::insert_post(conn, repo, cid, at_uri, record).await?; if let Some(labels) = labels { db::maintain_self_labels(conn, repo, Some(cid), at_uri, labels).await?; } + + deltas.incr(repo, AggregateType::ProfilePost).await; + if let Some(reply) = maybe_reply { + deltas.incr(&reply, AggregateType::Reply).await; + } + if let Some(embed) = maybe_embed { + deltas.incr(&embed, AggregateType::Embed).await; + } } RecordTypes::AppBskyFeedPostgate(record) => { let split_aturi = record.post.rsplitn(4, '/').collect::>(); @@ -452,6 +484,9 @@ pub async fn index_op( .await?; } RecordTypes::AppBskyFeedRepost(record) => { + deltas + .incr(&record.subject.uri, AggregateType::Repost) + .await; db::insert_repost(conn, repo, at_uri, record).await?; } RecordTypes::AppBskyFeedThreadgate(record) => { @@ -467,17 +502,27 @@ pub async fn index_op( db::insert_block(conn, repo, at_uri, record).await?; } RecordTypes::AppBskyGraphFollow(record) => { - db::insert_follow(conn, repo, at_uri, record).await?; + let subject = record.subject.clone(); + let count = db::insert_follow(conn, repo, at_uri, record).await?; + + deltas + .add_delta(repo, AggregateType::Follow, count as i32) + .await; + deltas + .add_delta(&subject, AggregateType::Follower, count as i32) + .await; } RecordTypes::AppBskyGraphList(record) => { let labels = record.labels.clone(); - db::upsert_list(conn, repo, at_uri, cid, record).await?; + let count = db::upsert_list(conn, repo, at_uri, cid, record).await?; if let Some(labels) = labels { db::maintain_self_labels(conn, repo, Some(cid), at_uri, labels).await?; } - // todo: when we have profile stats, update them. + deltas + .add_delta(repo, AggregateType::ProfileList, count as i32) + .await; } RecordTypes::AppBskyGraphListBlock(record) => { db::insert_list_block(conn, repo, at_uri, record).await?; @@ -493,7 +538,10 @@ pub async fn index_op( db::insert_list_item(conn, at_uri, record).await?; } RecordTypes::AppBskyGraphStarterPack(record) => { - db::upsert_starterpack(conn, repo, cid, at_uri, record).await?; + let count = db::upsert_starterpack(conn, repo, cid, at_uri, record).await?; + deltas + .add_delta(repo, AggregateType::ProfileStarterpack, count as i32) + .await; } RecordTypes::AppBskyGraphVerification(record) => { db::upsert_verification(conn, repo, cid, at_uri, record).await?; @@ -520,6 +568,7 @@ pub async fn index_op( pub async fn index_op_delete( conn: &mut AsyncPgConnection, + deltas: &mut impl AggregateDeltaStore, repo: &str, collection: CollectionType, at_uri: &str, @@ -527,23 +576,67 @@ pub async fn index_op_delete( match collection { CollectionType::BskyProfile => db::delete_profile(conn, repo).await?, CollectionType::BskyBlock => db::delete_block(conn, at_uri).await?, - CollectionType::BskyFeedGen => db::delete_feedgen(conn, at_uri).await?, - CollectionType::BskyFeedLike => db::delete_like(conn, at_uri).await?, - CollectionType::BskyFeedPost => db::delete_post(conn, at_uri).await?, + CollectionType::BskyFeedGen => { + let count = db::delete_feedgen(conn, at_uri).await?; + deltas + .add_delta(repo, AggregateType::ProfileFeed, -(count as i32)) + .await; + count + } + CollectionType::BskyFeedLike => { + if let Some(subject) = db::delete_like(conn, at_uri).await? { + deltas.decr(&subject, AggregateType::Like).await; + } + 0 + } + CollectionType::BskyFeedPost => { + let post_info = db::get_post_info_for_delete(conn, at_uri).await?; + + db::delete_post(conn, at_uri).await?; + + if let Some((reply_to, embed)) = post_info { + deltas.decr(repo, AggregateType::ProfilePost).await; + if let Some(reply_to) = reply_to { + deltas.decr(&reply_to, AggregateType::Reply).await; + } + if let Some(embed) = embed { + deltas.decr(&embed, AggregateType::Embed).await; + } + } + + 0 + } CollectionType::BskyFeedPostgate => db::delete_postgate(conn, at_uri).await?, - CollectionType::BskyFeedRepost => db::delete_repost(conn, at_uri).await?, + CollectionType::BskyFeedRepost => { + if let Some(subject) = db::delete_repost(conn, at_uri).await? { + deltas.decr(&subject, AggregateType::Repost).await; + } + 0 + } CollectionType::BskyFeedThreadgate => db::delete_threadgate(conn, at_uri).await?, CollectionType::BskyFollow => { - db::delete_follow(conn, at_uri).await?; + if let Some(followee) = db::delete_follow(conn, at_uri).await? { + deltas.decr(&followee, AggregateType::Follower).await; + deltas.decr(repo, AggregateType::Follow).await; + } 0 } CollectionType::BskyList => { - db::delete_list(conn, at_uri).await? - // todo: when we have profile stats, update them. + let count = db::delete_list(conn, at_uri).await?; + deltas + .add_delta(repo, AggregateType::ProfileList, -(count as i32)) + .await; + count } CollectionType::BskyListBlock => db::delete_list_block(conn, at_uri).await?, CollectionType::BskyListItem => db::delete_list_item(conn, at_uri).await?, - CollectionType::BskyStarterPack => db::delete_starterpack(conn, at_uri).await?, + CollectionType::BskyStarterPack => { + let count = db::delete_starterpack(conn, at_uri).await?; + deltas + .add_delta(repo, AggregateType::ProfileStarterpack, -(count as i32)) + .await; + count + } CollectionType::BskyVerification => db::delete_verification(conn, at_uri).await?, CollectionType::BskyLabelerService => db::delete_label_service(conn, at_uri).await?, CollectionType::ChatActorDecl => db::delete_chat_decl(conn, at_uri).await?, diff --git a/consumer/src/indexer/types.rs b/consumer/src/indexer/types.rs index 835175b8..05aa9d04 100644 --- a/consumer/src/indexer/types.rs +++ b/consumer/src/indexer/types.rs @@ -121,3 +121,36 @@ pub enum BackfillItemInner { Update(RecordTypes), Delete, } + +pub trait AggregateDeltaStore { + async fn add_delta(&mut self, uri: &str, typ: parakeet_index::AggregateType, delta: i32); + async fn incr(&mut self, uri: &str, typ: parakeet_index::AggregateType) { + self.add_delta(uri, typ, 1).await + } + async fn decr(&mut self, uri: &str, typ: parakeet_index::AggregateType) { + self.add_delta(uri, typ, -1).await + } +} + +impl AggregateDeltaStore for tokio::sync::mpsc::Sender { + async fn add_delta(&mut self, uri: &str, typ: parakeet_index::AggregateType, delta: i32) { + let res = self + .send(parakeet_index::AggregateDeltaReq { + typ: typ.into(), + uri: uri.to_string(), + delta, + }) + .await; + + if let Err(e) = res { + tracing::error!("failed to send aggregate delta: {e}"); + } + } +} + +impl AggregateDeltaStore for std::collections::HashMap<(String, i32), i32> { + async fn add_delta(&mut self, uri: &str, typ: parakeet_index::AggregateType, delta: i32) { + let key = (uri.to_string(), typ.into()); + self.entry(key).and_modify(|v| *v += delta).or_insert(delta); + } +} diff --git a/consumer/src/main.rs b/consumer/src/main.rs index f43902e3..eef2c02d 100644 --- a/consumer/src/main.rs +++ b/consumer/src/main.rs @@ -31,6 +31,8 @@ async fn main() -> eyre::Result<()> { ..Default::default() })?); + let index_client = parakeet_index::Client::connect(conf.index_uri).await?; + let (label_mgr, label_svc_tx) = label_indexer::LabelServiceManager::new( &conf.database_url, resolver.clone(), @@ -42,25 +44,30 @@ async fn main() -> eyre::Result<()> { let (backfiller, backfill_tx) = backfill::BackfillManager::new(pool.clone(), conf.history_mode, resolver.clone()).await?; + let (idxc_tx, idxc_rx) = tokio::sync::mpsc::channel(128); + let (relay_indexer, tx) = indexer::RelayIndexer::new( pool.clone(), backfill_tx, + idxc_tx, resolver.clone(), conf.history_mode, ) .await?; - let (firehose_res, indexer_res, backfill_res, label_res) = tokio::try_join! { + let (firehose_res, indexer_res, backfill_res, label_res, idxt_res) = tokio::try_join! { tokio::spawn(relay_consumer(relay_firehose, tx)), tokio::spawn(relay_indexer.run(conf.indexer_workers)), tokio::spawn(backfiller.run(conf.backfill_workers)), tokio::spawn(label_mgr.run(conf.initial_label_services)), + tokio::spawn(index_transport(index_client, idxc_rx)), }?; firehose_res .and(indexer_res) .and(backfill_res) .and(label_res) + .and(idxt_res) } async fn relay_consumer( @@ -87,6 +94,14 @@ async fn relay_consumer( Ok(()) } +async fn index_transport(mut idxc: parakeet_index::Client, rx: tokio::sync::mpsc::Receiver) -> eyre::Result<()> { + use tokio_stream::wrappers::ReceiverStream; + + idxc.submit_aggregate_delta_stream(ReceiverStream::new(rx)).await?; + + Ok(()) +} + fn build_ua(contact: &Option) -> String { let mut ua = format!("Parakeet {}", env!("CARGO_PKG_VERSION")); diff --git a/parakeet-index/Cargo.toml b/parakeet-index/Cargo.toml new file mode 100644 index 00000000..c9343cf8 --- /dev/null +++ b/parakeet-index/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "parakeet-index" +version = "0.1.0" +edition = "2024" + +[[bin]] +name = "parakeet-index" +required-features = ["server"] + +[dependencies] +tonic = "0.13.0" +prost = "0.13.5" + +eyre = { version = "0.6.12", optional = true } +figment = { version = "0.10.19", features = ["env", "toml"], optional = true } +itertools = { version = "0.14.0", optional = true } +serde = { version = "1.0.217", features = ["derive"], optional = true } +sled = { version = "0.34.7", optional = true } +tokio = { version = "1.42.0", features = ["full"], optional = true } +tracing = { version = "0.1.40", optional = true } +tracing-subscriber = { version = "0.3.18", optional = true } + +[build-dependencies] +tonic-build = "0.13.0" + +[features] +server = ["dep:eyre", "dep:figment", "dep:itertools", "dep:serde", "dep:sled", "dep:tokio", "dep:tracing", "dep:tracing-subscriber"] \ No newline at end of file diff --git a/parakeet-index/build.rs b/parakeet-index/build.rs new file mode 100644 index 00000000..ec77e8a7 --- /dev/null +++ b/parakeet-index/build.rs @@ -0,0 +1,5 @@ +fn main() -> Result<(), Box> { + tonic_build::configure().compile_protos(&["proto/parakeet.proto"], &[""])?; + + Ok(()) +} diff --git a/parakeet-index/proto/parakeet.proto b/parakeet-index/proto/parakeet.proto new file mode 100644 index 00000000..79cc5b19 --- /dev/null +++ b/parakeet-index/proto/parakeet.proto @@ -0,0 +1,96 @@ +syntax = "proto3"; +package parakeet; + +service Index { + rpc SubmitAggregateDelta(AggregateDeltaReq) returns (AggregateDeltaRes); + rpc SubmitAggregateDeltaBatch(AggregateDeltaBatchReq) returns (AggregateDeltaRes); + rpc SubmitAggregateDeltaStream(stream AggregateDeltaReq) returns (AggregateDeltaRes); + + rpc GetProfileStats(GetStatsReq) returns (GetProfileStatsRes); + rpc GetProfileStatsMany(GetStatsManyReq) returns (GetProfileStatsManyRes); + rpc GetPostStats(GetStatsReq) returns (GetPostStatsRes); + rpc GetPostStatsMany(GetStatsManyReq) returns (GetPostStatsManyRes); + rpc GetLikeCount(GetStatsReq) returns (GetLikeCountRes); + rpc GetLikeCountMany(GetStatsManyReq) returns (GetLikeCountManyRes); +} + +enum AggregateType { + UNKNOWN = 0; + FOLLOW = 1; + FOLLOWER = 2; + LIKE = 3; + REPLY = 4; + REPOST = 5; + // aka Quotes (in the context of posts) + EMBED = 6; + PROFILE_POST = 7; + PROFILE_LIST = 8; + PROFILE_FEED = 9; + PROFILE_STARTERPACK = 10; +} + +message AggregateDeltaReq { + // The type of aggregate to change + AggregateType typ = 1; + // The entry to change. Can be a full at:// uri for items or a did for actors/profiles + string uri = 2; + sint32 delta = 3; +} + +message AggregateDeltaBatchReq { + repeated AggregateDeltaReq deltas = 1; +} + +message AggregateDeltaRes {} + +message GetStatsReq { + string uri = 1; +} + +message GetStatsManyReq { + repeated string uris = 1; +} + +message ProfileStats { + int32 followers = 1; + int32 following = 2; + int32 posts = 3; + int32 lists = 4; + int32 feeds = 5; + int32 starterpacks = 6; +} + +message GetProfileStatsRes { + optional ProfileStats stats = 1; +} + +message GetProfileStatsManyRes { + map entries = 1; +} + +message PostStats { + int32 replies = 1; + int32 likes = 2; + int32 reposts = 3; + int32 quotes = 4; +} + +message GetPostStatsRes { + optional PostStats stats = 1; +} + +message GetPostStatsManyRes { + map entries = 1; +} + +message LikeCount { + int32 likes = 1; +} + +message GetLikeCountRes { + optional LikeCount likes = 1; +} + +message GetLikeCountManyRes { + map entries = 1; +} diff --git a/parakeet-index/run.sh b/parakeet-index/run.sh new file mode 100755 index 00000000..56ffa0ea --- /dev/null +++ b/parakeet-index/run.sh @@ -0,0 +1 @@ +cargo run --features server \ No newline at end of file diff --git a/parakeet-index/src/lib.rs b/parakeet-index/src/lib.rs new file mode 100644 index 00000000..9cc88ae3 --- /dev/null +++ b/parakeet-index/src/lib.rs @@ -0,0 +1,10 @@ +#[allow(clippy::all)] +pub mod index { + tonic::include_proto!("parakeet"); +} + +pub use index::*; +pub type Client = index_client::IndexClient; + +#[cfg(feature = "server")] +pub mod server; diff --git a/parakeet-index/src/main.rs b/parakeet-index/src/main.rs new file mode 100644 index 00000000..97ccbcaf --- /dev/null +++ b/parakeet-index/src/main.rs @@ -0,0 +1,24 @@ +use parakeet_index::index_server::IndexServer; +use parakeet_index::server::service::Service; +use parakeet_index::server::{GlobalState, config}; +use std::sync::Arc; +use tonic::transport::Server; + +#[tokio::main] +async fn main() -> eyre::Result<()> { + tracing_subscriber::fmt::init(); + + let conf = config::load_config()?; + + let db_root = conf.index_db_path.parse()?; + let addr = std::net::SocketAddr::new(conf.server.bind_address.parse()?, conf.server.port); + let state = Arc::new(GlobalState::new(db_root)?); + + let service = Service::new(state.clone()); + Server::builder() + .add_service(IndexServer::new(service)) + .serve(addr) + .await?; + + Ok(()) +} diff --git a/parakeet-index/src/server/config.rs b/parakeet-index/src/server/config.rs new file mode 100644 index 00000000..f4a4279a --- /dev/null +++ b/parakeet-index/src/server/config.rs @@ -0,0 +1,45 @@ +use figment::Figment; +use figment::providers::{Env, Format, Toml}; +use serde::Deserialize; + +pub fn load_config() -> eyre::Result { + let conf = Figment::new() + .merge(Toml::file("Config.toml")) + .merge(Env::prefixed("PKI_")) + .extract()?; + + Ok(conf) +} + +#[derive(Debug, Deserialize)] +pub struct Config { + pub database_url: String, + pub index_db_path: String, + #[serde(default)] + pub server: ConfigServer, +} + +#[derive(Debug, Deserialize)] +pub struct ConfigServer { + #[serde(default = "default_bind_address")] + pub bind_address: String, + #[serde(default = "default_port")] + pub port: u16, +} + +impl Default for ConfigServer { + fn default() -> Self { + ConfigServer { + bind_address: default_bind_address(), + port: default_port(), + } + } +} + +fn default_bind_address() -> String { + "0.0.0.0".to_string() +} + +fn default_port() -> u16 { + 6001 +} diff --git a/parakeet-index/src/server/db.rs b/parakeet-index/src/server/db.rs new file mode 100644 index 00000000..2fa86ced --- /dev/null +++ b/parakeet-index/src/server/db.rs @@ -0,0 +1,103 @@ +use crate::all_none; +use crate::server::utils::{ToIntExt, TreeExt, slice_as_i32}; +use sled::{Db, MergeOperator, Tree}; +use std::path::PathBuf; + +pub struct DbStore { + pub agg_db: Db, + pub label_db: Db, + + pub follows: Tree, + pub followers: Tree, + pub likes: Tree, + pub replies: Tree, + pub reposts: Tree, + pub embeds: Tree, + pub profile_posts: Tree, + pub profile_lists: Tree, + pub profile_feeds: Tree, + pub profile_starterpacks: Tree, +} + +impl DbStore { + pub fn new(db_root: PathBuf) -> eyre::Result { + let agg_db = sled::open(db_root.join("aggdb"))?; + let label_db = sled::open(db_root.join("labeldb"))?; + + Ok(DbStore { + follows: open_tree(&agg_db, "follows", merge_delta)?, + followers: open_tree(&agg_db, "followers", merge_delta)?, + likes: open_tree(&agg_db, "likes", merge_delta)?, + replies: open_tree(&agg_db, "replies", merge_delta)?, + reposts: open_tree(&agg_db, "reposts", merge_delta)?, + embeds: open_tree(&agg_db, "embeds", merge_delta)?, + profile_posts: open_tree(&agg_db, "profile_posts", merge_delta)?, + profile_lists: open_tree(&agg_db, "profile_lists", merge_delta)?, + profile_feeds: open_tree(&agg_db, "profile_feeds", merge_delta)?, + profile_starterpacks: open_tree(&agg_db, "profile_starterpacks", merge_delta)?, + + agg_db, + label_db, + }) + } + + pub fn get_post_stats(&self, post: &str) -> Option { + let replies = self.replies.get_i32(post); + let likes = self.likes.get_i32(post); + let reposts = self.reposts.get_i32(post); + let quotes = self.embeds.get_i32(post); + + if all_none![replies, likes, reposts, quotes] { + return None; + } + + Some(crate::PostStats { + replies: replies.unwrap_or_default(), + likes: likes.unwrap_or_default(), + reposts: reposts.unwrap_or_default(), + quotes: quotes.unwrap_or_default(), + }) + } + + pub fn get_profile_stats(&self, did: &str) -> Option { + let followers = self.followers.get_i32(did); + let following = self.follows.get_i32(did); + let posts = self.profile_posts.get_i32(did); + let lists = self.profile_lists.get_i32(did); + let feeds = self.profile_feeds.get_i32(did); + let starterpacks = self.profile_starterpacks.get_i32(did); + + if all_none![followers, following, posts, lists, feeds, starterpacks] { + return None; + } + + Some(crate::ProfileStats { + followers: followers.unwrap_or_default(), + following: following.unwrap_or_default(), + posts: posts.unwrap_or_default(), + lists: lists.unwrap_or_default(), + feeds: feeds.unwrap_or_default(), + starterpacks: starterpacks.unwrap_or_default(), + }) + } +} + +fn open_tree(db: &Db, name: &str, merge: impl MergeOperator + 'static) -> eyre::Result { + let tree = db.open_tree(name)?; + + tree.set_merge_operator(merge); + + Ok(tree) +} + +fn merge_delta(_key: &[u8], old: Option<&[u8]>, new: &[u8]) -> Option> { + let old = old.and_then(slice_as_i32); + let new = slice_as_i32(new)?; + + let res = match old { + Some(old) => old + new, + None => new, + }; + + Some(Vec::from_i32(res)) +} diff --git a/parakeet-index/src/server/mod.rs b/parakeet-index/src/server/mod.rs new file mode 100644 index 00000000..c19f5916 --- /dev/null +++ b/parakeet-index/src/server/mod.rs @@ -0,0 +1,18 @@ +use std::path::PathBuf; + +pub mod config; +pub mod db; +pub mod service; +mod utils; + +pub struct GlobalState { + pub dbs: db::DbStore, +} + +impl GlobalState { + pub fn new(db_root: PathBuf) -> eyre::Result { + let dbs = db::DbStore::new(db_root)?; + + Ok(GlobalState { dbs }) + } +} diff --git a/parakeet-index/src/server/service.rs b/parakeet-index/src/server/service.rs new file mode 100644 index 00000000..cd91963f --- /dev/null +++ b/parakeet-index/src/server/service.rs @@ -0,0 +1,201 @@ +use crate::index::*; +use crate::server::GlobalState; +use crate::server::utils::TreeExt; +use std::collections::HashMap; +use std::ops::Deref; +use std::sync::Arc; +use tonic::codegen::tokio_stream::StreamExt; +use tonic::{Request, Response, Status, Streaming, async_trait}; + +pub struct Service(Arc); + +impl Service { + pub fn new(state: Arc) -> Self { + Service(state) + } + + fn apply_delta( + &self, + uri: &str, + typ: AggregateType, + delta: i32, + ) -> sled::Result> { + let val = delta.to_le_bytes(); + + match typ { + AggregateType::Unknown => todo!(), + AggregateType::Follow => self.dbs.follows.merge(uri, val), + AggregateType::Follower => self.dbs.followers.merge(uri, val), + AggregateType::Like => self.dbs.likes.merge(uri, val), + AggregateType::Reply => self.dbs.replies.merge(uri, val), + AggregateType::Repost => self.dbs.reposts.merge(uri, val), + AggregateType::Embed => self.dbs.embeds.merge(uri, val), + AggregateType::ProfilePost => self.dbs.profile_posts.merge(uri, val), + AggregateType::ProfileList => self.dbs.profile_lists.merge(uri, val), + AggregateType::ProfileFeed => self.dbs.profile_feeds.merge(uri, val), + AggregateType::ProfileStarterpack => self.dbs.profile_starterpacks.merge(uri, val), + } + } +} + +impl Deref for Service { + type Target = Arc; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +#[async_trait] +impl index_server::Index for Service { + async fn submit_aggregate_delta( + &self, + request: Request, + ) -> Result, Status> { + let inner = request.into_inner(); + + let res = self.apply_delta(&inner.uri, inner.typ(), inner.delta); + + if let Err(e) = res { + tracing::error!("failed to update stats DB: {e}"); + return Err(Status::unknown("failed to update stats DB")); + } + + Ok(Response::new(AggregateDeltaRes {})) + } + + async fn submit_aggregate_delta_batch( + &self, + request: Request, + ) -> Result, Status> { + let inner = request.into_inner(); + + for data in inner.deltas { + let res = self.apply_delta(&data.uri, data.typ(), data.delta); + + if let Err(e) = res { + tracing::error!("failed to update stats DB: {e}"); + return Err(Status::unknown("failed to update stats DB")); + } + } + + Ok(Response::new(AggregateDeltaRes {})) + } + + async fn submit_aggregate_delta_stream( + &self, + request: Request>, + ) -> Result, Status> { + let mut inner = request.into_inner(); + + while let Some(req) = inner.next().await { + if let Ok(data) = req { + let res = self.apply_delta(&data.uri, data.typ(), data.delta); + + if let Err(e) = res { + tracing::error!("failed to update stats DB: {e}"); + return Err(Status::unknown("failed to update stats DB")); + } + } else { + tracing::error!("failed to read stream item") + } + } + + Ok(Response::new(AggregateDeltaRes {})) + } + + async fn get_profile_stats( + &self, + request: Request, + ) -> Result, Status> { + let inner = request.into_inner(); + + let stats = self.dbs.get_profile_stats(&inner.uri); + + Ok(Response::new(GetProfileStatsRes { stats })) + } + + async fn get_profile_stats_many( + &self, + request: Request, + ) -> Result, Status> { + let inner = request.into_inner(); + + // idk if this is the best way of doing this???? + let entries = inner + .uris + .into_iter() + .filter_map(|uri| { + let stats = self.dbs.get_profile_stats(&uri)?; + + Some((uri, stats)) + }) + .collect::>(); + + Ok(Response::new(GetProfileStatsManyRes { entries })) + } + + async fn get_post_stats( + &self, + request: Request, + ) -> Result, Status> { + let inner = request.into_inner(); + + let stats = self.dbs.get_post_stats(&inner.uri); + + Ok(Response::new(GetPostStatsRes { stats })) + } + + async fn get_post_stats_many( + &self, + request: Request, + ) -> Result, Status> { + let inner = request.into_inner(); + + let entries = inner + .uris + .into_iter() + .filter_map(|uri| { + let stats = self.dbs.get_post_stats(&uri)?; + + Some((uri, stats)) + }) + .collect::>(); + + Ok(Response::new(GetPostStatsManyRes { entries })) + } + + async fn get_like_count( + &self, + request: Request, + ) -> Result, Status> { + let inner = request.into_inner(); + + let likes = self + .dbs + .likes + .get_i32(inner.uri) + .map(|likes| LikeCount { likes }); + + Ok(Response::new(GetLikeCountRes { likes })) + } + + async fn get_like_count_many( + &self, + request: Request, + ) -> Result, Status> { + let inner = request.into_inner(); + + let entries = inner + .uris + .into_iter() + .filter_map(|uri| { + let likes = self.dbs.likes.get_i32(&uri)?; + + Some((uri, LikeCount { likes })) + }) + .collect(); + + Ok(Response::new(GetLikeCountManyRes { entries })) + } +} diff --git a/parakeet-index/src/server/utils.rs b/parakeet-index/src/server/utils.rs new file mode 100644 index 00000000..94a5c276 --- /dev/null +++ b/parakeet-index/src/server/utils.rs @@ -0,0 +1,59 @@ +use sled::{IVec, Tree}; + +pub trait ToIntExt { + fn as_i32(&self) -> Option; + fn from_i32(i: i32) -> Self; +} + +impl ToIntExt for IVec { + fn as_i32(&self) -> Option { + if self.len() == 4 { + let bytes = self[0..4].try_into().ok()?; + Some(i32::from_le_bytes(bytes)) + } else { + None + } + } + + fn from_i32(i: i32) -> Self { + IVec::from(&i.to_le_bytes()) + } +} + +impl ToIntExt for Vec { + fn as_i32(&self) -> Option { + if self.len() == 4 { + let bytes = self[0..4].try_into().ok()?; + Some(i32::from_le_bytes(bytes)) + } else { + None + } + } + + fn from_i32(i: i32) -> Self { + Vec::from(&i.to_le_bytes()) + } +} + +pub fn slice_as_i32(data: &[u8]) -> Option { + let bytes = data[0..4].try_into().ok()?; + + Some(i32::from_le_bytes(bytes)) +} + +pub trait TreeExt { + fn get_i32(&self, key: impl AsRef<[u8]>) -> Option; +} + +impl TreeExt for Tree { + fn get_i32(&self, key: impl AsRef<[u8]>) -> Option { + self.get(key).ok().flatten().and_then(|v| v.as_i32()) + } +} + +#[macro_export] +macro_rules! all_none { + ($var0:ident, $($var:ident),*) => { + $var0.is_none() $(&& $var.is_none())* + }; +} diff --git a/parakeet/Cargo.toml b/parakeet/Cargo.toml index c2275e8d..d40cc180 100644 --- a/parakeet/Cargo.toml +++ b/parakeet/Cargo.toml @@ -17,6 +17,7 @@ figment = { version = "0.10.19", features = ["env", "toml"] } itertools = "0.14.0" lexica = { path = "../lexica" } parakeet-db = { path = "../parakeet-db" } +parakeet-index = { path = "../parakeet-index" } serde = { version = "1.0.217", features = ["derive"] } serde_json = "1.0.134" tokio = { version = "1.42.0", features = ["full"] } diff --git a/parakeet/run.sh b/parakeet/run.sh new file mode 100755 index 00000000..0f4c3cf6 --- /dev/null +++ b/parakeet/run.sh @@ -0,0 +1 @@ +cargo run \ No newline at end of file diff --git a/parakeet/src/config.rs b/parakeet/src/config.rs index 4bc15cad..72095a76 100644 --- a/parakeet/src/config.rs +++ b/parakeet/src/config.rs @@ -13,6 +13,7 @@ pub(crate) fn load_config() -> eyre::Result { #[derive(Debug, Deserialize)] pub struct Config { + pub index_uri: String, pub database_url: String, #[serde(default)] pub server: ConfigServer, diff --git a/parakeet/src/hydration/feedgen.rs b/parakeet/src/hydration/feedgen.rs index 1d58f0b4..adf37a83 100644 --- a/parakeet/src/hydration/feedgen.rs +++ b/parakeet/src/hydration/feedgen.rs @@ -11,6 +11,7 @@ fn build_feedgen( feedgen: models::FeedGen, creator: ProfileView, labels: Vec, + likes: Option, ) -> GeneratorView { let content_mode = feedgen .content_mode @@ -31,7 +32,7 @@ fn build_feedgen( avatar: feedgen .avatar_cid .map(|v| format!("https://localhost/feedgen/{v}")), - like_count: 0, + like_count: likes.unwrap_or_default() as i64, accepts_interactions: feedgen.accepts_interactions, labels: map_labels(labels), content_mode, @@ -45,10 +46,10 @@ pub async fn hydrate_feedgen( apply_labelers: &[LabelConfigItem], ) -> Option { let labels = loaders.label.load(&feedgen, apply_labelers).await; - let feedgen = loaders.feedgen.load(feedgen).await?; + let (feedgen, likes) = loaders.feedgen.load(feedgen).await?; let profile = hydrate_profile(loaders, feedgen.owner.clone(), apply_labelers).await?; - Some(build_feedgen(feedgen, profile, labels)) + Some(build_feedgen(feedgen, profile, labels, likes)) } pub async fn hydrate_feedgens( @@ -61,17 +62,20 @@ pub async fn hydrate_feedgens( let creators = feedgens .values() - .map(|feedgen| feedgen.owner.clone()) + .map(|(feedgen, _)| feedgen.owner.clone()) .collect(); let creators = hydrate_profiles(loaders, creators, apply_labelers).await; feedgens .into_iter() - .filter_map(|(uri, feedgen)| { + .filter_map(|(uri, (feedgen, likes))| { let creator = creators.get(&feedgen.owner)?; let labels = labels.get(&uri).cloned().unwrap_or_default(); - Some((uri, build_feedgen(feedgen, creator.to_owned(), labels))) + Some(( + uri, + build_feedgen(feedgen, creator.to_owned(), labels, likes), + )) }) .collect() } diff --git a/parakeet/src/hydration/labeler.rs b/parakeet/src/hydration/labeler.rs index b380be34..33813abb 100644 --- a/parakeet/src/hydration/labeler.rs +++ b/parakeet/src/hydration/labeler.rs @@ -12,12 +12,13 @@ fn build_view( labeler: models::LabelerService, creator: ProfileView, labels: Vec, + likes: Option, ) -> LabelerView { LabelerView { uri: format!("at://{}/app.bsky.labeler.service/self", labeler.did), cid: labeler.cid, creator, - like_count: 0, + like_count: likes.unwrap_or_default() as i64, labels: map_labels(labels), indexed_at: labeler.indexed_at, } @@ -28,6 +29,7 @@ fn build_view_detailed( defs: Vec, creator: ProfileView, labels: Vec, + likes: Option, ) -> LabelerViewDetailed { let reason_types = labeler.reasons.map(|v| { v.into_iter() @@ -66,7 +68,7 @@ fn build_view_detailed( uri: format!("at://{}/app.bsky.labeler.service/self", labeler.did), cid: labeler.cid, creator, - like_count: 0, + like_count: likes.unwrap_or_default() as i64, policies: LabelerPolicy { label_values, label_value_definitions, @@ -85,10 +87,10 @@ pub async fn hydrate_labeler( apply_labelers: &[LabelConfigItem], ) -> Option { let labels = loaders.label.load(&labeler, apply_labelers).await; - let (labeler, _) = loaders.labeler.load(labeler).await?; + let (labeler, _, likes) = loaders.labeler.load(labeler).await?; let creator = hydrate_profile(loaders, labeler.did.clone(), apply_labelers).await?; - Some(build_view(labeler, creator, labels)) + Some(build_view(labeler, creator, labels, likes)) } pub async fn hydrate_labelers( @@ -101,17 +103,17 @@ pub async fn hydrate_labelers( let creators = labelers .values() - .map(|(labeler, _)| labeler.did.clone()) + .map(|(labeler, _, _)| labeler.did.clone()) .collect(); let creators = hydrate_profiles(loaders, creators, apply_labelers).await; labelers .into_iter() - .filter_map(|(k, (labeler, _))| { + .filter_map(|(k, (labeler, _, likes))| { let creator = creators.get(&labeler.did).cloned()?; let labels = labels.get(&k).cloned().unwrap_or_default(); - Some((k, build_view(labeler, creator, labels))) + Some((k, build_view(labeler, creator, labels, likes))) }) .collect() } @@ -122,10 +124,10 @@ pub async fn hydrate_labeler_detailed( apply_labelers: &[LabelConfigItem], ) -> Option { let labels = loaders.label.load(&labeler, apply_labelers).await; - let (labeler, defs) = loaders.labeler.load(labeler).await?; + let (labeler, defs, likes) = loaders.labeler.load(labeler).await?; let creator = hydrate_profile(loaders, labeler.did.clone(), apply_labelers).await?; - Some(build_view_detailed(labeler, defs, creator, labels)) + Some(build_view_detailed(labeler, defs, creator, labels, likes)) } pub async fn hydrate_labelers_detailed( @@ -138,17 +140,20 @@ pub async fn hydrate_labelers_detailed( let creators = labelers .values() - .map(|(labeler, _)| labeler.did.clone()) + .map(|(labeler, _, _)| labeler.did.clone()) .collect(); let creators = hydrate_profiles(loaders, creators, apply_labelers).await; labelers .into_iter() - .filter_map(|(k, (labeler, defs))| { + .filter_map(|(k, (labeler, defs, likes))| { let creator = creators.get(&labeler.did).cloned()?; let labels = labels.get(&k).cloned().unwrap_or_default(); - Some((k, build_view_detailed(labeler, defs, creator, labels))) + Some(( + k, + build_view_detailed(labeler, defs, creator, labels, likes), + )) }) .collect() } diff --git a/parakeet/src/hydration/posts.rs b/parakeet/src/hydration/posts.rs index 6007d5e1..58d56bcf 100644 --- a/parakeet/src/hydration/posts.rs +++ b/parakeet/src/hydration/posts.rs @@ -7,7 +7,9 @@ use lexica::app_bsky::actor::ProfileViewBasic; use lexica::app_bsky::embed::{AspectRatio, Embed}; use lexica::app_bsky::feed::{FeedViewPost, PostView, ReplyRef, ReplyRefPost, ThreadgateView}; use lexica::app_bsky::graph::ListViewBasic; +use lexica::app_bsky::RecordStats; use parakeet_db::models; +use parakeet_index::PostStats; use std::collections::HashMap; fn build_postview( @@ -16,14 +18,24 @@ fn build_postview( labels: Vec, embed: Option, threadgate: Option, + stats: Option, ) -> PostView { + let stats = stats + .map(|stats| RecordStats { + reply_count: stats.replies as i64, + repost_count: stats.reposts as i64, + like_count: stats.likes as i64, + quote_count: stats.quotes as i64, + }) + .unwrap_or_default(); + PostView { uri: post.at_uri, cid: post.cid, author, record: post.record, embed, - stats: Default::default(), + stats, labels: map_labels(labels), threadgate, indexed_at: post.created_at, @@ -96,13 +108,15 @@ pub async fn hydrate_post( post: String, apply_labelers: &[LabelConfigItem], ) -> Option { - let (post, threadgate) = loaders.posts.load(post).await?; + let (post, threadgate, stats) = loaders.posts.load(post).await?; let embed = hydrate_embed(loaders, post.at_uri.clone(), apply_labelers).await; let author = hydrate_profile_basic(loaders, post.did.clone(), apply_labelers).await?; let threadgate = hydrate_threadgate(loaders, threadgate, apply_labelers).await; let labels = loaders.label.load(&post.at_uri, apply_labelers).await; - Some(build_postview(post, author, labels, embed, threadgate)) + Some(build_postview( + post, author, labels, embed, threadgate, stats, + )) } pub async fn hydrate_posts( @@ -114,7 +128,7 @@ pub async fn hydrate_posts( 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 = hydrate_profiles_basic(loaders, authors, apply_labelers).await; @@ -122,7 +136,7 @@ pub async fn hydrate_posts( let threadgates = posts .values() - .filter_map(|(_, threadgate)| threadgate.clone()) + .filter_map(|(_, threadgate, _)| threadgate.clone()) .collect(); let threadgates = hydrate_threadgates(loaders, threadgates, apply_labelers).await; @@ -130,7 +144,7 @@ pub async fn hydrate_posts( posts .into_iter() - .filter_map(|(uri, (post, threadgate))| { + .filter_map(|(uri, (post, threadgate, stats))| { let author = authors.get(&post.did)?; let embed = embeds.get(&uri).cloned(); let threadgate = threadgate.and_then(|tg| threadgates.get(&tg.at_uri).cloned()); @@ -138,7 +152,7 @@ pub async fn hydrate_posts( Some(( uri, - build_postview(post, author.to_owned(), labels, embed, threadgate), + build_postview(post, author.to_owned(), labels, embed, threadgate, stats), )) }) .collect() @@ -153,7 +167,7 @@ pub async fn hydrate_feed_posts( 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 = hydrate_profiles_basic(loaders, authors, apply_labelers).await; @@ -163,7 +177,7 @@ pub async fn hydrate_feed_posts( 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::>(); @@ -171,7 +185,7 @@ pub async fn hydrate_feed_posts( posts .into_iter() - .filter_map(|(post_uri, (post, threadgate))| { + .filter_map(|(post_uri, (post, threadgate, stats))| { let author = authors.get(&post.did)?; let root = post.root_uri.as_ref().and_then(|uri| reply_posts.get(uri)); @@ -203,7 +217,7 @@ pub async fn hydrate_feed_posts( let embed = embeds.get(&post_uri).cloned(); let labels = post_labels.get(&post_uri).cloned().unwrap_or_default(); - let post = build_postview(post, author.to_owned(), labels, embed, None); + let post = build_postview(post, author.to_owned(), labels, embed, None, stats); Some(( post_uri, diff --git a/parakeet/src/hydration/profile.rs b/parakeet/src/hydration/profile.rs index e2885361..c22b97c7 100644 --- a/parakeet/src/hydration/profile.rs +++ b/parakeet/src/hydration/profile.rs @@ -2,17 +2,24 @@ use crate::hydration::{map_labels, LabelConfigItem}; use crate::loaders::Dataloaders; use lexica::app_bsky::actor::*; use parakeet_db::models; +use parakeet_index::ProfileStats; use std::collections::HashMap; use std::sync::OnceLock; pub static TRUSTED_VERIFIERS: OnceLock> = OnceLock::new(); -fn build_associated(chat: Option, labeler: bool) -> Option { - if chat.is_some() || labeler { +fn build_associated( + chat: Option, + labeler: bool, + stats: Option, +) -> Option { + if chat.is_some() || labeler || stats.is_some() { + let stats = stats.unwrap_or_default(); + Some(ProfileAssociated { - lists: 0, - feedgens: 0, - starter_packs: 0, + lists: stats.lists as i64, + feedgens: stats.feeds as i64, + starter_packs: stats.starterpacks as i64, labeler, chat: chat.map(|v| ProfileAssociatedChat { allow_incoming: v }), }) @@ -110,8 +117,9 @@ fn build_basic( is_labeler: bool, labels: Vec, verifications: Option>, + stats: Option, ) -> ProfileViewBasic { - let associated = build_associated(chat_decl, is_labeler); + let associated = build_associated(chat_decl, is_labeler, stats); let verification = build_verification(&profile, &handle, verifications); ProfileViewBasic { @@ -135,8 +143,9 @@ fn build_profile( is_labeler: bool, labels: Vec, verifications: Option>, + stats: Option, ) -> ProfileView { - let associated = build_associated(chat_decl, is_labeler); + let associated = build_associated(chat_decl, is_labeler, stats); let verification = build_verification(&profile, &handle, verifications); ProfileView { @@ -158,13 +167,13 @@ fn build_profile( fn build_detailed( handle: Option, profile: models::Profile, - follow_stats: Option, chat_decl: Option, is_labeler: bool, labels: Vec, verifications: Option>, + stats: Option, ) -> ProfileViewDetailed { - let associated = build_associated(chat_decl, is_labeler); + let associated = build_associated(chat_decl, is_labeler, stats); let verification = build_verification(&profile, &handle, verifications); ProfileViewDetailed { @@ -178,14 +187,8 @@ fn build_detailed( banner: profile .banner_cid .map(|v| format!("https://localhost/banner/{v}")), - followers_count: follow_stats - .as_ref() - .map(|v| v.followers as i64) - .unwrap_or_default(), - follows_count: follow_stats - .as_ref() - .map(|v| v.following as i64) - .unwrap_or_default(), + followers_count: stats.map(|v| v.followers as i64).unwrap_or_default(), + follows_count: stats.map(|v| v.following as i64).unwrap_or_default(), associated, labels: map_labels(labels), verification, @@ -201,10 +204,10 @@ pub async fn hydrate_profile_basic( ) -> Option { let labels = loaders.label.load(&did, apply_labelers).await; let verif = loaders.verification.load(did.clone()).await; - let (handle, profile, _, chat_decl, labeler) = loaders.profile.load(did).await?; + let (handle, profile, chat_decl, labeler, stats) = loaders.profile.load(did).await?; Some(build_basic( - handle, profile, chat_decl, labeler, labels, verif, + handle, profile, chat_decl, labeler, labels, verif, stats, )) } @@ -219,13 +222,13 @@ pub async fn hydrate_profiles_basic( profiles .into_iter() - .map(|(k, (handle, profile, _, chat_decl, labeler))| { + .map(|(k, (handle, profile, chat_decl, labeler, stats))| { let labels = labels.get(&k).cloned().unwrap_or_default(); let verif = verif.get(&k).cloned(); ( k, - build_basic(handle, profile, chat_decl, labeler, labels, verif), + build_basic(handle, profile, chat_decl, labeler, labels, verif, stats), ) }) .collect() @@ -238,10 +241,10 @@ pub async fn hydrate_profile( ) -> Option { let labels = loaders.label.load(&did, apply_labelers).await; let verif = loaders.verification.load(did.clone()).await; - let (handle, profile, _, chat_decl, labeler) = loaders.profile.load(did).await?; + let (handle, profile, chat_decl, labeler, stats) = loaders.profile.load(did).await?; Some(build_profile( - handle, profile, chat_decl, labeler, labels, verif, + handle, profile, chat_decl, labeler, labels, verif, stats, )) } @@ -256,13 +259,13 @@ pub async fn hydrate_profiles( profiles .into_iter() - .map(|(k, (handle, profile, _, chat_decl, labeler))| { + .map(|(k, (handle, profile, chat_decl, labeler, stats))| { let labels = labels.get(&k).cloned().unwrap_or_default(); let verif = verif.get(&k).cloned(); ( k, - build_profile(handle, profile, chat_decl, labeler, labels, verif), + build_profile(handle, profile, chat_decl, labeler, labels, verif, stats), ) }) .collect() @@ -275,16 +278,10 @@ pub async fn hydrate_profile_detailed( ) -> Option { let labels = loaders.label.load(&did, apply_labelers).await; let verif = loaders.verification.load(did.clone()).await; - let (handle, profile, follow_stats, chat_decl, labeler) = loaders.profile.load(did).await?; + let (handle, profile, chat_decl, labeler, stats) = loaders.profile.load(did).await?; Some(build_detailed( - handle, - profile, - follow_stats, - chat_decl, - labeler, - labels, - verif, + handle, profile, chat_decl, labeler, labels, verif, stats, )) } @@ -299,21 +296,13 @@ pub async fn hydrate_profiles_detailed( profiles .into_iter() - .map(|(k, (handle, profile, follow_stats, chat_decl, labeler))| { + .map(|(k, (handle, profile, chat_decl, labeler, stats))| { let labels = labels.get(&k).cloned().unwrap_or_default(); let verif = verif.get(&k).cloned(); ( k, - build_detailed( - handle, - profile, - follow_stats, - chat_decl, - labeler, - labels, - verif, - ), + build_detailed(handle, profile, chat_decl, labeler, labels, verif, stats), ) }) .collect() diff --git a/parakeet/src/loaders.rs b/parakeet/src/loaders.rs index 1281cb1e..69d2f7ef 100644 --- a/parakeet/src/loaders.rs +++ b/parakeet/src/loaders.rs @@ -26,16 +26,16 @@ pub struct Dataloaders { 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) -> Dataloaders { + pub fn new(pool: Pool, idxc: parakeet_index::Client) -> Dataloaders { Dataloaders { embed: Loader::new(EmbedLoader(pool.clone())), - feedgen: Loader::new(FeedGenLoader(pool.clone())), + feedgen: Loader::new(FeedGenLoader(pool.clone(), idxc.clone())), handle: Loader::new(HandleLoader(pool.clone())), label: LabelLoader(pool.clone()), // CARE: never cache this. - labeler: Loader::new(LabelServiceLoader(pool.clone())), + labeler: Loader::new(LabelServiceLoader(pool.clone(), idxc.clone())), list: Loader::new(ListLoader(pool.clone())), - posts: Loader::new(PostLoader(pool.clone())), - profile: Loader::new(ProfileLoader(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())), } @@ -66,13 +66,13 @@ impl BatchFn for HandleLoader { } } -pub struct ProfileLoader(Pool); +pub struct ProfileLoader(Pool, parakeet_index::Client); type ProfileLoaderRet = ( Option, models::Profile, - Option, Option, bool, + Option, ); impl BatchFn for ProfileLoader { async fn load(&mut self, keys: &[String]) -> HashMap { @@ -91,7 +91,6 @@ impl BatchFn for ProfileLoader { schema::actors::did, schema::actors::handle, models::Profile::as_select(), - Option::::as_select(), schema::chat_decls::allow_incoming.nullable(), schema::labelers::cid.nullable(), )) @@ -100,19 +99,30 @@ impl BatchFn for ProfileLoader { String, Option, models::Profile, - Option, Option, Option, )>(&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, follow_stats, chat_decl, labeler_cid)| { + |(did, handle, profile, chat_decl, labeler_cid)| { let chat_decl = chat_decl.and_then(|v| ChatAllowIncoming::from_str(&v).ok()); let is_labeler = labeler_cid.is_some(); + let maybe_stats = stats.remove(&did); - let val = (handle, profile, follow_stats, chat_decl, is_labeler); + let val = (handle, profile, chat_decl, is_labeler, maybe_stats); (did, val) }, @@ -157,8 +167,8 @@ impl BatchFn for ListLoader { } } -pub struct FeedGenLoader(Pool); -type FeedGenLoaderRet = models::FeedGen; //todo: when we have likes, we'll need the count here +pub struct FeedGenLoader(Pool, parakeet_index::Client); +type FeedGenLoaderRet = (models::FeedGen, Option); impl BatchFn for FeedGenLoader { async fn load(&mut self, keys: &[String]) -> HashMap { let mut conn = self.0.get().await.unwrap(); @@ -169,11 +179,23 @@ 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| (feedgen.at_uri.clone(), feedgen)), - ), + 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)) + })), Err(e) => { tracing::error!("feedgen load failed: {e}"); HashMap::new() @@ -182,8 +204,12 @@ impl BatchFn for FeedGenLoader { } } -pub struct PostLoader(Pool); -type PostLoaderRet = (models::Post, Option); +pub struct PostLoader(Pool, parakeet_index::Client); +type PostLoaderRet = ( + models::Post, + Option, + Option, +); impl BatchFn for PostLoader { async fn load(&mut self, keys: &[String]) -> HashMap { let mut conn = self.0.get().await.unwrap(); @@ -198,11 +224,23 @@ impl BatchFn for PostLoader { .load(&mut conn) .await; + let stats_req = parakeet_index::GetStatsManyReq { + uris: keys.to_vec(), + }; + let mut stats = self + .1 + .get_post_stats_many(stats_req) + .await + .unwrap() + .into_inner() + .entries; + match res { - Ok(res) => HashMap::from_iter( - res.into_iter() - .map(|(post, threadgate)| (post.at_uri.clone(), (post, threadgate))), - ), + 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() @@ -323,8 +361,12 @@ impl BatchFn for StarterPackLoader { } } -pub struct LabelServiceLoader(Pool); -type LabelServiceLoaderRet = (models::LabelerService, Vec); +pub struct LabelServiceLoader(Pool, parakeet_index::Client); +type LabelServiceLoaderRet = ( + models::LabelerService, + Vec, + Option, +); impl BatchFn for LabelServiceLoader { async fn load(&mut self, keys: &[String]) -> HashMap { let mut conn = self.0.get().await.unwrap(); @@ -343,10 +385,32 @@ 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)| (labeler.did.clone(), (labeler, 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)) + }) .collect() } } diff --git a/parakeet/src/main.rs b/parakeet/src/main.rs index 666cd307..fb68c729 100644 --- a/parakeet/src/main.rs +++ b/parakeet/src/main.rs @@ -14,6 +14,7 @@ mod xrpc; pub struct GlobalState { pub pool: Pool, pub dataloaders: Arc, + pub index_client: parakeet_index::Client, } #[tokio::main] @@ -25,7 +26,12 @@ async fn main() -> eyre::Result<()> { let db_mgr = AsyncDieselConnectionManager::::new(&conf.database_url); let pool = Pool::builder(db_mgr).build()?; - let dataloaders = Arc::new(loaders::Dataloaders::new(pool.clone())); + let index_client = parakeet_index::Client::connect(conf.index_uri).await?; + + let dataloaders = Arc::new(loaders::Dataloaders::new( + pool.clone(), + index_client.clone(), + )); #[allow(unused)] hydration::TRUSTED_VERIFIERS.set(conf.trusted_verifiers); @@ -44,7 +50,11 @@ async fn main() -> eyre::Result<()> { ) .layer(TraceLayer::new_for_http()) .layer(cors) - .with_state(GlobalState { pool, dataloaders }); + .with_state(GlobalState { + pool, + dataloaders, + index_client, + }); let addr = std::net::SocketAddr::new(conf.server.bind_address.parse()?, conf.server.port); let listener = tokio::net::TcpListener::bind(addr).await?;