init: parakeet-consumer

parakeet-consumer/Cargo.toml

···
+[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"

parakeet-consumer/src/core/actor_store.rs

···
+use crate::core::{ActorBackend, StorageError};
+use moka::future::Cache;
+use std::sync::Arc;
+use std::time::Duration;
+
+pub struct ActorIdStore<B: ActorBackend> {
+    backend: Arc<B>,
+    cache: Cache<String, i64>,
+}
+
+impl<B: ActorBackend> ActorIdStore<B> {
+    pub fn new(backend: Arc<B>, 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<i64, StorageError> {
+        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();
+    }
+}

parakeet-consumer/src/core/event_source.rs

···
+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<Self::Event, Self::Error>;
+    fn is_connected(&self) -> bool;
+    async fn disconnect(&mut self) -> Result<(), Self::Error>;
+}

parakeet-consumer/src/core/macros.rs

···
+#[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<Self, $crate::core::StorageError> {
+                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<DB: $crate::core::StorageBackend + ?Sized>(
+                &self,
+                db: &DB,
+            ) -> Result<(), $crate::core::StorageError> {
+                db.$db_method(self).await
+            }
+        }
+    };
+}

parakeet-consumer/src/core/storage.rs

···
+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<i64, StorageError>;
+}

parakeet-consumer/src/lib.rs

···
+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,
+};

parakeet-consumer/src/records/app_bsky.rs

···
+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
+);

parakeet-consumer/src/sources/tap/client.rs

···
+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<WebSocketStream<MaybeTlsStream<TcpStream>>>,
+}
+
+impl TapClient {
+    pub fn new(url: impl Into<String>) -> 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<TapEvent, TapError> {
+        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::<TapEvent>(&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<TapClient>,
+    max_retries: usize,
+    reconnect_delay: Duration,
+}
+
+impl ReconnectingTapClient {
+    pub fn new(url: impl Into<String>) -> 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<TapEvent, TapError> {
+        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(())
+        }
+    }
+}

parakeet-consumer/src/sources/tap/processor.rs

···
+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<Self, StorageError>;
+}
+
+#[async_trait]
+pub trait DatabaseWritable {
+    async fn write_to_db<DB: StorageBackend + ?Sized>(&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::<Post, _>($event, $db, $store).await,
+                "app.bsky.actor.profile" => process_record::<Profile, _>($event, $db, $store).await,
+                "app.bsky.graph.follow" => process_record::<Follow, _>($event, $db, $store).await,
+                "app.bsky.feed.like" => process_record::<Like, _>($event, $db, $store).await,
+                "app.bsky.feed.repost" => process_record::<Repost, _>($event, $db, $store).await,
+                _ => Ok(()),
+            },
+            _ => Ok(()),
+        }
+    };
+}
+
+pub fn spawn_worker<AB: ActorBackend + 'static>(
+    collection: &'static str,
+    mut rx: mpsc::Receiver<TapEvent>,
+    db: Arc<dyn StorageBackend>,
+    actor_store: Arc<ActorIdStore<AB>>,
+) -> 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<T, AB>(
+    event: &TapEvent,
+    db: &dyn StorageBackend,
+    actor_store: &Arc<ActorIdStore<AB>>,
+) -> 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<String, mpsc::Sender<TapEvent>>,
+    workers: Vec<JoinHandle<()>>,
+}
+
+impl Dispatcher {
+    pub fn new<AB: ActorBackend + 'static>(
+        db: Arc<dyn StorageBackend>,
+        actor_store: Arc<ActorIdStore<AB>>,
+        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: StorageBackend, AB: ActorBackend> {
+    db: Arc<DB>,
+    actor_store: Arc<ActorIdStore<AB>>,
+}
+
+impl<DB: StorageBackend, AB: ActorBackend> EventProcessor<DB, AB> {
+    pub fn new(db: Arc<DB>, actor_store: Arc<ActorIdStore<AB>>) -> 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)
+    }
+}

parakeet-consumer/src/sources/tap/types.rs

···
+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<String>,
+    pub record: Option<serde_json::Value>,
+}
+
+#[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<String>,
+    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),
+}