+40 -21

server/Cargo.lock

··· 88 88 ] 89 89 90 90 [[package]] 91 91 + name = "atomic-time" 92 92 + version = "0.1.5" 93 93 + source = "registry+https://github.com/rust-lang/crates.io-index" 94 94 + checksum = "9622f5c6fb50377516c70f65159e70b25465409760c6bd6d4e581318bf704e83" 95 95 + dependencies = [ 96 96 + "once_cell", 97 97 + "portable-atomic", 98 98 + ] 99 99 + 100 100 + [[package]] 91 101 name = "atomic-waker" 92 102 version = "1.1.2" 93 103 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 956 966 checksum = "d05e27ee213611ffe7d6348b942e8f942b37114c00cc03cec254295a4a17852e" 957 967 958 968 [[package]] 959 959 - name = "overload" 960 960 - version = "0.1.1" 969 969 + name = "ordered-varint" 970 970 + version = "2.0.0" 961 971 source = "registry+https://github.com/rust-lang/crates.io-index" 962 962 - checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" 972 972 + checksum = "e9cc9f18ab4bad1e01726bda1259feb8f11e5e76308708a966b4c0136e9db34c" 963 973 964 974 [[package]] 965 965 - name = "papaya" 966 966 - version = "0.2.3" 975 975 + name = "overload" 976 976 + version = "0.1.1" 967 977 source = "registry+https://github.com/rust-lang/crates.io-index" 968 968 - checksum = "f92dd0b07c53a0a0c764db2ace8c541dc47320dad97c2200c2a637ab9dd2328f" 969 969 - dependencies = [ 970 970 - "equivalent", 971 971 - "seize", 972 972 - ] 978 978 + checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" 973 979 974 980 [[package]] 975 981 name = "parking_lot" ··· 1044 1050 version = "0.3.32" 1045 1051 source = "registry+https://github.com/rust-lang/crates.io-index" 1046 1052 checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" 1053 1053 + 1054 1054 + [[package]] 1055 1055 + name = "portable-atomic" 1056 1056 + version = "1.11.1" 1057 1057 + source = "registry+https://github.com/rust-lang/crates.io-index" 1058 1058 + checksum = "f84267b20a16ea918e43c6a88433c2d54fa145c92a811b5b047ccbe153674483" 1047 1059 1048 1060 [[package]] 1049 1061 name = "proc-macro2" ··· 1335 1347 ] 1336 1348 1337 1349 [[package]] 1350 1350 + name = "scc" 1351 1351 + version = "2.3.4" 1352 1352 + source = "registry+https://github.com/rust-lang/crates.io-index" 1353 1353 + checksum = "22b2d775fb28f245817589471dd49c5edf64237f4a19d10ce9a92ff4651a27f4" 1354 1354 + dependencies = [ 1355 1355 + "sdd", 1356 1356 + ] 1357 1357 + 1358 1358 + [[package]] 1338 1359 name = "schannel" 1339 1360 version = "0.1.27" 1340 1361 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 1350 1371 checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" 1351 1372 1352 1373 [[package]] 1374 1374 + name = "sdd" 1375 1375 + version = "3.0.10" 1376 1376 + source = "registry+https://github.com/rust-lang/crates.io-index" 1377 1377 + checksum = "490dcfcbfef26be6800d11870ff2df8774fa6e86d047e3e8c8a76b25655e41ca" 1378 1378 + 1379 1379 + [[package]] 1353 1380 name = "security-framework" 1354 1381 version = "3.2.0" 1355 1382 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 1370 1397 dependencies = [ 1371 1398 "core-foundation-sys", 1372 1399 "libc", 1373 1373 - ] 1374 1374 - 1375 1375 - [[package]] 1376 1376 - name = "seize" 1377 1377 - version = "0.5.0" 1378 1378 - source = "registry+https://github.com/rust-lang/crates.io-index" 1379 1379 - checksum = "e4b8d813387d566f627f3ea1b914c068aac94c40ae27ec43f5f33bde65abefe7" 1380 1380 - dependencies = [ 1381 1381 - "libc", 1382 1382 - "windows-sys 0.52.0", 1383 1400 ] 1384 1401 1385 1402 [[package]] ··· 1448 1465 dependencies = [ 1449 1466 "anyhow", 1450 1467 "async-trait", 1468 1468 + "atomic-time", 1451 1469 "axum", 1452 1470 "axum-tws", 1453 1471 "fjall", 1454 1472 "futures-util", 1455 1455 - "papaya", 1473 1473 + "ordered-varint", 1456 1474 "pingora-limits", 1457 1475 "rkyv", 1458 1476 "rustls", 1477 1477 + "scc", 1459 1478 "serde", 1460 1479 "serde_json", 1461 1480 "smol_str",

+3 -1

server/Cargo.toml

··· 20 20 fjall = { version = "2", default-features = false, features = ["miniz", "lz4"] } 21 21 rkyv = {version = "0.8", features = ["unaligned"]} 22 22 smol_str = { version = "0.3", features = ["serde"] } 23 23 - papaya = "0.2" 24 23 serde = { version = "1", features = ["derive"] } 25 24 serde_json = "1.0.141" 25 25 + scc = "2.3.4" 26 26 + atomic-time = "0.1.5" 27 27 + ordered-varint = "2.0.0" 26 28 27 29 [target.'cfg(not(target_env = "msvc"))'.dependencies] 28 30 tikv-jemallocator = "0.6"

+8 -11

server/src/api.rs

··· 30 30 use crate::{ 31 31 db::Db, 32 32 error::{AppError, AppResult}, 33 33 + utils::time_now, 33 34 }; 34 35 35 36 struct LatencyMillis(u128); ··· 161 162 let maybe_hits = db 162 163 .get_hits( 163 164 &params.nsid, 164 164 - params.to.unwrap_or(0) 165 165 - ..params.from.unwrap_or( 166 166 - std::time::SystemTime::now() 167 167 - .duration_since(UNIX_EPOCH) 168 168 - .expect("oops") 169 169 - .as_micros() as u64, 170 170 - ), 171 171 - )? 165 165 + params.to.unwrap_or(0)..params.from.unwrap_or(time_now()), 166 166 + ) 172 167 .take(MAX_HITS); 173 168 let mut hits = Vec::with_capacity(maybe_hits.size_hint().0); 174 169 175 170 for maybe_hit in maybe_hits { 176 176 - let (timestamp, hit) = maybe_hit?; 171 171 + let hit = maybe_hit?; 172 172 + let hit_data = hit.access(); 173 173 + 177 174 hits.push(Hit { 178 178 - timestamp, 179 179 - deleted: hit.deleted, 175 175 + timestamp: hit.timestamp, 176 176 + deleted: hit_data.deleted, 180 177 }); 181 178 } 182 179

-244

server/src/db.rs

