rocks: hopefully shut down cleanly · microcosm.blue/microcosm-rs@7f35588

+30

Cargo.lock

··· 320 320 checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" 321 321 322 322 [[package]] 323 + name = "cfg_aliases" 324 + version = "0.2.1" 325 + source = "registry+https://github.com/rust-lang/crates.io-index" 326 + checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" 327 + 328 + [[package]] 323 329 name = "clang-sys" 324 330 version = "1.8.1" 325 331 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 424 430 dependencies = [ 425 431 "generic-array", 426 432 "typenum", 433 + ] 434 + 435 + [[package]] 436 + name = "ctrlc" 437 + version = "3.4.5" 438 + source = "registry+https://github.com/rust-lang/crates.io-index" 439 + checksum = "90eeab0aa92f3f9b4e87f258c72b139c207d251f9cbc1080a0086b86a8870dd3" 440 + dependencies = [ 441 + "nix", 442 + "windows-sys 0.59.0", 427 443 ] 428 444 429 445 [[package]] ··· 863 879 "axum-metrics", 864 880 "bincode", 865 881 "clap", 882 + "ctrlc", 866 883 "flume", 867 884 "fs4", 868 885 "links", ··· 874 891 "tempfile", 875 892 "tinyjson", 876 893 "tokio", 894 + "tokio-util", 877 895 "tungstenite", 878 896 "zstd", 879 897 ] ··· 1051 1069 "security-framework", 1052 1070 "security-framework-sys", 1053 1071 "tempfile", 1072 + ] 1073 + 1074 + [[package]] 1075 + name = "nix" 1076 + version = "0.29.0" 1077 + source = "registry+https://github.com/rust-lang/crates.io-index" 1078 + checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" 1079 + dependencies = [ 1080 + "bitflags", 1081 + "cfg-if", 1082 + "cfg_aliases", 1083 + "libc", 1054 1084 ] 1055 1085 1056 1086 [[package]]

+2

link_aggregator/Cargo.toml

··· 9 9 axum-metrics = "0.2" 10 10 bincode = "1.3.3" 11 11 clap = { version = "4.5.26", features = ["derive"] } 12 + ctrlc = "3.4.5" 12 13 flume = { version = "0.11.1", default-features = false } 13 14 fs4 = { version = "0.12.0", features = ["sync"] } 14 15 links = { path = "../links" } ··· 18 19 rocksdb = { version = "0.23.0", optional = true } 19 20 serde = { version = "1.0.215", features = ["derive"] } 20 21 tinyjson = "2.5.1" 22 + tokio-util = "0.7.13" 21 23 zstd = "0.13.2" 22 24 23 25 [dependencies.tokio]

+12 -5

link_aggregator/readme.md

··· 75 75 - [x] custom path for db storage 76 76 - [x] api server to look up backlink count 77 77 - [ ] other useful endpoints for the api server 78 - - [ ] show all nisd/path links to target 78 + - [x] show all nisd/path links to target 79 79 - [ ] get backlinking dids 80 80 - [ ] paging for all backlinking dids 81 81 - [ ] get count + most recent dids 82 82 - [ ] get count with any dids from provided set 83 - - [ ] write this readme 83 + - [~] write this readme 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 ··· 98 98 - [ ] deadletter queue of some kind for failed db writes 99 99 - [ ] also for valid json that was rejected? 100 100 - [x] get it running on raspi 101 - - [ ] get an estimate of disk usage per day after a few days of running 101 + - [~] get an estimate of disk usage per day after a few days of running 102 102 - [ ] make the did_init check only happen on test config (or remove it) 103 103 - [ ] actual error types (thiserror?) for lib-ish code 104 104 - [ ] clean up the main readme 105 105 - [x] web server metrics 106 106 - [ ] tokio metrics? 107 - - [ ] handle shutdown cleanly -- be nice to rocksdb 107 + - [x] handle shutdown cleanly -- be nice to rocksdb 108 108 - [x] add user-agent to jetstream request 109 - - [ ] possibly add tracing stuff also to complement metrics (at least for web reqs) 109 + 110 + data fixes 111 + - [ ] add rkey to linkers 🤦‍♀️ 112 + - [ ] links: 113 + - [ ] pull `$type`/`type` from object children of arrays (distinguish replies, quotes, etc) 114 + - [ ] actually define the format (deal with in-band dots etc) 115 + - [ ] _could_ throw cid neighbour into the target. probably should? but it's a lot of high volume uncompressible bytes 116 + - and it could be looked up from the linker's doc

