commit c5be8d39d2a4f79becc4c667cf8ed40c427977ac · nekomimi.pet/microcosm-rs

+5 -3

link_aggregator/readme.md

···

       84
       84
        
       - [ ] fix it sometimes getting stuck

     

       85
       85
        
         - seems to unstick in my possibly-different repro (letting laptop fall asleep) after a bit.

     

       86
       86
        
         - [ ] add a detection for no new links coming in after some period

     

       87
       87
       -
       - [~] handle jetstream restart: don't miss events (currently sketch: rewinds cursor by 1us so we will always double-count at least one event)

     

       88
       88
       -
         - [ ] especially: figure out what the risk is to rotating to another jetstream server in terms of gap/overlap from a different jetstream instance's cursor

     

       87
       87
       +
       - [x] handle jetstream restart: don't miss events (currently sketch: rewinds cursor by 1us so we will always double-count at least one event)

     

       88
       88
       +
         - [x] especially: figure out what the risk is to rotating to another jetstream server in terms of gap/overlap from a different jetstream instance's cursor (follow up separately)

     

       89
       89
       +
         - [x] jetstream: don't rotate servers, explicitly pass via cli

     

       89
       90
        
       - [x] metrics!

     

       90
       91
        
         - [x] event ts lag

     

       91
       92
        
       - [~] machine resource metrics

     

       92
       93
        
         - [x] disk consumption

     

       93
       94
        
         - [x] cpu usage

     

       94
       95
        
         - [x] mem usage

     

       95
       95
       -
         - [ ] network?

     

       96
       96
       +
         - [x] network?

     

       96
       97
        
       - [ ] make all storage apis return Result

     

       97
       98
        
       - [ ] handle all the unwraps

     

       98
       99
        
       - [ ] deadletter queue of some kind for failed db writes

     
···

       110
       111
        
       - [ ] serve html for browser requests

     

       111
       112
        
       - [ ] add a health check endpoint

     

       112
       113
        
       - [ ] add seq numbers to metrics

     

       114
       114
       +
       - [ ] persist the jetstream server url, error if started with a different one (maybe with --switch-streams or something)

     

       113
       115
        
       

     

       114
       116
        
       data fixes

     

       115
       117
        
       - [ ] add rkey to linkers 🤦‍♀️

+36 -37

link_aggregator/src/consumer/jetstream.rs

