atproto-jetstream 0.9.7

//! Async stream consumer for AT Protocol Jetstream events
//!
//! This module provides structures for consuming events from an async stream
//! and dispatching them to registered event handlers.

use anyhow::Result;
use async_trait::async_trait;
use futures::{SinkExt, StreamExt};
use http::Uri;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::{collections::HashMap, str::FromStr};
use tokio::sync::{RwLock, broadcast};
use tokio::time::{Instant, sleep};
use tokio_util::sync::CancellationToken;
use tokio_websockets::{ClientBuilder, Message};
use tracing::Instrument;

const MAX_MESSAGE_SIZE: usize = 56000;

/// Configuration for the Jetstream consumer task
#[cfg_attr(debug_assertions, derive(Debug))]
#[derive(Clone)]
pub struct ConsumerTaskConfig {
    /// User-Agent header value for WebSocket connections
    pub user_agent: String,
    /// Enable Zstandard compression for messages
    pub compression: bool,
    /// Path to Zstandard dictionary file (required if compression is enabled)
    pub zstd_dictionary_location: String,
    /// Hostname of the Jetstream instance to connect to
    pub jetstream_hostname: String,
    /// AT Protocol collections to subscribe to (empty for all)
    pub collections: Vec<String>,
    /// DIDs to filter events for (empty for all)
    pub dids: Vec<String>,
    /// Maximum message size in bytes (None for unlimited)
    pub max_message_size_bytes: Option<u64>,
    /// Optional cursor position to start streaming from
    pub cursor: Option<i64>,
    /// Whether to require a hello message before receiving events
    pub require_hello: bool,
}

/// Event data structure for Jetstream events
#[cfg_attr(debug_assertions, derive(Debug))]
#[derive(Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum JetstreamEvent {
    /// Repository commit event (create/update operations)
    Commit {
        /// DID of the repository that was updated
        did: String,
        /// Event timestamp in microseconds since Unix epoch
        time_us: u64,
        /// Event type identifier
        kind: String,

        #[serde(rename = "commit")]
        /// Commit operation details
        commit: JetstreamEventCommit,
    },

    /// Repository delete event
    Delete {
        /// DID of the repository that was updated
        did: String,
        /// Event timestamp in microseconds since Unix epoch
        time_us: u64,
        /// Event type identifier
        kind: String,

        #[serde(rename = "commit")]
        /// Delete operation details
        commit: JetstreamEventDelete,
    },

    /// Identity document update event
    Identity {
        /// DID whose identity was updated
        did: String,
        /// Event timestamp in microseconds since Unix epoch
        time_us: u64,
        /// Event type identifier
        kind: String,

        #[serde(rename = "identity")]
        /// Identity document data
        identity: serde_json::Value,
    },

    /// Account-related event
    Account {
        /// DID of the account
        did: String,
        /// Event timestamp in microseconds since Unix epoch
        time_us: u64,
        /// Event type identifier
        kind: String,

        #[serde(rename = "account")]
        /// Account data
        identity: serde_json::Value,
    },
}

/// Repository commit operation details
#[cfg_attr(debug_assertions, derive(Debug))]
#[derive(Clone, Serialize, Deserialize)]
pub struct JetstreamEventCommit {
    /// Repository revision identifier
    pub rev: String,
    /// Operation type (create, update)
    pub operation: String,
    /// AT Protocol collection name
    pub collection: String,
    /// Record key within the collection
    pub rkey: String,
    /// Content identifier (CID) of the record
    pub cid: String,
    /// Record data as JSON
    pub record: serde_json::Value,
}

/// Repository delete operation details
#[cfg_attr(debug_assertions, derive(Debug))]
#[derive(Clone, Serialize, Deserialize)]
pub struct JetstreamEventDelete {
    /// Repository revision identifier
    pub rev: String,
    /// Operation type (delete)
    pub operation: String,
    /// AT Protocol collection name
    pub collection: String,
    /// Record key that was deleted
    pub rkey: String,
}

