1use crate::errors::ConsumerError;
43use anyhow::Result;
44use async_trait::async_trait;
45use futures::{SinkExt, StreamExt};
46use http::Uri;
47use serde::{Deserialize, Serialize};
48use std::sync::Arc;
49use std::{collections::HashMap, str::FromStr};
50use tokio::sync::{RwLock, broadcast};
51use tokio::time::{Instant, sleep};
52use tokio_util::sync::CancellationToken;
53use tokio_websockets::{ClientBuilder, Message};
54use tracing::Instrument;
55
56const MAX_MESSAGE_SIZE: usize = 56000;
57
58#[cfg_attr(any(debug_assertions, test), derive(Debug))]
60#[derive(Clone)]
61pub struct ConsumerTaskConfig {
62 pub user_agent: String,
64 pub compression: bool,
66 pub zstd_dictionary_location: String,
68 pub jetstream_hostname: String,
70 pub collections: Vec<String>,
72 pub dids: Vec<String>,
74 pub max_message_size_bytes: Option<u64>,
76 pub cursor: Option<i64>,
78 pub require_hello: bool,
80}
81
82#[cfg_attr(any(debug_assertions, test), derive(Debug))]
84#[derive(Clone, Serialize, Deserialize)]
85#[serde(untagged)]
86pub enum JetstreamEvent {
87 Commit {
89 did: String,
91 time_us: u64,
93 kind: String,
95 commit: JetstreamEventCommit,
97 },
98
99 Delete {
101 did: String,
103 time_us: u64,
105 kind: String,
107 commit: JetstreamEventDelete,
109 },
110
111 Identity {
113 did: String,
115 time_us: u64,
117 kind: String,
119 identity: serde_json::Value,
121 },
122
123 Account {
125 did: String,
127 time_us: u64,
129 kind: String,
131 account: serde_json::Value,
133 },
134}
135
136#[cfg_attr(any(debug_assertions, test), derive(Debug))]
138#[derive(Clone, Serialize, Deserialize)]
139pub struct JetstreamEventCommit {
140 pub rev: String,
142 pub operation: String,
144 pub collection: String,
146 pub rkey: String,
148 pub cid: String,
150 pub record: serde_json::Value,
152}
153
154#[cfg_attr(any(debug_assertions, test), derive(Debug))]
156#[derive(Clone, Serialize, Deserialize)]
157pub struct JetstreamEventDelete {
158 pub rev: String,
160 pub operation: String,
162 pub collection: String,
164 pub rkey: String,
166}
167
168#[async_trait]
170pub trait EventHandler: Send + Sync {
171 async fn handle_event(&self, event: Arc<JetstreamEvent>) -> Result<()>;
176
177 fn handler_id(&self) -> &str;
181}
182
183#[cfg_attr(any(debug_assertions, test), derive(Debug))]
184#[derive(Clone, Serialize, Deserialize)]
185#[serde(tag = "type", content = "payload")]
186pub(crate) enum SubscriberSourcedMessage {
187 #[serde(rename = "options_update")]
188 Update {
189 #[serde(
190 rename = "wantedCollections",
191 skip_serializing_if = "Vec::is_empty",
192 default
193 )]
194 wanted_collections: Vec<String>,
195
196 #[serde(rename = "wantedDids", skip_serializing_if = "Vec::is_empty", default)]
197 wanted_dids: Vec<String>,
198
199 #[serde(rename = "maxMessageSizeBytes")]
200 max_message_size_bytes: u64,
201
202 #[serde(skip_serializing_if = "Option::is_none")]
203 cursor: Option<i64>,
204 },
205}
206
207pub struct Consumer {
209 config: ConsumerTaskConfig,
210 handlers: Arc<RwLock<HashMap<String, Arc<dyn EventHandler>>>>,
211 event_sender: Arc<RwLock<Option<broadcast::Sender<Arc<JetstreamEvent>>>>>,
212}
213
214impl Consumer {
215 pub fn new(config: ConsumerTaskConfig) -> Self {
217 Self {
218 config,
219 handlers: Arc::new(RwLock::new(HashMap::new())),
220 event_sender: Arc::new(RwLock::new(None)),
221 }
222 }
223
224 pub async fn register_handler(&self, handler: Arc<dyn EventHandler>) -> Result<()> {
226 let handler_id = handler.handler_id();
227 let mut handlers = self.handlers.write().await;
228
229 if handlers.contains_key(handler_id) {
230 return Err(ConsumerError::HandlerRegistrationFailed(format!(
231 "Handler with ID '{}' already registered",
232 handler_id
233 ))
234 .into());
235 }
236
237 handlers.insert(handler_id.to_string(), handler);
238 Ok(())
239 }
240
241 pub async fn unregister_handler(&self, handler_id: &str) -> Result<()> {
243 let mut handlers = self.handlers.write().await;
244 handlers.remove(handler_id);
245 Ok(())
246 }
247
248 pub async fn get_event_receiver(&self) -> Result<broadcast::Receiver<Arc<JetstreamEvent>>> {
252 let sender_guard = self.event_sender.read().await;
253 match sender_guard.as_ref() {
254 Some(sender) => Ok(sender.subscribe()),
255 None => Err(ConsumerError::EventSenderNotInitialized(
256 "consumer not running".to_string(),
257 )
258 .into()),
259 }
260 }
261
262 pub async fn run_background(&self, cancellation_token: CancellationToken) -> Result<()> {
292 tracing::info!("Starting Jetstream consumer");
293
294 let capacity = 50 + self.config.collections.len() * 30 + self.config.dids.len() * 60; let mut query_string = String::with_capacity(capacity);
300
301 query_string.push_str("compress=");
303 query_string.push_str(if self.config.compression {
304 "true"
305 } else {
306 "false"
307 });
308
309 query_string.push_str("&requireHello=");
311 query_string.push_str(if self.config.require_hello {
312 "true"
313 } else {
314 "false"
315 });
316
317 if !self.config.collections.is_empty() && !self.config.require_hello {
319 for collection in &self.config.collections {
320 query_string.push_str("&wantedCollections=");
321 query_string.push_str(&urlencoding::encode(collection));
322 }
323 }
324
325 if !self.config.dids.is_empty() && !self.config.require_hello {
327 for did in &self.config.dids {
328 query_string.push_str("&wantedDids=");
329 query_string.push_str(&urlencoding::encode(did));
330 }
331 }
332
333 if let Some(max_size) = self.config.max_message_size_bytes {
335 use std::fmt::Write;
336 write!(&mut query_string, "&maxMessageSizeBytes={}", max_size).unwrap();
337 }
338
339 if let Some(cursor) = self.config.cursor {
341 use std::fmt::Write;
342 write!(&mut query_string, "&cursor={}", cursor).unwrap();
343 }
344 let ws_url = Uri::from_str(&format!(
345 "wss://{}/subscribe?{}",
346 self.config.jetstream_hostname, query_string
347 ))?;
348
349 let (mut client, _) = ClientBuilder::from_uri(ws_url)
350 .add_header(
351 http::header::USER_AGENT,
352 http::HeaderValue::from_str(&self.config.user_agent)?,
353 )?
354 .connect()
355 .await?;
356
357 let update = SubscriberSourcedMessage::Update {
358 wanted_collections: self.config.collections.clone(),
359 wanted_dids: self.config.dids.clone(),
360 max_message_size_bytes: self
361 .config
362 .max_message_size_bytes
363 .unwrap_or(MAX_MESSAGE_SIZE as u64),
364 cursor: self.config.cursor,
365 };
366 let serialized_update = serde_json::to_string(&update)
367 .map_err(|err| ConsumerError::UpdateSerializationFailed(err.to_string()))?;
368
369 client
370 .send(Message::text(serialized_update))
371 .await
372 .map_err(|err| ConsumerError::UpdateSendFailed(err.to_string()))?;
373
374 let mut decompressor = if self.config.compression {
375 let data: Vec<u8> = std::fs::read(self.config.zstd_dictionary_location.clone())?;
377 zstd::bulk::Decompressor::with_dictionary(&data)
378 .map_err(|err| ConsumerError::DecompressorCreationFailed(err.to_string()))?
379 } else {
380 zstd::bulk::Decompressor::new()
381 .map_err(|err| ConsumerError::DecompressorCreationFailed(err.to_string()))?
382 };
383
384 let interval = std::time::Duration::from_secs(120);
385 let sleeper = sleep(interval);
386 tokio::pin!(sleeper);
387
388 loop {
389 tokio::select! {
390 () = cancellation_token.cancelled() => {
391 break;
392 },
393 () = &mut sleeper => {
394 sleeper.as_mut().reset(Instant::now() + interval);
395 },
396 item = client.next() => {
397 if item.is_none() {
398 tracing::warn!("jetstream connection closed");
399 break;
400 }
401 let item = item.unwrap();
402
403 if let Err(err) = item {
404 tracing::error!(error = ?err, "error processing jetstream message");
405 continue;
406 }
407 let item = item.unwrap();
408
409 let event = if self.config.compression {
410 if !item.is_binary() {
411 tracing::debug!("compression enabled but message from jetstream is not binary");
412 continue;
413 }
414 let payload = item.into_payload();
415
416 let decoded = decompressor.decompress(&payload, MAX_MESSAGE_SIZE * 3);
417 if let Err(err) = decoded {
418 tracing::debug!(err = ?err, "cannot decompress message");
419 continue;
420 }
421 let decoded = decoded.unwrap();
422 serde_json::from_slice::<JetstreamEvent>(&decoded)
423 .map_err(|err| ConsumerError::DeserializationFailed(err.to_string()))
424 } else {
425 if !item.is_text() {
426 tracing::debug!("compression disabled but message from jetstream is binary");
427 continue;
428 }
429 item.as_text()
430 .ok_or_else(|| ConsumerError::MessageConversionFailed("cannot convert message to text".to_string()))
431 .and_then(|value| {
432 serde_json::from_str::<JetstreamEvent>(value)
433 .map_err(|err| ConsumerError::DeserializationFailed(err.to_string()))
434 })
435 };
436 if let Err(err) = event {
437 tracing::error!(error = ?err, "error processing jetstream message");
438
439 continue;
440 }
441 let event = event.unwrap();
442
443 if let Err(err) = self.dispatch_to_handlers(event).await {
444 tracing::error!(error = ?err, "Failed to process message");
445 }
446
447 }
448 }
449 }
450
451 {
453 let mut sender_guard = self.event_sender.write().await;
454 *sender_guard = None;
455 }
456
457 Ok(())
458 }
459
460 async fn dispatch_to_handlers(&self, event: JetstreamEvent) -> Result<()> {
465 let handlers = self.handlers.read().await;
466 let event = Arc::new(event);
467
468 for (handler_id, handler) in handlers.iter() {
469 let handler_span = tracing::debug_span!("handler_dispatch", handler_id = %handler_id);
470 let event_ref = Arc::clone(&event);
471 async {
472 if let Err(err) = handler.handle_event(event_ref).await {
473 tracing::error!(
474 error = ?err,
475 handler_id = %handler_id,
476 "Handler failed to process event"
477 );
478 }
479 }
480 .instrument(handler_span)
481 .await;
482 }
483
484 Ok(())
485 }
486}
487
488pub struct LoggingHandler {
490 id: String,
491}
492
493impl LoggingHandler {
494 pub fn new(id: String) -> Self {
496 Self { id }
497 }
498}
499
500#[async_trait]
501impl EventHandler for LoggingHandler {
502 async fn handle_event(&self, _event: Arc<JetstreamEvent>) -> Result<()> {
503 Ok(())
504 }
505
506 fn handler_id(&self) -> &str {
507 &self.id
508 }
509}
510
511#[cfg(test)]
512mod tests {
513 use super::*;
514
515 #[test]
516 fn test_parse_account_event() {
517 let json_str = r#"{"did":"did:plc:yn72uqr4ihkjfbz7us7buqsq","time_us":1757517640675638,"kind":"account","account":{"active":false,"did":"did:plc:yn72uqr4ihkjfbz7us7buqsq","seq":13206502767,"status":"takendown","time":"2025-09-10T15:20:40.439Z"}}"#;
518
519 let event = serde_json::from_str::<JetstreamEvent>(json_str)
520 .expect("Failed to parse account event JSON");
521
522 match event {
523 JetstreamEvent::Account {
524 did,
525 time_us,
526 kind,
527 account: identity,
528 } => {
529 assert_eq!(did, "did:plc:yn72uqr4ihkjfbz7us7buqsq");
530 assert_eq!(time_us, 1757517640675638);
531 assert_eq!(kind, "account");
532
533 assert!(identity.is_object());
535 let account_obj = identity.as_object().unwrap();
536 assert_eq!(
537 account_obj.get("active").unwrap(),
538 &serde_json::json!(false)
539 );
540 assert_eq!(
541 account_obj.get("did").unwrap(),
542 &serde_json::json!("did:plc:yn72uqr4ihkjfbz7us7buqsq")
543 );
544 assert_eq!(
545 account_obj.get("seq").unwrap(),
546 &serde_json::json!(13206502767i64)
547 );
548 assert_eq!(
549 account_obj.get("status").unwrap(),
550 &serde_json::json!("takendown")
551 );
552 assert_eq!(
553 account_obj.get("time").unwrap(),
554 &serde_json::json!("2025-09-10T15:20:40.439Z")
555 );
556 }
557 _ => panic!("Expected JetstreamEvent::Account variant, got {:?}", event),
558 }
559 }
560
561 #[test]
562 fn test_parse_identity_event() {
563 let json_str = r#"{"did":"did:plc:mbuadp4xzlbmc2ncqp3pmtox","time_us":1757517628039893,"kind":"identity","identity":{"did":"did:plc:mbuadp4xzlbmc2ncqp3pmtox","handle":"nhieothv.bsky.social","seq":13206497272,"time":"2025-09-10T15:20:27.610Z"}}"#;
564
565 let event = serde_json::from_str::<JetstreamEvent>(json_str)
566 .expect("Failed to parse identity event JSON");
567
568 match event {
569 JetstreamEvent::Identity {
570 did,
571 time_us,
572 kind,
573 identity,
574 } => {
575 assert_eq!(did, "did:plc:mbuadp4xzlbmc2ncqp3pmtox");
576 assert_eq!(time_us, 1757517628039893);
577 assert_eq!(kind, "identity");
578
579 assert!(identity.is_object());
581 let identity_obj = identity.as_object().unwrap();
582 assert_eq!(
583 identity_obj.get("did").unwrap(),
584 &serde_json::json!("did:plc:mbuadp4xzlbmc2ncqp3pmtox")
585 );
586 assert_eq!(
587 identity_obj.get("handle").unwrap(),
588 &serde_json::json!("nhieothv.bsky.social")
589 );
590 assert_eq!(
591 identity_obj.get("seq").unwrap(),
592 &serde_json::json!(13206497272i64)
593 );
594 assert_eq!(
595 identity_obj.get("time").unwrap(),
596 &serde_json::json!("2025-09-10T15:20:27.610Z")
597 );
598 }
599 _ => panic!("Expected JetstreamEvent::Identity variant, got {:?}", event),
600 }
601 }
602
603 #[test]
604 fn test_parse_delete_event() {
605 let json_str = r#"{"did":"did:plc:5ozthefrqdo5kqnxzfgthhpp","time_us":1757519323847323,"kind":"commit","commit":{"rev":"3lyileto4q52k","operation":"delete","collection":"app.bsky.graph.follow","rkey":"3lxqxntaew32z"}}"#;
606
607 let event = serde_json::from_str::<JetstreamEvent>(json_str)
608 .expect("Failed to parse delete event JSON");
609
610 match event {
611 JetstreamEvent::Delete {
612 did,
613 time_us,
614 kind,
615 commit,
616 } => {
617 assert_eq!(did, "did:plc:5ozthefrqdo5kqnxzfgthhpp");
618 assert_eq!(time_us, 1757519323847323);
619 assert_eq!(kind, "commit");
620
621 assert_eq!(commit.rev, "3lyileto4q52k");
623 assert_eq!(commit.operation, "delete");
624 assert_eq!(commit.collection, "app.bsky.graph.follow");
625 assert_eq!(commit.rkey, "3lxqxntaew32z");
626 }
627 _ => panic!("Expected JetstreamEvent::Delete variant, got {:?}", event),
628 }
629 }
630
631 #[test]
632 fn test_parse_commit_event() {
633 let json_str = r#"{"did":"did:plc:suq5ijgyqmsawwf5tskf654x","time_us":1757519323848962,"kind":"commit","commit":{"rev":"3lyiletdopl2c","operation":"create","collection":"app.bsky.feed.like","rkey":"3lyiletddxt2c","record":{"$type":"app.bsky.feed.like","createdAt":"2025-09-10T15:47:13.086Z","subject":{"cid":"bafyreib2pygab7z5l7nkqf6bchcvgt4jwsqiaenpf3sr65lugum2uvzzf4","uri":"at://did:plc:yw65rktdby2chplqdytqzcao/app.bsky.feed.post/3lyildyjxgs2o"}},"cid":"bafyreigroo6vhxt62ufcndhaxzas6btq4jmniuz4egszbwuqgiyisqwqoy"}}"#;
634
635 let event = serde_json::from_str::<JetstreamEvent>(json_str)
636 .expect("Failed to parse commit event JSON");
637
638 match event {
639 JetstreamEvent::Commit {
640 did,
641 time_us,
642 kind,
643 commit,
644 } => {
645 assert_eq!(did, "did:plc:suq5ijgyqmsawwf5tskf654x");
646 assert_eq!(time_us, 1757519323848962);
647 assert_eq!(kind, "commit");
648
649 assert_eq!(commit.rev, "3lyiletdopl2c");
651 assert_eq!(commit.operation, "create");
652 assert_eq!(commit.collection, "app.bsky.feed.like");
653 assert_eq!(commit.rkey, "3lyiletddxt2c");
654 assert_eq!(
655 commit.cid,
656 "bafyreigroo6vhxt62ufcndhaxzas6btq4jmniuz4egszbwuqgiyisqwqoy"
657 );
658
659 assert!(commit.record.is_object());
661 let record_obj = commit.record.as_object().unwrap();
662 assert_eq!(
663 record_obj.get("$type").unwrap(),
664 &serde_json::json!("app.bsky.feed.like")
665 );
666 assert_eq!(
667 record_obj.get("createdAt").unwrap(),
668 &serde_json::json!("2025-09-10T15:47:13.086Z")
669 );
670
671 let subject = record_obj.get("subject").unwrap().as_object().unwrap();
673 assert_eq!(
674 subject.get("cid").unwrap(),
675 &serde_json::json!(
676 "bafyreib2pygab7z5l7nkqf6bchcvgt4jwsqiaenpf3sr65lugum2uvzzf4"
677 )
678 );
679 assert_eq!(
680 subject.get("uri").unwrap(),
681 &serde_json::json!(
682 "at://did:plc:yw65rktdby2chplqdytqzcao/app.bsky.feed.post/3lyildyjxgs2o"
683 )
684 );
685 }
686 _ => panic!("Expected JetstreamEvent::Commit variant, got {:?}", event),
687 }
688 }
689
690 #[test]
691 fn test_parse_commit_update_event() {
692 let json_str = r#"{"did":"did:plc:mek6cpladv2xrlu2zdykoxgz","time_us":1757519523286358,"kind":"commit","commit":{"rev":"3lyilmalk762z","operation":"update","collection":"app.bsky.actor.profile","rkey":"self","record":{"$type":"app.bsky.actor.profile","avatar":{"$type":"blob","ref":{"$link":"bafkreibmn7xi5iwugioov463wux62dg4m4w6qqrbsnileaobrzgxdwbsqy"},"mimeType":"image/jpeg","size":289838},"banner":{"$type":"blob","ref":{"$link":"bafkreicjgdlfs6fyyddjklfzrf6w2boychodkdebtjaiwavhcffjtuavsi"},"mimeType":"image/jpeg","size":676693},"description":"ela/dela | parte da fauna fantástica do céu azul | praticamente inofensiva","displayName":"la mucura mística","pinnedPost":{"cid":"bafyreihn2t4efvipbcignd6rlybmoecb7hx4jgntsojhpibjzxno3zhbuq","uri":"at://did:plc:mek6cpladv2xrlu2zdykoxgz/app.bsky.feed.post/3lxarfbd4ts2j"}},"cid":"bafyreifpmgw3podvvm4raq6zewn6jhoa73t7mlgf3f7hty2adb6f2ga7j4"}}"#;
693
694 let event = serde_json::from_str::<JetstreamEvent>(json_str)
695 .expect("Failed to parse commit update event JSON");
696
697 match event {
698 JetstreamEvent::Commit {
699 did,
700 time_us,
701 kind,
702 commit,
703 } => {
704 assert_eq!(did, "did:plc:mek6cpladv2xrlu2zdykoxgz");
705 assert_eq!(time_us, 1757519523286358);
706 assert_eq!(kind, "commit");
707
708 assert_eq!(commit.rev, "3lyilmalk762z");
710 assert_eq!(commit.operation, "update");
711 assert_eq!(commit.collection, "app.bsky.actor.profile");
712 assert_eq!(commit.rkey, "self");
713 assert_eq!(
714 commit.cid,
715 "bafyreifpmgw3podvvm4raq6zewn6jhoa73t7mlgf3f7hty2adb6f2ga7j4"
716 );
717
718 assert!(commit.record.is_object());
720 let record_obj = commit.record.as_object().unwrap();
721 assert_eq!(
722 record_obj.get("$type").unwrap(),
723 &serde_json::json!("app.bsky.actor.profile")
724 );
725 assert_eq!(
726 record_obj.get("description").unwrap(),
727 &serde_json::json!(
728 "ela/dela | parte da fauna fantástica do céu azul | praticamente inofensiva"
729 )
730 );
731 assert_eq!(
732 record_obj.get("displayName").unwrap(),
733 &serde_json::json!("la mucura mística")
734 );
735
736 let avatar = record_obj.get("avatar").unwrap().as_object().unwrap();
738 assert_eq!(avatar.get("$type").unwrap(), &serde_json::json!("blob"));
739 assert_eq!(
740 avatar.get("mimeType").unwrap(),
741 &serde_json::json!("image/jpeg")
742 );
743 assert_eq!(avatar.get("size").unwrap(), &serde_json::json!(289838));
744
745 let banner = record_obj.get("banner").unwrap().as_object().unwrap();
747 assert_eq!(banner.get("$type").unwrap(), &serde_json::json!("blob"));
748 assert_eq!(
749 banner.get("mimeType").unwrap(),
750 &serde_json::json!("image/jpeg")
751 );
752 assert_eq!(banner.get("size").unwrap(), &serde_json::json!(676693));
753
754 let pinned_post = record_obj.get("pinnedPost").unwrap().as_object().unwrap();
756 assert_eq!(
757 pinned_post.get("cid").unwrap(),
758 &serde_json::json!(
759 "bafyreihn2t4efvipbcignd6rlybmoecb7hx4jgntsojhpibjzxno3zhbuq"
760 )
761 );
762 assert_eq!(
763 pinned_post.get("uri").unwrap(),
764 &serde_json::json!(
765 "at://did:plc:mek6cpladv2xrlu2zdykoxgz/app.bsky.feed.post/3lxarfbd4ts2j"
766 )
767 );
768 }
769 _ => panic!("Expected JetstreamEvent::Commit variant, got {:?}", event),
770 }
771 }
772}