···

       12
       12
        
       

     

       13
       13
        
       const JETSTREAM_ZSTD_DICTIONARY: &[u8] = include_bytes!("../../zstd/dictionary");

     

       14
       14
        
       

     

       15
       15
       -
       const WS_URLS: [&str; 4] = [

     

       16
       16
       -
           "wss://jetstream2.us-east.bsky.network/subscribe",

     

       17
       17
       -
           "wss://jetstream1.us-east.bsky.network/subscribe",

     

       18
       18
       -
           "wss://jetstream1.us-west.bsky.network/subscribe",

     

       19
       19
       -
           "wss://jetstream2.us-west.bsky.network/subscribe",

     

       20
       20
       -
       ];

     

       21
       21
       -
       

     

       22
       15
        
       pub fn consume_jetstream(

     

       23
       16
        
           sender: flume::Sender<JsonValue>,

     

       24
       17
        
           cursor: Option<u64>,

     

       18
       18
       +
           stream: String,

     

       25
       19
        
           staying_alive: CancellationToken,

     

       26
       20
        
       ) -> Result<()> {

     

       27
       21
        
           describe_counter!(

     
···

       74
       68
        
           let mut connect_retries = 0;

     

       75
       69
        
           let mut latest_cursor = cursor;

     

       76
       70
        
           'outer: loop {

     

       77
       77
       -
               let url = WS_URLS[connect_retries % WS_URLS.len()];

     

       78
       71
        
               let stream_url = format!(

     

       79
       79
       -
                   "{url}?compress=true{}",

     

       72
       72
       +
                   "{stream}?compress=true{}",

     

       80
       73
        
                   latest_cursor

     

       81
       74
        
                       .map(|c| {

     

       82
       75
        
                           println!("starting with cursor from {:?} ago...", ts_age(c));

     
···

       88
       81
        
               let ua = format!("ucosm/link aggregator v{}", env!("CARGO_PKG_VERSION"));

     

       89
       82
        
               req.headers_mut().insert("user-agent", ua.parse()?);

     

       90
       83
        
       

     

       91
       91
       -
               counter!("jetstream_connect", "url" => url, "is_retry" => (connect_retries > 0).to_string()).increment(1);

     

       84
       84
       +
               counter!("jetstream_connect", "url" => stream.clone(), "is_retry" => (connect_retries > 0).to_string()).increment(1);

     

       92
       85
        
               println!("jetstream connecting, attempt #{connect_retries}, {stream_url:?} with user-agent: {ua:?}");

     

       93
       86
        
               let mut socket = match tungstenite::connect(req) {

     

       94
       87
        
                   Ok((socket, _)) => {

     
···

       128
       121
        
                   let b = match socket.read() {

     

       129
       122
        
                       Ok(Message::Binary(b)) => b,

     

       130
       123
        
                       Ok(Message::Text(_)) => {

     

       131
       131
       -
                           counter!("jetstream_read_fail", "url" => url, "reason" => "received text")

     

       124
       124
       +
                           counter!("jetstream_read_fail", "url" => stream.clone(), "reason" => "received text")

     

       132
       125
        
                               .increment(1);

     

       133
       126
        
                           eprintln!("jetstream: unexpected text message, should be binary for compressed (ignoring)");

     

       134
       127
        
                           continue;

     

       135
       128
        
                       }

     

       136
       129
        
                       Ok(Message::Close(f)) => {

     

       137
       137
       -
                           counter!("jetstream_read_fail", "url" => url, "reason" => "server closed")

     

       130
       130
       +
                           counter!("jetstream_read_fail", "url" => stream.clone(), "reason" => "server closed")

     

       138
       131
        
                               .increment(1);

     

       139
       132
        
                           println!("jetstream: closing the connection: {f:?}");

     

       140
       133
        
                           continue;

     

       141
       134
        
                       }

     

       142
       135
        
                       Ok(m) => {

     

       143
       143
       -
                           counter!("jetstream_read_fail", "url" => url, "reason" => "unexpected message", "message" => format!("{m:?}")).increment(1);

     

       136
       136
       +
                           counter!("jetstream_read_fail", "url" => stream.clone(), "reason" => "unexpected message", "message" => format!("{m:?}")).increment(1);

     

       144
       137
        
                           eprintln!("jetstream: unexpected from read (ignoring): {m:?}");

     

       145
       138
        
                           continue;

     

       146
       139
        
                       }

     

       147
       140
        
                       Err(TError::ConnectionClosed) => {

     

       148
       141
        
                           // clean exit

     

       149
       149
       -
                           counter!("jetstream_read_fail", "url" => url, "reason" => "clean close")

     

       142
       142
       +
                           counter!("jetstream_read_fail", "url" => stream.clone(), "reason" => "clean close")

     

       150
       143
        
                               .increment(1);

     

       151
       144
        
                           println!("jetstream closed the websocket cleanly.");

     

       152
       145
        
                           break;

     

       153
       146
        
                       }

     

       154
       147
        
                       Err(TError::AlreadyClosed) => {

     

       155
       148
        
                           // programming error

     

       156
       156
       -
                           counter!("jetstream_read_fail", "url" => url, "reason" => "already closed")

     

       149
       149
       +
                           counter!("jetstream_read_fail", "url" => stream.clone(), "reason" => "already closed")

     

       157
       150
        
                               .increment(1);

     

       158
       151
        
                           eprintln!(

     

       159
       152
        
                               "jetstream: got AlreadyClosed trying to .read() websocket. probably a bug."

     
···

       161
       154
        
                           break;

     

       162
       155
        
                       }

     

       163
       156
        
                       Err(TError::Capacity(e)) => {

     

       164
       164
       -
                           counter!("jetstream_read_fail", "url" => url, "reason" => "capacity error")

     

       157
       157
       +
                           counter!("jetstream_read_fail", "url" => stream.clone(), "reason" => "capacity error")

     

       165
       158
        
                               .increment(1);

     

       166
       159
        
                           eprintln!("jetstream: capacity error (ignoring): {e:?}");

     

       167
       160
        
                           continue;

     

       168
       161
        
                       }

     

       169
       162
        
                       Err(TError::Utf8) => {

     

       170
       170
       -
                           counter!("jetstream_read_fail", "url" => url, "reason" => "utf8 error")

     

       163
       163
       +
                           counter!("jetstream_read_fail", "url" => stream.clone(), "reason" => "utf8 error")

     

       171
       164
        
                               .increment(1);

     

       172
       165
        
                           eprintln!("jetstream: utf8 error (ignoring)");

     

       173
       166
        
                           continue;

     
···

       176
       169
        
                           eprintln!("jetstream: could not read message from socket. closing: {e:?}");

     

       177
       170
        
                           match socket.close(None) {

     

       178
       171
        
                               Err(TError::ConnectionClosed) => {

     

       179
       179
       -
                                   counter!("jetstream_read_fail", "url" => url, "reason" => "clean close").increment(1);

     

       172
       172
       +
                                   counter!("jetstream_read_fail", "url" => stream.clone(), "reason" => "clean close").increment(1);

     

       180
       173
        
                                   println!("jetstream closed the websocket cleanly.");

     

       181
       174
        
                                   break;

     

       182
       175
        
                               }

     

       183
       176
        
                               r => eprintln!("jetstream: close result after error: {r:?}"),

     

       184
       177
        
                           }

     

       185
       185
       -
                           counter!("jetstream_read_fail", "url" => url, "reason" => "read error")

     

       178
       178
       +
                           counter!("jetstream_read_fail", "url" => stream.clone(), "reason" => "read error")

     

       186
       179
        
                               .increment(1);

     

       187
       180
        
                           // if we didn't immediately get ConnectionClosed, we should keep polling read

     

       188
       181
        
                           // until we get it.

     
···

       190
       183
        
                       }

     

       191
       184
        
                   };

     

       192
       185
        
       

     

       193
       193
       -
                   counter!("jetstream_read_bytes", "url" => url).increment(b.len() as u64);

     

       186
       186
       +
                   counter!("jetstream_read_bytes", "url" => stream.clone()).increment(b.len() as u64);

     

       194
       187
        
                   let mut cursor = Cursor::new(b);

     

       195
       195
       -
                   let mut decoder =

     

       196
       196
       -
                       match zstd::stream::Decoder::with_prepared_dictionary(&mut cursor, &dict) {

     

       197
       197
       -
                           Ok(d) => d,

     

       198
       198
       -
                           Err(e) => {

     

       199
       199
       -
                               counter!("jetstream_read_fail", "url" => url, "reason" => "zstd decompress")

     

       188
       188
       +
                   let mut decoder = match zstd::stream::Decoder::with_prepared_dictionary(

     

       189
       189
       +
                       &mut cursor,

     

       190
       190
       +
                       &dict,

     

       191
       191
       +
                   ) {

     

       192
       192
       +
                       Ok(d) => d,

     

       193
       193
       +
                       Err(e) => {

     

       194
       194
       +
                           counter!("jetstream_read_fail", "url" => stream.clone(), "reason" => "zstd decompress")

     

       200
       195
        
                               .increment(1);

     

       201
       201
       -
                               eprintln!("jetstream: failed to decompress zstd message: {e:?}");

     

       202
       202
       -
                               continue;

     

       203
       203
       -
                           }

     

       204
       204
       -
                       };

     

       196
       196
       +
                           eprintln!("jetstream: failed to decompress zstd message: {e:?}");

     

       197
       197
       +
                           continue;

     

       198
       198
       +
                       }

     

       199
       199
       +
                   };

     

       205
       200
        
       

     

       206
       201
        
                   let mut s = String::new();

     

       207
       202
        
                   match decoder.read_to_string(&mut s) {

     

       208
       203
        
                       Ok(n) => {

     

       209
       209
       -
                           counter!("jetstream_read_bytes_decompressed", "url" => url).increment(n as u64);

     

       210
       210
       -
                           histogram!("jetstream_read_bytes_decompressed", "url" => url).record(n as f64);

     

       204
       204
       +
                           counter!("jetstream_read_bytes_decompressed", "url" => stream.clone())

     

       205
       205
       +
                               .increment(n as u64);

     

       206
       206
       +
                           histogram!("jetstream_read_bytes_decompressed", "url" => stream.clone())

     

       207
       207
       +
                               .record(n as f64);

     

       211
       208
        
                       }

     

       212
       209
        
                       Err(e) => {

     

       213
       213
       -
                           counter!("jetstream_read_fail", "url" => url, "reason" => "zstd string decode")

     

       210
       210
       +
                           counter!("jetstream_read_fail", "url" => stream.clone(), "reason" => "zstd string decode")

     

       214
       211
        
                               .increment(1);

     

       215
       212
        
                           eprintln!("jetstream: failed to decode zstd: {e:?}");

     

       216
       213
        
                           continue;

     
···

       220
       217
        
                   let v = match s.parse() {

     

       221
       218
        
                       Ok(v) => v,

     

       222
       219
        
                       Err(e) => {

     

       223
       223
       -
                           counter!("jetstream_read_fail", "url" => url, "reason" => "json parse")

     

       220
       220
       +
                           counter!("jetstream_read_fail", "url" => stream.clone(), "reason" => "json parse")

     

       224
       221
        
                               .increment(1);

     

       225
       222
        
                           eprintln!("jetstream: failed to parse message as json: {e:?}");

     

       226
       223
        
                           continue;

     
···

       231
       228
        
                   let ts = match get_event_time(&v) {

     

       232
       229
        
                       Some(ts) => ts,

     

       233
       230
        
                       None => {

     

       234
       234
       -
                           counter!("jetstream_read_fail", "url" => url, "reason" => "invalid event")

     

       231
       231
       +
                           counter!("jetstream_read_fail", "url" => stream.clone(), "reason" => "invalid event")

     

       235
       232
        
                               .increment(1);

     

       236
       233
        
                           eprintln!("jetstream: encountered an event without a timestamp: ignoring it.");

     

       237
       234
        
                           continue;

     
···

       239
       236
        
                   };

     

       240
       237
        
       

     

       241
       238
        
                   if let Err(flume::SendError(_rejected)) = sender.send(v) {

     

       242
       242
       -
                       counter!("jetstream_events", "url" => url).increment(1);

     

       239
       239
       +
                       counter!("jetstream_events", "url" => stream.clone()).increment(1);

     

       243
       240
        
                       if sender.is_disconnected() {

     

       244
       241
        
                           eprintln!("jetstream: send channel disconnected -- nothing to do, bye.");

     

       245
       242
        
                           bail!("jetstream: send channel disconnected");

     
···

       248
       245
        
                           "jetstream: failed to send on channel, dropping update! (FIXME / HANDLEME)"

     

       249
       246
        
                       );

     

       250
       247
        
                   }

     

       251
       251
       -
                   histogram!("jetstream_events_queued", "url" => url).record(sender.len() as f64);

     

       248
       248
       +
                   histogram!("jetstream_events_queued", "url" => stream.clone())

     

       249
       249
       +
                       .record(sender.len() as f64);

     

       252
       250
        
       

     

       253
       251
        
                   // only actually update our cursor after we've managed to queue the event

     

       254
       252
        
                   latest_cursor = Some(ts);

     

       255
       255
       -
                   gauge!("jetstream_cursor_age", "url" => url).set(ts_age(ts).as_micros() as f64);

     

       253
       253
       +
                   gauge!("jetstream_cursor_age", "url" => stream.clone())

     

       254
       254
       +
                       .set(ts_age(ts).as_micros() as f64);

     

       256
       255
        
               }

     

       257
       256
        
           }

     

       258
       257
        
           Ok(())

+2 -1

link_aggregator/src/consumer/mod.rs

···

       19
       19
        
           mut store: impl LinkStorage,

     

       20
       20
        
           qsize: Arc<AtomicU32>,

     

       21
       21
        
           fixture: Option<PathBuf>,

     

       22
       22
       +
           stream: String,

     

       22
       23
        
           staying_alive: CancellationToken,

     

       23
       24
        
       ) -> Result<()> {

     

       24
       25
        
           describe_counter!(

     
···

       53
       54
        
               let cursor = store.get_cursor().unwrap();

     

       54
       55
        
               (

     

       55
       56
        
                   receiver,

     

       56
       56
       -
                   thread::spawn(move || consume_jetstream(sender, cursor, staying_alive)),

     

       57
       57
       +
                   thread::spawn(move || consume_jetstream(sender, cursor, stream, staying_alive)),

     

       57
       58
        
               )

     

       58
       59
        
           };

     

       59
       60

+22 -3

link_aggregator/src/main.rs

···

       27
       27
        
       #[derive(Parser, Debug)]

     

       28
       28
        
       #[command(version, about, long_about = None)]

     

       29
       29
        
       struct Args {

     

       30
       30
       +
           #[arg(short, long)]

     

       31
       31
       +
           /// Jetstream server to connect to (exclusive with --fixture). Provide either a wss:// URL, or a shorhand value:

     

       32
       32
       +
           /// 'us-east-1', 'us-east-2', 'us-west-1', or 'us-west-2'

     

       33
       33
       +
           #[arg(short, long)]

     

       34
       34
       +
           jetstream: String,

     

       30
       35
        
           // TODO: make this part of rocks' own sub-config?

     

       31
       36
        
           /// Where to store data on disk, for backends that use disk storage

     

       32
       37
        
           #[arg(short, long)]

     
···

       47
       52
        
           Rocks,

     

       48
       53
        
       }

     

       49
       54
        
       

     

       55
       55
       +
       fn jetstream_url(provided: &str) -> String {

     

       56
       56
       +
           match provided {

     

       57
       57
       +
               "us-east-1" => "wss://jetstream1.us-east.bsky.network/subscribe".into(),

     

       58
       58
       +
               "us-east-2" => "wss://jetstream2.us-east.bsky.network/subscribe".into(),

     

       59
       59
       +
               "us-west-1" => "wss://jetstream1.us-west.bsky.network/subscribe".into(),

     

       60
       60
       +
               "us-west-2" => "wss://jetstream2.us-west.bsky.network/subscribe".into(),

     

       61
       61
       +
               custom => custom.into(),

     

       62
       62
       +
           }

     

       63
       63
       +
       }

     

       64
       64
       +
       

     

       50
       65
        
       fn main() -> Result<()> {

     

       51
       66
        
           let args = Args::parse();

     

       52
       67
        
       

     
···

       57
       72
        
               println!("using fixture at {p:?}...");

     

       58
       73
        
           }

     

       59
       74
        
       

     

       75
       75
       +
           let stream = jetstream_url(&args.jetstream);

     

       76
       76
       +
           println!("using jetstream server {stream:?}...",);

     

       77
       77
       +
       

     

       60
       78
        
           match args.backend {

     

       61
       61
       -
               StorageBackend::Memory => run(MemStorage::new(), fixture, None),

     

       79
       79
       +
               StorageBackend::Memory => run(MemStorage::new(), fixture, None, stream),

     

       62
       80
        
               #[cfg(feature = "rocks")]

     

       63
       81
        
               StorageBackend::Rocks => {

     

       64
       82
        
                   let storage_dir = args.data.clone().unwrap_or("rocks.test".into());

     

       65
       83
        
                   println!("starting rocksdb...");

     

       66
       84
        
                   let rocks = RocksStorage::new(storage_dir)?;

     

       67
       85
        
                   println!("rocks ready.");

     

       68
       68
       -
                   run(rocks, fixture, args.data)

     

       86
       86
       +
                   run(rocks, fixture, args.data, stream)

     

       69
       87
        
               }

     

       70
       88
        
           }

     

       71
       89
        
       }

     
···

       74
       92
        
           mut storage: impl LinkStorage,

     

       75
       93
        
           fixture: Option<PathBuf>,

     

       76
       94
        
           data_dir: Option<PathBuf>,

     

       95
       95
       +
           stream: String,

     

       77
       96
        
       ) -> Result<()> {

     

       78
       97
        
           let stay_alive = CancellationToken::new();

     

       79
       98
        
       

     
···

       102
       121
        
                   let stay_alive = stay_alive.clone();

     

       103
       122
        
                   let staying_alive = stay_alive.clone();

     

       104
       123
        
                   move || {

     

       105
       105
       -
                       if let Err(e) = consume(storage, qsize, fixture, staying_alive) {

     

       124
       124
       +
                       if let Err(e) = consume(storage, qsize, fixture, stream, staying_alive) {

     

       106
       125
        
                           eprintln!("jetstream finished with error: {e}");

     

       107
       126
        
                       }

     

       108
       127
        
                       stay_alive.drop_guard();

+1 -1

link_aggregator/src/storage/rocks_store.rs

···

       531
       531
        
                   let mut mini_batch = WriteBatch::default();

     

       532
       532
        
       

     

       533
       533
        
                   for (i, (record_link_key, links)) in chunk.iter().enumerate() {

     

       534
       534
       -
                       self.delete_record_link(&mut mini_batch, &record_link_key); // _could_ use delete range here instead of individual deletes, but since we have to scan anyway it's not obvious if it's better

     

       534
       534
       +
                       self.delete_record_link(&mut mini_batch, record_link_key); // _could_ use delete range here instead of individual deletes, but since we have to scan anyway it's not obvious if it's better

     

       535
       535
        
       

     

       536
       536
        
                       for (j, RecordLinkTarget(path, target_link_id)) in links.0.iter().enumerate() {

     

       537
       537
        
                           self.update_target_linkers(&mut mini_batch, target_link_id, |mut linkers| {