+14

Cargo.lock

reviewed

··· 3149 3149 ] 3150 3150 3151 3151 [[package]] 3152 3152 + name = "serde_qs" 3153 3153 + version = "1.0.0-rc.3" 3154 3154 + source = "registry+https://github.com/rust-lang/crates.io-index" 3155 3155 + checksum = "4cb0b9062a400c31442e67d1f2b1e7746bebd691110ebee1b7d0c7293b04fab1" 3156 3156 + dependencies = [ 3157 3157 + "itoa", 3158 3158 + "percent-encoding", 3159 3159 + "ryu", 3160 3160 + "serde", 3161 3161 + "thiserror 2.0.12", 3162 3162 + ] 3163 3163 + 3164 3164 + [[package]] 3152 3165 name = "serde_spanned" 3153 3166 version = "0.6.8" 3154 3167 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 3827 3840 "semver", 3828 3841 "serde", 3829 3842 "serde_json", 3843 3843 + "serde_qs", 3830 3844 "sha2", 3831 3845 "tempfile", 3832 3846 "thiserror 2.0.12",

+1

ufos/Cargo.toml

reviewed

··· 23 23 semver = "1.0.26" 24 24 serde = "1.0.219" 25 25 serde_json = "1.0.140" 26 26 + serde_qs = "1.0.0-rc.3" 26 27 sha2 = "0.10.9" 27 28 thiserror = "2.0.12" 28 29 tokio = { version = "1.44.2", features = ["full", "sync", "time"] }

+2 -1

ufos/src/lib.rs

reviewed

··· 6 6 pub mod server; 7 7 pub mod storage; 8 8 pub mod storage_fjall; 9 9 - pub mod storage_mem; 10 9 pub mod store_types; 11 10 12 11 use crate::error::BatchInsertError; ··· 287 286 #[derive(Debug, Serialize, JsonSchema)] 288 287 pub struct JustCount { 289 288 creates: u64, 289 289 + updates: u64, 290 290 + deletes: u64, 290 291 dids_estimate: u64, 291 292 } 292 293

+17 -57

ufos/src/main.rs

reviewed

··· 7 7 use ufos::server; 8 8 use ufos::storage::{StorageWhatever, StoreBackground, StoreReader, StoreWriter}; 9 9 use ufos::storage_fjall::FjallStorage; 10 10 - use ufos::storage_mem::MemStorage; 11 10 use ufos::store_types::SketchSecretPrefix; 12 11 use ufos::{nice_duration, ConsumerInfo}; 13 12 ··· 19 18 static GLOBAL: Jemalloc = Jemalloc; 20 19 21 20 /// Aggregate links in the at-mosphere 22 22 - #[derive(Parser, Debug)] 21 21 + #[derive(Parser, Debug, Clone)] 23 22 #[command(version, about, long_about = None)] 24 23 struct Args { 25 24 /// Jetstream server to connect to (exclusive with --fixture). Provide either a wss:// URL, or a shorhand value: ··· 47 46 /// todo: restore this 48 47 #[arg(long, action)] 49 48 pause_rw: bool, 50 50 - /// DEBUG: use an in-memory store instead of fjall 51 51 - #[arg(long, action)] 52 52 - in_mem: bool, 53 49 /// reset the rollup cursor, scrape through missed things in the past (backfill) 54 50 #[arg(long, action)] 55 51 reroll: bool, ··· 64 60 65 61 let args = Args::parse(); 66 62 let jetstream = args.jetstream.clone(); 67 67 - if args.in_mem { 68 68 - let (read_store, write_store, cursor, sketch_secret) = MemStorage::init( 69 69 - args.data, 70 70 - jetstream, 71 71 - args.jetstream_force, 72 72 - Default::default(), 73 73 - )?; 74 74 - go( 75 75 - args.jetstream, 76 76 - args.jetstream_fixture, 77 77 - args.pause_writer, 78 78 - args.backfill, 79 79 - args.reroll, 80 80 - read_store, 81 81 - write_store, 82 82 - cursor, 83 83 - sketch_secret, 84 84 - ) 85 85 - .await?; 86 86 - } else { 87 87 - let (read_store, write_store, cursor, sketch_secret) = FjallStorage::init( 88 88 - args.data, 89 89 - jetstream, 90 90 - args.jetstream_force, 91 91 - Default::default(), 92 92 - )?; 93 93 - go( 94 94 - args.jetstream, 95 95 - args.jetstream_fixture, 96 96 - args.pause_writer, 97 97 - args.backfill, 98 98 - args.reroll, 99 99 - read_store, 100 100 - write_store, 101 101 - cursor, 102 102 - sketch_secret, 103 103 - ) 104 104 - .await?; 105 105 - } 106 106 - 63 63 + let (read_store, write_store, cursor, sketch_secret) = FjallStorage::init( 64 64 + args.data.clone(), 65 65 + jetstream, 66 66 + args.jetstream_force, 67 67 + Default::default(), 68 68 + )?; 69 69 + go(args, read_store, write_store, cursor, sketch_secret).await?; 107 70 Ok(()) 108 71 } 109 72 110 110 - #[allow(clippy::too_many_arguments)] 111 73 async fn go<B: StoreBackground>( 112 112 - jetstream: String, 113 113 - jetstream_fixture: bool, 114 114 - pause_writer: bool, 115 115 - backfill: bool, 116 116 - reroll: bool, 74 74 + args: Args, 117 75 read_store: impl StoreReader + 'static + Clone, 118 76 mut write_store: impl StoreWriter<B> + 'static, 119 77 cursor: Option<Cursor>, ··· 122 80 println!("starting server with storage..."); 123 81 let serving = server::serve(read_store.clone()); 124 82 125 125 - if pause_writer { 83 83 + if args.pause_writer { 126 84 log::info!("not starting jetstream or the write loop."); 127 85 serving.await.map_err(|e| anyhow::anyhow!(e))?; 128 86 return Ok(()); 129 87 } 130 88 131 131 - let batches = if jetstream_fixture { 132 132 - log::info!("starting with jestream file fixture: {jetstream:?}"); 133 133 - file_consumer::consume(jetstream.into(), sketch_secret, cursor).await? 89 89 + let batches = if args.jetstream_fixture { 90 90 + log::info!("starting with jestream file fixture: {:?}", args.jetstream); 91 91 + file_consumer::consume(args.jetstream.into(), sketch_secret, cursor).await? 134 92 } else { 135 93 log::info!( 136 94 "starting consumer with cursor: {cursor:?} from {:?} ago", 137 95 cursor.map(|c| c.elapsed()) 138 96 ); 139 139 - consumer::consume(&jetstream, cursor, false, sketch_secret).await? 97 97 + consumer::consume(&args.jetstream, cursor, false, sketch_secret).await? 140 98 }; 141 99 142 142 - let rolling = write_store.background_tasks(reroll)?.run(backfill); 100 100 + let rolling = write_store 101 101 + .background_tasks(args.reroll)? 102 102 + .run(args.backfill); 143 103 let storing = write_store.receive_batches(batches); 144 104 145 105 let stating = do_update_stuff(read_store);

+58 -47

ufos/src/server.rs ufos/src/server/mod.rs

reviewed

··· 1 1 + mod collections_query; 2 2 + mod cors; 3 3 + 1 4 use crate::index_html::INDEX_HTML; 2 5 use crate::storage::StoreReader; 3 6 use crate::store_types::{HourTruncatedCursor, WeekTruncatedCursor}; 4 7 use crate::{ConsumerInfo, Cursor, JustCount, Nsid, NsidCount, OrderCollectionsBy, UFOsRecord}; 5 8 use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; 6 9 use chrono::{DateTime, Utc}; 10 10 + use collections_query::MultiCollectionQuery; 11 11 + use cors::{OkCors, OkCorsResponse}; 7 12 use dropshot::endpoint; 8 13 use dropshot::ApiDescription; 9 14 use dropshot::Body; ··· 11 16 use dropshot::ConfigLogging; 12 17 use dropshot::ConfigLoggingLevel; 13 18 use dropshot::HttpError; 14 14 - use dropshot::HttpResponseHeaders; 15 15 - use dropshot::HttpResponseOk; 16 19 use dropshot::Query; 17 20 use dropshot::RequestContext; 18 21 use dropshot::ServerBuilder; 22 22 + 19 23 use http::{Response, StatusCode}; 20 24 use schemars::JsonSchema; 21 25 use serde::{Deserialize, Serialize}; ··· 74 78 }] 75 79 async fn get_openapi(ctx: RequestContext<Context>) -> OkCorsResponse<serde_json::Value> { 76 80 let spec = (*ctx.context().spec).clone(); 77 77 - ok_cors(spec) 81 81 + OkCors(spec).into() 78 82 } 79 83 80 84 #[derive(Debug, Serialize, JsonSchema)] ··· 83 87 storage: serde_json::Value, 84 88 consumer: ConsumerInfo, 85 89 } 86 86 - /// Get meta information about UFOs itself 90 90 + /// UFOs meta-info 87 91 #[endpoint { 88 92 method = GET, 89 93 path = "/meta" ··· 103 107 .await 104 108 .map_err(failed_to_get("consumer info"))?; 105 109 106 106 - ok_cors(MetaInfo { 110 110 + OkCors(MetaInfo { 107 111 storage_name: storage.name(), 108 112 storage: storage_info, 109 113 consumer, 110 114 }) 115 115 + .into() 111 116 } 112 117 113 118 // TODO: replace with normal (🙃) multi-qs value somehow ··· 145 150 } 146 151 } 147 152 } 148 148 - /// Get recent records by collection 153 153 + /// Record samples 154 154 + /// 155 155 + /// Get most recent records seen in the firehose, by collection NSID 149 156 /// 150 157 /// Multiple collections are supported. They will be delivered in one big array with no 151 158 /// specified order. ··· 190 197 .map(|r| r.into()) 191 198 .collect(); 192 199 193 193 - ok_cors(records) 200 200 + OkCors(records).into() 194 201 } 195 202 196 203 #[derive(Debug, Deserialize, JsonSchema)] 197 197 - struct TotalSeenCollectionsQuery { 198 198 - collection: String, // JsonSchema not implemented for Nsid :( 204 204 + struct CollectionsStatsQuery { 205 205 + /// Limit stats to those seen after this UTC datetime 206 206 + /// 207 207 + /// default: 1 week ago 208 208 + since: Option<DateTime<Utc>>, 209 209 + /// Limit stats to those seen before this UTC datetime 210 210 + /// 211 211 + /// default: now 212 212 + until: Option<DateTime<Utc>>, 199 213 } 200 200 - #[derive(Debug, Serialize, JsonSchema)] 201 201 - struct TotalCounts { 202 202 - total_creates: u64, 203 203 - dids_estimate: u64, 204 204 - } 205 205 - /// Get total records seen by collection 214 214 + /// Collection stats 215 215 + /// 216 216 + /// Get record statistics for collections during a specific time period. 217 217 + /// 218 218 + /// Note: the statistics are "rolled up" into hourly buckets in the background, 219 219 + /// so the data here can be as stale as that background task is behind. See the 220 220 + /// meta info endpoint to find out how up-to-date the rollup currently is. (In 221 221 + /// general it sholud be pretty close to live) 206 222 #[endpoint { 207 223 method = GET, 208 208 - path = "/records/total-seen" 224 224 + path = "/collections/stats" 209 225 }] 210 210 - async fn get_records_total_seen( 226 226 + async fn get_collection_stats( 211 227 ctx: RequestContext<Context>, 212 212 - collection_query: Query<TotalSeenCollectionsQuery>, 213 213 - ) -> OkCorsResponse<HashMap<String, TotalCounts>> { 228 228 + collections_query: MultiCollectionQuery, 229 229 + query: Query<CollectionsStatsQuery>, 230 230 + ) -> OkCorsResponse<HashMap<String, JustCount>> { 214 231 let Context { storage, .. } = ctx.context(); 232 232 + let q = query.into_inner(); 233 233 + let collections: HashSet<Nsid> = collections_query.try_into()?; 215 234 216 216 - let query = collection_query.into_inner(); 217 217 - let collections = to_multiple_nsids(&query.collection) 218 218 - .map_err(|reason| HttpError::for_bad_request(None, reason))?; 235 235 + let since = q.since.map(dt_to_cursor).transpose()?.unwrap_or_else(|| { 236 236 + let week_ago_secs = 7 * 86_400; 237 237 + let week_ago = SystemTime::now() - Duration::from_secs(week_ago_secs); 238 238 + Cursor::at(week_ago).into() 239 239 + }); 240 240 + 241 241 + let until = q.until.map(dt_to_cursor).transpose()?; 219 242 220 243 let mut seen_by_collection = HashMap::with_capacity(collections.len()); 221 244 222 245 for collection in &collections { 223 223 - let (total_creates, dids_estimate) = storage 224 224 - .get_counts_by_collection(collection) 246 246 + let counts = storage 247 247 + .get_collection_counts(collection, since, until) 225 248 .await 226 249 .map_err(|e| HttpError::for_internal_error(format!("boooo: {e:?}")))?; 227 250 228 228 - seen_by_collection.insert( 229 229 - collection.to_string(), 230 230 - TotalCounts { 231 231 - total_creates, 232 232 - dids_estimate, 233 233 - }, 234 234 - ); 251 251 + seen_by_collection.insert(collection.to_string(), counts); 235 252 } 236 253 237 237 - ok_cors(seen_by_collection) 254 254 + OkCors(seen_by_collection).into() 238 255 } 239 256 240 257 #[derive(Debug, Serialize, JsonSchema)] ··· 283 300 order: Option<CollectionsQueryOrder>, 284 301 } 285 302 286 286 - /// Get collection with statistics 303 303 + /// List collections 304 304 + /// 305 305 + /// With statistics. 287 306 /// 288 307 /// ## To fetch a full list: 289 308 /// ··· 353 372 354 373 let next_cursor = next_cursor.map(|c| URL_SAFE_NO_PAD.encode(c)); 355 374 356 356 - ok_cors(CollectionsResponse { 375 375 + OkCors(CollectionsResponse { 357 376 collections, 358 377 cursor: next_cursor, 359 378 }) 379 379 + .into() 360 380 } 361 381 362 382 #[derive(Debug, Deserialize, JsonSchema)] ··· 384 404 range: Vec<DateTime<Utc>>, 385 405 series: HashMap<String, Vec<JustCount>>, 386 406 } 387 387 - /// Get timeseries data 407 407 + /// Collection timeseries stats 388 408 #[endpoint { 389 409 method = GET, 390 410 path = "/timeseries" ··· 407 427 let step = if let Some(secs) = q.step { 408 428 if secs < 3600 { 409 429 let msg = format!("step is too small: {}", secs); 410 410 - return Err(HttpError::for_bad_request(None, msg)); 430 430 + Err(HttpError::for_bad_request(None, msg))?; 411 431 } 412 432 (secs / 3600) * 3600 // trucate to hour 413 433 } else { ··· 433 453 .map(|(k, v)| (k.to_string(), v.iter().map(Into::into).collect())) 434 454 .collect(); 435 455 436 436 - ok_cors(CollectionTimeseriesResponse { range, series }) 456 456 + OkCors(CollectionTimeseriesResponse { range, series }).into() 437 457 } 438 458 439 459 pub async fn serve(storage: impl StoreReader + 'static) -> Result<(), String> { ··· 449 469 api.register(get_openapi).unwrap(); 450 470 api.register(get_meta_info).unwrap(); 451 471 api.register(get_records_by_collections).unwrap(); 452 452 - api.register(get_records_total_seen).unwrap(); 472 472 + api.register(get_collection_stats).unwrap(); 453 473 api.register(get_collections).unwrap(); 454 474 api.register(get_timeseries).unwrap(); 455 475 ··· 482 502 .map_err(|error| format!("failed to start server: {}", error))? 483 503 .await 484 504 } 485 485 - 486 486 - /// awkward helpers 487 487 - type OkCorsResponse<T> = Result<HttpResponseHeaders<HttpResponseOk<T>>, HttpError>; 488 488 - fn ok_cors<T: Send + Sync + Serialize + JsonSchema>(t: T) -> OkCorsResponse<T> { 489 489 - let mut res = HttpResponseHeaders::new_unnamed(HttpResponseOk(t)); 490 490 - res.headers_mut() 491 491 - .insert("access-control-allow-origin", "*".parse().unwrap()); 492 492 - Ok(res) 493 493 - }

