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#[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 "app.bsky.feed.post".to_string(),
49 ])
50 .timeout_time_sec(20) .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 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 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 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 jetstream
143 .connect(cursor.clone())
144 .await
145 .map_err(|e| anyhow::anyhow!("error running ingest: {}", e))
146}
147
148pub struct ProfileIngestor;
150
151#[async_trait]
152impl LexiconIngestor for ProfileIngestor {
153 async fn ingest(&self, message: Event<Value>) -> Result<()> {
154 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}