/// Trait for handling Jetstream events
#[async_trait]
pub trait EventHandler: Send + Sync {
    /// Handle a received event
    async fn handle_event(&self, event: JetstreamEvent) -> Result<()>;

    /// Get the handler's identifier
    fn handler_id(&self) -> String;
}

/// Errors specific to the consumer module
#[derive(thiserror::Error, Debug)]
pub enum ConsumerError {
    /// WebSocket connection establishment failed
    #[error("error-atproto-jetstream-consumer-1 WebSocket connection failed: {0}")]
    ConnectionFailed(String),
    /// Message decompression operation failed
    #[error("error-atproto-jetstream-consumer-2 Message decompression failed: {0}")]
    DecompressionFailed(String),
    /// JSON deserialization of event data failed
    #[error("error-atproto-jetstream-consumer-3 Event deserialization failed: {0}")]
    DeserializationFailed(String),
    /// Event handler registration failed (e.g., duplicate handler ID)
    #[error("error-atproto-jetstream-consumer-4 Handler registration failed: {0}")]
    HandlerRegistrationFailed(String),
    /// Event broadcast sender not initialized (consumer not running)
    #[error("error-atproto-jetstream-consumer-5 Event sender not initialized: {0}")]
    EventSenderNotInitialized(String),
    /// WebSocket message format conversion failed
    #[error("error-atproto-jetstream-consumer-6 Message conversion failed: {0}")]
    MessageConversionFailed(String),
    /// Serialization of subscription update message failed
    #[error("error-atproto-jetstream-consumer-7 Update serialization failed: {0}")]
    UpdateSerializationFailed(String),
    /// Sending subscription update message failed
    #[error("error-atproto-jetstream-consumer-8 Update send failed: {0}")]
    UpdateSendFailed(String),
    /// Zstandard decompressor initialization failed
    #[error("error-atproto-jetstream-consumer-9 Decompressor creation failed: {0}")]
    DecompressorCreationFailed(String),
}

#[cfg_attr(debug_assertions, derive(Debug))]
#[derive(Clone, Serialize, Deserialize)]
#[serde(tag = "type", content = "payload")]
pub(crate) enum SubscriberSourcedMessage {
    #[serde(rename = "options_update")]
    Update {
        #[serde(rename = "wantedCollections")]
        wanted_collections: Vec<String>,

        #[serde(rename = "wantedDids", skip_serializing_if = "Vec::is_empty", default)]
        wanted_dids: Vec<String>,

        #[serde(rename = "maxMessageSizeBytes")]
        max_message_size_bytes: u64,

        #[serde(skip_serializing_if = "Option::is_none")]
        cursor: Option<i64>,
    },
}

/// Main consumer structure for handling async streams and event dispatching
pub struct Consumer {
    config: ConsumerTaskConfig,
    handlers: Arc<RwLock<HashMap<String, Arc<dyn EventHandler>>>>,
    event_sender: Arc<RwLock<Option<broadcast::Sender<JetstreamEvent>>>>,
}

impl Consumer {
    /// Create a new consumer with the given configuration
    pub fn new(config: ConsumerTaskConfig) -> Self {
        Self {
            config,
            handlers: Arc::new(RwLock::new(HashMap::new())),
            event_sender: Arc::new(RwLock::new(None)),
        }
    }

    /// Register an event handler
    pub async fn register_handler(&self, handler: Arc<dyn EventHandler>) -> Result<()> {
        let handler_id = handler.handler_id();
        let mut handlers = self.handlers.write().await;

        if handlers.contains_key(&handler_id) {
            return Err(ConsumerError::HandlerRegistrationFailed(format!(
                "Handler with ID '{}' already registered",
                handler_id
            ))
            .into());
        }

        handlers.insert(handler_id.clone(), handler);
        Ok(())
    }

    /// Unregister an event handler
    pub async fn unregister_handler(&self, handler_id: &str) -> Result<()> {
        let mut handlers = self.handlers.write().await;
        handlers.remove(handler_id);
        Ok(())
    }