+12 -3

link_aggregator/src/consumer/jetstream.rs

··· 1 + use anyhow::{bail, Result}; 1 2 use metrics::{ 2 3 counter, describe_counter, describe_gauge, describe_histogram, gauge, histogram, Unit, 3 4 }; ··· 5 6 use std::thread; 6 7 use std::time; 7 8 use tinyjson::JsonValue; 9 + use tokio_util::sync::CancellationToken; 8 10 use tungstenite::{client::IntoClientRequest, Error as TError, Message}; 9 11 use zstd::dict::DecoderDictionary; 10 12 ··· 20 22 pub fn consume_jetstream( 21 23 sender: flume::Sender<JsonValue>, 22 24 cursor: Option<u64>, 23 - ) -> anyhow::Result<()> { 25 + staying_alive: CancellationToken, 26 + ) -> Result<()> { 24 27 describe_counter!( 25 28 "jetstream_connnect", 26 29 Unit::Count, ··· 115 118 break; 116 119 } 117 120 121 + if staying_alive.is_cancelled() { 122 + eprintln!("jetstream: cancelling"); 123 + // TODO: cleanly close the connection? 124 + break 'outer; 125 + } 126 + 118 127 counter!("jetstream_read").increment(1); 119 128 let b = match socket.read() { 120 129 Ok(Message::Binary(b)) => b, ··· 233 242 counter!("jetstream_events", "url" => url).increment(1); 234 243 if sender.is_disconnected() { 235 244 eprintln!("jetstream: send channel disconnected -- nothing to do, bye."); 236 - break 'outer; 245 + bail!("jetstream: send channel disconnected"); 237 246 } 238 247 eprintln!( 239 248 "jetstream: failed to send on channel, dropping update! (FIXME / HANDLEME)" ··· 246 255 gauge!("jetstream_cursor_age", "url" => url).set(ts_age(ts).as_micros() as f64); 247 256 } 248 257 } 249 - Err(anyhow::anyhow!("broke out of jetstream loop")) 258 + Ok(()) 250 259 } 251 260 252 261 fn get_event_time(v: &JsonValue) -> Option<u64> {

+10 -3

link_aggregator/src/consumer/mod.rs

··· 2 2 mod jsonl_file; 3 3 4 4 use crate::storage::LinkStorage; 5 + use anyhow::Result; 5 6 use jetstream::consume_jetstream; 6 7 use jsonl_file::consume_jsonl_file; 7 8 use link_aggregator::{ActionableEvent, RecordId}; ··· 12 13 use std::sync::Arc; 13 14 use std::thread; 14 15 use tinyjson::JsonValue; 16 + use tokio_util::sync::CancellationToken; 15 17 16 - pub fn consume(mut store: impl LinkStorage, qsize: Arc<AtomicU32>, fixture: Option<PathBuf>) { 18 + pub fn consume( 19 + mut store: impl LinkStorage, 20 + qsize: Arc<AtomicU32>, 21 + fixture: Option<PathBuf>, 22 + staying_alive: CancellationToken, 23 + ) -> Result<()> { 17 24 describe_counter!( 18 25 "consumer_events_non_actionable", 19 26 Unit::Count, ··· 46 53 let cursor = store.get_cursor().unwrap(); 47 54 ( 48 55 receiver, 49 - thread::spawn(move || consume_jetstream(sender, cursor)), 56 + thread::spawn(move || consume_jetstream(sender, cursor, staying_alive)), 50 57 ) 51 58 }; 52 59 ··· 61 68 } 62 69 } 63 70 64 - consumer_handle.join().unwrap().unwrap() 71 + consumer_handle.join().unwrap() 65 72 } 66 73 67 74 pub fn get_actionable(event: &JsonValue) -> Option<(ActionableEvent, u64)> {

+55 -12

link_aggregator/src/main.rs

··· 6 6 use clap::{Parser, ValueEnum}; 7 7 use metrics_exporter_prometheus::PrometheusBuilder; 8 8 use std::path::PathBuf; 9 - use std::sync::atomic::{AtomicU32, Ordering}; 10 - use std::sync::Arc; 9 + use std::sync::{ 10 + atomic::{AtomicU32, Ordering}, 11 + Arc, 12 + }; 11 13 use std::thread; 12 14 use std::time; 13 15 use tokio::runtime; 14 - use tokio::sync::oneshot; 16 + use tokio_util::sync::CancellationToken; 15 17 16 18 use consumer::consume; 17 19 use server::serve; 18 20 #[cfg(feature = "rocks")] 19 21 use storage::RocksStorage; 20 22 use storage::{LinkReader, LinkStorage, MemStorage}; 23 + 24 + const MONITOR_INTERVAL: time::Duration = time::Duration::from_secs(2); 21 25 22 26 /// Aggregate links in the at-mosphere 23 27 #[derive(Parser, Debug)] ··· 71 75 fixture: Option<PathBuf>, 72 76 data_dir: Option<PathBuf>, 73 77 ) -> Result<()> { 78 + let stay_alive = CancellationToken::new(); 79 + 80 + ctrlc::set_handler({ 81 + let mut desperation: u8 = 0; 82 + let stay_alive = stay_alive.clone(); 83 + move || { 84 + match desperation { 85 + 0 => { 86 + println!("ok, shutting down..."); 87 + stay_alive.cancel(); 88 + } 89 + 1.. => panic!("fine, panicking!"), 90 + } 91 + desperation += 1; 92 + } 93 + })?; 94 + 74 95 let qsize = Arc::new(AtomicU32::new(0)); 75 96 76 97 thread::scope(|s| { 77 98 let readable = storage.to_readable(); 78 99 79 - let consumer = s.spawn({ 100 + s.spawn({ 80 101 let qsize = qsize.clone(); 81 - move || consume(storage, qsize, fixture) 102 + let stay_alive = stay_alive.clone(); 103 + let staying_alive = stay_alive.clone(); 104 + move || { 105 + if let Err(e) = consume(storage, qsize, fixture, staying_alive) { 106 + eprintln!("jetstream finished with error: {e}"); 107 + } 108 + stay_alive.drop_guard(); 109 + } 82 110 }); 83 111 84 - let (stop_server, shutdown) = oneshot::channel::<()>(); 85 112 s.spawn({ 86 113 let readable = readable.clone(); 114 + let stay_alive = stay_alive.clone(); 115 + let staying_alive = stay_alive.clone(); 87 116 || { 88 117 runtime::Builder::new_multi_thread() 89 118 .worker_threads(1) 90 119 .max_blocking_threads(2) 91 120 .enable_all() 92 - .build()? 121 + .build() 122 + .expect("axum startup") 93 123 .block_on(async { 94 124 install_metrics_server()?; 95 - serve(readable, "0.0.0.0:6789", shutdown).await 125 + serve(readable, "0.0.0.0:6789", staying_alive).await 96 126 }) 127 + .unwrap(); 128 + stay_alive.drop_guard(); 97 129 } 98 130 }); 99 131 100 - s.spawn(move || { 132 + s.spawn(move || { // monitor thread 133 + let stay_alive = stay_alive.clone(); 134 + let check_alive = stay_alive.clone(); 135 + 101 136 let process_collector = metrics_process::Collector::default(); 102 137 process_collector.describe(); 103 138 metrics::describe_gauge!( ··· 117 152 println!("disk space monitoring should work, watching at {p:?}"); 118 153 } 119 154 } 120 - while !consumer.is_finished() { 155 + 156 + 'monitor: loop { 121 157 readable.summarize(qsize.load(Ordering::Relaxed)); 122 158 process_collector.collect(); 123 159 if let Some(ref p) = data_dir { ··· 128 164 metrics::gauge!("storage.free").set(free as f64); 129 165 } 130 166 } 131 - thread::sleep(time::Duration::from_secs(3)); 167 + let wait = time::Instant::now(); 168 + while wait.elapsed() < MONITOR_INTERVAL { 169 + thread::sleep(time::Duration::from_millis(100)); 170 + if check_alive.is_cancelled() { 171 + break 'monitor 172 + } 173 + } 132 174 } 133 - let _ = stop_server.send(()); 175 + stay_alive.drop_guard(); 134 176 }); 135 177 }); 178 + 136 179 println!("byeeee"); 137 180 138 181 Ok(())

+3 -5

link_aggregator/src/server.rs

··· 2 2 use serde::Deserialize; 3 3 use std::collections::HashMap; 4 4 use tokio::net::{TcpListener, ToSocketAddrs}; 5 - use tokio::sync::oneshot::Receiver; 6 5 use tokio::task::block_in_place; 6 + use tokio_util::sync::CancellationToken; 7 7 8 8 use crate::storage::LinkReader; 9 9 10 - pub async fn serve<S, A>(store: S, addr: A, cancel: Receiver<()>) -> anyhow::Result<()> 10 + pub async fn serve<S, A>(store: S, addr: A, stay_alive: CancellationToken) -> anyhow::Result<()> 11 11 where 12 12 S: LinkReader, 13 13 A: ToSocketAddrs, ··· 33 33 let listener = TcpListener::bind(addr).await?; 34 34 println!("api: listening at http://{:?}", listener.local_addr()?); 35 35 axum::serve(listener, app) 36 - .with_graceful_shutdown(async { 37 - cancel.await.ok(); 38 - }) 36 + .with_graceful_shutdown(async move { stay_alive.cancelled().await }) 39 37 .await?; 40 38 41 39 Ok(())

+24 -1

link_aggregator/src/storage/rocks_store.rs

··· 44 44 db: Arc<DBWithThreadMode<MultiThreaded>>, // TODO: mov seqs here (concat merge op will be fun) 45 45 did_id_table: IdTable<Did, DidIdValue>, 46 46 target_id_table: IdTable<TargetKey, TargetId>, 47 + is_writer: bool, 47 48 } 48 49 49 50 trait IdTableValue: ValueFromRocks + Clone { ··· 213 214 db, 214 215 did_id_table, 215 216 target_id_table, 217 + is_writer: true, 216 218 }) 217 219 } 218 220 ··· 557 559 } 558 560 } 559 561 562 + impl Drop for RocksStorage { 563 + fn drop(&mut self) { 564 + if self.is_writer { 565 + println!("rocksdb writer: cleaning up for shutdown..."); 566 + if let Err(e) = self.db.flush_wal(true) { 567 + eprintln!("rocks: flushing wal failed: {e:?}"); 568 + } 569 + if let Err(e) = self.db.flush_opt(&{ 570 + let mut opt = rocksdb::FlushOptions::default(); 571 + opt.set_wait(true); 572 + opt 573 + }) { 574 + eprintln!("rocks: flushing memtables failed: {e:?}"); 575 + } 576 + self.db.cancel_all_background_work(true); 577 + } 578 + } 579 + } 580 + 560 581 impl AsRocksValue for u64 {} 561 582 impl ValueFromRocks for u64 {} 562 583 ··· 602 623 } 603 624 604 625 fn to_readable(&mut self) -> impl LinkReader { 605 - self.clone() 626 + let mut readable = self.clone(); 627 + readable.is_writer = false; 628 + readable 606 629 } 607 630 } 608 631

Configure Feed

Configure Feed