Skip to main content

bsky_posts/
bsky-posts.rs

1use std::{
2    collections::HashMap,
3    sync::{Arc, Mutex},
4};
5
6use anyhow::Result;
7use async_trait::async_trait;
8use rocketman::{
9    connection::JetstreamConnection,
10    handler::{self, Ingestors},
11    ingestion::LexiconIngestor,
12    options::JetstreamOptions,
13    types::event::Event,
14};
15use serde_json::Value;
16use tracing::{error, info, warn};
17
18/// This example tests the reconnection fix for idle connections.
19///
20/// When listening to a custom collection that gets no traffic,
21/// the server closes the connection after timeout and
22/// rocketman should reconnect automatically.
23///
24/// Run with: cargo run --example bsky-posts
25///
26/// Expected behavior:
27/// 1. Initial connection
28/// 2. No messages for timeout period (20s)
29/// 3. "No messages received in X seconds, reconnecting" or "Stream closed by server"
30/// 4. Automatic reconnection (should repeat indefinitely)
31/// 5. No infinite loops or hanging
32
33#[tokio::main]
34async fn main() -> Result<()> {
35    tracing_subscriber::fmt()
36        .with_max_level(tracing::Level::INFO)
37        .init();
38
39    start_idle_test().await
40}
41
42pub async fn start_idle_test() -> Result<()> {
43    info!("Testing reconnection fix with idle connection...");
44
45    let opts = JetstreamOptions::builder()
46        .wanted_collections(vec![
47            // This is a custom collection that will get zero traffic
48            "app.bsky.feed.post".to_string(),
49        ])
50        .timeout_time_sec(20) // Shorter timeout for faster testing
51        .bound(65536)
52        .build();
53
54    let jetstream = JetstreamConnection::new(opts);
55
56    let mut ingestors: HashMap<String, Box<dyn LexiconIngestor + Send + Sync>> = HashMap::new();
57    ingestors.insert(
58        "fm.teal.alpha.feed.play".to_string(),
59        Box::new(ProfileIngestor),
60    );
61
62    let cursor: Arc<Mutex<Option<u64>>> = Arc::new(Mutex::new(None));
63
64    let msg_rx = jetstream.get_msg_rx();
65    let reconnect_tx = jetstream.get_reconnect_tx();
66
67    let monitor_rx = msg_rx.clone();
68    tokio::spawn(async move {
69        let mut interval = tokio::time::interval(std::time::Duration::from_secs(10));
70
71        loop {
72            interval.tick().await;
73            let queue_len = monitor_rx.len();
74            if queue_len > 0 {
75                warn!("Queue has {} messages pending", queue_len);
76            }
77        }
78    });
79
80    // Spawn task to process messages (there should be none)
81    let c_cursor = cursor.clone();
82    tokio::spawn(async move {
83        info!("Message processing task started");
84        let mut message_count = 0u64;
85        let mut last_log = std::time::Instant::now();
86
87        let ing = Ingestors {
88            commits: ingestors,
89            identity: None,
90            account: None,
91        };
92
93        while let Ok(message) = msg_rx.recv().await {
94            message_count += 1;
95
96            // Log every 10 messages to see if we're actually processing
97            if message_count % 10 == 0 || last_log.elapsed() > std::time::Duration::from_secs(5) {
98                info!(
99                    "Processing message #{} (queue len: {}/{})",
100                    message_count,
101                    msg_rx.len(),
102                    msg_rx.capacity().unwrap_or(0)
103                );
104                last_log = std::time::Instant::now();
105            }
106
107            match handler::handle_message(message, &ing, reconnect_tx.clone(), c_cursor.clone())
108                .await
109            {
110                Ok(_) => {}
111                Err(e) => {
112                    error!("Error processing message #{}: {}", message_count, e);
113                }
114            }
115        }
116
117        error!("Message processing task ended unexpectedly!");
118    });
119
120    // Add a monitoring task to track connection health
121    let start_time = std::time::Instant::now();
122    tokio::spawn(async move {
123        let mut interval = tokio::time::interval(std::time::Duration::from_secs(15));
124        let mut tick_count = 0;
125
126        loop {
127            interval.tick().await;
128            tick_count += 1;
129            let elapsed = start_time.elapsed().as_secs();
130            info!(
131                "Health check #{} - uptime: {}s (expecting reconnects every ~20s)",
132                tick_count, elapsed
133            );
134        }
135    });
136
137    info!("Connecting to jetstream (testing reconnection fix)...");
138    info!("Watch for 'Stream closed by server' or 'No messages received' followed by reconnection");
139
140    // This should connect, sit idle, get disconnected, then reconnect automatically
141    // The fix should prevent infinite loops when reconnecting
142    jetstream
143        .connect(cursor.clone())
144        .await
145        .map_err(|e| anyhow::anyhow!("error running ingest: {}", e))
146}
147
148/// Simple ingestor that should never be called
149pub struct ProfileIngestor;
150
151#[async_trait]
152impl LexiconIngestor for ProfileIngestor {
153    async fn ingest(&self, message: Event<Value>) -> Result<()> {
154        // If we get messages, they're probably deletions with no record
155        if let Some(ref commit) = message.commit {
156            if commit.record.is_some() {
157                info!("got message with record: {:?}", message);
158            } else {
159                info!(
160                    "ProfileIngestor got deletion/empty record for DID: {}",
161                    message.did
162                );
163            }
164        } else {
165            info!("ProfileIngestor got message with no commit: {:?}", message);
166        }
167        Ok(())
168    }
169}