this repo has no description
lewis.moe
1use crate::state::AppState;
2use crate::sync::firehose::SequencedEvent;
3use crate::sync::util::{format_event_for_sending, format_event_with_prefetched_blocks, prefetch_blocks_for_events};
4use axum::{
5 extract::{ws::Message, ws::WebSocket, ws::WebSocketUpgrade, Query, State},
6 response::Response,
7};
8use futures::{sink::SinkExt, stream::StreamExt};
9use serde::Deserialize;
10use std::sync::atomic::{AtomicUsize, Ordering};
11use tokio::sync::broadcast::error::RecvError;
12use tracing::{error, info, warn};
13
14const BACKFILL_BATCH_SIZE: i64 = 1000;
15static SUBSCRIBER_COUNT: AtomicUsize = AtomicUsize::new(0);
16
17#[derive(Deserialize)]
18pub struct SubscribeReposParams {
19 pub cursor: Option<i64>,
20}
21
22#[axum::debug_handler]
23pub async fn subscribe_repos(
24 ws: WebSocketUpgrade,
25 State(state): State<AppState>,
26 Query(params): Query<SubscribeReposParams>,
27) -> Response {
28 ws.on_upgrade(move |socket| handle_socket(socket, state, params))
29}
30
31async fn send_event(
32 socket: &mut WebSocket,
33 state: &AppState,
34 event: SequencedEvent,
35) -> Result<(), anyhow::Error> {
36 let bytes = format_event_for_sending(state, event).await?;
37 socket.send(Message::Binary(bytes.into())).await?;
38 Ok(())
39}
40
41pub fn get_subscriber_count() -> usize {
42 SUBSCRIBER_COUNT.load(Ordering::SeqCst)
43}
44
45async fn handle_socket(mut socket: WebSocket, state: AppState, params: SubscribeReposParams) {
46 let count = SUBSCRIBER_COUNT.fetch_add(1, Ordering::SeqCst) + 1;
47 crate::metrics::set_firehose_subscribers(count);
48 info!(cursor = ?params.cursor, subscribers = count, "New firehose subscriber");
49
50 let _ = handle_socket_inner(&mut socket, &state, params).await;
51
52 let count = SUBSCRIBER_COUNT.fetch_sub(1, Ordering::SeqCst) - 1;
53 crate::metrics::set_firehose_subscribers(count);
54 info!(subscribers = count, "Firehose subscriber disconnected");
55}
56
57async fn handle_socket_inner(socket: &mut WebSocket, state: &AppState, params: SubscribeReposParams) -> Result<(), ()> {
58 if let Some(cursor) = params.cursor {
59 let mut current_cursor = cursor;
60 loop {
61 let events = sqlx::query_as!(
62 SequencedEvent,
63 r#"
64 SELECT seq, did, created_at, event_type, commit_cid, prev_cid, ops, blobs, blocks_cids
65 FROM repo_seq
66 WHERE seq > $1
67 ORDER BY seq ASC
68 LIMIT $2
69 "#,
70 current_cursor,
71 BACKFILL_BATCH_SIZE
72 )
73 .fetch_all(&state.db)
74 .await;
75
76 match events {
77 Ok(events) => {
78 if events.is_empty() {
79 break;
80 }
81
82 let events_count = events.len();
83
84 let prefetched = match prefetch_blocks_for_events(state, &events).await {
85 Ok(blocks) => blocks,
86 Err(e) => {
87 error!("Failed to prefetch blocks for backfill: {}", e);
88 socket.close().await.ok();
89 return Err(());
90 }
91 };
92
93 for event in events {
94 current_cursor = event.seq;
95 let bytes = match format_event_with_prefetched_blocks(event, &prefetched).await {
96 Ok(b) => b,
97 Err(e) => {
98 warn!("Failed to format backfill event: {}", e);
99 return Err(());
100 }
101 };
102 if let Err(e) = socket.send(Message::Binary(bytes.into())).await {
103 warn!("Failed to send backfill event: {}", e);
104 return Err(());
105 }
106 crate::metrics::record_firehose_event();
107 }
108 if (events_count as i64) < BACKFILL_BATCH_SIZE {
109 break;
110 }
111 }
112 Err(e) => {
113 error!("Failed to fetch backfill events: {}", e);
114 socket.close().await.ok();
115 return Err(());
116 }
117 }
118 }
119 }
120
121 let mut rx = state.firehose_tx.subscribe();
122 let max_lag_before_disconnect: u64 = std::env::var("FIREHOSE_MAX_LAG")
123 .ok()
124 .and_then(|v| v.parse().ok())
125 .unwrap_or(5000);
126
127 loop {
128 tokio::select! {
129 result = rx.recv() => {
130 match result {
131 Ok(event) => {
132 if let Err(e) = send_event(socket, state, event).await {
133 warn!("Failed to send event: {}", e);
134 break;
135 }
136 crate::metrics::record_firehose_event();
137 }
138 Err(RecvError::Lagged(skipped)) => {
139 warn!(skipped = skipped, "Firehose subscriber lagged behind");
140 if skipped > max_lag_before_disconnect {
141 warn!(skipped = skipped, max_lag = max_lag_before_disconnect,
142 "Disconnecting slow firehose consumer");
143 break;
144 }
145 }
146 Err(RecvError::Closed) => {
147 info!("Firehose channel closed");
148 break;
149 }
150 }
151 }
152 Some(Ok(msg)) = socket.next() => {
153 if let Message::Close(_) = msg {
154 info!("Client closed connection");
155 break;
156 }
157 }
158 else => {
159 break;
160 }
161 }
162 }
163 Ok(())
164}