+1

Cargo.lock

reviewed

··· 2655 2655 "time", 2656 2656 "tokio", 2657 2657 "tokio-tungstenite", 2658 2658 + "tokio-util", 2658 2659 "tracing", 2659 2660 "tracing-subscriber", 2660 2661 "url",

+1

crates/jetstream/Cargo.toml

reviewed

··· 22 22 futures-util = "0.3.31" 23 23 tokio = { version = "1.48.0", features = ["macros", "rt", "sync", "time"] } 24 24 tokio-tungstenite = { version = "0.28.0", features = ["native-tls"] } 25 25 + tokio-util = "0.7.17" 25 26 tracing-subscriber = "0.3.20" 26 27 27 28 zstd = { version = "0.13.3", optional = true }

+4

crates/jetstream/src/client.rs

reviewed

··· 7 7 use bytes::Bytes; 8 8 use std::sync::Arc; 9 9 use tokio::{sync::Notify, task::JoinHandle}; 10 10 + use tokio_util::sync::{CancellationToken, DropGuard}; 10 11 11 12 #[derive(Debug)] 12 13 pub struct JetstreamClient { 13 14 handle: JoinHandle<()>, 14 15 client_tx: flume::Sender<ClientCommand>, 15 16 metrics: Metrics, 17 17 + shutdown: DropGuard, 16 18 } 17 19 18 20 impl JetstreamClient { ··· 22 20 handle: tokio::task::JoinHandle<()>, 23 21 client_tx: flume::Sender<ClientCommand>, 24 22 metrics: Metrics, 23 23 + shutdown: CancellationToken, 25 24 ) -> Self { 26 25 Self { 27 26 handle, 28 27 client_tx, 29 28 metrics, 29 29 + shutdown: shutdown.drop_guard(), 30 30 } 31 31 } 32 32

+7 -4

crates/jetstream/src/client_builder.rs

reviewed

··· 1 1 use super::JetstreamClient; 2 2 - use crate::{Receiver, metrics::Metrics, types::Options}; 2 2 + use crate::{Receiver, metrics::Metrics, subscriber_options::Options}; 3 3 use atproto::{Did, Nsid}; 4 4 + use tokio_util::sync::CancellationToken; 4 5 use url::Url; 5 6 6 7 #[derive(Default)] ··· 15 14 let (event_tx, event_rx) = flume::bounded(8); 16 15 let (client_tx, client_rx) = flume::bounded(8); 17 16 17 17 + let shutdown = CancellationToken::new(); 18 18 let metrics = Metrics::new(); 19 19 let handle = tokio::task::spawn(crate::task::jetstream_subscriber( 20 20 event_tx, ··· 24 22 instance, 25 23 self.options, 26 24 self.cursor, 25 25 + shutdown.child_token(), 27 26 )); 28 27 29 29 - let client = JetstreamClient::new(handle, client_tx, metrics); 28 28 + let client = JetstreamClient::new(handle, client_tx, metrics, shutdown); 30 29 let receiver = Receiver::new(event_rx); 31 30 32 31 (client, receiver) 33 32 } 34 33 35 34 pub fn add_collection(&mut self, nsid: Box<Nsid>) -> &mut Self { 36 36 - self.options.insert_collection(nsid); 35 35 + self.options.add_collection(nsid); 37 36 self 38 37 } 39 38 ··· 51 48 } 52 49 53 50 pub fn add_did(&mut self, did: Box<Did>) -> &mut Self { 54 54 - self.options.insert_did(did); 51 51 + self.options.add_did(did); 55 52 self 56 53 } 57 54

+3 -1

crates/jetstream/src/de.rs

reviewed

··· 6 6 // @NOTE 7 7 // 8 8 // Using `serde_json::value::RawValue` breaks if any part of the deserialization 9 9 - // features a tagged enum or if `#[serde(flatten)]` is used. 9 9 + // features a tagged enum or `#[serde(flatten)]` is used. 10 10 // 11 11 + // To get around this limitation we first deserialize into a less ideal type, 12 12 + // then manually transform into the types we want. 11 13 12 14 #[derive(Debug, Deserialize, Serialize)] 13 15 struct InnerEvent<'a> {