+72

ufos/src/server/collections_query.rs

reviewed

··· 1 1 + use crate::Nsid; 2 2 + use async_trait::async_trait; 3 3 + use dropshot::{ 4 4 + ApiEndpointBodyContentType, ExtractorMetadata, HttpError, Query, RequestContext, ServerContext, 5 5 + SharedExtractor, 6 6 + }; 7 7 + use schemars::JsonSchema; 8 8 + use serde::Deserialize; 9 9 + use std::collections::HashSet; 10 10 + 11 11 + /// The real type that gets deserialized 12 12 + #[derive(Debug, Deserialize, JsonSchema)] 13 13 + pub struct MultiCollectionQuery { 14 14 + pub collection: Vec<String>, 15 15 + } 16 16 + 17 17 + /// The fake corresponding type for docs that dropshot won't freak out about a 18 18 + /// vec for 19 19 + #[derive(Deserialize, JsonSchema)] 20 20 + #[allow(dead_code)] 21 21 + struct MultiCollectionQueryForDocs { 22 22 + /// One or more collection [NSID](https://atproto.com/specs/nsid)s 23 23 + /// 24 24 + /// Pass this parameter multiple times to specify multiple collections, like 25 25 + /// `collection=app.bsky.feed.like&collection=app.bsky.feed.post` 26 26 + collection: String, 27 27 + } 28 28 + 29 29 + impl TryFrom<MultiCollectionQuery> for HashSet<Nsid> { 30 30 + type Error = HttpError; 31 31 + fn try_from(mcq: MultiCollectionQuery) -> Result<Self, Self::Error> { 32 32 + let mut out = HashSet::with_capacity(mcq.collection.len()); 33 33 + for c in mcq.collection { 34 34 + let nsid = Nsid::new(c).map_err(|e| { 35 35 + HttpError::for_bad_request( 36 36 + None, 37 37 + format!("failed to convert collection to an NSID: {e:?}"), 38 38 + ) 39 39 + })?; 40 40 + out.insert(nsid); 41 41 + } 42 42 + Ok(out) 43 43 + } 44 44 + } 45 45 + 46 46 + // The `SharedExtractor` implementation for Query<QueryType> describes how to 47 47 + // construct an instance of `Query<QueryType>` from an HTTP request: namely, by 48 48 + // parsing the query string to an instance of `QueryType`. 49 49 + #[async_trait] 50 50 + impl SharedExtractor for MultiCollectionQuery { 51 51 + async fn from_request<Context: ServerContext>( 52 52 + ctx: &RequestContext<Context>, 53 53 + ) -> Result<MultiCollectionQuery, HttpError> { 54 54 + let raw_query = ctx.request.uri().query().unwrap_or(""); 55 55 + let q = serde_qs::from_str(raw_query).map_err(|e| { 56 56 + HttpError::for_bad_request(None, format!("unable to parse query string: {}", e)) 57 57 + })?; 58 58 + Ok(q) 59 59 + } 60 60 + 61 61 + fn metadata(body_content_type: ApiEndpointBodyContentType) -> ExtractorMetadata { 62 62 + // HACK: query type switcheroo: passing MultiCollectionQuery to 63 63 + // `metadata` would "helpfully" panic because dropshot believes we can 64 64 + // only have scalar types in a query. 65 65 + // 66 66 + // so instead we have a fake second type whose only job is to look the 67 67 + // same as MultiCollectionQuery exept that it has `String` instead of 68 68 + // `Vec<String>`, which dropshot will accept, and generate ~close-enough 69 69 + // docs for. 70 70 + <Query<MultiCollectionQueryForDocs> as SharedExtractor>::metadata(body_content_type) 71 71 + } 72 72 + }

+23

ufos/src/server/cors.rs

reviewed

··· 1 1 + use dropshot::{HttpError, HttpResponseHeaders, HttpResponseOk}; 2 2 + use schemars::JsonSchema; 3 3 + use serde::Serialize; 4 4 + 5 5 + pub type OkCorsResponse<T> = Result<HttpResponseHeaders<HttpResponseOk<T>>, HttpError>; 6 6 + 7 7 + /// Helper for constructing Ok responses: return OkCors(T).into() 8 8 + /// (not happy with this yet) 9 9 + pub struct OkCors<T: Serialize + JsonSchema + Send + Sync>(pub T); 10 10 + 11 11 + impl<T> From<OkCors<T>> for OkCorsResponse<T> 12 12 + where 13 13 + T: Serialize + JsonSchema + Send + Sync, 14 14 + { 15 15 + fn from(ok: OkCors<T>) -> OkCorsResponse<T> { 16 16 + let mut res = HttpResponseHeaders::new_unnamed(HttpResponseOk(ok.0)); 17 17 + res.headers_mut() 18 18 + .insert("access-control-allow-origin", "*".parse().unwrap()); 19 19 + Ok(res) 20 20 + } 21 21 + } 22 22 + 23 23 + // TODO: cors for HttpError

+8 -3

ufos/src/storage.rs

reviewed