    /// Get a broadcast receiver for events
    pub async fn get_event_receiver(&self) -> Result<broadcast::Receiver<JetstreamEvent>> {
        let sender_guard = self.event_sender.read().await;
        match sender_guard.as_ref() {
            Some(sender) => Ok(sender.subscribe()),
            None => Err(ConsumerError::EventSenderNotInitialized(
                "consumer not running".to_string(),
            )
            .into()),
        }
    }

    /// Run the consumer in the background
    ///
    /// # Example
    /// ```rust,no_run
    /// use atproto_jetstream::{Consumer, ConsumerTaskConfig, CancellationToken};
    ///
    /// # async fn example() -> anyhow::Result<()> {
    /// let config = ConsumerTaskConfig {
    ///     user_agent: "my-app/1.0".to_string(),
    ///     compression: false,
    ///     zstd_dictionary_location: String::new(),
    ///     jetstream_hostname: "jetstream1.us-east.bsky.network".to_string(),
    ///     collections: vec!["app.bsky.feed.post".to_string()],
    ///     dids: vec![], // Subscribe to all DIDs
    ///     max_message_size_bytes: None, // No limit
    ///     cursor: None, // Live-tail from current time
    ///     require_hello: true, // Wait for initial options update
    /// };
    ///
    /// let consumer = Consumer::new(config);
    /// let cancellation_token = CancellationToken::new();
    ///
    /// // To cancel the consumer later:
    /// // cancellation_token.cancel();
    ///
    /// consumer.run_background(cancellation_token).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn run_background(&self, cancellation_token: CancellationToken) -> Result<()> {
        tracing::info!("Starting Jetstream consumer");

        // Build WebSocket URL with query parameters
        let mut query_params = vec![];

        // Add compression parameter
        query_params.push(format!("compress={}", self.config.compression));

        // Add requireHello parameter
        query_params.push(format!("requireHello={}", self.config.require_hello));

        // Add wantedCollections if specified
        if !self.config.collections.is_empty() {
            let collections = self
                .config
                .collections
                .iter()
                .map(|c| urlencoding::encode(c))
                .collect::<Vec<_>>()
                .join(",");
            query_params.push(format!("wantedCollections={}", collections));
        }

        // Add wantedDids if specified
        if !self.config.dids.is_empty() {
            let dids = self
                .config
                .dids
                .iter()
                .map(|d| urlencoding::encode(d))
                .collect::<Vec<_>>()
                .join(",");
            query_params.push(format!("wantedDids={}", dids));
        }

        // Add maxMessageSizeBytes if specified
        if let Some(max_size) = self.config.max_message_size_bytes {
            query_params.push(format!("maxMessageSizeBytes={}", max_size));
        }

        // Add cursor if specified
        if let Some(cursor) = self.config.cursor {
            query_params.push(format!("cursor={}", cursor));
        }

        let query_string = query_params.join("&");
        let ws_url = Uri::from_str(&format!(
            "wss://{}/subscribe?{}",
            self.config.jetstream_hostname, query_string
        ))?;

        let (mut client, _) = ClientBuilder::from_uri(ws_url)
            .add_header(
                http::header::USER_AGENT,
                http::HeaderValue::from_str(&self.config.user_agent)?,
            )?
            .connect()
            .await?;

        let update = SubscriberSourcedMessage::Update {
            wanted_collections: self.config.collections.clone(),
            wanted_dids: self.config.dids.clone(),
            max_message_size_bytes: self
                .config
                .max_message_size_bytes
                .unwrap_or(MAX_MESSAGE_SIZE as u64),
            cursor: self.config.cursor,
        };
        let serialized_update = serde_json::to_string(&update)
            .map_err(|err| ConsumerError::UpdateSerializationFailed(err.to_string()))?;

        client
            .send(Message::text(serialized_update))
            .await
            .map_err(|err| ConsumerError::UpdateSendFailed(err.to_string()))?;