+1 -1

crates/jetstream/src/lib.rs

reviewed

··· 2 2 mod client_builder; 3 3 mod de; 4 4 pub mod metrics; 5 5 + pub mod subscriber_options; 5 6 mod task; 6 6 - pub mod types; 7 7 8 8 pub use atproto::{Did, Nsid}; 9 9 pub use client::{EventRef, JetstreamClient, JetstreamClientError, Receiver};

+3 -15

crates/jetstream/src/main.rs

reviewed

··· 91 91 writeln!(log, "{msg}").expect("Failed to write to log file"); 92 92 } 93 93 94 94 - let event = match message.deserialize() { 95 95 - Ok(message) => message, 96 96 - Err(error) => { 97 97 - tracing::error!(?error, %msg, "failed to deserialize event"); 98 98 - continue; 99 99 - } 100 100 - }; 101 101 - 102 102 - match event { 103 103 - Event::Account { .. } if arguments.hide_account => continue, 104 104 - Event::Identity { .. } if arguments.hide_identity => continue, 105 105 - _ => {} 94 94 + if let Ok(Event::Commit(commit)) = message.deserialize() { 95 95 + println!("{commit:#?}"); 96 96 + eprintln!("{:?}", client.metrics()); 106 97 } 107 107 - 108 108 - println!("{event:#?}"); 109 109 - eprintln!("{:?}", client.metrics()); 110 98 } 111 99 112 100 client.join().await.unwrap();

+11 -7

crates/jetstream/src/metrics.rs

reviewed

··· 1 1 use std::sync::{Arc, Mutex, MutexGuard}; 2 2 3 3 + /// Jetstream client metrics. 3 4 #[derive(Debug, Default)] 4 5 pub struct Metrics { 5 6 inner: Arc<Mutex<MetricsData>>, ··· 36 35 Self::default() 37 36 } 38 37 38 38 + /// Create a clone of the metrics data. 39 39 pub fn export(&self) -> MetricsData { 40 40 - self.lock().clone() 40 40 + self.inner.lock().unwrap().clone() 41 41 } 42 42 43 43 pub(crate) fn modify<F>(&self, f: F) 44 44 where 45 45 F: Fn(MutexGuard<MetricsData>), 46 46 { 47 47 - let guard = self.lock(); 47 47 + let guard = self.inner.lock().unwrap(); 48 48 f(guard); 49 49 } 50 50 51 51 - pub(crate) fn lock(&self) -> MutexGuard<'_, MetricsData> { 52 52 - self.inner.lock().unwrap_or_else(|poison| { 53 53 - self.inner.clear_poison(); 54 54 - poison.into_inner() 55 55 - }) 51 51 + pub(crate) fn increment_message_kind(&self, kind: &str) { 52 52 + match kind { 53 53 + "account" => self.inner.lock().unwrap().account_messages += 1, 54 54 + "commit" => self.inner.lock().unwrap().commit_messages += 1, 55 55 + "identity" => self.inner.lock().unwrap().identity_messages += 1, 56 56 + _ => self.inner.lock().unwrap().unknown_messages += 1, 57 57 + } 56 58 } 57 59 }

+220

crates/jetstream/src/subscriber_options.rs

reviewed