··· 1 1 use crate::store_types::{CountsValue, HourTruncatedCursor, SketchSecretPrefix}; 2 2 use crate::{ 3 3 - error::StorageError, ConsumerInfo, Cursor, EventBatch, NsidCount, OrderCollectionsBy, 4 4 - UFOsRecord, 3 3 + error::StorageError, ConsumerInfo, Cursor, EventBatch, JustCount, NsidCount, 4 4 + OrderCollectionsBy, UFOsRecord, 5 5 }; 6 6 use async_trait::async_trait; 7 7 use jetstream::exports::{Did, Nsid}; ··· 92 92 step: u64, 93 93 ) -> StorageResult<(Vec<HourTruncatedCursor>, HashMap<Nsid, Vec<CountsValue>>)>; 94 94 95 95 - async fn get_counts_by_collection(&self, collection: &Nsid) -> StorageResult<(u64, u64)>; 95 95 + async fn get_collection_counts( 96 96 + &self, 97 97 + collection: &Nsid, 98 98 + since: HourTruncatedCursor, 99 99 + until: Option<HourTruncatedCursor>, 100 100 + ) -> StorageResult<JustCount>; 96 101 97 102 async fn get_records_by_collections( 98 103 &self,

+129 -67

ufos/src/storage_fjall.rs

reviewed

··· 13 13 WEEK_IN_MICROS, 14 14 }; 15 15 use crate::{ 16 16 - nice_duration, CommitAction, ConsumerInfo, Did, EventBatch, Nsid, NsidCount, 16 16 + nice_duration, CommitAction, ConsumerInfo, Did, EventBatch, JustCount, Nsid, NsidCount, 17 17 OrderCollectionsBy, UFOsRecord, 18 18 }; 19 19 use async_trait::async_trait; ··· 715 715 Ok((output_hours, output_series)) 716 716 } 717 717 718 718 - fn get_counts_by_collection(&self, collection: &Nsid) -> StorageResult<(u64, u64)> { 719 719 - // 0. grab a snapshot in case rollups happen while we're working 720 720 - let instant = self.keyspace.instant(); 721 721 - let global = self.global.snapshot_at(instant); 722 722 - let rollups = self.rollups.snapshot_at(instant); 718 718 + fn get_collection_counts( 719 719 + &self, 720 720 + collection: &Nsid, 721 721 + since: HourTruncatedCursor, 722 722 + until: Option<HourTruncatedCursor>, 723 723 + ) -> StorageResult<JustCount> { 724 724 + // grab snapshots in case rollups happen while we're working 725 725 + let rollups = self.rollups.snapshot(); 723 726 724 724 - // 1. all-time counts 725 725 - let all_time_key = AllTimeRollupKey::new(collection).to_db_bytes()?; 726 726 - let mut total_counts = rollups 727 727 - .get(&all_time_key)? 728 728 - .as_deref() 729 729 - .map(db_complete::<CountsValue>) 730 730 - .transpose()? 731 731 - .unwrap_or_default(); 727 727 + let until = until.unwrap_or_else(|| Cursor::at(SystemTime::now()).into()); 728 728 + let buckets = CursorBucket::buckets_spanning(since, until); 729 729 + let mut total_counts = CountsValue::default(); 732 730 733 733 - // 2. live counts that haven't been rolled into all-time yet. 734 734 - let rollup_cursor = 735 735 - get_snapshot_static_neu::<NewRollupCursorKey, NewRollupCursorValue>(&global)?.ok_or( 736 736 - StorageError::BadStateError("Could not find current rollup cursor".to_string()), 737 737 - )?; 738 738 - 739 739 - let full_range = LiveCountsKey::range_from_cursor(rollup_cursor)?; 740 740 - for kv in rollups.range(full_range) { 741 741 - let (key_bytes, val_bytes) = kv?; 742 742 - let key = db_complete::<LiveCountsKey>(&key_bytes)?; 743 743 - if key.collection() == collection { 744 744 - let counts = db_complete::<CountsValue>(&val_bytes)?; 745 745 - total_counts.merge(&counts); 746 746 - } 731 731 + for bucket in buckets { 732 732 + let key = match bucket { 733 733 + CursorBucket::Hour(t) => HourlyRollupKey::new(t, collection).to_db_bytes()?, 734 734 + CursorBucket::Week(t) => WeeklyRollupKey::new(t, collection).to_db_bytes()?, 735 735 + CursorBucket::AllTime => unreachable!(), // TODO: fall back on this if the time span spans the whole dataset? 736 736 + }; 737 737 + let count = rollups 738 738 + .get(&key)? 739 739 + .as_deref() 740 740 + .map(db_complete::<CountsValue>) 741 741 + .transpose()? 742 742 + .unwrap_or_default(); 743 743 + total_counts.merge(&count); 747 744 } 748 748 - Ok(( 749 749 - total_counts.counts().creates, 750 750 - total_counts.dids().estimate() as u64, 751 751 - )) 745 745 + 746 746 + Ok((&total_counts).into()) 752 747 } 753 748 754 749 fn get_records_by_collections( ··· 840 835 }) 841 836 .await? 842 837 } 843 843 - async fn get_counts_by_collection(&self, collection: &Nsid) -> StorageResult<(u64, u64)> { 838 838 + async fn get_collection_counts( 839 839 + &self, 840 840 + collection: &Nsid, 841 841 + since: HourTruncatedCursor, 842 842 + until: Option<HourTruncatedCursor>, 843 843 + ) -> StorageResult<JustCount> { 844 844 let s = self.clone(); 845 845 let collection = collection.clone(); 846 846 - tokio::task::spawn_blocking(move || FjallReader::get_counts_by_collection(&s, &collection)) 847 847 - .await? 846 846 + tokio::task::spawn_blocking(move || { 847 847 + FjallReader::get_collection_counts(&s, &collection, since, until) 848 848 + }) 849 849 + .await? 848 850 } 849 851 async fn get_records_by_collections( 850 852 &self, ··· 1485 1487 } 1486 1488 1487 1489 const TEST_BATCH_LIMIT: usize = 16; 1490 1490 + fn beginning() -> HourTruncatedCursor { 1491 1491 + Cursor::from_start().into() 1492 1492 + } 1488 1493 1489 1494 #[derive(Debug, Default)] 1490 1495 struct TestBatch { ··· 1622 1627 fn test_hello() -> anyhow::Result<()> { 1623 1628 let (read, mut write) = fjall_db(); 1624 1629 write.insert_batch::<TEST_BATCH_LIMIT>(EventBatch::default())?; 1625 1625 - let (records, dids) = 1626 1626 - read.get_counts_by_collection(&Nsid::new("a.b.c".to_string()).unwrap())?; 1627 1627 - assert_eq!(records, 0); 1628 1628 - assert_eq!(dids, 0); 1630 1630 + let JustCount { 1631 1631 + creates, 1632 1632 + dids_estimate, 1633 1633 + .. 1634 1634 + } = read.get_collection_counts( 1635 1635 + &Nsid::new("a.b.c".to_string()).unwrap(), 1636 1636 + beginning(), 1637 1637 + None, 1638 1638 + )?; 1639 1639 + assert_eq!(creates, 0); 1640 1640 + assert_eq!(dids_estimate, 0); 1629 1641 Ok(()) 1630 1642 } 1631 1643 ··· 1644 1656 100, 1645 1657 ); 1646 1658 write.insert_batch(batch.batch)?; 1659 1659 + write.step_rollup()?; 1647 1660 1648 1648 - let (records, dids) = read.get_counts_by_collection(&collection)?; 1649 1649 - assert_eq!(records, 1); 1650 1650 - assert_eq!(dids, 1); 1651 1651 - let (records, dids) = 1652 1652 - read.get_counts_by_collection(&Nsid::new("d.e.f".to_string()).unwrap())?; 1653 1653 - assert_eq!(records, 0); 1654 1654 - assert_eq!(dids, 0); 1661 1661 + let JustCount { 1662 1662 + creates, 1663 1663 + dids_estimate, 1664 1664 + .. 1665 1665 + } = read.get_collection_counts(&collection, beginning(), None)?; 1666 1666 + assert_eq!(creates, 1); 1667 1667 + assert_eq!(dids_estimate, 1); 1668 1668 + let JustCount { 1669 1669 + creates, 1670 1670 + dids_estimate, 1671 1671 + .. 1672 1672 + } = read.get_collection_counts( 1673 1673 + &Nsid::new("d.e.f".to_string()).unwrap(), 1674 1674 + beginning(), 1675 1675 + None, 1676 1676 + )?; 1677 1677 + assert_eq!(creates, 0); 1678 1678 + assert_eq!(dids_estimate, 0); 1655 1679 1656 1680 let records = read.get_records_by_collections([collection].into(), 2, false)?; 1657 1681 assert_eq!(records.len(), 1); ··· 1815 1839 101, 1816 1840 ); 1817 1841 write.insert_batch(batch.batch)?; 1842 1842 + write.step_rollup()?; 1818 1843 1819 1819 - let (records, dids) = read.get_counts_by_collection(&collection)?; 1820 1820 - assert_eq!(records, 1); 1821 1821 - assert_eq!(dids, 1); 1844 1844 + let JustCount { 1845 1845 + creates, 1846 1846 + dids_estimate, 1847 1847 + .. 1848 1848 + } = read.get_collection_counts(&collection, beginning(), None)?; 1849 1849 + assert_eq!(creates, 1); 1850 1850 + assert_eq!(dids_estimate, 1); 1822 1851 1823 1852 let records = read.get_records_by_collections([collection].into(), 2, false)?; 1824 1853 assert_eq!(records.len(), 1); ··· 1853 1882 101, 1854 1883 ); 1855 1884 write.insert_batch(batch.batch)?; 1885 1885 + write.step_rollup()?; 1856 1886 1857 1857 - let (creates, dids) = read.get_counts_by_collection(&collection)?; 1887 1887 + let JustCount { 1888 1888 + creates, 1889 1889 + dids_estimate, 1890 1890 + .. 1891 1891 + } = read.get_collection_counts(&collection, beginning(), None)?; 1858 1892 assert_eq!(creates, 1); 1859 1859 - assert_eq!(dids, 1); 1893 1893 + assert_eq!(dids_estimate, 1); 1860 1894 1861 1895 let records = read.get_records_by_collections([collection].into(), 2, false)?; 1862 1896 assert_eq!(records.len(), 0); ··· 2170 2204 write.insert_batch(batch.batch)?; 2171 2205 2172 2206 // before any rollup 2173 2173 - let (records, dids) = 2174 2174 - read.get_counts_by_collection(&Nsid::new("a.a.a".to_string()).unwrap())?; 2175 2175 - assert_eq!(records, 3); 2176 2176 - assert_eq!(dids, 2); 2207 2207 + let JustCount { 2208 2208 + creates, 2209 2209 + dids_estimate, 2210 2210 + .. 2211 2211 + } = read.get_collection_counts( 2212 2212 + &Nsid::new("a.a.a".to_string()).unwrap(), 2213 2213 + beginning(), 2214 2214 + None, 2215 2215 + )?; 2216 2216 + assert_eq!(creates, 0); 2217 2217 + assert_eq!(dids_estimate, 0); 2177 2218 2178 2219 // first batch rolled up 2179 2220 let (n, _) = write.step_rollup()?; 2180 2221 assert_eq!(n, 1); 2181 2222 2182 2182 - let (records, dids) = 2183 2183 - read.get_counts_by_collection(&Nsid::new("a.a.a".to_string()).unwrap())?; 2184 2184 - assert_eq!(records, 3); 2185 2185 - assert_eq!(dids, 2); 2223 2223 + let JustCount { 2224 2224 + creates, 2225 2225 + dids_estimate, 2226 2226 + .. 2227 2227 + } = read.get_collection_counts( 2228 2228 + &Nsid::new("a.a.a".to_string()).unwrap(), 2229 2229 + beginning(), 2230 2230 + None, 2231 2231 + )?; 2232 2232 + assert_eq!(creates, 2); 2233 2233 + assert_eq!(dids_estimate, 2); 2186 2234 2187 2235 // delete account rolled up 2188 2236 let (n, _) = write.step_rollup()?; 2189 2237 assert_eq!(n, 1); 2190 2238 2191 2191 - let (records, dids) = 2192 2192 - read.get_counts_by_collection(&Nsid::new("a.a.a".to_string()).unwrap())?; 2193 2193 - assert_eq!(records, 3); 2194 2194 - assert_eq!(dids, 2); 2239 2239 + let JustCount { 2240 2240 + creates, 2241 2241 + dids_estimate, 2242 2242 + .. 2243 2243 + } = read.get_collection_counts( 2244 2244 + &Nsid::new("a.a.a".to_string()).unwrap(), 2245 2245 + beginning(), 2246 2246 + None, 2247 2247 + )?; 2248 2248 + assert_eq!(creates, 2); 2249 2249 + assert_eq!(dids_estimate, 2); 2195 2250 2196 2251 // second batch rolled up 2197 2252 let (n, _) = write.step_rollup()?; 2198 2253 assert_eq!(n, 1); 2199 2254 2200 2200 - let (records, dids) = 2201 2201 - read.get_counts_by_collection(&Nsid::new("a.a.a".to_string()).unwrap())?; 2202 2202 - assert_eq!(records, 3); 2203 2203 - assert_eq!(dids, 2); 2255 2255 + let JustCount { 2256 2256 + creates, 2257 2257 + dids_estimate, 2258 2258 + .. 2259 2259 + } = read.get_collection_counts( 2260 2260 + &Nsid::new("a.a.a".to_string()).unwrap(), 2261 2261 + beginning(), 2262 2262 + None, 2263 2263 + )?; 2264 2264 + assert_eq!(creates, 3); 2265 2265 + assert_eq!(dids_estimate, 2); 2204 2266 2205 2267 // no more rollups left 2206 2268 let (n, _) = write.step_rollup()?;

-1757

ufos/src/storage_mem.rs

reviewed

··· 1 1 - use std::ops::Bound; 2 2 - use std::sync::Arc; 3 3 - 4 4 - use crate::db_types::{db_complete, DbBytes, DbStaticStr, StaticStr}; 5 5 - use crate::error::StorageError; 6 6 - use crate::storage::{StorageResult, StorageWhatever, StoreBackground, StoreReader, StoreWriter}; 7 7 - use crate::store_types::{ 8 8 - AllTimeRollupKey, CommitCounts, CountsValue, DeleteAccountQueueKey, DeleteAccountQueueVal, 9 9 - HourTruncatedCursor, HourlyRollupKey, JetstreamCursorKey, JetstreamCursorValue, 10 10 - JetstreamEndpointKey, JetstreamEndpointValue, LiveCountsKey, NewRollupCursorKey, 11 11 - NewRollupCursorValue, NsidRecordFeedKey, NsidRecordFeedVal, RecordLocationKey, 12 12 - RecordLocationMeta, RecordLocationVal, RecordRawValue, SketchSecretPrefix, TakeoffKey, 13 13 - TakeoffValue, WeekTruncatedCursor, WeeklyRollupKey, WithCollection, 14 14 - }; 15 15 - use crate::{ 16 16 - CommitAction, ConsumerInfo, Did, EventBatch, Nsid, NsidCount, OrderCollectionsBy, UFOsRecord, 17 17 - }; 18 18 - use async_trait::async_trait; 19 19 - use jetstream::events::Cursor; 20 20 - use lsm_tree::range::prefix_to_range; 21 21 - use std::collections::{BTreeMap, HashMap, HashSet}; 22 22 - use std::path::Path; 23 23 - use std::sync::{Mutex, RwLock}; 24 24 - use std::time::SystemTime; 25 25 - 26 26 - const MAX_BATCHED_CLEANUP_SIZE: usize = 1024; // try to commit progress for longer feeds 27 27 - const MAX_BATCHED_ACCOUNT_DELETE_RECORDS: usize = 1024; 28 28 - const MAX_BATCHED_ROLLUP_COUNTS: usize = 256; 29 29 - 30 30 - /// 31 31 - /// new data format, roughly: 32 32 - /// 33 33 - /// Partion: 'global' 34 34 - /// 35 35 - /// - Global sequence counter (is the jetstream cursor -- monotonic with many gaps) 36 36 - /// - key: "js_cursor" (literal) 37 37 - /// - val: u64 38 38 - /// 39 39 - /// - Jetstream server endpoint (persisted because the cursor can't be used on another instance without data loss) 40 40 - /// - key: "js_endpoint" (literal) 41 41 - /// - val: string (URL of the instance) 42 42 - /// 43 43 - /// - Launch date 44 44 - /// - key: "takeoff" (literal) 45 45 - /// - val: u64 (micros timestamp, not from jetstream for now so not precise) 46 46 - /// 47 47 - /// - Rollup cursor (bg work: roll stats into hourlies, delete accounts, old record deletes) 48 48 - /// - key: "rollup_cursor" (literal) 49 49 - /// - val: u64 (tracks behind js_cursor) 50 50 - /// 51 51 - /// 52 52 - /// Partition: 'feed' 53 53 - /// 54 54 - /// - Per-collection list of record references ordered by jetstream cursor 55 55 - /// - key: nullstr || u64 (collection nsid null-terminated, jetstream cursor) 56 56 - /// - val: nullstr || nullstr || nullstr (did, rkey, rev. rev is mostly a sanity-check for now.) 57 57 - /// 58 58 - /// 59 59 - /// Partition: 'records' 60 60 - /// 61 61 - /// - Actual records by their atproto location 62 62 - /// - key: nullstr || nullstr || nullstr (did, collection, rkey) 63 63 - /// - val: u64 || bool || nullstr || rawval (js_cursor, is_update, rev, actual record) 64 64 - /// 65 65 - /// 66 66 - /// Partition: 'rollups' 67 67 - /// 68 68 - /// - Live (batched) records counts and dids estimate per collection 69 69 - /// - key: "live_counts" || u64 || nullstr (js_cursor, nsid) 70 70 - /// - val: u64 || HLL (count (not cursor), estimator) 71 71 - /// 72 72 - /// - Hourly total record counts and dids estimate per collection 73 73 - /// - key: "hourly_counts" || u64 || nullstr (hour, nsid) 74 74 - /// - val: u64 || HLL (count (not cursor), estimator) 75 75 - /// 76 76 - /// - Weekly total record counts and dids estimate per collection 77 77 - /// - key: "weekly_counts" || u64 || nullstr (hour, nsid) 78 78 - /// - val: u64 || HLL (count (not cursor), estimator) 79 79 - /// 80 80 - /// - All-time total record counts and dids estimate per collection 81 81 - /// - key: "ever_counts" || nullstr (nsid) 82 82 - /// - val: u64 || HLL (count (not cursor), estimator) 83 83 - /// 84 84 - /// - TODO: sorted indexes for all-times? 85 85 - /// 86 86 - /// 87 87 - /// Partition: 'queues' 88 88 - /// 89 89 - /// - Delete account queue 90 90 - /// - key: "delete_acount" || u64 (js_cursor) 91 91 - /// - val: nullstr (did) 92 92 - /// 93 93 - /// 94 94 - /// TODO: moderation actions 95 95 - /// TODO: account privacy preferences. Might wait for the protocol-level (PDS-level?) stuff to land. Will probably do lazy fetching + caching on read. 96 96 - #[derive(Debug)] 97 97 - pub struct MemStorage {} 98 98 - 99 99 - #[derive(Debug, Default)] 100 100 - pub struct MemConfig { 101 101 - /// drop the db when the storage is dropped 102 102 - /// 103 103 - /// this is only meant for tests 104 104 - #[cfg(test)] 105 105 - pub temp: bool, 106 106 - } 107 107 - 108 108 - //////////// 109 109 - //////////// 110 110 - //////////// 111 111 - //////////// 112 112 - //////////// 113 113 - //////////// 114 114 - 115 115 - struct BatchSentinel {} 116 116 - 117 117 - #[derive(Clone)] 118 118 - struct MemKeyspace { 119 119 - keyspace_guard: Arc<RwLock<BatchSentinel>>, 120 120 - } 121 121 - 122 122 - impl MemKeyspace { 123 123 - pub fn open() -> Self { 124 124 - Self { 125 125 - keyspace_guard: Arc::new(RwLock::new(BatchSentinel {})), 126 126 - } 127 127 - } 128 128 - pub fn open_partition(&self, _name: &str) -> StorageResult<MemPartion> { 129 129 - Ok(MemPartion { 130 130 - // name: name.to_string(), 131 131 - keyspace_guard: self.keyspace_guard.clone(), 132 132 - contents: Default::default(), 133 133 - }) 134 134 - } 135 135 - pub fn batch(&self) -> MemBatch { 136 136 - MemBatch { 137 137 - keyspace_guard: self.keyspace_guard.clone(), 138 138 - tasks: Vec::new(), 139 139 - } 140 140 - } 141 141 - pub fn instant(&self) -> u64 { 142 142 - 1 143 143 - } 144 144 - } 145 145 - 146 146 - enum BatchTask { 147 147 - Insert { 148 148 - p: MemPartion, 149 149 - key: Vec<u8>, 150 150 - val: Vec<u8>, 151 151 - }, 152 152 - Remove { 153 153 - p: MemPartion, 154 154 - key: Vec<u8>, 155 155 - }, 156 156 - } 157 157 - struct MemBatch { 158 158 - keyspace_guard: Arc<RwLock<BatchSentinel>>, 159 159 - tasks: Vec<BatchTask>, 160 160 - } 161 161 - impl MemBatch { 162 162 - pub fn insert(&mut self, p: &MemPartion, key: &[u8], val: &[u8]) { 163 163 - self.tasks.push(BatchTask::Insert { 164 164 - p: p.clone(), 165 165 - key: key.to_vec(), 166 166 - val: val.to_vec(), 167 167 - }); 168 168 - } 169 169 - pub fn remove(&mut self, p: &MemPartion, key: &[u8]) { 170 170 - self.tasks.push(BatchTask::Remove { 171 171 - p: p.clone(), 172 172 - key: key.to_vec(), 173 173 - }); 174 174 - } 175 175 - pub fn len(&self) -> usize { 176 176 - self.tasks.len() 177 177 - } 178 178 - pub fn commit(&mut self) -> StorageResult<()> { 179 179 - let _guard = self.keyspace_guard.write().unwrap(); 180 180 - for task in &mut self.tasks { 181 181 - match task { 182 182 - BatchTask::Insert { p, key, val } => p 183 183 - .contents 184 184 - .try_lock() 185 185 - .unwrap() 186 186 - .insert(key.to_vec(), val.to_vec()), 187 187 - BatchTask::Remove { p, key } => p.contents.try_lock().unwrap().remove(key), 188 188 - }; 189 189 - } 190 190 - Ok(()) 191 191 - } 192 192 - } 193 193 - 194 194 - #[derive(Clone)] 195 195 - struct MemPartion { 196 196 - // name: String, 197 197 - keyspace_guard: Arc<RwLock<BatchSentinel>>, 198 198 - contents: Arc<Mutex<BTreeMap<Vec<u8>, Vec<u8>>>>, 199 199 - } 200 200 - impl MemPartion { 201 201 - pub fn get(&self, key: &[u8]) -> StorageResult<Option<Vec<u8>>> { 202 202 - let _guard = self.keyspace_guard.read().unwrap(); 203 203 - Ok(self.contents.lock().unwrap().get(key).cloned()) 204 204 - } 205 205 - pub fn prefix(&self, pre: &[u8]) -> Vec<StorageResult<(Vec<u8>, Vec<u8>)>> { 206 206 - // let prefix_bytes = prefix.to_db_bytes()?; 207 207 - let (_, Bound::Excluded(range_end)) = prefix_to_range(pre) else { 208 208 - panic!("bad range thing"); 209 209 - }; 210 210 - 211 211 - return self.range(pre.to_vec()..range_end.to_vec()); 212 212 - } 213 213 - pub fn range(&self, r: std::ops::Range<Vec<u8>>) -> Vec<StorageResult<(Vec<u8>, Vec<u8>)>> { 214 214 - let _guard = self.keyspace_guard.read().unwrap(); 215 215 - self.contents 216 216 - .lock() 217 217 - .unwrap() 218 218 - .range(r) 219 219 - .map(|(k, v)| Ok((k.clone(), v.clone()))) 220 220 - .collect() 221 221 - } 222 222 - pub fn insert(&self, key: &[u8], val: &[u8]) -> StorageResult<()> { 223 223 - let _guard = self.keyspace_guard.read().unwrap(); 224 224 - self.contents 225 225 - .lock() 226 226 - .unwrap() 227 227 - .insert(key.to_vec(), val.to_vec()); 228 228 - Ok(()) 229 229 - } 230 230 - // pub fn remove(&self, key: &[u8]) -> StorageResult<()> { 231 231 - // let _guard = self.keyspace_guard.read().unwrap(); 232 232 - // self.contents 233 233 - // .lock() 234 234 - // .unwrap() 235 235 - // .remove(key); 236 236 - // Ok(()) 237 237 - // } 238 238 - pub fn snapshot_at(&self, _instant: u64) -> Self { 239 239 - self.clone() 240 240 - } 241 241 - pub fn snapshot(&self) -> Self { 242 242 - self.clone() 243 243 - } 244 244 - } 245 245 - 246 246 - //////////// 247 247 - //////////// 248 248 - //////////// 249 249 - //////////// 250 250 - //////////// 251 251 - //////////// 252 252 - 253 253 - impl StorageWhatever<MemReader, MemWriter, MemBackground, MemConfig> for MemStorage { 254 254 - fn init( 255 255 - _path: impl AsRef<Path>, 256 256 - endpoint: String, 257 257 - force_endpoint: bool, 258 258 - _config: MemConfig, 259 259 - ) -> StorageResult<(MemReader, MemWriter, Option<Cursor>, SketchSecretPrefix)> { 260 260 - let keyspace = MemKeyspace::open(); 261 261 - 262 262 - let global = keyspace.open_partition("global")?; 263 263 - let feeds = keyspace.open_partition("feeds")?; 264 264 - let records = keyspace.open_partition("records")?; 265 265 - let rollups = keyspace.open_partition("rollups")?; 266 266 - let queues = keyspace.open_partition("queues")?; 267 267 - 268 268 - let js_cursor = get_static_neu::<JetstreamCursorKey, JetstreamCursorValue>(&global)?; 269 269 - 270 270 - if js_cursor.is_some() { 271 271 - let stored_endpoint = 272 272 - get_static_neu::<JetstreamEndpointKey, JetstreamEndpointValue>(&global)?; 273 273 - 274 274 - let JetstreamEndpointValue(stored) = stored_endpoint.ok_or(StorageError::InitError( 275 275 - "found cursor but missing js_endpoint, refusing to start.".to_string(), 276 276 - ))?; 277 277 - 278 278 - if stored != endpoint { 279 279 - if force_endpoint { 280 280 - log::warn!("forcing a jetstream switch from {stored:?} to {endpoint:?}"); 281 281 - insert_static_neu::<JetstreamEndpointKey>( 282 282 - &global, 283 283 - JetstreamEndpointValue(endpoint.to_string()), 284 284 - )?; 285 285 - } else { 286 286 - return Err(StorageError::InitError(format!( 287 287 - "stored js_endpoint {stored:?} differs from provided {endpoint:?}, refusing to start."))); 288 288 - } 289 289 - } 290 290 - } else { 291 291 - insert_static_neu::<JetstreamEndpointKey>( 292 292 - &global, 293 293 - JetstreamEndpointValue(endpoint.to_string()), 294 294 - )?; 295 295 - insert_static_neu::<TakeoffKey>(&global, Cursor::at(SystemTime::now()))?; 296 296 - insert_static_neu::<NewRollupCursorKey>(&global, Cursor::from_start())?; 297 297 - } 298 298 - 299 299 - let reader = MemReader { 300 300 - keyspace: keyspace.clone(), 301 301 - global: global.clone(), 302 302 - feeds: feeds.clone(), 303 303 - records: records.clone(), 304 304 - rollups: rollups.clone(), 305 305 - }; 306 306 - let writer = MemWriter { 307 307 - keyspace, 308 308 - global, 309 309 - feeds, 310 310 - records, 311 311 - rollups, 312 312 - queues, 313 313 - }; 314 314 - let secret_prefix = [0u8; 16]; // in-mem store is always deterministic: no secret 315 315 - Ok((reader, writer, js_cursor, secret_prefix)) 316 316 - } 317 317 - } 318 318 - 319 319 - type MemRKV = StorageResult<(Vec<u8>, Vec<u8>)>; 320 320 - 321 321 - #[derive(Clone)] 322 322 - pub struct MemReader { 323 323 - keyspace: MemKeyspace, 324 324 - global: MemPartion, 325 325 - feeds: MemPartion, 326 326 - records: MemPartion, 327 327 - rollups: MemPartion, 328 328 - } 329 329 - 330 330 - /// An iterator that knows how to skip over deleted/invalidated records 331 331 - struct RecordIterator { 332 332 - db_iter: Box<dyn Iterator<Item = MemRKV>>, 333 333 - records: MemPartion, 334 334 - limit: usize, 335 335 - fetched: usize, 336 336 - } 337 337 - impl RecordIterator { 338 338 - pub fn new( 339 339 - feeds: &MemPartion, 340 340 - records: MemPartion, 341 341 - collection: &Nsid, 342 342 - limit: usize, 343 343 - ) -> StorageResult<Self> { 344 344 - let prefix = NsidRecordFeedKey::from_prefix_to_db_bytes(collection)?; 345 345 - let db_iter = feeds.prefix(&prefix).into_iter().rev(); 346 346 - Ok(Self { 347 347 - db_iter: Box::new(db_iter), 348 348 - records, 349 349 - limit, 350 350 - fetched: 0, 351 351 - }) 352 352 - } 353 353 - fn get_record(&self, db_next: MemRKV) -> StorageResult<Option<UFOsRecord>> { 354 354 - let (key_bytes, val_bytes) = db_next?; 355 355 - let feed_key = db_complete::<NsidRecordFeedKey>(&key_bytes)?; 356 356 - let feed_val = db_complete::<NsidRecordFeedVal>(&val_bytes)?; 357 357 - let location_key: RecordLocationKey = (&feed_key, &feed_val).into(); 358 358 - 359 359 - let Some(location_val_bytes) = self.records.get(&location_key.to_db_bytes()?)? else { 360 360 - // record was deleted (hopefully) 361 361 - return Ok(None); 362 362 - }; 363 363 - 364 364 - let (meta, n) = RecordLocationMeta::from_db_bytes(&location_val_bytes)?; 365 365 - 366 366 - if meta.cursor() != feed_key.cursor() { 367 367 - // older/different version 368 368 - return Ok(None); 369 369 - } 370 370 - if meta.rev != feed_val.rev() { 371 371 - // weird... 372 372 - log::warn!("record lookup: cursor match but rev did not...? excluding."); 373 373 - return Ok(None); 374 374 - } 375 375 - let Some(raw_value_bytes) = location_val_bytes.get(n..) else { 376 376 - log::warn!( 377 377 - "record lookup: found record but could not get bytes to decode the record??" 378 378 - ); 379 379 - return Ok(None); 380 380 - }; 381 381 - let rawval = db_complete::<RecordRawValue>(raw_value_bytes)?; 382 382 - Ok(Some(UFOsRecord { 383 383 - collection: feed_key.collection().clone(), 384 384 - cursor: feed_key.cursor(), 385 385 - did: feed_val.did().clone(), 386 386 - rkey: feed_val.rkey().clone(), 387 387 - rev: meta.rev.to_string(), 388 388 - record: rawval.try_into()?, 389 389 - is_update: meta.is_update, 390 390 - })) 391 391 - } 392 392 - } 393 393 - impl Iterator for RecordIterator { 394 394 - type Item = StorageResult<Option<UFOsRecord>>; 395 395 - fn next(&mut self) -> Option<Self::Item> { 396 396 - if self.fetched == self.limit { 397 397 - return Some(Ok(None)); 398 398 - } 399 399 - let record = loop { 400 400 - let db_next = self.db_iter.next()?; // None short-circuits here 401 401 - match self.get_record(db_next) { 402 402 - Err(e) => return Some(Err(e)), 403 403 - Ok(Some(record)) => break record, 404 404 - Ok(None) => continue, 405 405 - } 406 406 - }; 407 407 - self.fetched += 1; 408 408 - Some(Ok(Some(record))) 409 409 - } 410 410 - } 411 411 - 412 412 - impl MemReader { 413 413 - fn get_storage_stats(&self) -> StorageResult<serde_json::Value> { 414 414 - let rollup_cursor = 415 415 - get_static_neu::<NewRollupCursorKey, NewRollupCursorValue>(&self.global)? 416 416 - .map(|c| c.to_raw_u64()); 417 417 - 418 418 - Ok(serde_json::json!({ 419 419 - "rollup_cursor": rollup_cursor, 420 420 - })) 421 421 - } 422 422 - 423 423 - fn get_consumer_info(&self) -> StorageResult<ConsumerInfo> { 424 424 - let global = self.global.snapshot(); 425 425 - 426 426 - let endpoint = 427 427 - get_snapshot_static_neu::<JetstreamEndpointKey, JetstreamEndpointValue>(&global)? 428 428 - .ok_or(StorageError::BadStateError( 429 429 - "Could not find jetstream endpoint".to_string(), 430 430 - ))? 431 431 - .0; 432 432 - 433 433 - let started_at = get_snapshot_static_neu::<TakeoffKey, TakeoffValue>(&global)? 434 434 - .ok_or(StorageError::BadStateError( 435 435 - "Could not find jetstream takeoff time".to_string(), 436 436 - ))? 437 437 - .to_raw_u64(); 438 438 - 439 439 - let latest_cursor = 440 440 - get_snapshot_static_neu::<JetstreamCursorKey, JetstreamCursorValue>(&global)? 441 441 - .map(|c| c.to_raw_u64()); 442 442 - 443 443 - let rollup_cursor = 444 444 - get_snapshot_static_neu::<NewRollupCursorKey, JetstreamCursorValue>(&global)? 445 445 - .map(|c| c.to_raw_u64()); 446 446 - 447 447 - Ok(ConsumerInfo::Jetstream { 448 448 - endpoint, 449 449 - started_at, 450 450 - latest_cursor, 451 451 - rollup_cursor, 452 452 - }) 453 453 - } 454 454 - 455 455 - fn get_counts_by_collection(&self, collection: &Nsid) -> StorageResult<(u64, u64)> { 456 456 - // 0. grab a snapshot in case rollups happen while we're working 457 457 - let instant = self.keyspace.instant(); 458 458 - let global = self.global.snapshot_at(instant); 459 459 - let rollups = self.rollups.snapshot_at(instant); 460 460 - 461 461 - // 1. all-time counts 462 462 - let all_time_key = AllTimeRollupKey::new(collection).to_db_bytes()?; 463 463 - let mut total_counts = rollups 464 464 - .get(&all_time_key)? 465 465 - .as_deref() 466 466 - .map(db_complete::<CountsValue>) 467 467 - .transpose()? 468 468 - .unwrap_or_default(); 469 469 - 470 470 - // 2. live counts that haven't been rolled into all-time yet. 471 471 - let rollup_cursor = 472 472 - get_snapshot_static_neu::<NewRollupCursorKey, NewRollupCursorValue>(&global)?.ok_or( 473 473 - StorageError::BadStateError("Could not find current rollup cursor".to_string()), 474 474 - )?; 475 475 - 476 476 - let full_range = LiveCountsKey::range_from_cursor(rollup_cursor)?; 477 477 - for kv in rollups.range(full_range) { 478 478 - let (key_bytes, val_bytes) = kv?; 479 479 - let key = db_complete::<LiveCountsKey>(&key_bytes)?; 480 480 - if key.collection() == collection { 481 481 - let counts = db_complete::<CountsValue>(&val_bytes)?; 482 482 - total_counts.merge(&counts); 483 483 - } 484 484 - } 485 485 - Ok(( 486 486 - total_counts.counts().creates, 487 487 - total_counts.dids().estimate() as u64, 488 488 - )) 489 489 - } 490 490 - 491 491 - fn get_records_by_collections( 492 492 - &self, 493 493 - collections: HashSet<Nsid>, 494 494 - limit: usize, 495 495 - _expand_each_collection: bool, 496 496 - ) -> StorageResult<Vec<UFOsRecord>> { 497 497 - if collections.is_empty() { 498 498 - return Ok(vec![]); 499 499 - } 500 500 - let mut record_iterators = Vec::new(); 501 501 - for collection in collections { 502 502 - let iter = RecordIterator::new(&self.feeds, self.records.clone(), &collection, limit)?; 503 503 - record_iterators.push(iter.peekable()); 504 504 - } 505 505 - let mut merged = Vec::new(); 506 506 - loop { 507 507 - let mut latest: Option<(Cursor, usize)> = None; // ugh 508 508 - for (i, iter) in record_iterators.iter_mut().enumerate() { 509 509 - let Some(it) = iter.peek_mut() else { 510 510 - continue; 511 511 - }; 512 512 - let it = match it { 513 513 - Ok(v) => v, 514 514 - Err(e) => Err(std::mem::replace(e, StorageError::Stolen))?, 515 515 - }; 516 516 - let Some(rec) = it else { 517 517 - break; 518 518 - }; 519 519 - if let Some((cursor, _)) = latest { 520 520 - if rec.cursor > cursor { 521 521 - latest = Some((rec.cursor, i)) 522 522 - } 523 523 - } else { 524 524 - latest = Some((rec.cursor, i)); 525 525 - } 526 526 - } 527 527 - let Some((_, idx)) = latest else { 528 528 - break; 529 529 - }; 530 530 - // yeah yeah whateverrrrrrrrrrrrrrrr 531 531 - merged.push(record_iterators[idx].next().unwrap().unwrap().unwrap()); 532 532 - } 533 533 - Ok(merged) 534 534 - } 535 535 - } 536 536 - 537 537 - #[async_trait] 538 538 - impl StoreReader for MemReader { 539 539 - fn name(&self) -> String { 540 540 - "in-memory store".into() 541 541 - } 542 542 - async fn get_storage_stats(&self) -> StorageResult<serde_json::Value> { 543 543 - let s = self.clone(); 544 544 - tokio::task::spawn_blocking(move || MemReader::get_storage_stats(&s)).await? 545 545 - } 546 546 - async fn get_consumer_info(&self) -> StorageResult<ConsumerInfo> { 547 547 - let s = self.clone(); 548 548 - tokio::task::spawn_blocking(move || MemReader::get_consumer_info(&s)).await? 549 549 - } 550 550 - async fn get_collections( 551 551 - &self, 552 552 - _: usize, 553 553 - _: OrderCollectionsBy, 554 554 - _: Option<HourTruncatedCursor>, 555 555 - _: Option<HourTruncatedCursor>, 556 556 - ) -> StorageResult<(Vec<NsidCount>, Option<Vec<u8>>)> { 557 557 - todo!() 558 558 - } 559 559 - async fn get_timeseries( 560 560 - &self, 561 561 - _: Vec<Nsid>, 562 562 - _: HourTruncatedCursor, 563 563 - _: Option<HourTruncatedCursor>, 564 564 - _: u64, 565 565 - ) -> StorageResult<(Vec<HourTruncatedCursor>, HashMap<Nsid, Vec<CountsValue>>)> { 566 566 - todo!() 567 567 - } 568 568 - async fn get_counts_by_collection(&self, collection: &Nsid) -> StorageResult<(u64, u64)> { 569 569 - let s = self.clone(); 570 570 - let collection = collection.clone(); 571 571 - tokio::task::spawn_blocking(move || MemReader::get_counts_by_collection(&s, &collection)) 572 572 - .await? 573 573 - } 574 574 - async fn get_records_by_collections( 575 575 - &self, 576 576 - collections: HashSet<Nsid>, 577 577 - limit: usize, 578 578 - expand_each_collection: bool, 579 579 - ) -> StorageResult<Vec<UFOsRecord>> { 580 580 - let s = self.clone(); 581 581 - tokio::task::spawn_blocking(move || { 582 582 - MemReader::get_records_by_collections(&s, collections, limit, expand_each_collection) 583 583 - }) 584 584 - .await? 585 585 - } 586 586 - } 587 587 - 588 588 - pub struct MemWriter { 589 589 - keyspace: MemKeyspace, 590 590 - global: MemPartion, 591 591 - feeds: MemPartion, 592 592 - records: MemPartion, 593 593 - rollups: MemPartion, 594 594 - queues: MemPartion, 595 595 - } 596 596 - 597 597 - impl MemWriter { 598 598 - fn rollup_delete_account( 599 599 - &mut self, 600 600 - cursor: Cursor, 601 601 - key_bytes: &[u8], 602 602 - val_bytes: &[u8], 603 603 - ) -> StorageResult<usize> { 604 604 - let did = db_complete::<DeleteAccountQueueVal>(val_bytes)?; 605 605 - self.delete_account(&did)?; 606 606 - let mut batch = self.keyspace.batch(); 607 607 - batch.remove(&self.queues, key_bytes); 608 608 - insert_batch_static_neu::<NewRollupCursorKey>(&mut batch, &self.global, cursor)?; 609 609 - batch.commit()?; 610 610 - Ok(1) 611 611 - } 612 612 - 613 613 - fn rollup_live_counts( 614 614 - &mut self, 615 615 - timelies: impl Iterator<Item = Result<(Vec<u8>, Vec<u8>), StorageError>>, 616 616 - cursor_exclusive_limit: Option<Cursor>, 617 617 - rollup_limit: usize, 618 618 - ) -> StorageResult<usize> { 619 619 - // current strategy is to buffer counts in mem before writing the rollups 620 620 - // we *could* read+write every single batch to rollup.. but their merge is associative so 621 621 - // ...so save the db some work up front? is this worth it? who knows... 622 622 - 623 623 - log::warn!("sup!!!"); 624 624 - 625 625 - #[derive(Eq, Hash, PartialEq)] 626 626 - enum Rollup { 627 627 - Hourly(HourTruncatedCursor), 628 628 - Weekly(WeekTruncatedCursor), 629 629 - AllTime, 630 630 - } 631 631 - 632 632 - let mut batch = self.keyspace.batch(); 633 633 - let mut cursors_advanced = 0; 634 634 - let mut last_cursor = Cursor::from_start(); 635 635 - let mut counts_by_rollup: HashMap<(Nsid, Rollup), CountsValue> = HashMap::new(); 636 636 - 637 637 - log::warn!("about to loop...."); 638 638 - for (i, kv) in timelies.enumerate() { 639 639 - log::warn!("loop {i} {kv:?}..."); 640 640 - if i >= rollup_limit { 641 641 - break; 642 642 - } 643 643 - 644 644 - let (key_bytes, val_bytes) = kv?; 645 645 - let key = db_complete::<LiveCountsKey>(&key_bytes) 646 646 - .inspect_err(|e| log::warn!("rlc: key: {e:?}"))?; 647 647 - 648 648 - if cursor_exclusive_limit 649 649 - .map(|limit| key.cursor() > limit) 650 650 - .unwrap_or(false) 651 651 - { 652 652 - break; 653 653 - } 654 654 - 655 655 - batch.remove(&self.rollups, &key_bytes); 656 656 - let val = db_complete::<CountsValue>(&val_bytes) 657 657 - .inspect_err(|e| log::warn!("rlc: val: {e:?}"))?; 658 658 - counts_by_rollup 659 659 - .entry(( 660 660 - key.collection().clone(), 661 661 - Rollup::Hourly(key.cursor().into()), 662 662 - )) 663 663 - .or_default() 664 664 - .merge(&val); 665 665 - counts_by_rollup 666 666 - .entry(( 667 667 - key.collection().clone(), 668 668 - Rollup::Weekly(key.cursor().into()), 669 669 - )) 670 670 - .or_default() 671 671 - .merge(&val); 672 672 - counts_by_rollup 673 673 - .entry((key.collection().clone(), Rollup::AllTime)) 674 674 - .or_default() 675 675 - .merge(&val); 676 676 - 677 677 - cursors_advanced += 1; 678 678 - last_cursor = key.cursor(); 679 679 - } 680 680 - log::warn!("done looping. looping cbr counts(?).."); 681 681 - 682 682 - for ((nsid, rollup), counts) in counts_by_rollup { 683 683 - log::warn!( 684 684 - "######################## cbr loop {nsid:?} {counts:?} ########################" 685 685 - ); 686 686 - let key_bytes = match rollup { 687 687 - Rollup::Hourly(hourly_cursor) => { 688 688 - let k = HourlyRollupKey::new(hourly_cursor, &nsid); 689 689 - log::info!("hrly k: {k:?}"); 690 690 - k.to_db_bytes()? 691 691 - } 692 692 - Rollup::Weekly(weekly_cursor) => { 693 693 - let k = WeeklyRollupKey::new(weekly_cursor, &nsid); 694 694 - log::info!("weekly k: {k:?}"); 695 695 - k.to_db_bytes()? 696 696 - } 697 697 - Rollup::AllTime => { 698 698 - let k = AllTimeRollupKey::new(&nsid); 699 699 - log::info!("alltime k: {k:?}"); 700 700 - k.to_db_bytes()? 701 701 - } 702 702 - }; 703 703 - // log::info!("key bytes: {key_bytes:?}"); 704 704 - let mut rolled: CountsValue = self 705 705 - .rollups 706 706 - .get(&key_bytes)? 707 707 - .inspect(|v| { 708 708 - let lax = CountsValue::from_db_bytes(v); 709 709 - log::info!( 710 710 - "val: len={}, lax={lax:?} first32={:?}", 711 711 - v.len(), 712 712 - v.get(..32) 713 713 - ); 714 714 - }) 715 715 - .as_deref() 716 716 - .map(db_complete::<CountsValue>) 717 717 - .transpose() 718 718 - .inspect_err(|e| log::warn!("oooh did we break on the rolled thing? {e:?}"))? 719 719 - .unwrap_or_default(); 720 720 - 721 721 - // try to round-trip before inserting, for funsies 722 722 - let tripppin = counts.to_db_bytes()?; 723 723 - let (and_back, n) = CountsValue::from_db_bytes(&tripppin)?; 724 724 - assert_eq!(n, tripppin.len()); 725 725 - assert_eq!(counts.prefix, and_back.prefix); 726 726 - assert_eq!(counts.dids().estimate(), and_back.dids().estimate()); 727 727 - if counts.counts().creates > 20000000 { 728 728 - panic!("COUNTS maybe wtf? {counts:?}") 729 729 - } 730 730 - // assert_eq!(rolled, and_back); 731 731 - 732 732 - rolled.merge(&counts); 733 733 - 734 734 - // try to round-trip before inserting, for funsies 735 735 - let tripppin = rolled.to_db_bytes()?; 736 736 - let (and_back, n) = CountsValue::from_db_bytes(&tripppin)?; 737 737 - assert_eq!(n, tripppin.len()); 738 738 - assert_eq!(rolled.prefix, and_back.prefix); 739 739 - assert_eq!(rolled.dids().estimate(), and_back.dids().estimate()); 740 740 - if rolled.counts().creates > 20000000 { 741 741 - panic!("maybe wtf? {rolled:?}") 742 742 - } 743 743 - // assert_eq!(rolled, and_back); 744 744 - 745 745 - batch.insert(&self.rollups, &key_bytes, &rolled.to_db_bytes()?); 746 746 - } 747 747 - 748 748 - log::warn!("done cbr loop."); 749 749 - 750 750 - insert_batch_static_neu::<NewRollupCursorKey>(&mut batch, &self.global, last_cursor) 751 751 - .inspect_err(|e| log::warn!("insert neu: {e:?}"))?; 752 752 - 753 753 - batch.commit()?; 754 754 - 755 755 - log::warn!("ok finished rlc stuff. huh."); 756 756 - Ok(cursors_advanced) 757 757 - } 758 758 - } 759 759 - 760 760 - impl StoreWriter<MemBackground> for MemWriter { 761 761 - fn background_tasks(&mut self, _reroll: bool) -> StorageResult<MemBackground> { 762 762 - Ok(MemBackground {}) 763 763 - } 764 764 - 765 765 - fn insert_batch<const LIMIT: usize>( 766 766 - &mut self, 767 767 - event_batch: EventBatch<LIMIT>, 768 768 - ) -> StorageResult<()> { 769 769 - if event_batch.is_empty() { 770 770 - return Ok(()); 771 771 - } 772 772 - 773 773 - let mut batch = self.keyspace.batch(); 774 774 - 775 775 - // would be nice not to have to iterate everything at once here 776 776 - let latest = event_batch.latest_cursor().unwrap(); 777 777 - 778 778 - for (nsid, commits) in event_batch.commits_by_nsid { 779 779 - for commit in commits.commits { 780 780 - let location_key: RecordLocationKey = (&commit, &nsid).into(); 781 781 - 782 782 - match commit.action { 783 783 - CommitAction::Cut => { 784 784 - batch.remove(&self.records, &location_key.to_db_bytes()?); 785 785 - } 786 786 - CommitAction::Put(put_action) => { 787 787 - let feed_key = NsidRecordFeedKey::from_pair(nsid.clone(), commit.cursor); 788 788 - let feed_val: NsidRecordFeedVal = 789 789 - (&commit.did, &commit.rkey, commit.rev.as_str()).into(); 790 790 - batch.insert( 791 791 - &self.feeds, 792 792 - &feed_key.to_db_bytes()?, 793 793 - &feed_val.to_db_bytes()?, 794 794 - ); 795 795 - 796 796 - let location_val: RecordLocationVal = 797 797 - (commit.cursor, commit.rev.as_str(), put_action).into(); 798 798 - batch.insert( 799 799 - &self.records, 800 800 - &location_key.to_db_bytes()?, 801 801 - &location_val.to_db_bytes()?, 802 802 - ); 803 803 - } 804 804 - } 805 805 - } 806 806 - let live_counts_key: LiveCountsKey = (latest, &nsid).into(); 807 807 - let counts_value = CountsValue::new( 808 808 - CommitCounts { 809 809 - creates: commits.creates as u64, 810 810 - updates: commits.updates as u64, 811 811 - deletes: commits.deletes as u64, 812 812 - }, 813 813 - commits.dids_estimate, 814 814 - ); 815 815 - batch.insert( 816 816 - &self.rollups, 817 817 - &live_counts_key.to_db_bytes()?, 818 818 - &counts_value.to_db_bytes()?, 819 819 - ); 820 820 - } 821 821 - 822 822 - for remove in event_batch.account_removes { 823 823 - let queue_key = DeleteAccountQueueKey::new(remove.cursor); 824 824 - let queue_val: DeleteAccountQueueVal = remove.did; 825 825 - batch.insert( 826 826 - &self.queues, 827 827 - &queue_key.to_db_bytes()?, 828 828 - &queue_val.to_db_bytes()?, 829 829 - ); 830 830 - } 831 831 - 832 832 - batch.insert( 833 833 - &self.global, 834 834 - &DbStaticStr::<JetstreamCursorKey>::default().to_db_bytes()?, 835 835 - &latest.to_db_bytes()?, 836 836 - ); 837 837 - 838 838 - batch.commit()?; 839 839 - Ok(()) 840 840 - } 841 841 - 842 842 - fn step_rollup(&mut self) -> StorageResult<(usize, HashSet<Nsid>)> { 843 843 - let mut dirty_nsids = HashSet::new(); 844 844 - 845 845 - let rollup_cursor = 846 846 - get_static_neu::<NewRollupCursorKey, NewRollupCursorValue>(&self.global)? 847 847 - .ok_or(StorageError::BadStateError( 848 848 - "Could not find current rollup cursor".to_string(), 849 849 - )) 850 850 - .inspect_err(|e| log::warn!("failed getting rollup cursor: {e:?}"))?; 851 851 - 852 852 - // timelies 853 853 - let live_counts_range = LiveCountsKey::range_from_cursor(rollup_cursor) 854 854 - .inspect_err(|e| log::warn!("live counts range: {e:?}"))?; 855 855 - let mut timely_iter = self.rollups.range(live_counts_range).into_iter().peekable(); 856 856 - 857 857 - let timely_next = timely_iter 858 858 - .peek_mut() 859 859 - .map(|kv| -> StorageResult<LiveCountsKey> { 860 860 - match kv { 861 861 - Err(e) => Err(std::mem::replace(e, StorageError::Stolen))?, 862 862 - Ok((key_bytes, _)) => { 863 863 - let key = db_complete::<LiveCountsKey>(key_bytes).inspect_err(|e| { 864 864 - log::warn!("failed getting key for next timely: {e:?}") 865 865 - })?; 866 866 - Ok(key) 867 867 - } 868 868 - } 869 869 - }) 870 870 - .transpose() 871 871 - .inspect_err(|e| log::warn!("something about timely: {e:?}"))?; 872 872 - 873 873 - // delete accounts 874 874 - let delete_accounts_range = 875 875 - DeleteAccountQueueKey::new(rollup_cursor).range_to_prefix_end()?; 876 876 - 877 877 - let next_delete = self 878 878 - .queues 879 879 - .range(delete_accounts_range) 880 880 - .into_iter() 881 881 - .next() 882 882 - .transpose() 883 883 - .inspect_err(|e| log::warn!("range for next delete: {e:?}"))? 884 884 - .map(|(key_bytes, val_bytes)| { 885 885 - db_complete::<DeleteAccountQueueKey>(&key_bytes) 886 886 - .inspect_err(|e| log::warn!("failed inside next delete thing????: {e:?}")) 887 887 - .map(|k| (k.suffix, key_bytes, val_bytes)) 888 888 - }) 889 889 - .transpose() 890 890 - .inspect_err(|e| log::warn!("failed getting next delete: {e:?}"))?; 891 891 - 892 892 - let cursors_stepped = match (timely_next, next_delete) { 893 893 - (Some(timely), Some((delete_cursor, delete_key_bytes, delete_val_bytes))) => { 894 894 - if timely.cursor() < delete_cursor { 895 895 - let n = self 896 896 - .rollup_live_counts( 897 897 - timely_iter, 898 898 - Some(delete_cursor), 899 899 - MAX_BATCHED_ROLLUP_COUNTS, 900 900 - ) 901 901 - .inspect_err(|e| log::warn!("rolling up live counts: {e:?}"))?; 902 902 - dirty_nsids.insert(timely.collection().clone()); 903 903 - n 904 904 - } else { 905 905 - self.rollup_delete_account(delete_cursor, &delete_key_bytes, &delete_val_bytes) 906 906 - .inspect_err(|e| log::warn!("deleting acocunt: {e:?}"))? 907 907 - } 908 908 - } 909 909 - (Some(timely), None) => { 910 910 - let n = self 911 911 - .rollup_live_counts(timely_iter, None, MAX_BATCHED_ROLLUP_COUNTS) 912 912 - .inspect_err(|e| log::warn!("rolling up (lasjdflkajs): {e:?}"))?; 913 913 - dirty_nsids.insert(timely.collection().clone()); 914 914 - n 915 915 - } 916 916 - (None, Some((delete_cursor, delete_key_bytes, delete_val_bytes))) => self 917 917 - .rollup_delete_account(delete_cursor, &delete_key_bytes, &delete_val_bytes) 918 918 - .inspect_err(|e| log::warn!("deleting acocunt other branch: {e:?}"))?, 919 919 - (None, None) => 0, 920 920 - }; 921 921 - 922 922 - Ok((cursors_stepped, dirty_nsids)) 923 923 - } 924 924 - 925 925 - fn trim_collection( 926 926 - &mut self, 927 927 - collection: &Nsid, 928 928 - limit: usize, 929 929 - _full_scan: bool, 930 930 - // TODO: could add a start cursor limit to avoid iterating deleted stuff at the start (/end) 931 931 - ) -> StorageResult<(usize, usize, bool)> { 932 932 - let mut dangling_feed_keys_cleaned = 0; 933 933 - let mut records_deleted = 0; 934 934 - 935 935 - let mut batch = self.keyspace.batch(); 936 936 - 937 937 - let prefix = NsidRecordFeedKey::from_prefix_to_db_bytes(collection)?; 938 938 - let mut found = 0; 939 939 - for kv in self.feeds.prefix(&prefix).into_iter().rev() { 940 940 - let (key_bytes, val_bytes) = kv?; 941 941 - let feed_key = db_complete::<NsidRecordFeedKey>(&key_bytes)?; 942 942 - let feed_val = db_complete::<NsidRecordFeedVal>(&val_bytes)?; 943 943 - let location_key: RecordLocationKey = (&feed_key, &feed_val).into(); 944 944 - let location_key_bytes = location_key.to_db_bytes()?; 945 945 - 946 946 - let Some(location_val_bytes) = self.records.get(&location_key_bytes)? else { 947 947 - // record was deleted (hopefully) 948 948 - batch.remove(&self.feeds, &location_key_bytes); 949 949 - dangling_feed_keys_cleaned += 1; 950 950 - continue; 951 951 - }; 952 952 - 953 953 - let (meta, _) = RecordLocationMeta::from_db_bytes(&location_val_bytes)?; 954 954 - 955 955 - if meta.cursor() != feed_key.cursor() { 956 956 - // older/different version 957 957 - batch.remove(&self.feeds, &location_key_bytes); 958 958 - dangling_feed_keys_cleaned += 1; 959 959 - continue; 960 960 - } 961 961 - if meta.rev != feed_val.rev() { 962 962 - // weird... 963 963 - log::warn!("record lookup: cursor match but rev did not...? removing."); 964 964 - batch.remove(&self.feeds, &location_key_bytes); 965 965 - dangling_feed_keys_cleaned += 1; 966 966 - continue; 967 967 - } 968 968 - 969 969 - if batch.len() >= MAX_BATCHED_CLEANUP_SIZE { 970 970 - batch.commit()?; 971 971 - batch = self.keyspace.batch(); 972 972 - } 973 973 - 974 974 - found += 1; 975 975 - if found <= limit { 976 976 - continue; 977 977 - } 978 978 - 979 979 - batch.remove(&self.feeds, &location_key_bytes); 980 980 - batch.remove(&self.records, &location_key_bytes); 981 981 - records_deleted += 1; 982 982 - } 983 983 - 984 984 - batch.commit()?; 985 985 - 986 986 - log::info!("trim_collection ({collection:?}) removed {dangling_feed_keys_cleaned} dangling feed entries and {records_deleted} records"); 987 987 - Ok((dangling_feed_keys_cleaned, records_deleted, false)) 988 988 - } 989 989 - 990 990 - fn delete_account(&mut self, did: &Did) -> Result<usize, StorageError> { 991 991 - let mut records_deleted = 0; 992 992 - let mut batch = self.keyspace.batch(); 993 993 - let prefix = RecordLocationKey::from_prefix_to_db_bytes(did)?; 994 994 - for kv in self.records.prefix(&prefix) { 995 995 - let (key_bytes, _) = kv?; 996 996 - batch.remove(&self.records, &key_bytes); 997 997 - records_deleted += 1; 998 998 - if batch.len() >= MAX_BATCHED_ACCOUNT_DELETE_RECORDS { 999 999 - batch.commit()?; 1000 1000 - batch = self.keyspace.batch(); 1001 1001 - } 1002 1002 - } 1003 1003 - batch.commit()?; 1004 1004 - Ok(records_deleted) 1005 1005 - } 1006 1006 - } 1007 1007 - 1008 1008 - pub struct MemBackground; 1009 1009 - 1010 1010 - #[async_trait] 1011 1011 - impl StoreBackground for MemBackground { 1012 1012 - async fn run(mut self, _backfill: bool) -> StorageResult<()> { 1013 1013 - // noop for mem (is there a nicer way to do this?) 1014 1014 - loop { 1015 1015 - tokio::time::sleep(std::time::Duration::from_secs_f64(10.)).await; 1016 1016 - } 1017 1017 - } 1018 1018 - } 1019 1019 - 1020 1020 - /// Get a value from a fixed key 1021 1021 - fn get_static_neu<K: StaticStr, V: DbBytes>(global: &MemPartion) -> StorageResult<Option<V>> { 1022 1022 - let key_bytes = DbStaticStr::<K>::default().to_db_bytes()?; 1023 1023 - let value = global 1024 1024 - .get(&key_bytes)? 1025 1025 - .map(|value_bytes| db_complete(&value_bytes)) 1026 1026 - .transpose()?; 1027 1027 - Ok(value) 1028 1028 - } 1029 1029 - 1030 1030 - /// Get a value from a fixed key 1031 1031 - fn get_snapshot_static_neu<K: StaticStr, V: DbBytes>( 1032 1032 - global: &MemPartion, 1033 1033 - ) -> StorageResult<Option<V>> { 1034 1034 - let key_bytes = DbStaticStr::<K>::default().to_db_bytes()?; 1035 1035 - let value = global 1036 1036 - .get(&key_bytes)? 1037 1037 - .map(|value_bytes| db_complete(&value_bytes)) 1038 1038 - .transpose()?; 1039 1039 - Ok(value) 1040 1040 - } 1041 1041 - 1042 1042 - /// Set a value to a fixed key 1043 1043 - fn insert_static_neu<K: StaticStr>(global: &MemPartion, value: impl DbBytes) -> StorageResult<()> { 1044 1044 - let key_bytes = DbStaticStr::<K>::default().to_db_bytes()?; 1045 1045 - let value_bytes = value.to_db_bytes()?; 1046 1046 - global.insert(&key_bytes, &value_bytes)?; 1047 1047 - Ok(()) 1048 1048 - } 1049 1049 - 1050 1050 - /// Set a value to a fixed key 1051 1051 - fn insert_batch_static_neu<K: StaticStr>( 1052 1052 - batch: &mut MemBatch, 1053 1053 - global: &MemPartion, 1054 1054 - value: impl DbBytes, 1055 1055 - ) -> StorageResult<()> { 1056 1056 - let key_bytes = DbStaticStr::<K>::default().to_db_bytes()?; 1057 1057 - let value_bytes = value.to_db_bytes()?; 1058 1058 - batch.insert(global, &key_bytes, &value_bytes); 1059 1059 - Ok(()) 1060 1060 - } 1061 1061 - 1062 1062 - #[derive(Debug, serde::Serialize, schemars::JsonSchema)] 1063 1063 - pub struct StorageInfo { 1064 1064 - pub keyspace_disk_space: u64, 1065 1065 - pub keyspace_journal_count: usize, 1066 1066 - pub keyspace_sequence: u64, 1067 1067 - pub global_approximate_len: usize, 1068 1068 - } 1069 1069 - 1070 1070 - #[cfg(test)] 1071 1071 - mod tests { 1072 1072 - use super::*; 1073 1073 - use crate::{DeleteAccount, RecordKey, UFOsCommit}; 1074 1074 - use jetstream::events::{CommitEvent, CommitOp}; 1075 1075 - use jetstream::exports::Cid; 1076 1076 - use serde_json::value::RawValue; 1077 1077 - 1078 1078 - fn fjall_db() -> (MemReader, MemWriter) { 1079 1079 - let (read, write, _, _) = MemStorage::init( 1080 1080 - tempfile::tempdir().unwrap(), 1081 1081 - "offline test (no real jetstream endpoint)".to_string(), 1082 1082 - false, 1083 1083 - MemConfig { temp: true }, 1084 1084 - ) 1085 1085 - .unwrap(); 1086 1086 - (read, write) 1087 1087 - } 1088 1088 - 1089 1089 - const TEST_BATCH_LIMIT: usize = 16; 1090 1090 - 1091 1091 - #[derive(Debug, Default)] 1092 1092 - struct TestBatch { 1093 1093 - pub batch: EventBatch<TEST_BATCH_LIMIT>, 1094 1094 - } 1095 1095 - 1096 1096 - impl TestBatch { 1097 1097 - #[allow(clippy::too_many_arguments)] 1098 1098 - pub fn create( 1099 1099 - &mut self, 1100 1100 - did: &str, 1101 1101 - collection: &str, 1102 1102 - rkey: &str, 1103 1103 - record: &str, 1104 1104 - rev: Option<&str>, 1105 1105 - cid: Option<Cid>, 1106 1106 - cursor: u64, 1107 1107 - ) -> Nsid { 1108 1108 - let did = Did::new(did.to_string()).unwrap(); 1109 1109 - let collection = Nsid::new(collection.to_string()).unwrap(); 1110 1110 - let record = RawValue::from_string(record.to_string()).unwrap(); 1111 1111 - let cid = cid.unwrap_or( 1112 1112 - "bafyreidofvwoqvd2cnzbun6dkzgfucxh57tirf3ohhde7lsvh4fu3jehgy" 1113 1113 - .parse() 1114 1114 - .unwrap(), 1115 1115 - ); 1116 1116 - 1117 1117 - let event = CommitEvent { 1118 1118 - collection, 1119 1119 - rkey: RecordKey::new(rkey.to_string()).unwrap(), 1120 1120 - rev: rev.unwrap_or("asdf").to_string(), 1121 1121 - operation: CommitOp::Create, 1122 1122 - record: Some(record), 1123 1123 - cid: Some(cid), 1124 1124 - }; 1125 1125 - 1126 1126 - let (commit, collection) = 1127 1127 - UFOsCommit::from_commit_info(event, did.clone(), Cursor::from_raw_u64(cursor)) 1128 1128 - .unwrap(); 1129 1129 - 1130 1130 - self.batch 1131 1131 - .commits_by_nsid 1132 1132 - .entry(collection.clone()) 1133 1133 - .or_default() 1134 1134 - .truncating_insert(commit, &[0u8; 16]) 1135 1135 - .unwrap(); 1136 1136 - 1137 1137 - collection 1138 1138 - } 1139 1139 - #[allow(clippy::too_many_arguments)] 1140 1140 - pub fn update( 1141 1141 - &mut self, 1142 1142 - did: &str, 1143 1143 - collection: &str, 1144 1144 - rkey: &str, 1145 1145 - record: &str, 1146 1146 - rev: Option<&str>, 1147 1147 - cid: Option<Cid>, 1148 1148 - cursor: u64, 1149 1149 - ) -> Nsid { 1150 1150 - let did = Did::new(did.to_string()).unwrap(); 1151 1151 - let collection = Nsid::new(collection.to_string()).unwrap(); 1152 1152 - let record = RawValue::from_string(record.to_string()).unwrap(); 1153 1153 - let cid = cid.unwrap_or( 1154 1154 - "bafyreidofvwoqvd2cnzbun6dkzgfucxh57tirf3ohhde7lsvh4fu3jehgy" 1155 1155 - .parse() 1156 1156 - .unwrap(), 1157 1157 - ); 1158 1158 - 1159 1159 - let event = CommitEvent { 1160 1160 - collection, 1161 1161 - rkey: RecordKey::new(rkey.to_string()).unwrap(), 1162 1162 - rev: rev.unwrap_or("asdf").to_string(), 1163 1163 - operation: CommitOp::Update, 1164 1164 - record: Some(record), 1165 1165 - cid: Some(cid), 1166 1166 - }; 1167 1167 - 1168 1168 - let (commit, collection) = 1169 1169 - UFOsCommit::from_commit_info(event, did.clone(), Cursor::from_raw_u64(cursor)) 1170 1170 - .unwrap(); 1171 1171 - 1172 1172 - self.batch 1173 1173 - .commits_by_nsid 1174 1174 - .entry(collection.clone()) 1175 1175 - .or_default() 1176 1176 - .truncating_insert(commit, &[0u8; 16]) 1177 1177 - .unwrap(); 1178 1178 - 1179 1179 - collection 1180 1180 - } 1181 1181 - #[allow(clippy::too_many_arguments)] 1182 1182 - pub fn delete( 1183 1183 - &mut self, 1184 1184 - did: &str, 1185 1185 - collection: &str, 1186 1186 - rkey: &str, 1187 1187 - rev: Option<&str>, 1188 1188 - cursor: u64, 1189 1189 - ) -> Nsid { 1190 1190 - let did = Did::new(did.to_string()).unwrap(); 1191 1191 - let collection = Nsid::new(collection.to_string()).unwrap(); 1192 1192 - let event = CommitEvent { 1193 1193 - collection, 1194 1194 - rkey: RecordKey::new(rkey.to_string()).unwrap(), 1195 1195 - rev: rev.unwrap_or("asdf").to_string(), 1196 1196 - operation: CommitOp::Delete, 1197 1197 - record: None, 1198 1198 - cid: None, 1199 1199 - }; 1200 1200 - 1201 1201 - let (commit, collection) = 1202 1202 - UFOsCommit::from_commit_info(event, did, Cursor::from_raw_u64(cursor)).unwrap(); 1203 1203 - 1204 1204 - self.batch 1205 1205 - .commits_by_nsid 1206 1206 - .entry(collection.clone()) 1207 1207 - .or_default() 1208 1208 - .truncating_insert(commit, &[0u8; 16]) 1209 1209 - .unwrap(); 1210 1210 - 1211 1211 - collection 1212 1212 - } 1213 1213 - pub fn delete_account(&mut self, did: &str, cursor: u64) -> Did { 1214 1214 - let did = Did::new(did.to_string()).unwrap(); 1215 1215 - self.batch.account_removes.push(DeleteAccount { 1216 1216 - did: did.clone(), 1217 1217 - cursor: Cursor::from_raw_u64(cursor), 1218 1218 - }); 1219 1219 - did 1220 1220 - } 1221 1221 - } 1222 1222 - 1223 1223 - #[test] 1224 1224 - fn test_hello() -> anyhow::Result<()> { 1225 1225 - let (read, mut write) = fjall_db(); 1226 1226 - write.insert_batch::<TEST_BATCH_LIMIT>(EventBatch::default())?; 1227 1227 - let (records, dids) = 1228 1228 - read.get_counts_by_collection(&Nsid::new("a.b.c".to_string()).unwrap())?; 1229 1229 - assert_eq!(records, 0); 1230 1230 - assert_eq!(dids, 0); 1231 1231 - Ok(()) 1232 1232 - } 1233 1233 - 1234 1234 - #[test] 1235 1235 - fn test_insert_one() -> anyhow::Result<()> { 1236 1236 - let (read, mut write) = fjall_db(); 1237 1237 - 1238 1238 - let mut batch = TestBatch::default(); 1239 1239 - let collection = batch.create( 1240 1240 - "did:plc:inze6wrmsm7pjl7yta3oig77", 1241 1241 - "a.b.c", 1242 1242 - "asdf", 1243 1243 - "{}", 1244 1244 - Some("rev-z"), 1245 1245 - None, 1246 1246 - 100, 1247 1247 - ); 1248 1248 - write.insert_batch(batch.batch)?; 1249 1249 - 1250 1250 - let (records, dids) = read.get_counts_by_collection(&collection)?; 1251 1251 - assert_eq!(records, 1); 1252 1252 - assert_eq!(dids, 1); 1253 1253 - let (records, dids) = 1254 1254 - read.get_counts_by_collection(&Nsid::new("d.e.f".to_string()).unwrap())?; 1255 1255 - assert_eq!(records, 0); 1256 1256 - assert_eq!(dids, 0); 1257 1257 - 1258 1258 - let records = read.get_records_by_collections(HashSet::from([collection]), 2, false)?; 1259 1259 - assert_eq!(records.len(), 1); 1260 1260 - let rec = &records[0]; 1261 1261 - assert_eq!(rec.record.get(), "{}"); 1262 1262 - assert!(!rec.is_update); 1263 1263 - 1264 1264 - let records = read.get_records_by_collections( 1265 1265 - HashSet::from([Nsid::new("d.e.f".to_string()).unwrap()]), 1266 1266 - 2, 1267 1267 - false, 1268 1268 - )?; 1269 1269 - assert_eq!(records.len(), 0); 1270 1270 - 1271 1271 - Ok(()) 1272 1272 - } 1273 1273 - 1274 1274 - #[test] 1275 1275 - fn test_get_multi_collection() -> anyhow::Result<()> { 1276 1276 - let (read, mut write) = fjall_db(); 1277 1277 - 1278 1278 - let mut batch = TestBatch::default(); 1279 1279 - batch.create( 1280 1280 - "did:plc:inze6wrmsm7pjl7yta3oig77", 1281 1281 - "a.a.a", 1282 1282 - "aaa", 1283 1283 - r#""earliest""#, 1284 1284 - Some("rev-a"), 1285 1285 - None, 1286 1286 - 100, 1287 1287 - ); 1288 1288 - batch.create( 1289 1289 - "did:plc:inze6wrmsm7pjl7yta3oig77", 1290 1290 - "a.a.b", 1291 1291 - "aab", 1292 1292 - r#""in between""#, 1293 1293 - Some("rev-ab"), 1294 1294 - None, 1295 1295 - 101, 1296 1296 - ); 1297 1297 - batch.create( 1298 1298 - "did:plc:inze6wrmsm7pjl7yta3oig77", 1299 1299 - "a.a.a", 1300 1300 - "aaa-2", 1301 1301 - r#""last""#, 1302 1302 - Some("rev-a-2"), 1303 1303 - None, 1304 1304 - 102, 1305 1305 - ); 1306 1306 - write.insert_batch(batch.batch)?; 1307 1307 - 1308 1308 - let records = read.get_records_by_collections( 1309 1309 - HashSet::from([ 1310 1310 - Nsid::new("a.a.a".to_string()).unwrap(), 1311 1311 - Nsid::new("a.a.b".to_string()).unwrap(), 1312 1312 - Nsid::new("a.a.c".to_string()).unwrap(), 1313 1313 - ]), 1314 1314 - 100, 1315 1315 - false, 1316 1316 - )?; 1317 1317 - assert_eq!(records.len(), 3); 1318 1318 - assert_eq!(records[0].record.get(), r#""last""#); 1319 1319 - assert_eq!( 1320 1320 - records[0].collection, 1321 1321 - Nsid::new("a.a.a".to_string()).unwrap() 1322 1322 - ); 1323 1323 - assert_eq!(records[1].record.get(), r#""in between""#); 1324 1324 - assert_eq!( 1325 1325 - records[1].collection, 1326 1326 - Nsid::new("a.a.b".to_string()).unwrap() 1327 1327 - ); 1328 1328 - assert_eq!(records[2].record.get(), r#""earliest""#); 1329 1329 - assert_eq!( 1330 1330 - records[2].collection, 1331 1331 - Nsid::new("a.a.a".to_string()).unwrap() 1332 1332 - ); 1333 1333 - 1334 1334 - Ok(()) 1335 1335 - } 1336 1336 - 1337 1337 - #[test] 1338 1338 - fn test_update_one() -> anyhow::Result<()> { 1339 1339 - let (read, mut write) = fjall_db(); 1340 1340 - 1341 1341 - let mut batch = TestBatch::default(); 1342 1342 - let collection = batch.create( 1343 1343 - "did:plc:inze6wrmsm7pjl7yta3oig77", 1344 1344 - "a.b.c", 1345 1345 - "rkey-asdf", 1346 1346 - "{}", 1347 1347 - Some("rev-a"), 1348 1348 - None, 1349 1349 - 100, 1350 1350 - ); 1351 1351 - write.insert_batch(batch.batch)?; 1352 1352 - 1353 1353 - let mut batch = TestBatch::default(); 1354 1354 - batch.update( 1355 1355 - "did:plc:inze6wrmsm7pjl7yta3oig77", 1356 1356 - "a.b.c", 1357 1357 - "rkey-asdf", 1358 1358 - r#"{"ch": "ch-ch-ch-changes"}"#, 1359 1359 - Some("rev-z"), 1360 1360 - None, 1361 1361 - 101, 1362 1362 - ); 1363 1363 - write.insert_batch(batch.batch)?; 1364 1364 - 1365 1365 - let (records, dids) = read.get_counts_by_collection(&collection)?; 1366 1366 - assert_eq!(records, 1); 1367 1367 - assert_eq!(dids, 1); 1368 1368 - 1369 1369 - let records = read.get_records_by_collections(HashSet::from([collection]), 2, false)?; 1370 1370 - assert_eq!(records.len(), 1); 1371 1371 - let rec = &records[0]; 1372 1372 - assert_eq!(rec.record.get(), r#"{"ch": "ch-ch-ch-changes"}"#); 1373 1373 - assert!(rec.is_update); 1374 1374 - Ok(()) 1375 1375 - } 1376 1376 - 1377 1377 - #[test] 1378 1378 - fn test_delete_one() -> anyhow::Result<()> { 1379 1379 - let (read, mut write) = fjall_db(); 1380 1380 - 1381 1381 - let mut batch = TestBatch::default(); 1382 1382 - let collection = batch.create( 1383 1383 - "did:plc:inze6wrmsm7pjl7yta3oig77", 1384 1384 - "a.b.c", 1385 1385 - "rkey-asdf", 1386 1386 - "{}", 1387 1387 - Some("rev-a"), 1388 1388 - None, 1389 1389 - 100, 1390 1390 - ); 1391 1391 - write.insert_batch(batch.batch)?; 1392 1392 - 1393 1393 - let mut batch = TestBatch::default(); 1394 1394 - batch.delete( 1395 1395 - "did:plc:inze6wrmsm7pjl7yta3oig77", 1396 1396 - "a.b.c", 1397 1397 - "rkey-asdf", 1398 1398 - Some("rev-z"), 1399 1399 - 101, 1400 1400 - ); 1401 1401 - write.insert_batch(batch.batch)?; 1402 1402 - 1403 1403 - let (records, dids) = read.get_counts_by_collection(&collection)?; 1404 1404 - assert_eq!(records, 1); 1405 1405 - assert_eq!(dids, 1); 1406 1406 - 1407 1407 - let records = read.get_records_by_collections(HashSet::from([collection]), 2, false)?; 1408 1408 - assert_eq!(records.len(), 0); 1409 1409 - 1410 1410 - Ok(()) 1411 1411 - } 1412 1412 - 1413 1413 - #[test] 1414 1414 - fn test_collection_trim() -> anyhow::Result<()> { 1415 1415 - let (read, mut write) = fjall_db(); 1416 1416 - 1417 1417 - let mut batch = TestBatch::default(); 1418 1418 - batch.create( 1419 1419 - "did:plc:inze6wrmsm7pjl7yta3oig77", 1420 1420 - "a.a.a", 1421 1421 - "rkey-aaa", 1422 1422 - "{}", 1423 1423 - Some("rev-aaa"), 1424 1424 - None, 1425 1425 - 10_000, 1426 1426 - ); 1427 1427 - let mut last_b_cursor; 1428 1428 - for i in 1..=10 { 1429 1429 - last_b_cursor = 11_000 + i; 1430 1430 - batch.create( 1431 1431 - &format!("did:plc:inze6wrmsm7pjl7yta3oig7{}", i % 3), 1432 1432 - "a.a.b", 1433 1433 - &format!("rkey-bbb-{i}"), 1434 1434 - &format!(r#"{{"n": {i}}}"#), 1435 1435 - Some(&format!("rev-bbb-{i}")), 1436 1436 - None, 1437 1437 - last_b_cursor, 1438 1438 - ); 1439 1439 - } 1440 1440 - batch.create( 1441 1441 - "did:plc:inze6wrmsm7pjl7yta3oig77", 1442 1442 - "a.a.c", 1443 1443 - "rkey-ccc", 1444 1444 - "{}", 1445 1445 - Some("rev-ccc"), 1446 1446 - None, 1447 1447 - 12_000, 1448 1448 - ); 1449 1449 - 1450 1450 - write.insert_batch(batch.batch)?; 1451 1451 - 1452 1452 - let records = read.get_records_by_collections( 1453 1453 - HashSet::from([Nsid::new("a.a.a".to_string()).unwrap()]), 1454 1454 - 100, 1455 1455 - false, 1456 1456 - )?; 1457 1457 - assert_eq!(records.len(), 1); 1458 1458 - let records = read.get_records_by_collections( 1459 1459 - HashSet::from([Nsid::new("a.a.b".to_string()).unwrap()]), 1460 1460 - 100, 1461 1461 - false, 1462 1462 - )?; 1463 1463 - assert_eq!(records.len(), 10); 1464 1464 - let records = read.get_records_by_collections( 1465 1465 - HashSet::from([Nsid::new("a.a.c".to_string()).unwrap()]), 1466 1466 - 100, 1467 1467 - false, 1468 1468 - )?; 1469 1469 - assert_eq!(records.len(), 1); 1470 1470 - let records = read.get_records_by_collections( 1471 1471 - HashSet::from([Nsid::new("a.a.d".to_string()).unwrap()]), 1472 1472 - 100, 1473 1473 - false, 1474 1474 - )?; 1475 1475 - assert_eq!(records.len(), 0); 1476 1476 - 1477 1477 - write.trim_collection(&Nsid::new("a.a.a".to_string()).unwrap(), 6, false)?; 1478 1478 - write.trim_collection(&Nsid::new("a.a.b".to_string()).unwrap(), 6, false)?; 1479 1479 - write.trim_collection(&Nsid::new("a.a.c".to_string()).unwrap(), 6, false)?; 1480 1480 - write.trim_collection(&Nsid::new("a.a.d".to_string()).unwrap(), 6, false)?; 1481 1481 - 1482 1482 - let records = read.get_records_by_collections( 1483 1483 - HashSet::from([Nsid::new("a.a.a".to_string()).unwrap()]), 1484 1484 - 100, 1485 1485 - false, 1486 1486 - )?; 1487 1487 - assert_eq!(records.len(), 1); 1488 1488 - let records = read.get_records_by_collections( 1489 1489 - HashSet::from([Nsid::new("a.a.b".to_string()).unwrap()]), 1490 1490 - 100, 1491 1491 - false, 1492 1492 - )?; 1493 1493 - assert_eq!(records.len(), 6); 1494 1494 - let records = read.get_records_by_collections( 1495 1495 - HashSet::from([Nsid::new("a.a.c".to_string()).unwrap()]), 1496 1496 - 100, 1497 1497 - false, 1498 1498 - )?; 1499 1499 - assert_eq!(records.len(), 1); 1500 1500 - let records = read.get_records_by_collections( 1501 1501 - HashSet::from([Nsid::new("a.a.d".to_string()).unwrap()]), 1502 1502 - 100, 1503 1503 - false, 1504 1504 - )?; 1505 1505 - assert_eq!(records.len(), 0); 1506 1506 - 1507 1507 - Ok(()) 1508 1508 - } 1509 1509 - 1510 1510 - #[test] 1511 1511 - fn test_delete_account() -> anyhow::Result<()> { 1512 1512 - let (read, mut write) = fjall_db(); 1513 1513 - 1514 1514 - let mut batch = TestBatch::default(); 1515 1515 - batch.create( 1516 1516 - "did:plc:person-a", 1517 1517 - "a.a.a", 1518 1518 - "rkey-aaa", 1519 1519 - "{}", 1520 1520 - Some("rev-aaa"), 1521 1521 - None, 1522 1522 - 10_000, 1523 1523 - ); 1524 1524 - for i in 1..=2 { 1525 1525 - batch.create( 1526 1526 - "did:plc:person-b", 1527 1527 - "a.a.a", 1528 1528 - &format!("rkey-bbb-{i}"), 1529 1529 - &format!(r#"{{"n": {i}}}"#), 1530 1530 - Some(&format!("rev-bbb-{i}")), 1531 1531 - None, 1532 1532 - 11_000 + i, 1533 1533 - ); 1534 1534 - } 1535 1535 - write.insert_batch(batch.batch)?; 1536 1536 - 1537 1537 - let records = read.get_records_by_collections( 1538 1538 - HashSet::from([Nsid::new("a.a.a".to_string()).unwrap()]), 1539 1539 - 100, 1540 1540 - false, 1541 1541 - )?; 1542 1542 - assert_eq!(records.len(), 3); 1543 1543 - 1544 1544 - let records_deleted = 1545 1545 - write.delete_account(&Did::new("did:plc:person-b".to_string()).unwrap())?; 1546 1546 - assert_eq!(records_deleted, 2); 1547 1547 - 1548 1548 - let records = read.get_records_by_collections( 1549 1549 - HashSet::from([Nsid::new("a.a.a".to_string()).unwrap()]), 1550 1550 - 100, 1551 1551 - false, 1552 1552 - )?; 1553 1553 - assert_eq!(records.len(), 1); 1554 1554 - 1555 1555 - Ok(()) 1556 1556 - } 1557 1557 - 1558 1558 - #[test] 1559 1559 - fn rollup_delete_account_removes_record() -> anyhow::Result<()> { 1560 1560 - let (read, mut write) = fjall_db(); 1561 1561 - 1562 1562 - let mut batch = TestBatch::default(); 1563 1563 - batch.create( 1564 1564 - "did:plc:person-a", 1565 1565 - "a.a.a", 1566 1566 - "rkey-aaa", 1567 1567 - "{}", 1568 1568 - Some("rev-aaa"), 1569 1569 - None, 1570 1570 - 10_000, 1571 1571 - ); 1572 1572 - write.insert_batch(batch.batch)?; 1573 1573 - 1574 1574 - let mut batch = TestBatch::default(); 1575 1575 - batch.delete_account("did:plc:person-a", 9_999); // queue it before the rollup 1576 1576 - write.insert_batch(batch.batch)?; 1577 1577 - 1578 1578 - write.step_rollup()?; 1579 1579 - 1580 1580 - let records = read.get_records_by_collections( 1581 1581 - HashSet::from([Nsid::new("a.a.a".to_string()).unwrap()]), 1582 1582 - 1, 1583 1583 - false, 1584 1584 - )?; 1585 1585 - assert_eq!(records.len(), 0); 1586 1586 - 1587 1587 - Ok(()) 1588 1588 - } 1589 1589 - 1590 1590 - #[test] 1591 1591 - fn rollup_delete_live_count_step() -> anyhow::Result<()> { 1592 1592 - let (read, mut write) = fjall_db(); 1593 1593 - 1594 1594 - let mut batch = TestBatch::default(); 1595 1595 - batch.create( 1596 1596 - "did:plc:person-a", 1597 1597 - "a.a.a", 1598 1598 - "rkey-aaa", 1599 1599 - "{}", 1600 1600 - Some("rev-aaa"), 1601 1601 - None, 1602 1602 - 10_000, 1603 1603 - ); 1604 1604 - write.insert_batch(batch.batch)?; 1605 1605 - 1606 1606 - let (n, _) = write.step_rollup()?; 1607 1607 - assert_eq!(n, 1); 1608 1608 - 1609 1609 - let mut batch = TestBatch::default(); 1610 1610 - batch.delete_account("did:plc:person-a", 10_001); 1611 1611 - write.insert_batch(batch.batch)?; 1612 1612 - 1613 1613 - let records = read.get_records_by_collections( 1614 1614 - HashSet::from([Nsid::new("a.a.a".to_string()).unwrap()]), 1615 1615 - 1, 1616 1616 - false, 1617 1617 - )?; 1618 1618 - assert_eq!(records.len(), 1); 1619 1619 - 1620 1620 - let (n, _) = write.step_rollup()?; 1621 1621 - assert_eq!(n, 1); 1622 1622 - 1623 1623 - let records = read.get_records_by_collections( 1624 1624 - HashSet::from([Nsid::new("a.a.a".to_string()).unwrap()]), 1625 1625 - 1, 1626 1626 - false, 1627 1627 - )?; 1628 1628 - assert_eq!(records.len(), 0); 1629 1629 - 1630 1630 - let mut batch = TestBatch::default(); 1631 1631 - batch.delete_account("did:plc:person-a", 9_999); 1632 1632 - write.insert_batch(batch.batch)?; 1633 1633 - 1634 1634 - let (n, _) = write.step_rollup()?; 1635 1635 - assert_eq!(n, 0); 1636 1636 - 1637 1637 - Ok(()) 1638 1638 - } 1639 1639 - 1640 1640 - #[test] 1641 1641 - fn rollup_multiple_count_batches() -> anyhow::Result<()> { 1642 1642 - let (_read, mut write) = fjall_db(); 1643 1643 - 1644 1644 - let mut batch = TestBatch::default(); 1645 1645 - batch.create( 1646 1646 - "did:plc:person-a", 1647 1647 - "a.a.a", 1648 1648 - "rkey-aaa", 1649 1649 - "{}", 1650 1650 - Some("rev-aaa"), 1651 1651 - None, 1652 1652 - 10_000, 1653 1653 - ); 1654 1654 - write.insert_batch(batch.batch)?; 1655 1655 - 1656 1656 - let mut batch = TestBatch::default(); 1657 1657 - batch.create( 1658 1658 - "did:plc:person-a", 1659 1659 - "a.a.a", 1660 1660 - "rkey-aab", 1661 1661 - "{}", 1662 1662 - Some("rev-aab"), 1663 1663 - None, 1664 1664 - 10_001, 1665 1665 - ); 1666 1666 - write.insert_batch(batch.batch)?; 1667 1667 - 1668 1668 - let (n, _) = write.step_rollup()?; 1669 1669 - assert_eq!(n, 2); 1670 1670 - 1671 1671 - let (n, _) = write.step_rollup()?; 1672 1672 - assert_eq!(n, 0); 1673 1673 - 1674 1674 - Ok(()) 1675 1675 - } 1676 1676 - 1677 1677 - #[test] 1678 1678 - fn counts_before_and_after_rollup() -> anyhow::Result<()> { 1679 1679 - let (read, mut write) = fjall_db(); 1680 1680 - 1681 1681 - let mut batch = TestBatch::default(); 1682 1682 - batch.create( 1683 1683 - "did:plc:person-a", 1684 1684 - "a.a.a", 1685 1685 - "rkey-aaa", 1686 1686 - "{}", 1687 1687 - Some("rev-aaa"), 1688 1688 - None, 1689 1689 - 10_000, 1690 1690 - ); 1691 1691 - batch.create( 1692 1692 - "did:plc:person-b", 1693 1693 - "a.a.a", 1694 1694 - "rkey-bbb", 1695 1695 - "{}", 1696 1696 - Some("rev-bbb"), 1697 1697 - None, 1698 1698 - 10_001, 1699 1699 - ); 1700 1700 - write.insert_batch(batch.batch)?; 1701 1701 - 1702 1702 - let mut batch = TestBatch::default(); 1703 1703 - batch.delete_account("did:plc:person-a", 11_000); 1704 1704 - write.insert_batch(batch.batch)?; 1705 1705 - 1706 1706 - let mut batch = TestBatch::default(); 1707 1707 - batch.create( 1708 1708 - "did:plc:person-a", 1709 1709 - "a.a.a", 1710 1710 - "rkey-aac", 1711 1711 - "{}", 1712 1712 - Some("rev-aac"), 1713 1713 - None, 1714 1714 - 12_000, 1715 1715 - ); 1716 1716 - write.insert_batch(batch.batch)?; 1717 1717 - 1718 1718 - // before any rollup 1719 1719 - let (records, dids) = 1720 1720 - read.get_counts_by_collection(&Nsid::new("a.a.a".to_string()).unwrap())?; 1721 1721 - assert_eq!(records, 3); 1722 1722 - assert_eq!(dids, 2); 1723 1723 - 1724 1724 - // first batch rolled up 1725 1725 - let (n, _) = write.step_rollup()?; 1726 1726 - assert_eq!(n, 1); 1727 1727 - 1728 1728 - let (records, dids) = 1729 1729 - read.get_counts_by_collection(&Nsid::new("a.a.a".to_string()).unwrap())?; 1730 1730 - assert_eq!(records, 3); 1731 1731 - assert_eq!(dids, 2); 1732 1732 - 1733 1733 - // delete account rolled up 1734 1734 - let (n, _) = write.step_rollup()?; 1735 1735 - assert_eq!(n, 1); 1736 1736 - 1737 1737 - let (records, dids) = 1738 1738 - read.get_counts_by_collection(&Nsid::new("a.a.a".to_string()).unwrap())?; 1739 1739 - assert_eq!(records, 3); 1740 1740 - assert_eq!(dids, 2); 1741 1741 - 1742 1742 - // second batch rolled up 1743 1743 - let (n, _) = write.step_rollup()?; 1744 1744 - assert_eq!(n, 1); 1745 1745 - 1746 1746 - let (records, dids) = 1747 1747 - read.get_counts_by_collection(&Nsid::new("a.a.a".to_string()).unwrap())?; 1748 1748 - assert_eq!(records, 3); 1749 1749 - assert_eq!(dids, 2); 1750 1750 - 1751 1751 - // no more rollups left 1752 1752 - let (n, _) = write.step_rollup()?; 1753 1753 - assert_eq!(n, 0); 1754 1754 - 1755 1755 - Ok(()) 1756 1756 - } 1757 1757 - }

+8 -1

ufos/src/store_types.rs

reviewed

··· 269 269 } 270 270 impl From<&CountsValue> for JustCount { 271 271 fn from(cv: &CountsValue) -> Self { 272 272 + let CommitCounts { 273 273 + creates, 274 274 + updates, 275 275 + deletes, 276 276 + } = cv.counts(); 272 277 Self { 273 273 - creates: cv.counts().creates, 278 278 + creates, 279 279 + updates, 280 280 + deletes, 274 281 dids_estimate: cv.dids().estimate() as u64, 275 282 } 276 283 }