··· 1 1 - use std::{ 2 2 - ops::{Bound, Deref, RangeBounds}, 3 3 - path::Path, 4 4 - time::Duration, 5 5 - }; 6 6 - 7 7 - use fjall::{Config, Keyspace, Partition, PartitionCreateOptions}; 8 8 - use pingora_limits::rate::Rate; 9 9 - use rkyv::{Archive, Deserialize, Serialize, rancor::Error}; 10 10 - use smol_str::SmolStr; 11 11 - use tokio::sync::broadcast; 12 12 - 13 13 - use crate::{ 14 14 - error::{AppError, AppResult}, 15 15 - jetstream::JetstreamEvent, 16 16 - }; 17 17 - 18 18 - #[derive(Clone, Debug, Default, Archive, Deserialize, Serialize, PartialEq)] 19 19 - #[rkyv(compare(PartialEq), derive(Debug))] 20 20 - pub struct NsidCounts { 21 21 - pub count: u128, 22 22 - pub deleted_count: u128, 23 23 - pub last_seen: u64, 24 24 - } 25 25 - 26 26 - #[derive(Debug, Default, Archive, Deserialize, Serialize, PartialEq)] 27 27 - #[rkyv(compare(PartialEq), derive(Debug))] 28 28 - pub struct NsidHit { 29 29 - pub deleted: bool, 30 30 - } 31 31 - 32 32 - pub struct EventRecord { 33 33 - pub nsid: SmolStr, 34 34 - pub timestamp: u64, 35 35 - pub deleted: bool, 36 36 - } 37 37 - 38 38 - impl EventRecord { 39 39 - pub fn from_jetstream(event: JetstreamEvent) -> Option<Self> { 40 40 - match event { 41 41 - JetstreamEvent::Commit { 42 42 - time_us, commit, .. 43 43 - } => Some(Self { 44 44 - nsid: commit.collection.into(), 45 45 - timestamp: time_us, 46 46 - deleted: false, 47 47 - }), 48 48 - JetstreamEvent::Delete { 49 49 - time_us, commit, .. 50 50 - } => Some(Self { 51 51 - nsid: commit.collection.into(), 52 52 - timestamp: time_us, 53 53 - deleted: true, 54 54 - }), 55 55 - _ => None, 56 56 - } 57 57 - } 58 58 - } 59 59 - 60 60 - // counts is nsid -> NsidCounts 61 61 - // hits is tree per nsid: timestamp -> NsidHit 62 62 - pub struct Db { 63 63 - inner: Keyspace, 64 64 - hits: papaya::HashMap<SmolStr, Partition>, 65 65 - counts: Partition, 66 66 - event_broadcaster: broadcast::Sender<(SmolStr, NsidCounts)>, 67 67 - eps: Rate, 68 68 - } 69 69 - 70 70 - impl Db { 71 71 - pub fn new(path: impl AsRef<Path>) -> AppResult<Self> { 72 72 - tracing::info!("opening db..."); 73 73 - let ks = Config::new(path) 74 74 - .cache_size(8 * 1024 * 1024) // from talna 75 75 - .open()?; 76 76 - Ok(Self { 77 77 - hits: Default::default(), 78 78 - counts: ks.open_partition( 79 79 - "_counts", 80 80 - PartitionCreateOptions::default().compression(fjall::CompressionType::None), 81 81 - )?, 82 82 - inner: ks, 83 83 - event_broadcaster: broadcast::channel(1000).0, 84 84 - eps: Rate::new(Duration::from_secs(1)), 85 85 - }) 86 86 - } 87 87 - 88 88 - pub fn eps(&self) -> usize { 89 89 - self.eps.rate(&()) as usize 90 90 - } 91 91 - 92 92 - pub fn new_listener(&self) -> broadcast::Receiver<(SmolStr, NsidCounts)> { 93 93 - self.event_broadcaster.subscribe() 94 94 - } 95 95 - 96 96 - #[inline(always)] 97 97 - fn run_in_nsid_tree<T>( 98 98 - &self, 99 99 - nsid: &str, 100 100 - f: impl FnOnce(&Partition) -> AppResult<T>, 101 101 - ) -> AppResult<T> { 102 102 - f(self.hits.pin().get_or_insert_with(SmolStr::new(nsid), || { 103 103 - let opts = PartitionCreateOptions::default() 104 104 - .compression(fjall::CompressionType::Miniz(9)) 105 105 - .compaction_strategy(fjall::compaction::Strategy::Fifo(fjall::compaction::Fifo { 106 106 - limit: 5 * 1024 * 1024 * 1024, // 5 gb 107 107 - ttl_seconds: Some(60 * 60 * 24 * 30), // 30 days 108 108 - })); 109 109 - self.inner.open_partition(nsid, opts).unwrap() 110 110 - })) 111 111 - } 112 112 - 113 113 - pub fn record_event(&self, e: EventRecord) -> AppResult<()> { 114 114 - let EventRecord { 115 115 - nsid, 116 116 - timestamp, 117 117 - deleted, 118 118 - } = e; 119 119 - 120 120 - self.insert_event(&nsid, timestamp, deleted)?; 121 121 - // increment count 122 122 - let mut counts = self.get_count(&nsid)?; 123 123 - counts.last_seen = timestamp; 124 124 - if deleted { 125 125 - counts.deleted_count += 1; 126 126 - } else { 127 127 - counts.count += 1; 128 128 - } 129 129 - self.insert_count(&nsid, counts.clone())?; 130 130 - if self.event_broadcaster.receiver_count() > 0 { 131 131 - let _ = self.event_broadcaster.send((SmolStr::new(&nsid), counts)); 132 132 - } 133 133 - self.eps.observe(&(), 1); 134 134 - Ok(()) 135 135 - } 136 136 - 137 137 - #[inline(always)] 138 138 - fn insert_event(&self, nsid: &str, timestamp: u64, deleted: bool) -> AppResult<()> { 139 139 - self.run_in_nsid_tree(nsid, |tree| { 140 140 - tree.insert( 141 141 - timestamp.to_be_bytes(), 142 142 - unsafe { rkyv::to_bytes::<Error>(&NsidHit { deleted }).unwrap_unchecked() } 143 143 - .as_slice(), 144 144 - ) 145 145 - .map_err(AppError::from) 146 146 - }) 147 147 - } 148 148 - 149 149 - #[inline(always)] 150 150 - fn insert_count(&self, nsid: &str, counts: NsidCounts) -> AppResult<()> { 151 151 - self.counts 152 152 - .insert( 153 153 - nsid, 154 154 - unsafe { rkyv::to_bytes::<Error>(&counts).unwrap_unchecked() }.as_slice(), 155 155 - ) 156 156 - .map_err(AppError::from) 157 157 - } 158 158 - 159 159 - pub fn get_count(&self, nsid: &str) -> AppResult<NsidCounts> { 160 160 - let Some(raw) = self.counts.get(nsid)? else { 161 161 - return Ok(NsidCounts::default()); 162 162 - }; 163 163 - Ok(unsafe { rkyv::from_bytes_unchecked::<_, Error>(&raw).unwrap_unchecked() }) 164 164 - } 165 165 - 166 166 - pub fn get_counts(&self) -> impl Iterator<Item = AppResult<(SmolStr, NsidCounts)>> { 167 167 - self.counts.iter().map(|res| { 168 168 - res.map_err(AppError::from).map(|(key, val)| { 169 169 - ( 170 170 - SmolStr::new(unsafe { str::from_utf8_unchecked(&key) }), 171 171 - unsafe { rkyv::from_bytes_unchecked::<_, Error>(&val).unwrap_unchecked() }, 172 172 - ) 173 173 - }) 174 174 - }) 175 175 - } 176 176 - 177 177 - pub fn get_nsids(&self) -> impl Iterator<Item = impl Deref<Target = str>> { 178 178 - self.inner 179 179 - .list_partitions() 180 180 - .into_iter() 181 181 - .filter(|k| k.deref() != "_counts") 182 182 - } 183 183 - 184 184 - pub fn get_hits( 185 185 - &self, 186 186 - nsid: &str, 187 187 - range: impl RangeBounds<u64>, 188 188 - ) -> AppResult<Box<dyn Iterator<Item = AppResult<(u64, NsidHit)>>>> { 189 189 - let start = range.start_bound().cloned().map(u64::to_be_bytes); 190 190 - let end = range.end_bound().cloned().map(u64::to_be_bytes); 191 191 - 192 192 - let _guard = self.hits.guard(); 193 193 - let Some(tree) = self.hits.get(nsid, &_guard) else { 194 194 - return Ok(Box::new(std::iter::empty())); 195 195 - }; 196 196 - 197 197 - Ok(Box::new(tree.range(TimestampRange { start, end }).map( 198 198 - |res| { 199 199 - res.map_err(AppError::from).map(|(key, val)| { 200 200 - ( 201 201 - u64::from_be_bytes(key.as_ref().try_into().unwrap()), 202 202 - unsafe { rkyv::from_bytes_unchecked::<_, Error>(&val).unwrap_unchecked() }, 203 203 - ) 204 204 - }) 205 205 - }, 206 206 - ))) 207 207 - } 208 208 - 209 209 - pub fn tracking_since(&self) -> AppResult<u64> { 210 210 - let _guard = self.hits.guard(); 211 211 - // HACK: we should actually store when we started tracking but im lazy 212 212 - // should be accurate enough 213 213 - let Some(tree) = self.hits.get("app.bsky.feed.like", &_guard) else { 214 214 - return Ok(0); 215 215 - }; 216 216 - 217 217 - let Some((timestamp_raw, _)) = tree.first_key_value()? else { 218 218 - return Ok(0); 219 219 - }; 220 220 - 221 221 - Ok(u64::from_be_bytes( 222 222 - timestamp_raw.as_ref().try_into().unwrap(), 223 223 - )) 224 224 - } 225 225 - } 226 226 - 227 227 - type TimestampRepr = [u8; 8]; 228 228 - 229 229 - struct TimestampRange { 230 230 - start: Bound<TimestampRepr>, 231 231 - end: Bound<TimestampRepr>, 232 232 - } 233 233 - 234 234 - impl RangeBounds<TimestampRepr> for TimestampRange { 235 235 - #[inline(always)] 236 236 - fn start_bound(&self) -> Bound<&TimestampRepr> { 237 237 - self.start.as_ref() 238 238 - } 239 239 - 240 240 - #[inline(always)] 241 241 - fn end_bound(&self) -> Bound<&TimestampRepr> { 242 242 - self.end.as_ref() 243 243 - } 244 244 - }

