From 227de7c78529daf2370f5c59cc2b67afcb4e1991 Mon Sep 17 00:00:00 2001 From: Timothy Quilling Date: Mon, 29 Dec 2025 20:20:13 -0500 Subject: [PATCH] init: parakeet-consumer --- parakeet-consumer/Cargo.toml | 51 ++++++ parakeet-consumer/src/core/actor_store.rs | 32 ++++ parakeet-consumer/src/core/event_source.rs | 18 ++ parakeet-consumer/src/core/macros.rs | 55 ++++++ parakeet-consumer/src/core/storage.rs | 38 ++++ parakeet-consumer/src/lib.rs | 44 +++++ parakeet-consumer/src/records/app_bsky.rs | 43 +++++ parakeet-consumer/src/sources/tap/client.rs | 166 ++++++++++++++++++ .../src/sources/tap/processor.rs | 149 ++++++++++++++++ parakeet-consumer/src/sources/tap/types.rs | 91 ++++++++++ 10 files changed, 687 insertions(+) create mode 100644 parakeet-consumer/Cargo.toml create mode 100644 parakeet-consumer/src/core/actor_store.rs create mode 100644 parakeet-consumer/src/core/event_source.rs create mode 100644 parakeet-consumer/src/core/macros.rs create mode 100644 parakeet-consumer/src/core/storage.rs create mode 100644 parakeet-consumer/src/lib.rs create mode 100644 parakeet-consumer/src/records/app_bsky.rs create mode 100644 parakeet-consumer/src/sources/tap/client.rs create mode 100644 parakeet-consumer/src/sources/tap/processor.rs create mode 100644 parakeet-consumer/src/sources/tap/types.rs diff --git a/parakeet-consumer/Cargo.toml b/parakeet-consumer/Cargo.toml new file mode 100644 index 00000000..e485077d --- /dev/null +++ b/parakeet-consumer/Cargo.toml @@ -0,0 +1,51 @@ +[package] +name = "parakeet-consumer" +version = "0.1.0" +edition = "2021" + +[dependencies] +# Workspace dependencies +parakeet-db = { path = "../parakeet-db" } + +# AT Protocol +jacquard-api = { workspace = true } +jacquard-common = { workspace = true } + +# Async runtime +tokio = { workspace = true, features = ["full"] } +async-trait = { workspace = true } + +# WebSocket +tokio-tungstenite = { version = "0.21", features = ["native-tls"] } + +# Database +diesel = { workspace = true, features = ["postgres", "chrono", "uuid", "serde_json"] } +diesel-async = { workspace = true, features = ["deadpool", "postgres"] } +deadpool-diesel = { workspace = true, features = ["postgres"] } + +# Serialization +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } + +# Time handling +chrono = { workspace = true, features = ["serde"] } + +# Caching +moka = { version = "0.12", features = ["future"] } + +# Error handling +thiserror = { workspace = true } + +# Logging +tracing = { workspace = true } + +# Utilities +futures = { workspace = true } +# bytes = { workspace = true } +# uuid = { workspace = true, features = ["v4", "serde"] } + +[dev-dependencies] +tokio-test = "0.4" +mockall = "0.12" +proptest = "1.4" +tracing-subscriber = "0.3" diff --git a/parakeet-consumer/src/core/actor_store.rs b/parakeet-consumer/src/core/actor_store.rs new file mode 100644 index 00000000..c628f093 --- /dev/null +++ b/parakeet-consumer/src/core/actor_store.rs @@ -0,0 +1,32 @@ +use crate::core::{ActorBackend, StorageError}; +use moka::future::Cache; +use std::sync::Arc; +use std::time::Duration; + +pub struct ActorIdStore { + backend: Arc, + cache: Cache, +} + +impl ActorIdStore { + pub fn new(backend: Arc, cache_capacity: u64, cache_ttl: Duration) -> Self { + Self { + backend, + cache: Cache::builder() + .max_capacity(cache_capacity) + .time_to_live(cache_ttl) + .build(), + } + } + pub async fn get(&self, did: &str) -> Result { + if let Some(id) = self.cache.get(did).await { + return Ok(id); + } + let id = self.backend.get_actor_id(did).await?; + self.cache.insert(did.to_string(), id).await; + Ok(id) + } + pub async fn clear_cache(&self) { + self.cache.invalidate_all(); + } +} diff --git a/parakeet-consumer/src/core/event_source.rs b/parakeet-consumer/src/core/event_source.rs new file mode 100644 index 00000000..3c121d54 --- /dev/null +++ b/parakeet-consumer/src/core/event_source.rs @@ -0,0 +1,18 @@ +use async_trait::async_trait; +use std::fmt::Debug; + +pub trait Event: Send + Sync + Debug { + fn id(&self) -> i64; + fn event_type(&self) -> &str; +} + +#[async_trait] +pub trait EventSource: Send + Sync { + type Event: Event; + type Error: std::error::Error + Send + Sync + 'static; + + async fn connect(&mut self) -> Result<(), Self::Error>; + async fn next_event(&mut self) -> Result; + fn is_connected(&self) -> bool; + async fn disconnect(&mut self) -> Result<(), Self::Error>; +} diff --git a/parakeet-consumer/src/core/macros.rs b/parakeet-consumer/src/core/macros.rs new file mode 100644 index 00000000..8a3fbdee --- /dev/null +++ b/parakeet-consumer/src/core/macros.rs @@ -0,0 +1,55 @@ +#[macro_export] +macro_rules! define_record { + ( + struct_name: $name:ident, + field_name: $field:ident, + field_type: $field_type:ty, + deserialize_type: $deser_type:ty, + db_method: $db_method:ident + ) => { + #[derive(Debug, Clone)] + pub struct $name<'a> { + pub $field: $field_type, + pub actor_id: i64, + pub cid: String, + pub uri: String, + } + + impl $crate::sources::tap::processor::FromTapRecord for $name<'static> { + fn from_tap_record( + record: &$crate::sources::tap::types::TapRecord, + actor_id: i64, + ) -> Result { + let val = record + .record + .as_ref() + .ok_or_else(|| $crate::core::StorageError::Parse("Missing record".into()))?; + let json_str = serde_json::to_string(val) + .map_err(|e| $crate::core::StorageError::Parse(e.to_string()))?; + let $field = serde_json::from_str::<$deser_type>(&json_str) + .map_err(|e| $crate::core::StorageError::Parse(e.to_string()))? + .into_static(); + Ok($name { + $field, + actor_id, + cid: record + .cid + .as_ref() + .ok_or_else(|| $crate::core::StorageError::Parse("Missing CID".into()))? + .clone(), + uri: format!("at://{}/{}/{}", record.did, record.collection, record.rkey), + }) + } + } + + #[async_trait::async_trait] + impl $crate::sources::tap::processor::DatabaseWritable for $name<'static> { + async fn write_to_db( + &self, + db: &DB, + ) -> Result<(), $crate::core::StorageError> { + db.$db_method(self).await + } + } + }; +} diff --git a/parakeet-consumer/src/core/storage.rs b/parakeet-consumer/src/core/storage.rs new file mode 100644 index 00000000..fd42c863 --- /dev/null +++ b/parakeet-consumer/src/core/storage.rs @@ -0,0 +1,38 @@ +use crate::records::{Follow, Like, Post, Profile, Repost}; +use async_trait::async_trait; + +#[derive(Debug, thiserror::Error)] +pub enum StorageError { + #[error("Parse error: {0}")] + Parse(String), + + #[error("Connection error: {0}")] + Connection(String), + + #[error("Query error: {0}")] + Query(String), + + #[error("Transaction error: {0}")] + Transaction(String), + + #[error("Not found")] + NotFound, + + #[error("Constraint violation: {0}")] + ConstraintViolation(String), +} + +#[async_trait] +pub trait StorageBackend: Send + Sync { + async fn upsert_post(&self, post: &Post<'static>) -> Result<(), StorageError>; + async fn upsert_profile(&self, profile: &Profile<'static>) -> Result<(), StorageError>; + async fn create_follow(&self, follow: &Follow<'static>) -> Result<(), StorageError>; + async fn create_like(&self, like: &Like<'static>) -> Result<(), StorageError>; + async fn create_repost(&self, repost: &Repost<'static>) -> Result<(), StorageError>; + async fn delete_record(&self, uri: &str) -> Result<(), StorageError>; +} + +#[async_trait] +pub trait ActorBackend: Send + Sync { + async fn get_actor_id(&self, did: &str) -> Result; +} diff --git a/parakeet-consumer/src/lib.rs b/parakeet-consumer/src/lib.rs new file mode 100644 index 00000000..37a20528 --- /dev/null +++ b/parakeet-consumer/src/lib.rs @@ -0,0 +1,44 @@ +pub mod core { + pub mod actor_store; + pub mod event_source; + pub mod macros; + pub mod storage; + + pub use event_source::{Event, EventSource}; + pub use storage::{ActorBackend, StorageBackend, StorageError}; +} + +pub mod sources { + pub mod tap { + pub mod client; + pub mod processor; + pub mod types; + + pub use client::{ReconnectingTapClient, TapClient}; + pub use processor::{spawn_worker, Dispatcher, EventProcessor}; + pub use types::{ + IdentityData, LabelData, RepoStatus, TapAction, TapError, TapEvent, TapRecord, + }; + } +} + +pub mod records { + pub mod app_bsky; + + pub use app_bsky::{Follow, Like, Post, Profile, Repost}; +} + +pub mod storage { + // pub mod postgres; + + // pub use postgres::PostgresBackend; +} + +pub use core::actor_store::ActorIdStore; +pub use core::{ActorBackend, Event, EventSource, StorageBackend, StorageError}; +pub use records::{Follow, Like, Post, Profile, Repost}; + +pub use sources::tap::{ + spawn_worker, Dispatcher, EventProcessor, ReconnectingTapClient, TapAction, TapClient, + TapError, TapEvent, TapRecord, +}; diff --git a/parakeet-consumer/src/records/app_bsky.rs b/parakeet-consumer/src/records/app_bsky.rs new file mode 100644 index 00000000..09e31c51 --- /dev/null +++ b/parakeet-consumer/src/records/app_bsky.rs @@ -0,0 +1,43 @@ +use crate::define_record; +use jacquard_api::app_bsky; +use jacquard_common::IntoStatic; + +define_record!( + struct_name: Post, + field_name: post, + field_type: app_bsky::feed::post::Post<'a>, + deserialize_type: app_bsky::feed::post::Post<'_>, + db_method: upsert_post +); + +define_record!( + struct_name: Profile, + field_name: profile, + field_type: app_bsky::actor::profile::Profile<'a>, + deserialize_type: app_bsky::actor::profile::Profile<'_>, + db_method: upsert_profile +); + +define_record!( + struct_name: Follow, + field_name: follow, + field_type: app_bsky::graph::follow::Follow<'a>, + deserialize_type: app_bsky::graph::follow::Follow<'_>, + db_method: create_follow +); + +define_record!( + struct_name: Like, + field_name: like, + field_type: app_bsky::feed::like::Like<'a>, + deserialize_type: app_bsky::feed::like::Like<'_>, + db_method: create_like +); + +define_record!( + struct_name: Repost, + field_name: repost, + field_type: app_bsky::feed::repost::Repost<'a>, + deserialize_type: app_bsky::feed::repost::Repost<'_>, + db_method: create_repost +); diff --git a/parakeet-consumer/src/sources/tap/client.rs b/parakeet-consumer/src/sources/tap/client.rs new file mode 100644 index 00000000..e983844a --- /dev/null +++ b/parakeet-consumer/src/sources/tap/client.rs @@ -0,0 +1,166 @@ +use super::types::{TapError, TapEvent}; +use crate::core::{Event, EventSource}; +use async_trait::async_trait; +use futures::{SinkExt, StreamExt}; +use std::time::Duration; +use tokio::net::TcpStream; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream}; + +pub struct TapClient { + url: String, + ws: Option>>, +} + +impl TapClient { + pub fn new(url: impl Into) -> Self { + let mut url = url.into(); + url = url + .replacen("http://", "ws://", 1) + .replacen("https://", "wss://", 1); + if !url.ends_with("/channel") { + if !url.ends_with('/') { + url.push('/'); + } + url.push_str("channel"); + } + Self { url, ws: None } + } + + async fn send_ack(&mut self, event_id: i64) -> Result<(), TapError> { + match &mut self.ws { + Some(ws) => ws + .send(Message::Text(format!( + r#"{{"type":"ack","id":{}}}"#, + event_id + ))) + .await + .map_err(|e| TapError::Connection(e.to_string())), + None => Err(TapError::Disconnected), + } + } +} + +#[async_trait] +impl EventSource for TapClient { + type Event = TapEvent; + type Error = TapError; + + async fn connect(&mut self) -> Result<(), TapError> { + self.ws = Some( + connect_async(&self.url) + .await + .map_err(|e| TapError::Connection(e.to_string()))? + .0, + ); + Ok(()) + } + + async fn next_event(&mut self) -> Result { + loop { + let ws = self.ws.as_mut().ok_or(TapError::Disconnected)?; + match ws.next().await { + Some(Ok(Message::Text(text))) => { + if let Ok(event) = serde_json::from_str::(&text) { + let _ = self.send_ack(event.id()).await; + return Ok(event); + } + } + Some(Ok(Message::Ping(data))) => { + ws.send(Message::Pong(data)) + .await + .map_err(|e| TapError::Connection(e.to_string()))?; + } + Some(Ok(Message::Close(_))) | None => { + self.ws = None; + return Err(TapError::Disconnected); + } + Some(Err(e)) => { + self.ws = None; + return Err(TapError::Connection(e.to_string())); + } + _ => {} + } + } + } + + fn is_connected(&self) -> bool { + self.ws.is_some() + } + async fn disconnect(&mut self) -> Result<(), TapError> { + self.ws.take(); + Ok(()) + } +} + +pub struct ReconnectingTapClient { + url: String, + inner: Option, + max_retries: usize, + reconnect_delay: Duration, +} + +impl ReconnectingTapClient { + pub fn new(url: impl Into) -> Self { + let url = url.into(); + Self { + url: url.clone(), + inner: Some(TapClient::new(url)), + max_retries: 10, + reconnect_delay: Duration::from_secs(1), + } + } + async fn ensure_connected(&mut self) -> Result<(), TapError> { + if self.inner.as_ref().is_some_and(|c| c.is_connected()) { + return Ok(()); + } + let (mut retries, mut delay) = (0, self.reconnect_delay); + while retries < self.max_retries { + let mut client = TapClient::new(self.url.clone()); + if client.connect().await.is_ok() { + self.inner = Some(client); + return Ok(()); + } + retries += 1; + if retries < self.max_retries { + tokio::time::sleep(delay).await; + delay = (delay * 2).min(Duration::from_secs(60)); + } + } + Err(TapError::Connection(format!( + "Failed after {} attempts", + self.max_retries + ))) + } +} + +#[async_trait] +impl EventSource for ReconnectingTapClient { + type Event = TapEvent; + type Error = TapError; + async fn connect(&mut self) -> Result<(), TapError> { + self.ensure_connected().await + } + async fn next_event(&mut self) -> Result { + loop { + self.ensure_connected().await?; + if let Some(client) = &mut self.inner { + match client.next_event().await { + Ok(event) => return Ok(event), + Err(TapError::Disconnected) => self.inner = None, + Err(e) => return Err(e), + } + } + } + } + fn is_connected(&self) -> bool { + self.inner.as_ref().is_some_and(|c| c.is_connected()) + } + async fn disconnect(&mut self) -> Result<(), TapError> { + if let Some(mut client) = self.inner.take() { + client.disconnect().await + } else { + Ok(()) + } + } +} diff --git a/parakeet-consumer/src/sources/tap/processor.rs b/parakeet-consumer/src/sources/tap/processor.rs new file mode 100644 index 00000000..8c6ac306 --- /dev/null +++ b/parakeet-consumer/src/sources/tap/processor.rs @@ -0,0 +1,149 @@ +use super::types::{TapAction, TapEvent, TapRecord}; +use crate::core::actor_store::ActorIdStore; +use crate::core::{ActorBackend, Event, StorageBackend, StorageError}; +use crate::records::{Follow, Like, Post, Profile, Repost}; +use async_trait::async_trait; +use std::sync::Arc; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; +use tracing::{error, info}; + +pub trait FromTapRecord: Sized { + fn from_tap_record(record: &TapRecord, actor_id: i64) -> Result; +} + +#[async_trait] +pub trait DatabaseWritable { + async fn write_to_db(&self, db: &DB) -> Result<(), StorageError>; +} + +const COLLECTIONS: &[&str] = &[ + "app.bsky.feed.post", + "app.bsky.actor.profile", + "app.bsky.graph.follow", + "app.bsky.feed.like", + "app.bsky.feed.repost", +]; + +macro_rules! process_collection { + ($event:expr, $db:expr, $store:expr) => { + match $event { + TapEvent::Record { record, .. } => match record.collection.as_str() { + "app.bsky.feed.post" => process_record::($event, $db, $store).await, + "app.bsky.actor.profile" => process_record::($event, $db, $store).await, + "app.bsky.graph.follow" => process_record::($event, $db, $store).await, + "app.bsky.feed.like" => process_record::($event, $db, $store).await, + "app.bsky.feed.repost" => process_record::($event, $db, $store).await, + _ => Ok(()), + }, + _ => Ok(()), + } + }; +} + +pub fn spawn_worker( + collection: &'static str, + mut rx: mpsc::Receiver, + db: Arc, + actor_store: Arc>, +) -> JoinHandle<()> { + tokio::spawn(async move { + info!("Worker started for collection: {}", collection); + while let Some(event) = rx.recv().await { + if let Err(e) = process_collection!(&event, db.as_ref(), &actor_store) { + error!( + "Failed to process {} event {}: {}", + collection, + event.id(), + e + ); + } + } + info!("Worker stopped for collection: {}", collection); + }) +} + +async fn process_record( + event: &TapEvent, + db: &dyn StorageBackend, + actor_store: &Arc>, +) -> Result<(), StorageError> +where + T: FromTapRecord + DatabaseWritable + Send + 'static, + AB: ActorBackend, +{ + let TapEvent::Record { record, .. } = event else { + return Ok(()); + }; + match record.action { + TapAction::Create | TapAction::Update => { + T::from_tap_record(record, actor_store.get(&record.did).await?)? + .write_to_db(db) + .await + } + TapAction::Delete => { + db.delete_record(&format!( + "at://{}/{}/{}", + record.did, record.collection, record.rkey + )) + .await + } + } +} + +pub struct Dispatcher { + channels: std::collections::HashMap>, + workers: Vec>, +} + +impl Dispatcher { + pub fn new( + db: Arc, + actor_store: Arc>, + channel_size: usize, + ) -> Self { + let (channels, workers) = COLLECTIONS.iter().fold( + (std::collections::HashMap::new(), Vec::new()), + |(mut ch, mut w), c| { + let (tx, rx) = mpsc::channel(channel_size); + ch.insert(c.to_string(), tx); + w.push(spawn_worker(c, rx, db.clone(), actor_store.clone())); + (ch, w) + }, + ); + Self { channels, workers } + } + + pub async fn dispatch(&self, event: TapEvent) -> Result<(), StorageError> { + let TapEvent::Record { ref record, .. } = event else { + return Ok(()); + }; + if let Some(tx) = self.channels.get(&record.collection) { + tx.send(event) + .await + .map_err(|_| StorageError::Query("Worker channel closed".into()))?; + } + Ok(()) + } + + pub async fn shutdown(self) { + drop(self.channels); + for worker in self.workers { + let _ = worker.await; + } + } +} + +pub struct EventProcessor { + db: Arc, + actor_store: Arc>, +} + +impl EventProcessor { + pub fn new(db: Arc, actor_store: Arc>) -> Self { + Self { db, actor_store } + } + pub async fn process_event(&self, event: &TapEvent) -> Result<(), StorageError> { + process_collection!(event, self.db.as_ref(), &self.actor_store) + } +} diff --git a/parakeet-consumer/src/sources/tap/types.rs b/parakeet-consumer/src/sources/tap/types.rs new file mode 100644 index 00000000..106df47b --- /dev/null +++ b/parakeet-consumer/src/sources/tap/types.rs @@ -0,0 +1,91 @@ +use crate::core::Event; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(tag = "type", rename_all = "lowercase")] +pub enum TapEvent { + Record { id: i64, record: TapRecord }, + Identity { id: i64, identity: IdentityData }, + Label { id: i64, label: LabelData }, +} + +impl Event for TapEvent { + fn id(&self) -> i64 { + match self { + TapEvent::Record { id, .. } + | TapEvent::Identity { id, .. } + | TapEvent::Label { id, .. } => *id, + } + } + fn event_type(&self) -> &str { + match self { + TapEvent::Record { .. } => "record", + TapEvent::Identity { .. } => "identity", + TapEvent::Label { .. } => "label", + } + } +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct TapRecord { + pub live: bool, + pub rev: String, + pub did: String, + pub collection: String, + pub rkey: String, + pub action: TapAction, + pub cid: Option, + pub record: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct IdentityData { + pub did: String, + pub handle: String, + #[serde(rename = "isActive")] + pub is_active: bool, + pub status: RepoStatus, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct LabelData { + pub live: bool, + #[serde(rename = "labelerDID")] + pub labeler_did: String, + pub uri: String, + pub val: String, + pub cts: String, + pub src: String, + pub cid: Option, + pub neg: bool, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum TapAction { + Create, + Update, + Delete, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum RepoStatus { + Active, + Takendown, + Suspended, + Deactivated, + Deleted, +} + +#[derive(Debug, thiserror::Error)] +pub enum TapError { + #[error("Connection error: {0}")] + Connection(String), + #[error("Disconnected from tap")] + Disconnected, + #[error("Deserialization error: {0}")] + Deserialization(String), + #[error("IO error: {0}")] + Io(#[from] std::io::Error), +} -- 2.51.2