··· 1 1 + use core::{error, fmt}; 2 2 + use std::collections::HashSet; 3 3 + 4 4 + use serde::{Serialize, Serializer}; 5 5 + 6 6 + use crate::{Did, Nsid}; 7 7 + 8 8 + pub const MAX_WANTED_COLLECTIONS: usize = 100; 9 9 + 10 10 + pub const MAX_WANTED_DIDS: usize = 10_000; 11 11 + 12 12 + // @TODO Review 13 13 + pub const MAX_URL_LENGTH: usize = 4000; 14 14 + 15 15 + /// Subscriber sourced message. 16 16 + /// 17 17 + /// Ref: <https://github.com/bluesky-social/jetstream?tab=readme-ov-file#subscriber-sourced-messages> 18 18 + /// 19 19 + #[derive(Debug, Serialize)] 20 20 + #[serde(tag = "type", content = "payload", rename_all = "snake_case")] 21 21 + pub enum SubscriberSourcedMessage<'a> { 22 22 + OptionsUpdate(&'a Options), 23 23 + } 24 24 + 25 25 + impl<'a> SubscriberSourcedMessage<'a> { 26 26 + /// Serialize the [`SubscriberSourcedMessage`] to JSON. 27 27 + pub fn to_json(&self) -> String { 28 28 + serde_json::to_string(self).expect("SubscriberSourcedMessage should be serializable") 29 29 + } 30 30 + } 31 31 + 32 32 + /// Jetstream subscription options. 33 33 + /// 34 34 + /// Can either be appended to the `/subscribe` URL on connection to the Jetstream instance 35 35 + /// or sent as an options update message after connection. 36 36 + /// 37 37 + /// Ref: <https://github.com/bluesky-social/jetstream?tab=readme-ov-file#options-updates> 38 38 + /// 39 39 + #[derive(Debug, Default, Serialize)] 40 40 + #[serde(rename_all = "camelCase")] 41 41 + pub struct Options { 42 42 + /// Collection NSIDs to filter which records are received. 43 43 + /// 44 44 + /// Maximum: 100 45 45 + wanted_collections: HashSet<Box<Nsid>>, 46 46 + 47 47 + /// Repository DIDs to filter which records are received. 48 48 + /// 49 49 + /// Maximum: 10_000 50 50 + wanted_dids: HashSet<Box<Did>>, 51 51 + 52 52 + /// Maximum message size in bytes the subscriber wants to receive. 53 53 + /// 54 54 + /// Zero means no limit, negative values are treated as zero by Jetstream, and 55 55 + /// will be normalized to zero when serialized. 56 56 + #[serde(serialize_with = "serialize_max_message_size")] 57 57 + pub max_message_size_bytes: i64, 58 58 + } 59 59 + 60 60 + fn normalize_max_message_size(value: i64) -> i64 { 61 61 + value.abs() 62 62 + } 63 63 + 64 64 + fn serialize_max_message_size<S>(value: &i64, serializer: S) -> Result<S::Ok, S::Error> 65 65 + where 66 66 + S: Serializer, 67 67 + { 68 68 + serializer.serialize_i64(normalize_max_message_size(*value)) 69 69 + } 70 70 + 71 71 + impl Options { 72 72 + /// Add a collection NSID to the subscription options. 73 73 + /// 74 74 + /// Returns an error if the maximum number of subscribed collections has been reached; `Ok(true)` 75 75 + /// if the collection was newly added to the set, or `Ok(false)` if the colletion was already in the 76 76 + /// the set. 77 77 + pub fn add_collection(&mut self, collection: Box<Nsid>) -> Result<bool, Box<Nsid>> { 78 78 + if self.wanted_collections.len() == MAX_WANTED_COLLECTIONS 79 79 + && !self.wanted_collections.contains(&collection) 80 80 + { 81 81 + return Err(collection); 82 82 + } 83 83 + 84 84 + Ok(self.wanted_collections.insert(collection)) 85 85 + } 86 86 + 87 87 + pub fn remove_collection(&mut self, collection: &Nsid) -> bool { 88 88 + self.wanted_collections.remove(collection) 89 89 + } 90 90 + 91 91 + /// Add a DID to the subscription options. 92 92 + /// 93 93 + /// Returns an error if the maximum number of subscribed DIDs has been reached; `Ok(true)` 94 94 + /// if the DID was newly added to the set, or `Ok(false)` if the DID was already in the 95 95 + /// the set. 96 96 + pub fn add_did(&mut self, did: Box<Did>) -> Result<bool, Box<Did>> { 97 97 + if self.wanted_dids.len() == MAX_WANTED_DIDS && !self.wanted_dids.contains(&did) { 98 98 + return Err(did); 99 99 + } 100 100 + 101 101 + Ok(self.wanted_dids.insert(did)) 102 102 + } 103 103 + 104 104 + pub fn remove_did(&mut self, did: &Did) -> bool { 105 105 + self.wanted_dids.remove(did) 106 106 + } 107 107 + 108 108 + /// Get the normalized maximum message size. 109 109 + pub fn max_message_size(&self) -> i64 { 110 110 + normalize_max_message_size(self.max_message_size_bytes) 111 111 + } 112 112 + 113 113 + /// Compute the length, in bytes, of the query string required to express 114 114 + /// this set of options in a Jetstream `/subscribe` request. 115 115 + pub fn subscribe_url_len(&self, base: &url::Url) -> usize { 116 116 + const WANTED_DIDS_LEN: usize = "wantedDids=".len(); 117 117 + const WANTED_COLLECTIONS_LEN: usize = "wantedCollections=".len(); 118 118 + 119 119 + let (wdl, wdc): (_, usize) = self.wanted_dids.iter().fold((0, 0), |(sl, sc), v| { 120 120 + (sl + WANTED_DIDS_LEN + v.len(), sc + 1) 121 121 + }); 122 122 + 123 123 + let (wcl, wcc): (_, usize) = self.wanted_collections.iter().fold((0, 0), |(sl, sc), v| { 124 124 + (sl + WANTED_COLLECTIONS_LEN + v.len(), sc + 1) 125 125 + }); 126 126 + 127 127 + let mml = match self.max_message_size() { 128 128 + 0 => 0, 129 129 + n if wdc > 0 || wcc > 0 => 1 + n.to_string().len() + "maxMessageSizeBytes=".len(), 130 130 + n => n.to_string().len() + "maxMessageSizeBytes=".len(), 131 131 + }; 132 132 + 133 133 + mml + (wdl + wdc.saturating_sub(1)) + wcl + wcc.saturating_sub(1) + base.as_str().len() 134 134 + } 135 135 + 136 136 + /// Appends the subscription options to the specified URL. 137 137 + pub fn subscribe_url(&self, url: &url::Url, cursor: &Option<u128>) -> SubscribeMethod { 138 138 + let mut url = url.to_owned(); 139 139 + url.set_path("/subscribe"); 140 140 + url.query_pairs_mut().clear(); 141 141 + 142 142 + if let Some(cursor) = cursor { 143 143 + url.query_pairs_mut() 144 144 + .append_pair("cursor", &cursor.to_string()); 145 145 + } 146 146 + 147 147 + if self.subscribe_url_len(&url) > MAX_URL_LENGTH { 148 148 + url.query_pairs_mut().append_pair("requireHello", "true"); 149 149 + return SubscribeMethod::Hello(url); 150 150 + } 151 151 + 152 152 + { 153 153 + let mut query = url.query_pairs_mut(); 154 154 + for collection in &self.wanted_collections { 155 155 + query.append_pair("wantedCollections", collection); 156 156 + } 157 157 + for did in &self.wanted_dids { 158 158 + query.append_pair("wantedDids", did.as_str()); 159 159 + } 160 160 + if self.max_message_size() > 0 { 161 161 + query.append_pair( 162 162 + "maxMessageSizeBytes", 163 163 + &self.max_message_size_bytes.to_string(), 164 164 + ); 165 165 + } 166 166 + } 167 167 + 168 168 + SubscribeMethod::Query(url) 169 169 + } 170 170 + } 171 171 + 172 172 + #[derive(Debug)] 173 173 + pub struct UrlTooLong; 174 174 + 175 175 + impl fmt::Display for UrlTooLong { 176 176 + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 177 177 + f.write_str("URL too long") 178 178 + } 179 179 + } 180 180 + 181 181 + impl error::Error for UrlTooLong {} 182 182 + 183 183 + pub enum SubscribeMethod { 184 184 + Query(url::Url), 185 185 + Hello(url::Url), 186 186 + } 187 187 + 188 188 + #[cfg(test)] 189 189 + mod tests { 190 190 + use atproto::Nsid; 191 191 + 192 192 + use super::Options; 193 193 + 194 194 + #[test] 195 195 + fn query_len() { 196 196 + let mut options = Options::default(); 197 197 + options 198 198 + .add_collection(Nsid::from_static("sh.tangled.*").into()) 199 199 + .unwrap(); 200 200 + 201 201 + assert_eq!( 202 202 + options.subscribe_url_len(&"https://example.url/subscribe".parse().unwrap()), 203 203 + "wantedCollections=sh.tangled.*".len() 204 204 + ); 205 205 + 206 206 + options 207 207 + .add_collection(Nsid::from_static("app.bsky.*").into()) 208 208 + .unwrap(); 209 209 + assert_eq!( 210 210 + options.subscribe_url_len(&"https://example.url/subscribe".parse().unwrap()), 211 211 + "wantedCollections=sh.tangled.*&wantedCollections=app.bsky.*".len() 212 212 + ); 213 213 + 214 214 + options.max_message_size_bytes = 1_000_000; 215 215 + assert_eq!( 216 216 + options.subscribe_url_len(&"https://example.url/subscribe".parse().unwrap()), 217 217 + "wantedCollections=sh.tangled.*&wantedCollections=app.bsky.*&maxMessageSizeBytes=1000000".len() 218 218 + ); 219 219 + } 220 220 + }