        let mut decompressor = if self.config.compression {
            // mkdir -p data/ && curl -o data/zstd_dictionary https://github.com/bluesky-social/jetstream/raw/refs/heads/main/pkg/models/zstd_dictionary
            let data: Vec<u8> = std::fs::read(self.config.zstd_dictionary_location.clone())?;
            zstd::bulk::Decompressor::with_dictionary(&data)
                .map_err(|err| ConsumerError::DecompressorCreationFailed(err.to_string()))?
        } else {
            zstd::bulk::Decompressor::new()
                .map_err(|err| ConsumerError::DecompressorCreationFailed(err.to_string()))?
        };

        let interval = std::time::Duration::from_secs(120);
        let sleeper = sleep(interval);
        tokio::pin!(sleeper);

        loop {
            tokio::select! {
                () = cancellation_token.cancelled() => {
                    break;
                },
                () = &mut sleeper => {
                        // consumer_control_insert(&self.pool, &self.config.jetstream_hostname, time_usec).await?;

                        sleeper.as_mut().reset(Instant::now() + interval);
                },
                item = client.next() => {
                    if item.is_none() {
                        tracing::warn!("jetstream connection closed");
                        break;
                    }
                    let item = item.unwrap();

                    if let Err(err) = item {
                        tracing::error!(error = ?err, "error processing jetstream message");
                        continue;
                    }
                    let item = item.unwrap();

                    let event = if self.config.compression {
                        if !item.is_binary() {
                            tracing::debug!("compression enabled but message from jetstream is not binary");
                            continue;
                        }
                        let payload = item.into_payload();

                        let decoded = decompressor.decompress(&payload, MAX_MESSAGE_SIZE * 3);
                        if let Err(err) = decoded {
                            tracing::debug!(err = ?err, "cannot decompress message");
                            continue;
                        }
                        let decoded = decoded.unwrap();
                        serde_json::from_slice::<JetstreamEvent>(&decoded)
                        .map_err(|err| ConsumerError::DeserializationFailed(err.to_string()))
                    } else {
                        if !item.is_text() {
                            tracing::debug!("compression disabled but message from jetstream is binary");
                            continue;
                        }
                        item.as_text()
                            .ok_or_else(|| ConsumerError::MessageConversionFailed("cannot convert message to text".to_string()))
                            .and_then(|value| {
                                serde_json::from_str::<JetstreamEvent>(value)
                                .map_err(|err| ConsumerError::DeserializationFailed(err.to_string()))
                            })
                    };
                    if let Err(err) = event {
                        tracing::error!(error = ?err, "error processing jetstream message");

                        continue;
                    }
                    let event = event.unwrap();

                    if let Err(err) = self.dispatch_to_handlers(event).await {
                        tracing::error!(error = ?err, "Failed to process message");
                    }

                }
            }
        }

        // Clean up
        {
            let mut sender_guard = self.event_sender.write().await;
            *sender_guard = None;
        }

        Ok(())
    }

    /// Dispatch event to all registered handlers
    async fn dispatch_to_handlers(&self, event: JetstreamEvent) -> Result<()> {
        let handlers = self.handlers.read().await;

        for (handler_id, handler) in handlers.iter() {
            let handler_span = tracing::debug_span!("handler_dispatch", handler_id = %handler_id);
            async {
                if let Err(err) = handler.handle_event(event.clone()).await {
                    tracing::error!(
                        error = ?err,
                        handler_id = %handler_id,
                        "Handler failed to process event"
                    );
                }
            }
            .instrument(handler_span)
            .await;
        }

        Ok(())
    }
}

/// Example event handler implementation
pub struct LoggingHandler {
    id: String,
}

impl LoggingHandler {
    /// Create a new logging handler with the specified ID
    pub fn new(id: String) -> Self {
        Self { id }
    }
}

#[async_trait]
impl EventHandler for LoggingHandler {
    async fn handle_event(&self, _event: JetstreamEvent) -> Result<()> {
        Ok(())
    }

    fn handler_id(&self) -> String {
        self.id.clone()
    }
}