+485

server/src/db/block.rs

··· 1 1 + use ordered_varint::Variable; 2 2 + use rkyv::{ 3 3 + Archive, Serialize, api::high::HighSerializer, rancor, ser::allocator::ArenaHandle, 4 4 + util::AlignedVec, 5 5 + }; 6 6 + use std::{ 7 7 + io::{self, Read, Write}, 8 8 + marker::PhantomData, 9 9 + }; 10 10 + 11 11 + use crate::error::AppResult; 12 12 + 13 13 + pub struct Item<T> { 14 14 + pub timestamp: u64, 15 15 + data: AlignedVec, 16 16 + phantom: PhantomData<T>, 17 17 + } 18 18 + 19 19 + impl<T: Archive> Item<T> { 20 20 + pub fn access(&self) -> &T::Archived { 21 21 + unsafe { rkyv::access_unchecked::<T::Archived>(&self.data) } 22 22 + } 23 23 + } 24 24 + 25 25 + impl<T: for<'a> Serialize<HighSerializer<AlignedVec, ArenaHandle<'a>, rancor::Error>>> Item<T> { 26 26 + pub fn new(timestamp: u64, data: &T) -> Self { 27 27 + Item { 28 28 + timestamp, 29 29 + data: unsafe { rkyv::to_bytes(data).unwrap_unchecked() }, 30 30 + phantom: PhantomData, 31 31 + } 32 32 + } 33 33 + } 34 34 + 35 35 + pub struct ItemEncoder<W: Write, T> { 36 36 + writer: W, 37 37 + prev_timestamp: u64, 38 38 + prev_delta: i64, 39 39 + _item: PhantomData<T>, 40 40 + } 41 41 + 42 42 + impl<W: Write, T> ItemEncoder<W, T> { 43 43 + pub fn new(writer: W) -> Self { 44 44 + ItemEncoder { 45 45 + writer, 46 46 + prev_timestamp: 0, 47 47 + prev_delta: 0, 48 48 + _item: PhantomData, 49 49 + } 50 50 + } 51 51 + 52 52 + pub fn encode(&mut self, item: &Item<T>) -> AppResult<()> { 53 53 + if self.prev_timestamp == 0 { 54 54 + // self.writer.write_varint(item.timestamp)?; 55 55 + self.prev_timestamp = item.timestamp; 56 56 + self.write_data(&item.data)?; 57 57 + return Ok(()); 58 58 + } 59 59 + 60 60 + let delta = (item.timestamp as i128 - self.prev_timestamp as i128) as i64; 61 61 + 62 62 + self.writer.write_varint(delta - self.prev_delta)?; 63 63 + self.prev_timestamp = item.timestamp; 64 64 + self.prev_delta = delta; 65 65 + 66 66 + self.write_data(&item.data)?; 67 67 + 68 68 + Ok(()) 69 69 + } 70 70 + 71 71 + fn write_data(&mut self, data: &[u8]) -> AppResult<()> { 72 72 + self.writer.write_varint(data.len())?; 73 73 + self.writer.write_all(data)?; 74 74 + Ok(()) 75 75 + } 76 76 + 77 77 + pub fn finish(mut self) -> AppResult<W> { 78 78 + self.writer.flush()?; 79 79 + Ok(self.writer) 80 80 + } 81 81 + } 82 82 + 83 83 + pub struct ItemDecoder<R, T> { 84 84 + reader: R, 85 85 + current_timestamp: u64, 86 86 + current_delta: i64, 87 87 + first_item: bool, 88 88 + _item: PhantomData<T>, 89 89 + } 90 90 + 91 91 + impl<R: Read, T: Archive> ItemDecoder<R, T> { 92 92 + pub fn new(reader: R, start_timestamp: u64) -> AppResult<Self> { 93 93 + Ok(ItemDecoder { 94 94 + reader, 95 95 + current_timestamp: start_timestamp, 96 96 + current_delta: 0, 97 97 + first_item: true, 98 98 + _item: PhantomData, 99 99 + }) 100 100 + } 101 101 + 102 102 + pub fn decode(&mut self) -> AppResult<Option<Item<T>>> { 103 103 + if self.first_item { 104 104 + // read the first timestamp 105 105 + // let timestamp = match self.reader.read_varint::<u64>() { 106 106 + // Ok(timestamp) => timestamp, 107 107 + // Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None), 108 108 + // Err(e) => return Err(e.into()), 109 109 + // }; 110 110 + // self.current_timestamp = timestamp; 111 111 + 112 112 + let Some(data_raw) = self.read_item()? else { 113 113 + return Ok(None); 114 114 + }; 115 115 + self.first_item = false; 116 116 + return Ok(Some(Item { 117 117 + timestamp: self.current_timestamp, 118 118 + data: data_raw, 119 119 + phantom: PhantomData, 120 120 + })); 121 121 + } 122 122 + 123 123 + let Some(_delta) = self.read_timestamp()? else { 124 124 + return Ok(None); 125 125 + }; 126 126 + 127 127 + // read data 128 128 + let data_raw = match self.read_item()? { 129 129 + Some(data_raw) => data_raw, 130 130 + None => { 131 131 + return Err(io::Error::new( 132 132 + io::ErrorKind::UnexpectedEof, 133 133 + "expected data after delta", 134 134 + ) 135 135 + .into()); 136 136 + } 137 137 + }; 138 138 + 139 139 + Ok(Some(Item { 140 140 + timestamp: self.current_timestamp, 141 141 + data: data_raw, 142 142 + phantom: PhantomData, 143 143 + })) 144 144 + } 145 145 + 146 146 + // [10, 11, 12, 14] -> [1, 1, 2] -> [0, 1] 147 147 + fn read_timestamp(&mut self) -> AppResult<Option<u64>> { 148 148 + let delta = match self.reader.read_varint::<i64>() { 149 149 + Ok(delta) => delta, 150 150 + Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None), 151 151 + Err(e) => return Err(e.into()), 152 152 + }; 153 153 + self.current_delta += delta; 154 154 + self.current_timestamp = 155 155 + (self.current_timestamp as i128 + self.current_delta as i128) as u64; 156 156 + Ok(Some(self.current_timestamp)) 157 157 + } 158 158 + 159 159 + fn read_item(&mut self) -> AppResult<Option<AlignedVec>> { 160 160 + let data_len = match self.reader.read_varint::<usize>() { 161 161 + Ok(data_len) => data_len, 162 162 + Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None), 163 163 + Err(e) => return Err(e.into()), 164 164 + }; 165 165 + let mut data_raw = AlignedVec::with_capacity(data_len); 166 166 + for _ in 0..data_len { 167 167 + data_raw.push(0); 168 168 + } 169 169 + self.reader.read_exact(data_raw.as_mut_slice())?; 170 170 + Ok(Some(data_raw)) 171 171 + } 172 172 + } 173 173 + 174 174 + impl<R: Read, T: Archive> Iterator for ItemDecoder<R, T> { 175 175 + type Item = AppResult<Item<T>>; 176 176 + 177 177 + fn next(&mut self) -> Option<Self::Item> { 178 178 + self.decode().transpose() 179 179 + } 180 180 + } 181 181 + 182 182 + pub trait WriteVariableExt: Write { 183 183 + fn write_varint(&mut self, value: impl Variable) -> io::Result<usize> { 184 184 + value.encode_variable(self) 185 185 + } 186 186 + } 187 187 + impl<W: Write> WriteVariableExt for W {} 188 188 + 189 189 + pub trait ReadVariableExt: Read { 190 190 + fn read_varint<T: Variable>(&mut self) -> io::Result<T> { 191 191 + T::decode_variable(self) 192 192 + } 193 193 + } 194 194 + impl<R: Read> ReadVariableExt for R {} 195 195 + 196 196 + #[cfg(test)] 197 197 + mod test { 198 198 + use super::*; 199 199 + use rkyv::{Archive, Deserialize, Serialize}; 200 200 + use std::io::Cursor; 201 201 + 202 202 + #[derive(Archive, Deserialize, Serialize, Debug, PartialEq)] 203 203 + #[rkyv(compare(PartialEq))] 204 204 + struct TestData { 205 205 + id: u32, 206 206 + value: String, 207 207 + } 208 208 + 209 209 + #[test] 210 210 + fn test_encoder_decoder_single_item() { 211 211 + let data = TestData { 212 212 + id: 123, 213 213 + value: "test".to_string(), 214 214 + }; 215 215 + 216 216 + let item = Item::new(1000, &data); 217 217 + 218 218 + // encode 219 219 + let mut buffer = Vec::new(); 220 220 + let mut encoder = ItemEncoder::new(&mut buffer); 221 221 + encoder.encode(&item).unwrap(); 222 222 + encoder.finish().unwrap(); 223 223 + 224 224 + // decode 225 225 + let cursor = Cursor::new(buffer); 226 226 + let mut decoder = ItemDecoder::<_, TestData>::new(cursor, 1000).unwrap(); 227 227 + 228 228 + let decoded_item = decoder.decode().unwrap().unwrap(); 229 229 + assert_eq!(decoded_item.timestamp, 1000); 230 230 + 231 231 + let decoded_data = decoded_item.access(); 232 232 + assert_eq!(decoded_data.id, 123); 233 233 + assert_eq!(decoded_data.value.as_str(), "test"); 234 234 + } 235 235 + 236 236 + #[test] 237 237 + fn test_encoder_decoder_multiple_items() { 238 238 + let items = vec![ 239 239 + Item::new( 240 240 + 1000, 241 241 + &TestData { 242 242 + id: 1, 243 243 + value: "first".to_string(), 244 244 + }, 245 245 + ), 246 246 + Item::new( 247 247 + 1010, 248 248 + &TestData { 249 249 + id: 2, 250 250 + value: "second".to_string(), 251 251 + }, 252 252 + ), 253 253 + Item::new( 254 254 + 1015, 255 255 + &TestData { 256 256 + id: 3, 257 257 + value: "third".to_string(), 258 258 + }, 259 259 + ), 260 260 + Item::new( 261 261 + 1025, 262 262 + &TestData { 263 263 + id: 4, 264 264 + value: "fourth".to_string(), 265 265 + }, 266 266 + ), 267 267 + ]; 268 268 + 269 269 + // encode 270 270 + let mut buffer = Vec::new(); 271 271 + let mut encoder = ItemEncoder::new(&mut buffer); 272 272 + 273 273 + for item in &items { 274 274 + encoder.encode(item).unwrap(); 275 275 + } 276 276 + encoder.finish().unwrap(); 277 277 + 278 278 + // decode 279 279 + let cursor = Cursor::new(buffer); 280 280 + let mut decoder = ItemDecoder::<_, TestData>::new(cursor, 1000).unwrap(); 281 281 + 282 282 + let mut decoded_items = Vec::new(); 283 283 + while let Some(item) = decoder.decode().unwrap() { 284 284 + decoded_items.push(item); 285 285 + } 286 286 + 287 287 + assert_eq!(decoded_items.len(), 4); 288 288 + 289 289 + for (original, decoded) in items.iter().zip(decoded_items.iter()) { 290 290 + assert_eq!(original.timestamp, decoded.timestamp); 291 291 + assert_eq!(original.access().id, decoded.access().id); 292 292 + assert_eq!( 293 293 + original.access().value.as_str(), 294 294 + decoded.access().value.as_str() 295 295 + ); 296 296 + } 297 297 + } 298 298 + 299 299 + #[test] 300 300 + fn test_encoder_decoder_with_iterator() { 301 301 + let items = vec![ 302 302 + Item::new( 303 303 + 2000, 304 304 + &TestData { 305 305 + id: 10, 306 306 + value: "a".to_string(), 307 307 + }, 308 308 + ), 309 309 + Item::new( 310 310 + 2005, 311 311 + &TestData { 312 312 + id: 20, 313 313 + value: "b".to_string(), 314 314 + }, 315 315 + ), 316 316 + Item::new( 317 317 + 2012, 318 318 + &TestData { 319 319 + id: 30, 320 320 + value: "c".to_string(), 321 321 + }, 322 322 + ), 323 323 + ]; 324 324 + 325 325 + // encode 326 326 + let mut buffer = Vec::new(); 327 327 + let mut encoder = ItemEncoder::new(&mut buffer); 328 328 + 329 329 + for item in &items { 330 330 + encoder.encode(item).unwrap(); 331 331 + } 332 332 + encoder.finish().unwrap(); 333 333 + 334 334 + // decode 335 335 + let cursor = Cursor::new(buffer); 336 336 + let decoder = ItemDecoder::<_, TestData>::new(cursor, 2000).unwrap(); 337 337 + 338 338 + let decoded_items: Result<Vec<_>, _> = decoder.collect(); 339 339 + let decoded_items = decoded_items.unwrap(); 340 340 + 341 341 + assert_eq!(decoded_items.len(), 3); 342 342 + assert_eq!(decoded_items[0].timestamp, 2000); 343 343 + assert_eq!(decoded_items[1].timestamp, 2005); 344 344 + assert_eq!(decoded_items[2].timestamp, 2012); 345 345 + 346 346 + assert_eq!(decoded_items[0].access().id, 10); 347 347 + assert_eq!(decoded_items[1].access().id, 20); 348 348 + assert_eq!(decoded_items[2].access().id, 30); 349 349 + } 350 350 + 351 351 + #[test] 352 352 + fn test_delta_compression() { 353 353 + let items = vec![ 354 354 + Item::new( 355 355 + 1000, 356 356 + &TestData { 357 357 + id: 1, 358 358 + value: "a".to_string(), 359 359 + }, 360 360 + ), 361 361 + Item::new( 362 362 + 1010, 363 363 + &TestData { 364 364 + id: 2, 365 365 + value: "b".to_string(), 366 366 + }, 367 367 + ), // delta = 10 368 368 + Item::new( 369 369 + 1020, 370 370 + &TestData { 371 371 + id: 3, 372 372 + value: "c".to_string(), 373 373 + }, 374 374 + ), // delta = 10, delta-of-delta = 0 375 375 + Item::new( 376 376 + 1025, 377 377 + &TestData { 378 378 + id: 4, 379 379 + value: "d".to_string(), 380 380 + }, 381 381 + ), // delta = 5, delta-of-delta = -5 382 382 + ]; 383 383 + 384 384 + let mut buffer = Vec::new(); 385 385 + let mut encoder = ItemEncoder::new(&mut buffer); 386 386 + 387 387 + for item in &items { 388 388 + encoder.encode(item).unwrap(); 389 389 + } 390 390 + encoder.finish().unwrap(); 391 391 + 392 392 + // decode and verify 393 393 + let cursor = Cursor::new(buffer); 394 394 + let decoder = ItemDecoder::<_, TestData>::new(cursor, 1000).unwrap(); 395 395 + 396 396 + let decoded_items: Result<Vec<_>, _> = decoder.collect(); 397 397 + let decoded_items = decoded_items.unwrap(); 398 398 + 399 399 + for (original, decoded) in items.iter().zip(decoded_items.iter()) { 400 400 + assert_eq!(original.timestamp, decoded.timestamp); 401 401 + assert_eq!(original.access().id, decoded.access().id); 402 402 + } 403 403 + } 404 404 + 405 405 + #[test] 406 406 + fn test_empty_decode() { 407 407 + let buffer = Vec::new(); 408 408 + let cursor = Cursor::new(buffer); 409 409 + let mut decoder = ItemDecoder::<_, TestData>::new(cursor, 1000).unwrap(); 410 410 + 411 411 + let result = decoder.decode().unwrap(); 412 412 + assert!(result.is_none()); 413 413 + } 414 414 + 415 415 + #[test] 416 416 + fn test_backwards_timestamp() { 417 417 + let items = vec![ 418 418 + Item::new( 419 419 + 1000, 420 420 + &TestData { 421 421 + id: 1, 422 422 + value: "first".to_string(), 423 423 + }, 424 424 + ), 425 425 + Item::new( 426 426 + 900, 427 427 + &TestData { 428 428 + id: 2, 429 429 + value: "second".to_string(), 430 430 + }, 431 431 + ), 432 432 + ]; 433 433 + 434 434 + let mut buffer = Vec::new(); 435 435 + let mut encoder = ItemEncoder::new(&mut buffer); 436 436 + 437 437 + for item in &items { 438 438 + encoder.encode(item).unwrap(); 439 439 + } 440 440 + encoder.finish().unwrap(); 441 441 + 442 442 + let cursor = Cursor::new(buffer); 443 443 + let decoder = ItemDecoder::<_, TestData>::new(cursor, 1000).unwrap(); 444 444 + 445 445 + let decoded_items: Result<Vec<_>, _> = decoder.collect(); 446 446 + let decoded_items = decoded_items.unwrap(); 447 447 + 448 448 + assert_eq!(decoded_items.len(), 2); 449 449 + assert_eq!(decoded_items[0].timestamp, 1000); 450 450 + assert_eq!(decoded_items[1].timestamp, 900); 451 451 + } 452 452 + 453 453 + #[test] 454 454 + fn test_different_data_sizes() { 455 455 + let small_data = TestData { 456 456 + id: 1, 457 457 + value: "x".to_string(), 458 458 + }; 459 459 + let large_data = TestData { 460 460 + id: 2, 461 461 + value: "a".repeat(1000), 462 462 + }; 463 463 + 464 464 + let items = vec![Item::new(1000, &small_data), Item::new(1001, &large_data)]; 465 465 + 466 466 + let mut buffer = Vec::new(); 467 467 + let mut encoder = ItemEncoder::new(&mut buffer); 468 468 + 469 469 + for item in &items { 470 470 + encoder.encode(item).unwrap(); 471 471 + } 472 472 + encoder.finish().unwrap(); 473 473 + 474 474 + let cursor = Cursor::new(buffer); 475 475 + let decoder = ItemDecoder::<_, TestData>::new(cursor, 1000).unwrap(); 476 476 + 477 477 + let decoded_items: Result<Vec<_>, _> = decoder.collect(); 478 478 + let decoded_items = decoded_items.unwrap(); 479 479 + 480 480 + assert_eq!(decoded_items.len(), 2); 481 481 + assert_eq!(decoded_items[0].access().value.as_str(), "x"); 482 482 + assert_eq!(decoded_items[1].access().value.len(), 1000); 483 483 + assert_eq!(decoded_items[1].access().value.as_str(), "a".repeat(1000)); 484 484 + } 485 485 + }

+589

server/src/db/mod.rs

··· 1 1 + use std::{ 2 2 + io::Cursor, 3 3 + ops::{Bound, Deref, RangeBounds}, 4 4 + path::Path, 5 5 + sync::{ 6 6 + Arc, 7 7 + atomic::{AtomicUsize, Ordering as AtomicOrdering}, 8 8 + }, 9 9 + time::{Duration, Instant}, 10 10 + }; 11 11 + 12 12 + use atomic_time::AtomicInstant; 13 13 + use fjall::{Config, Keyspace, Partition, PartitionCreateOptions, Slice}; 14 14 + use ordered_varint::Variable; 15 15 + use pingora_limits::rate::Rate; 16 16 + use rkyv::{Archive, Deserialize, Serialize, rancor::Error}; 17 17 + use smol_str::SmolStr; 18 18 + use tokio::sync::broadcast; 19 19 + 20 20 + use crate::{ 21 21 + db::block::{ReadVariableExt, WriteVariableExt}, 22 22 + error::{AppError, AppResult}, 23 23 + jetstream::JetstreamEvent, 24 24 + utils::time_now, 25 25 + }; 26 26 + 27 27 + mod block; 28 28 + 29 29 + #[derive(Clone, Debug, Default, Archive, Deserialize, Serialize, PartialEq)] 30 30 + #[rkyv(compare(PartialEq), derive(Debug))] 31 31 + pub struct NsidCounts { 32 32 + pub count: u128, 33 33 + pub deleted_count: u128, 34 34 + pub last_seen: u64, 35 35 + } 36 36 + 37 37 + #[derive(Debug, Default, Archive, Deserialize, Serialize, PartialEq)] 38 38 + #[rkyv(compare(PartialEq), derive(Debug))] 39 39 + pub struct NsidHit { 40 40 + pub deleted: bool, 41 41 + } 42 42 + 43 43 + #[derive(Clone)] 44 44 + pub struct EventRecord { 45 45 + pub nsid: SmolStr, 46 46 + pub timestamp: u64, // seconds 47 47 + pub deleted: bool, 48 48 + } 49 49 + 50 50 + impl EventRecord { 51 51 + pub fn from_jetstream(event: JetstreamEvent) -> Option<Self> { 52 52 + match event { 53 53 + JetstreamEvent::Commit { 54 54 + time_us, commit, .. 55 55 + } => Some(Self { 56 56 + nsid: commit.collection.into(), 57 57 + timestamp: time_us / 1_000_000, 58 58 + deleted: false, 59 59 + }), 60 60 + JetstreamEvent::Delete { 61 61 + time_us, commit, .. 62 62 + } => Some(Self { 63 63 + nsid: commit.collection.into(), 64 64 + timestamp: time_us / 1_000_000, 65 65 + deleted: true, 66 66 + }), 67 67 + _ => None, 68 68 + } 69 69 + } 70 70 + } 71 71 + 72 72 + // counts is nsid -> NsidCounts 73 73 + // hits is tree per nsid: timestamp -> NsidHit 74 74 + pub struct DbOld { 75 75 + inner: Keyspace, 76 76 + hits: scc::HashIndex<SmolStr, Partition>, 77 77 + counts: Partition, 78 78 + event_broadcaster: broadcast::Sender<(SmolStr, NsidCounts)>, 79 79 + eps: Rate, 80 80 + } 81 81 + 82 82 + impl DbOld { 83 83 + pub fn new(path: impl AsRef<Path>) -> AppResult<Self> { 84 84 + tracing::info!("opening db..."); 85 85 + let ks = Config::new(path) 86 86 + .cache_size(8 * 1024 * 1024) // from talna 87 87 + .open()?; 88 88 + Ok(Self { 89 89 + hits: Default::default(), 90 90 + counts: ks.open_partition( 91 91 + "_counts", 92 92 + PartitionCreateOptions::default().compression(fjall::CompressionType::None), 93 93 + )?, 94 94 + inner: ks, 95 95 + event_broadcaster: broadcast::channel(1000).0, 96 96 + eps: Rate::new(Duration::from_secs(1)), 97 97 + }) 98 98 + } 99 99 + 100 100 + pub fn eps(&self) -> usize { 101 101 + self.eps.rate(&()) as usize 102 102 + } 103 103 + 104 104 + pub fn new_listener(&self) -> broadcast::Receiver<(SmolStr, NsidCounts)> { 105 105 + self.event_broadcaster.subscribe() 106 106 + } 107 107 + 108 108 + #[inline(always)] 109 109 + fn run_in_nsid_tree<T>( 110 110 + &self, 111 111 + nsid: &str, 112 112 + f: impl FnOnce(&Partition) -> AppResult<T>, 113 113 + ) -> AppResult<T> { 114 114 + f(self 115 115 + .hits 116 116 + .entry(SmolStr::new(nsid)) 117 117 + .or_insert_with(|| { 118 118 + let opts = PartitionCreateOptions::default() 119 119 + .compression(fjall::CompressionType::Miniz(9)) 120 120 + .compaction_strategy(fjall::compaction::Strategy::Fifo( 121 121 + fjall::compaction::Fifo { 122 122 + limit: 5 * 1024 * 1024 * 1024, // 5 gb 123 123 + ttl_seconds: Some(60 * 60 * 24 * 30), // 30 days 124 124 + }, 125 125 + )); 126 126 + self.inner.open_partition(nsid, opts).unwrap() 127 127 + }) 128 128 + .get()) 129 129 + } 130 130 + 131 131 + pub fn record_event(&self, e: EventRecord) -> AppResult<()> { 132 132 + let EventRecord { 133 133 + nsid, 134 134 + timestamp, 135 135 + deleted, 136 136 + } = e; 137 137 + 138 138 + self.insert_event(&nsid, timestamp, deleted)?; 139 139 + // increment count 140 140 + let mut counts = self.get_count(&nsid)?; 141 141 + counts.last_seen = timestamp; 142 142 + if deleted { 143 143 + counts.deleted_count += 1; 144 144 + } else { 145 145 + counts.count += 1; 146 146 + } 147 147 + self.insert_count(&nsid, counts.clone())?; 148 148 + if self.event_broadcaster.receiver_count() > 0 { 149 149 + let _ = self.event_broadcaster.send((SmolStr::new(&nsid), counts)); 150 150 + } 151 151 + self.eps.observe(&(), 1); 152 152 + Ok(()) 153 153 + } 154 154 + 155 155 + #[inline(always)] 156 156 + fn insert_event(&self, nsid: &str, timestamp: u64, deleted: bool) -> AppResult<()> { 157 157 + self.run_in_nsid_tree(nsid, |tree| { 158 158 + tree.insert( 159 159 + timestamp.to_be_bytes(), 160 160 + unsafe { rkyv::to_bytes::<Error>(&NsidHit { deleted }).unwrap_unchecked() } 161 161 + .as_slice(), 162 162 + ) 163 163 + .map_err(AppError::from) 164 164 + }) 165 165 + } 166 166 + 167 167 + #[inline(always)] 168 168 + fn insert_count(&self, nsid: &str, counts: NsidCounts) -> AppResult<()> { 169 169 + self.counts 170 170 + .insert( 171 171 + nsid, 172 172 + unsafe { rkyv::to_bytes::<Error>(&counts).unwrap_unchecked() }.as_slice(), 173 173 + ) 174 174 + .map_err(AppError::from) 175 175 + } 176 176 + 177 177 + pub fn get_count(&self, nsid: &str) -> AppResult<NsidCounts> { 178 178 + let Some(raw) = self.counts.get(nsid)? else { 179 179 + return Ok(NsidCounts::default()); 180 180 + }; 181 181 + Ok(unsafe { rkyv::from_bytes_unchecked::<_, Error>(&raw).unwrap_unchecked() }) 182 182 + } 183 183 + 184 184 + pub fn get_counts(&self) -> impl Iterator<Item = AppResult<(SmolStr, NsidCounts)>> { 185 185 + self.counts.iter().map(|res| { 186 186 + res.map_err(AppError::from).map(|(key, val)| { 187 187 + ( 188 188 + SmolStr::new(unsafe { str::from_utf8_unchecked(&key) }), 189 189 + unsafe { rkyv::from_bytes_unchecked::<_, Error>(&val).unwrap_unchecked() }, 190 190 + ) 191 191 + }) 192 192 + }) 193 193 + } 194 194 + 195 195 + pub fn get_nsids(&self) -> impl Iterator<Item = impl Deref<Target = str>> { 196 196 + self.inner 197 197 + .list_partitions() 198 198 + .into_iter() 199 199 + .filter(|k| k.deref() != "_counts") 200 200 + } 201 201 + 202 202 + pub fn get_hits( 203 203 + &self, 204 204 + nsid: &str, 205 205 + range: impl RangeBounds<u64>, 206 206 + ) -> AppResult<Box<dyn Iterator<Item = AppResult<(u64, NsidHit)>>>> { 207 207 + let start = range.start_bound().cloned().map(u64::to_be_bytes); 208 208 + let end = range.end_bound().cloned().map(u64::to_be_bytes); 209 209 + 210 210 + let _guard = scc::ebr::Guard::new(); 211 211 + let Some(tree) = self.hits.peek(nsid, &_guard) else { 212 212 + return Ok(Box::new(std::iter::empty())); 213 213 + }; 214 214 + 215 215 + Ok(Box::new(tree.range(TimestampRangeOld { start, end }).map( 216 216 + |res| { 217 217 + res.map_err(AppError::from).map(|(key, val)| { 218 218 + ( 219 219 + u64::from_be_bytes(key.as_ref().try_into().unwrap()), 220 220 + unsafe { rkyv::from_bytes_unchecked::<_, Error>(&val).unwrap_unchecked() }, 221 221 + ) 222 222 + }) 223 223 + }, 224 224 + ))) 225 225 + } 226 226 + 227 227 + pub fn tracking_since(&self) -> AppResult<u64> { 228 228 + let _guard = scc::ebr::Guard::new(); 229 229 + // HACK: we should actually store when we started tracking but im lazy 230 230 + // should be accurate enough 231 231 + let Some(tree) = self.hits.peek("app.bsky.feed.like", &_guard) else { 232 232 + return Ok(0); 233 233 + }; 234 234 + let Some((timestamp_raw, _)) = tree.first_key_value()? else { 235 235 + return Ok(0); 236 236 + }; 237 237 + drop(_guard); 238 238 + 239 239 + Ok(u64::from_be_bytes( 240 240 + timestamp_raw.as_ref().try_into().unwrap(), 241 241 + )) 242 242 + } 243 243 + } 244 244 + 245 245 + type ItemDecoder = block::ItemDecoder<Cursor<Slice>, NsidHit>; 246 246 + type ItemEncoder = block::ItemEncoder<Vec<u8>, NsidHit>; 247 247 + type Item = block::Item<NsidHit>; 248 248 + 249 249 + pub struct LexiconHandle { 250 250 + tree: Partition, 251 251 + buf: Arc<scc::Queue<EventRecord>>, 252 252 + buf_len: AtomicUsize, 253 253 + last_insert: AtomicInstant, 254 254 + eps: Rate, 255 255 + block_size: AtomicUsize, 256 256 + } 257 257 + 258 258 + impl LexiconHandle { 259 259 + fn new(keyspace: &Keyspace, nsid: &str) -> Self { 260 260 + let opts = PartitionCreateOptions::default().compression(fjall::CompressionType::Miniz(9)); 261 261 + Self { 262 262 + tree: keyspace.open_partition(nsid, opts).unwrap(), 263 263 + buf: Default::default(), 264 264 + buf_len: AtomicUsize::new(0), 265 265 + last_insert: AtomicInstant::now(), 266 266 + eps: Rate::new(Duration::from_secs(5)), 267 267 + block_size: AtomicUsize::new(1000), 268 268 + } 269 269 + } 270 270 + 271 271 + fn item_count(&self) -> usize { 272 272 + self.buf_len.load(AtomicOrdering::Acquire) 273 273 + } 274 274 + 275 275 + fn last_insert(&self) -> Instant { 276 276 + self.last_insert.load(AtomicOrdering::Acquire) 277 277 + } 278 278 + 279 279 + fn suggested_block_size(&self) -> usize { 280 280 + self.block_size.load(AtomicOrdering::Relaxed) 281 281 + } 282 282 + 283 283 + fn insert(&self, event: EventRecord) { 284 284 + self.buf.push(event); 285 285 + self.buf_len.fetch_add(1, AtomicOrdering::Release); 286 286 + self.last_insert 287 287 + .store(Instant::now(), AtomicOrdering::Release); 288 288 + self.eps.observe(&(), 1); 289 289 + let rate = self.eps.rate(&()) as usize; 290 290 + if rate != 0 { 291 291 + self.block_size.store(rate * 60, AtomicOrdering::Relaxed); 292 292 + } 293 293 + } 294 294 + 295 295 + fn sync(&self) -> AppResult<()> { 296 296 + let mut writer = ItemEncoder::new(Vec::with_capacity( 297 297 + size_of::<u64>() + self.item_count() * size_of::<(u64, NsidHit)>(), 298 298 + )); 299 299 + let mut start_timestamp = None; 300 300 + let mut end_timestamp = None; 301 301 + while let Some(event) = self.buf.pop() { 302 302 + let item = Item::new( 303 303 + event.timestamp, 304 304 + &NsidHit { 305 305 + deleted: event.deleted, 306 306 + }, 307 307 + ); 308 308 + writer.encode(&item)?; 309 309 + if start_timestamp.is_none() { 310 310 + start_timestamp = Some(event.timestamp); 311 311 + } 312 312 + end_timestamp = Some(event.timestamp); 313 313 + } 314 314 + if let (Some(start_timestamp), Some(end_timestamp)) = (start_timestamp, end_timestamp) { 315 315 + self.buf_len.store(0, AtomicOrdering::Release); 316 316 + let value = writer.finish()?; 317 317 + let mut key = Vec::with_capacity(size_of::<u64>() * 2); 318 318 + key.write_varint(start_timestamp)?; 319 319 + key.write_varint(end_timestamp)?; 320 320 + self.tree.insert(key, value)?; 321 321 + } 322 322 + Ok(()) 323 323 + } 324 324 + } 325 325 + 326 326 + type BoxedIter<T> = Box<dyn Iterator<Item = T>>; 327 327 + 328 328 + // counts is nsid -> NsidCounts 329 329 + // hits is tree per nsid: varint start time + varint end time -> block of hits 330 330 + pub struct Db { 331 331 + inner: Keyspace, 332 332 + hits: scc::HashIndex<SmolStr, Arc<LexiconHandle>>, 333 333 + counts: Partition, 334 334 + event_broadcaster: broadcast::Sender<(SmolStr, NsidCounts)>, 335 335 + eps: Rate, 336 336 + min_block_size: usize, 337 337 + max_last_activity: Duration, 338 338 + } 339 339 + 340 340 + impl Db { 341 341 + pub fn new(path: impl AsRef<Path>) -> AppResult<Self> { 342 342 + tracing::info!("opening db..."); 343 343 + let ks = Config::new(path) 344 344 + .cache_size(8 * 1024 * 1024) // from talna 345 345 + .open()?; 346 346 + Ok(Self { 347 347 + hits: Default::default(), 348 348 + counts: ks.open_partition( 349 349 + "_counts", 350 350 + PartitionCreateOptions::default().compression(fjall::CompressionType::None), 351 351 + )?, 352 352 + inner: ks, 353 353 + event_broadcaster: broadcast::channel(1000).0, 354 354 + eps: Rate::new(Duration::from_secs(1)), 355 355 + min_block_size: 512, 356 356 + max_last_activity: Duration::from_secs(10), 357 357 + }) 358 358 + } 359 359 + 360 360 + pub fn sync(&self, all: bool) -> AppResult<()> { 361 361 + let _guard = scc::ebr::Guard::new(); 362 362 + for (nsid, tree) in self.hits.iter(&_guard) { 363 363 + let count = tree.item_count(); 364 364 + let is_max_block_size = count > self.min_block_size.max(tree.suggested_block_size()); 365 365 + let is_too_old = tree.last_insert().elapsed() > self.max_last_activity; 366 366 + if count > 0 && (all || is_max_block_size || is_too_old) { 367 367 + tracing::info!("syncing {count} of {nsid} to db"); 368 368 + tree.sync()?; 369 369 + } 370 370 + } 371 371 + Ok(()) 372 372 + } 373 373 + 374 374 + #[inline(always)] 375 375 + pub fn eps(&self) -> usize { 376 376 + self.eps.rate(&()) as usize 377 377 + } 378 378 + 379 379 + #[inline(always)] 380 380 + pub fn new_listener(&self) -> broadcast::Receiver<(SmolStr, NsidCounts)> { 381 381 + self.event_broadcaster.subscribe() 382 382 + } 383 383 + 384 384 + #[inline(always)] 385 385 + fn maybe_run_in_nsid_tree<T>( 386 386 + &self, 387 387 + nsid: &str, 388 388 + f: impl FnOnce(&LexiconHandle) -> T, 389 389 + ) -> Option<T> { 390 390 + let _guard = scc::ebr::Guard::new(); 391 391 + let handle = match self.hits.peek(nsid, &_guard) { 392 392 + Some(handle) => handle.clone(), 393 393 + None => { 394 394 + if self.inner.partition_exists(nsid) { 395 395 + let handle = Arc::new(LexiconHandle::new(&self.inner, nsid)); 396 396 + let _ = self.hits.insert(SmolStr::new(nsid), handle.clone()); 397 397 + handle 398 398 + } else { 399 399 + return None; 400 400 + } 401 401 + } 402 402 + }; 403 403 + Some(f(&handle)) 404 404 + } 405 405 + 406 406 + #[inline(always)] 407 407 + fn run_in_nsid_tree<T>( 408 408 + &self, 409 409 + nsid: SmolStr, 410 410 + f: impl FnOnce(&LexiconHandle) -> AppResult<T>, 411 411 + ) -> AppResult<T> { 412 412 + f(self 413 413 + .hits 414 414 + .entry(nsid.clone()) 415 415 + .or_insert_with(move || Arc::new(LexiconHandle::new(&self.inner, &nsid))) 416 416 + .get()) 417 417 + } 418 418 + 419 419 + pub fn record_event(&self, e: EventRecord) -> AppResult<()> { 420 420 + let EventRecord { 421 421 + nsid, 422 422 + timestamp, 423 423 + deleted, 424 424 + } = e.clone(); 425 425 + 426 426 + // insert event 427 427 + self.run_in_nsid_tree(nsid.clone(), move |tree| Ok(tree.insert(e)))?; 428 428 + // increment count 429 429 + let mut counts = self.get_count(&nsid)?; 430 430 + counts.last_seen = timestamp; 431 431 + if deleted { 432 432 + counts.deleted_count += 1; 433 433 + } else { 434 434 + counts.count += 1; 435 435 + } 436 436 + self.insert_count(&nsid, counts.clone())?; 437 437 + if self.event_broadcaster.receiver_count() > 0 { 438 438 + let _ = self.event_broadcaster.send((SmolStr::new(&nsid), counts)); 439 439 + } 440 440 + self.eps.observe(&(), 1); 441 441 + Ok(()) 442 442 + } 443 443 + 444 444 + #[inline(always)] 445 445 + fn insert_count(&self, nsid: &str, counts: NsidCounts) -> AppResult<()> { 446 446 + self.counts 447 447 + .insert( 448 448 + nsid, 449 449 + unsafe { rkyv::to_bytes::<Error>(&counts).unwrap_unchecked() }.as_slice(), 450 450 + ) 451 451 + .map_err(AppError::from) 452 452 + } 453 453 + 454 454 + pub fn get_count(&self, nsid: &str) -> AppResult<NsidCounts> { 455 455 + let Some(raw) = self.counts.get(nsid)? else { 456 456 + return Ok(NsidCounts::default()); 457 457 + }; 458 458 + Ok(unsafe { rkyv::from_bytes_unchecked::<_, Error>(&raw).unwrap_unchecked() }) 459 459 + } 460 460 + 461 461 + pub fn get_counts(&self) -> impl Iterator<Item = AppResult<(SmolStr, NsidCounts)>> { 462 462 + self.counts.iter().map(|res| { 463 463 + res.map_err(AppError::from).map(|(key, val)| { 464 464 + ( 465 465 + SmolStr::new(unsafe { str::from_utf8_unchecked(&key) }), 466 466 + unsafe { rkyv::from_bytes_unchecked::<_, Error>(&val).unwrap_unchecked() }, 467 467 + ) 468 468 + }) 469 469 + }) 470 470 + } 471 471 + 472 472 + pub fn get_nsids(&self) -> impl Iterator<Item = impl Deref<Target = str>> { 473 473 + self.inner 474 474 + .list_partitions() 475 475 + .into_iter() 476 476 + .filter(|k| k.deref() != "_counts") 477 477 + } 478 478 + 479 479 + pub fn get_hits_debug(&self, nsid: &str) -> BoxedIter<AppResult<(Slice, Slice)>> { 480 480 + self.maybe_run_in_nsid_tree(nsid, |handle| -> BoxedIter<AppResult<(Slice, Slice)>> { 481 481 + Box::new( 482 482 + handle 483 483 + .tree 484 484 + .iter() 485 485 + .rev() 486 486 + .map(|res| res.map_err(AppError::from)), 487 487 + ) 488 488 + }) 489 489 + .unwrap_or_else(|| Box::new(std::iter::empty())) 490 490 + } 491 491 + 492 492 + pub fn get_hits( 493 493 + &self, 494 494 + nsid: &str, 495 495 + range: impl RangeBounds<u64> + std::fmt::Debug, 496 496 + ) -> BoxedIter<AppResult<Item>> { 497 497 + let start = range 498 498 + .start_bound() 499 499 + .cloned() 500 500 + .map(|t| unsafe { t.to_variable_vec().unwrap_unchecked() }); 501 501 + let end = range 502 502 + .end_bound() 503 503 + .cloned() 504 504 + .map(|t| unsafe { t.to_variable_vec().unwrap_unchecked() }); 505 505 + let limit = match range.end_bound().cloned() { 506 506 + Bound::Included(end) => end, 507 507 + Bound::Excluded(end) => end.saturating_sub(1), 508 508 + Bound::Unbounded => u64::MAX, 509 509 + }; 510 510 + 511 511 + self.maybe_run_in_nsid_tree(nsid, move |handle| -> BoxedIter<AppResult<Item>> { 512 512 + let map_block = move |(key, val)| { 513 513 + let mut key_reader = Cursor::new(key); 514 514 + let start_timestamp = key_reader.read_varint::<u64>()?; 515 515 + let items = 516 516 + ItemDecoder::new(Cursor::new(val), start_timestamp)?.take_while(move |item| { 517 517 + item.as_ref().map_or(true, |item| item.timestamp <= limit) 518 518 + }); 519 519 + Ok(items) 520 520 + }; 521 521 + 522 522 + Box::new( 523 523 + handle 524 524 + .tree 525 525 + .range(TimestampRange { start, end }) 526 526 + .map(move |res| res.map_err(AppError::from).and_then(map_block)) 527 527 + .flatten() 528 528 + .flatten(), 529 529 + ) 530 530 + }) 531 531 + .unwrap_or_else(|| Box::new(std::iter::empty())) 532 532 + } 533 533 + 534 534 + pub fn tracking_since(&self) -> AppResult<u64> { 535 535 + let _guard = scc::ebr::Guard::new(); 536 536 + // HACK: we should actually store when we started tracking but im lazy 537 537 + // should be accurate enough 538 538 + let Some(handle) = self.hits.peek("app.bsky.feed.like", &_guard) else { 539 539 + return Ok(0); 540 540 + }; 541 541 + let Some((timestamps_raw, _)) = handle.tree.first_key_value()? else { 542 542 + return Ok(0); 543 543 + }; 544 544 + drop(_guard); 545 545 + 546 546 + let mut timestamp_reader = Cursor::new(timestamps_raw); 547 547 + timestamp_reader 548 548 + .read_varint::<u64>() 549 549 + .map_err(AppError::from) 550 550 + } 551 551 + } 552 552 + 553 553 + type TimestampRepr = Vec<u8>; 554 554 + 555 555 + struct TimestampRange { 556 556 + start: Bound<TimestampRepr>, 557 557 + end: Bound<TimestampRepr>, 558 558 + } 559 559 + 560 560 + impl RangeBounds<TimestampRepr> for TimestampRange { 561 561 + #[inline(always)] 562 562 + fn start_bound(&self) -> Bound<&TimestampRepr> { 563 563 + self.start.as_ref() 564 564 + } 565 565 + 566 566 + #[inline(always)] 567 567 + fn end_bound(&self) -> Bound<&TimestampRepr> { 568 568 + self.end.as_ref() 569 569 + } 570 570 + } 571 571 + 572 572 + type TimestampReprOld = [u8; 8]; 573 573 + 574 574 + struct TimestampRangeOld { 575 575 + start: Bound<TimestampReprOld>, 576 576 + end: Bound<TimestampReprOld>, 577 577 + } 578 578 + 579 579 + impl RangeBounds<TimestampReprOld> for TimestampRangeOld { 580 580 + #[inline(always)] 581 581 + fn start_bound(&self) -> Bound<&TimestampReprOld> { 582 582 + self.start.as_ref() 583 583 + } 584 584 + 585 585 + #[inline(always)] 586 586 + fn end_bound(&self) -> Bound<&TimestampReprOld> { 587 587 + self.end.as_ref() 588 588 + } 589 589 + }

+41 -21

server/src/main.rs

··· 9 9 10 10 use crate::{ 11 11 api::serve, 12 12 - db::{Db, EventRecord}, 12 12 + db::{Db, DbOld, EventRecord}, 13 13 error::AppError, 14 14 jetstream::JetstreamClient, 15 15 }; ··· 18 18 mod db; 19 19 mod error; 20 20 mod jetstream; 21 21 + mod utils; 21 22 22 23 #[cfg(not(target_env = "msvc"))] 23 24 #[global_allocator] ··· 34 35 .compact() 35 36 .init(); 36 37 37 37 - if std::env::args() 38 38 - .nth(1) 39 39 - .map_or(false, |arg| arg == "migrate") 40 40 - { 41 41 - migrate_to_miniz(); 42 42 - return; 38 38 + match std::env::args().nth(1).as_deref() { 39 39 + Some("migrate") => { 40 40 + migrate(); 41 41 + return; 42 42 + } 43 43 + Some("debug") => { 44 44 + debug(); 45 45 + return; 46 46 + } 47 47 + _ => {} 43 48 } 44 49 45 50 let db = Arc::new(Db::new(".fjall_data").expect("couldnt create db")); ··· 58 63 }; 59 64 60 65 let cancel_token = CancellationToken::new(); 61 61 - let (event_tx, mut event_rx) = tokio::sync::mpsc::channel(1000); 62 66 63 67 let consume_events = tokio::spawn({ 64 68 let consume_cancel = cancel_token.child_token(); 69 69 + let db = db.clone(); 65 70 async move { 66 71 jetstream.connect().await?; 67 72 loop { ··· 71 76 let Some(record) = EventRecord::from_jetstream(event) else { 72 77 continue; 73 78 }; 74 74 - let _ = event_tx.send(record).await; 79 79 + let db = db.clone(); 80 80 + tokio::task::spawn_blocking(move || { 81 81 + if let Err(err) = db.record_event(record) { 82 82 + tracing::error!("failed to record event: {}", err); 83 83 + } 84 84 + }); 75 85 } 76 86 Err(err) => return Err(err), 77 87 }, ··· 81 91 } 82 92 }); 83 93 84 84 - let ingest_events = std::thread::spawn({ 94 94 + std::thread::spawn({ 85 95 let db = db.clone(); 86 96 move || { 87 87 - tracing::info!("starting ingest events thread..."); 88 88 - while let Some(e) = event_rx.blocking_recv() { 89 89 - if let Err(e) = db.record_event(e) { 90 90 - tracing::error!("failed to record event: {}", e); 97 97 + loop { 98 98 + match db.sync(false) { 99 99 + Ok(_) => (), 100 100 + Err(e) => tracing::error!("failed to sync db: {}", e), 91 101 } 102 102 + std::thread::sleep(std::time::Duration::from_secs(1)); 92 103 } 93 104 } 94 105 }); 95 106 96 107 tokio::select! { 97 97 - res = serve(db, cancel_token.child_token()) => { 108 108 + res = serve(db.clone(), cancel_token.child_token()) => { 98 109 if let Err(e) = res { 99 110 tracing::error!("serve failed: {}", e); 100 111 } ··· 114 125 } 115 126 116 127 tracing::info!("shutting down..."); 117 117 - ingest_events 118 118 - .join() 119 119 - .expect("failed to join ingest events thread"); 128 128 + db.sync(true).expect("couldnt sync db"); 129 129 + } 130 130 + 131 131 + fn debug() { 132 132 + let db = Db::new(".fjall_data").expect("couldnt create db"); 133 133 + for nsid in db.get_nsids() { 134 134 + let nsid = nsid.deref(); 135 135 + for hit in db.get_hits(nsid, ..) { 136 136 + let hit = hit.expect("cant read event"); 137 137 + println!("{nsid} {}", hit.timestamp); 138 138 + } 139 139 + } 120 140 } 121 141 122 122 - fn migrate_to_miniz() { 123 123 - let from = Db::new(".fjall_data").expect("couldnt create db"); 124 124 - let to = Db::new(".fjall_data_miniz").expect("couldnt create db"); 142 142 + fn migrate() { 143 143 + let from = DbOld::new(".fjall_data").expect("couldnt create db"); 144 144 + let to = Db::new(".fjall_data_migrated").expect("couldnt create db"); 125 145 126 146 let mut total_count = 0_u64; 127 147 for nsid in from.get_nsids() {

+8

server/src/utils.rs

··· 1 1 + use std::time::UNIX_EPOCH; 2 2 + 3 3 + pub fn time_now() -> u64 { 4 4 + std::time::SystemTime::now() 5 5 + .duration_since(UNIX_EPOCH) 6 6 + .expect("oops") 7 7 + .as_micros() as u64 8 8 + }