+60 -53

crates/jetstream/src/task.rs

reviewed

··· 1 1 - use crate::{client::ClientCommand, metrics::Metrics, types::Options}; 1 1 + use crate::{ 2 2 + client::ClientCommand, 3 3 + metrics::Metrics, 4 4 + subscriber_options::{Options, SubscribeMethod}, 5 5 + }; 2 6 use bytes::Bytes; 3 7 use futures_util::{SinkExt, StreamExt as _}; 4 8 use serde::Deserialize; 5 9 use std::{ops::ControlFlow, sync::Arc, time::Duration}; 6 10 use tokio::{sync::Notify, time::Instant}; 7 7 - use tokio_tungstenite::tungstenite::{Error as TungsteniteError, Message}; 11 11 + use tokio_tungstenite::{ 12 12 + connect_async, 13 13 + tungstenite::{Error as TungsteniteError, Message}, 14 14 + }; 15 15 + use tokio_util::sync::CancellationToken; 8 16 9 17 #[cfg(feature = "zstd")] 10 18 const ZSTD_DICTIONARY: &[u8] = include_bytes!("dictionary"); 11 11 - 12 12 - const MAX_URL_LENGTH: usize = 2000; 13 19 14 20 /// Duration since last receipt of a message to consider a subscription broken. 15 21 const THRESHOLD: Duration = Duration::from_secs(35); ··· 30 24 mut instance: url::Url, 31 25 mut options: Options, 32 26 initial_cursor: Option<u128>, 27 27 + shutdown: CancellationToken, 33 28 ) { 34 29 // Prepare the URL 35 30 instance.set_path("/subscribe"); ··· 52 45 // timeout.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); 53 46 54 47 'outer: loop { 55 55 - let mut url = instance.clone(); 48 48 + let (subscribe_url, require_hello) = match options.subscribe_url(&instance, &cursor) { 49 49 + SubscribeMethod::Query(url) => (url, false), 50 50 + SubscribeMethod::Hello(url) => (url, true), 51 51 + }; 52 52 + 53 53 + tracing::debug!(%subscribe_url, "connecting to jetstream"); 54 54 + let (socket, _) = match shutdown 55 55 + .run_until_cancelled(connect_async(subscribe_url.as_str())) 56 56 + .await 56 57 { 57 57 - let mut query = url.query_pairs_mut(); 58 58 - if let Some(cursor) = cursor { 59 59 - query.append_pair("cursor", &format!("{cursor}")); 60 60 - } 61 61 - } 62 62 - 63 63 - let mut require_hello = true; 64 64 - if url.as_str().len() + options.subscribe_query_len() < MAX_URL_LENGTH { 65 65 - tracing::debug!("appending subscribe options to URL"); 66 66 - options.append_to_url(&mut url); 67 67 - require_hello = false; 68 68 - } 69 69 - 70 70 - if require_hello { 71 71 - url.query_pairs_mut().append_pair("requireHello", "true"); 72 72 - } 73 73 - 74 74 - tracing::debug!(%url, "connecting to jetstream"); 75 75 - let (socket, _) = match tokio_tungstenite::connect_async(url.as_str()).await { 76 76 - Ok(socket) => socket, 77 77 - Err(error) => { 58 58 + Some(Ok(socket)) => socket, 59 59 + Some(Err(error)) => { 78 60 match error { 79 61 tokio_tungstenite::tungstenite::Error::Http(response) => { 80 62 let body = response.body().as_deref().unwrap_or_default(); ··· 75 79 tokio::time::sleep(Duration::from_secs(5)).await; 76 80 continue; 77 81 } 82 82 + None => break, 78 83 }; 79 84 80 85 metrics.modify(|mut data| data.connects += 1); 81 86 let (mut write, mut read) = socket.split(); 82 87 83 88 if require_hello { 84 84 - // Subscribe! 85 89 if let Err(e) = send_options_update::<_, TungsteniteError>(&mut write, &options).await { 86 90 tracing::error!(?e, "failed to send subscibe options update"); 87 91 continue; ··· 141 145 payload.into() 142 146 } 143 147 #[cfg(feature = "zstd")] 144 144 - Message::Binary(raw_payload) => { 148 148 + Message::Binary(compressed_payload) => { 145 149 use bytes::Buf as _; 146 150 use std::io::Read as _; 147 151 148 148 - let raw_bytes = raw_payload.len(); 149 149 - let mut payload = Vec::with_capacity(raw_payload.len()); 150 150 - let mut decoder = 151 151 - zstd::Decoder::with_prepared_dictionary(raw_payload.reader(), &dictionary) 152 152 - .expect("Failed to initialise zstd dictionary"); 152 152 + let compressed_bytes = compressed_payload.len(); 153 153 + let mut payload = Vec::with_capacity(compressed_payload.len()); 154 154 + let mut decoder = zstd::Decoder::with_prepared_dictionary( 155 155 + compressed_payload.reader(), 156 156 + &dictionary, 157 157 + ) 158 158 + .expect("prepared zstd dictionary should be valid"); 153 159 154 160 let Ok(_) = decoder.read_to_end(&mut payload) else { 155 161 tracing::error!("failed to dezstd message with zstd"); ··· 159 161 }; 160 162 161 163 metrics.modify(|mut data| { 162 162 - data.bytes_received_raw += raw_bytes; 164 164 + data.bytes_received_raw += compressed_bytes; 163 165 data.bytes_received += payload.len(); 164 166 }); 165 167 Bytes::from(payload) ··· 200 202 kind: &'a str, 201 203 } 202 204 203 203 - // Deserialize just the event timestamp. 205 205 + let mut new_cursor = cursor; 206 206 + 207 207 + // Deserialize just the event timestamp and kind. 204 208 match serde_json::from_slice::<PartialEvent>(&bytes) { 205 209 Ok(event) => { 206 206 - _ = { 207 207 - match event.kind { 208 208 - "identity" => metrics.modify(|mut data| data.identity_messages += 1), 209 209 - "account" => metrics.modify(|mut data| data.account_messages += 1), 210 210 - "commit" => metrics.modify(|mut data| data.commit_messages += 1), 211 211 - _ => metrics.modify(|mut data| data.unknown_messages += 1), 212 212 - } 213 213 - cursor.replace(event.time_us.into()) 214 214 - } 210 210 + metrics.increment_message_kind(event.kind); 211 211 + new_cursor.replace(event.time_us.into()); 215 212 } 216 213 Err(error) => { 217 214 match std::str::from_utf8(&bytes) { ··· 215 222 } 216 223 Err(_) => tracing::error!(?error, ?bytes, "failed to deserialize event"), 217 224 } 218 218 - continue; 225 225 + break; 219 226 } 220 227 } 221 228 ··· 228 235 } 229 236 break 'outer; 230 237 } 238 238 + 239 239 + // Update the cursor since the message has been dispatched. 240 240 + cursor = new_cursor; 231 241 } 232 242 233 243 metrics.modify(|mut data| data.disconnects += 1); ··· 244 248 S: SinkExt<Message> + Unpin, 245 249 E: From<S::Error>, 246 250 { 247 247 - use crate::types::SubscriberOptionsUpdate; 251 251 + use crate::subscriber_options::SubscriberSourcedMessage; 248 252 249 249 - let update = SubscriberOptionsUpdate::serialize_options_update(options); 253 253 + let update = SubscriberSourcedMessage::OptionsUpdate(options).to_json(); 250 254 tracing::debug!(%update, "sending options update"); 251 255 sink.send(Message::Text(update.into())).await?; 252 256 Ok(()) ··· 262 266 options: &mut Options, 263 267 ) -> ControlFlow<(), (Outcome, Arc<Notify>)> { 264 268 match command { 265 265 - ClientCommand::AddDid { complete, did } => match options.insert_did(did) { 266 266 - true => ControlFlow::Continue((Outcome::UpdateOptions, complete)), 267 267 - false => ControlFlow::Continue((Outcome::NoAction, complete)), 269 269 + ClientCommand::AddDid { complete, did } => match options.add_did(did) { 270 270 + Ok(true) => ControlFlow::Continue((Outcome::UpdateOptions, complete)), 271 271 + Ok(false) => ControlFlow::Continue((Outcome::NoAction, complete)), 272 272 + Err(did) => { 273 273 + tracing::error!(?did, "too many DIDs encountered when adding DID"); 274 274 + ControlFlow::Break(()) 275 275 + } 268 276 }, 269 277 ClientCommand::RemoveDid { complete, did } => match options.remove_did(&did) { 270 278 true => ControlFlow::Continue((Outcome::UpdateOptions, complete)), ··· 277 277 ClientCommand::AddCollection { 278 278 complete, 279 279 collection, 280 280 - } => match options.insert_collection(collection) { 281 281 - true => ControlFlow::Continue((Outcome::UpdateOptions, complete)), 282 282 - false => ControlFlow::Continue((Outcome::NoAction, complete)), 280 280 + } => match options.add_collection(collection) { 281 281 + Ok(true) => ControlFlow::Continue((Outcome::UpdateOptions, complete)), 282 282 + Ok(false) => ControlFlow::Continue((Outcome::NoAction, complete)), 283 283 + Err(collection) => { 284 284 + tracing::error!( 285 285 + ?collection, 286 286 + "too many collections encountered when adding collection" 287 287 + ); 288 288 + ControlFlow::Break(()) 289 289 + } 283 290 }, 284 291 ClientCommand::RemoveCollection { 285 292 complete,

-110

crates/jetstream/src/types.rs

reviewed

··· 1 1 - use crate::{Did, Nsid}; 2 2 - use serde::Serialize; 3 3 - use std::collections::HashSet; 4 4 - 5 5 - #[derive(Debug, Serialize)] 6 6 - #[serde(tag = "type", content = "payload", rename_all = "snake_case")] 7 7 - pub(crate) enum SubscriberOptionsUpdate<'a> { 8 8 - OptionsUpdate(&'a Options), 9 9 - } 10 10 - 11 11 - impl<'a> SubscriberOptionsUpdate<'a> { 12 12 - pub fn serialize_options_update(options: &'a Options) -> String { 13 13 - serde_json::to_string(&Self::OptionsUpdate(options)) 14 14 - .expect("OptionsUpdate should be serialized") 15 15 - } 16 16 - } 17 17 - 18 18 - #[derive(Debug, Default, Serialize)] 19 19 - #[serde(rename_all = "camelCase")] 20 20 - pub(crate) struct Options { 21 21 - pub wanted_dids: HashSet<Box<Did>>, 22 22 - pub wanted_collections: HashSet<Box<Nsid>>, 23 23 - pub max_message_size_bytes: i64, 24 24 - } 25 25 - 26 26 - impl Options { 27 27 - pub fn insert_did(&mut self, did: impl Into<Box<Did>>) -> bool { 28 28 - self.wanted_dids.insert(did.into()) 29 29 - } 30 30 - 31 31 - pub fn remove_did(&mut self, did: &Did) -> bool { 32 32 - self.wanted_dids.remove(did) 33 33 - } 34 34 - 35 35 - pub fn insert_collection(&mut self, collection: impl Into<Box<Nsid>>) -> bool { 36 36 - self.wanted_collections.insert(collection.into()) 37 37 - } 38 38 - 39 39 - pub fn remove_collection(&mut self, collection: &Nsid) -> bool { 40 40 - self.wanted_collections.remove(collection) 41 41 - } 42 42 - 43 43 - /// Compute the length, in bytes, of the query string required to express 44 44 - /// this set of options in a Jetstream `/subscribe` request. 45 45 - pub fn subscribe_query_len(&self) -> usize { 46 46 - const WANTED_DIDS_LEN: usize = "wantedDids=".len(); 47 47 - const WANTED_COLLECTIONS_LEN: usize = "wantedCollections=".len(); 48 48 - 49 49 - let (wdl, wdc): (_, usize) = self.wanted_dids.iter().fold((0, 0), |(sl, sc), v| { 50 50 - (sl + WANTED_DIDS_LEN + v.len(), sc + 1) 51 51 - }); 52 52 - 53 53 - let (wcl, wcc): (_, usize) = self.wanted_collections.iter().fold((0, 0), |(sl, sc), v| { 54 54 - (sl + WANTED_COLLECTIONS_LEN + v.len(), sc + 1) 55 55 - }); 56 56 - 57 57 - let mml = match self.max_message_size_bytes { 58 58 - 0 => 0, 59 59 - n if wdc > 0 || wcc > 0 => 1 + n.to_string().len() + "maxMessageSizeBytes=".len(), 60 60 - n => n.to_string().len() + "maxMessageSizeBytes=".len(), 61 61 - }; 62 62 - 63 63 - mml + (wdl + wdc.saturating_sub(1)) + wcl + wcc.saturating_sub(1) 64 64 - } 65 65 - 66 66 - pub fn append_to_url(&self, url: &mut url::Url) { 67 67 - let mut query = url.query_pairs_mut(); 68 68 - for collection in &self.wanted_collections { 69 69 - query.append_pair("wantedCollections", collection); 70 70 - } 71 71 - for did in &self.wanted_dids { 72 72 - query.append_pair("wantedDids", did.as_str()); 73 73 - } 74 74 - if self.max_message_size_bytes > 0 { 75 75 - query.append_pair( 76 76 - "maxMessageSizeBytes", 77 77 - &self.max_message_size_bytes.to_string(), 78 78 - ); 79 79 - } 80 80 - } 81 81 - } 82 82 - 83 83 - #[cfg(test)] 84 84 - mod tests { 85 85 - use super::Options; 86 86 - use atproto::Nsid; 87 87 - 88 88 - #[test] 89 89 - fn query_len() { 90 90 - let mut options = Options::default(); 91 91 - options.insert_collection(Nsid::from_static("sh.tangled.*")); 92 92 - 93 93 - assert_eq!( 94 94 - options.subscribe_query_len(), 95 95 - "wantedCollections=sh.tangled.*".len() 96 96 - ); 97 97 - 98 98 - options.insert_collection(Nsid::from_static("app.bsky.*")); 99 99 - assert_eq!( 100 100 - options.subscribe_query_len(), 101 101 - "wantedCollections=sh.tangled.*&wantedCollections=app.bsky.*".len() 102 102 - ); 103 103 - 104 104 - options.max_message_size_bytes = 1_000_000; 105 105 - assert_eq!( 106 106 - options.subscribe_query_len(), 107 107 - "wantedCollections=sh.tangled.*&wantedCollections=app.bsky.*&maxMessageSizeBytes=1000000".len() 108 108 - ); 109 109 - } 110 110 - }