very fast at protocol indexer with flexible filtering, xrpc queries, cursor-backed event stream, and more, built on fjall
rust fjall at-protocol atproto indexer
58
fork

Configure Feed

Select the types of activity you want to include in your feed.

[db] implement zstd compression with dict support

dawn 65621074 4803e5ac

+233 -89
+46 -80
Cargo.lock
··· 285 285 286 286 [[package]] 287 287 name = "bon" 288 - version = "3.9.0" 288 + version = "3.9.1" 289 289 source = "registry+https://github.com/rust-lang/crates.io-index" 290 - checksum = "2d13a61f2963b88eef9c1be03df65d42f6996dfeac1054870d950fcf66686f83" 290 + checksum = "f47dbe92550676ee653353c310dfb9cf6ba17ee70396e1f7cf0a2020ad49b2fe" 291 291 dependencies = [ 292 292 "bon-macros", 293 293 "rustversion", ··· 295 295 296 296 [[package]] 297 297 name = "bon-macros" 298 - version = "3.9.0" 298 + version = "3.9.1" 299 299 source = "registry+https://github.com/rust-lang/crates.io-index" 300 - checksum = "d314cc62af2b6b0c65780555abb4d02a03dd3b799cd42419044f0c38d99738c0" 300 + checksum = "519bd3116aeeb42d5372c29d982d16d0170d3d4a5ed85fc7dd91642ffff3c67c" 301 301 dependencies = [ 302 - "darling 0.23.0", 302 + "darling", 303 303 "ident_case", 304 304 "prettyplease", 305 305 "proc-macro2", ··· 310 310 311 311 [[package]] 312 312 name = "borsh" 313 - version = "1.6.0" 313 + version = "1.6.1" 314 314 source = "registry+https://github.com/rust-lang/crates.io-index" 315 - checksum = "d1da5ab77c1437701eeff7c88d968729e7766172279eab0676857b3d63af7a6f" 315 + checksum = "cfd1e3f8955a5d7de9fab72fc8373fade9fb8a703968cb200ae3dc6cf08e185a" 316 316 dependencies = [ 317 + "bytes", 317 318 "cfg_aliases", 318 319 ] 319 320 ··· 382 383 383 384 [[package]] 384 385 name = "cc" 385 - version = "1.2.56" 386 + version = "1.2.57" 386 387 source = "registry+https://github.com/rust-lang/crates.io-index" 387 - checksum = "aebf35691d1bfb0ac386a69bac2fde4dd276fb618cf8bf4f5318fe285e821bb2" 388 + checksum = "7a0dd1ca384932ff3641c8718a02769f1698e7563dc6974ffd03346116310423" 388 389 dependencies = [ 389 390 "find-msvc-tools", 390 391 "jobserver", ··· 710 711 711 712 [[package]] 712 713 name = "darling" 713 - version = "0.21.3" 714 - source = "registry+https://github.com/rust-lang/crates.io-index" 715 - checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0" 716 - dependencies = [ 717 - "darling_core 0.21.3", 718 - "darling_macro 0.21.3", 719 - ] 720 - 721 - [[package]] 722 - name = "darling" 723 714 version = "0.23.0" 724 715 source = "registry+https://github.com/rust-lang/crates.io-index" 725 716 checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" 726 717 dependencies = [ 727 - "darling_core 0.23.0", 728 - "darling_macro 0.23.0", 729 - ] 730 - 731 - [[package]] 732 - name = "darling_core" 733 - version = "0.21.3" 734 - source = "registry+https://github.com/rust-lang/crates.io-index" 735 - checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4" 736 - dependencies = [ 737 - "fnv", 738 - "ident_case", 739 - "proc-macro2", 740 - "quote", 741 - "strsim", 742 - "syn", 718 + "darling_core", 719 + "darling_macro", 743 720 ] 744 721 745 722 [[package]] ··· 757 734 758 735 [[package]] 759 736 name = "darling_macro" 760 - version = "0.21.3" 761 - source = "registry+https://github.com/rust-lang/crates.io-index" 762 - checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" 763 - dependencies = [ 764 - "darling_core 0.21.3", 765 - "quote", 766 - "syn", 767 - ] 768 - 769 - [[package]] 770 - name = "darling_macro" 771 737 version = "0.23.0" 772 738 source = "registry+https://github.com/rust-lang/crates.io-index" 773 739 checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" 774 740 dependencies = [ 775 - "darling_core 0.23.0", 741 + "darling_core", 776 742 "quote", 777 743 "syn", 778 744 ] ··· 1041 1007 1042 1008 [[package]] 1043 1009 name = "fjall" 1044 - version = "3.1.0" 1045 - source = "registry+https://github.com/rust-lang/crates.io-index" 1046 - checksum = "40cb1eb0cef3792900897b32c8282f6417bc978f6af46400a2f14bf0e649ae30" 1010 + version = "3.1.2" 1011 + source = "git+https://github.com/90-008/fjall.git#f21f459c79315cb1ac93762ffa8c05e038473c7a" 1047 1012 dependencies = [ 1048 1013 "byteorder-lite", 1049 1014 "byteview", ··· 1054 1019 "lz4_flex", 1055 1020 "tempfile", 1056 1021 "xxhash-rust", 1022 + "zstd", 1057 1023 ] 1058 1024 1059 1025 [[package]] ··· 2151 2117 2152 2118 [[package]] 2153 2119 name = "lsm-tree" 2154 - version = "3.1.0" 2155 - source = "registry+https://github.com/rust-lang/crates.io-index" 2156 - checksum = "fc5fa40c207eed45c811085aaa1b0a25fead22e298e286081cd4b98785fe759b" 2120 + version = "3.1.2" 2121 + source = "git+https://github.com/90-008/lsm-tree.git#85e22bb1875634bd25d19284d26c93d9edda4b4d" 2157 2122 dependencies = [ 2158 2123 "byteorder-lite", 2159 2124 "byteview", ··· 2169 2134 "tempfile", 2170 2135 "varint-rs", 2171 2136 "xxhash-rust", 2137 + "zstd", 2172 2138 ] 2173 2139 2174 2140 [[package]] 2175 2141 name = "lz4_flex" 2176 - version = "0.11.5" 2142 + version = "0.13.0" 2177 2143 source = "registry+https://github.com/rust-lang/crates.io-index" 2178 - checksum = "08ab2867e3eeeca90e844d1940eab391c9dc5228783db2ed999acbc0a9ed375a" 2144 + checksum = "db9a0d582c2874f68138a16ce1867e0ffde6c0bb0a0df85e1f36d04146db488a" 2179 2145 dependencies = [ 2180 2146 "twox-hash", 2181 2147 ] ··· 2401 2367 2402 2368 [[package]] 2403 2369 name = "once_cell" 2404 - version = "1.21.3" 2370 + version = "1.21.4" 2405 2371 source = "registry+https://github.com/rust-lang/crates.io-index" 2406 - checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" 2372 + checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" 2407 2373 dependencies = [ 2408 2374 "critical-section", 2409 2375 "portable-atomic", ··· 2644 2610 2645 2611 [[package]] 2646 2612 name = "quick_cache" 2647 - version = "0.6.18" 2613 + version = "0.6.21" 2648 2614 source = "registry+https://github.com/rust-lang/crates.io-index" 2649 - checksum = "7ada44a88ef953a3294f6eb55d2007ba44646015e18613d2f213016379203ef3" 2615 + checksum = "5a70b1b8b47e31d0498ecbc3c5470bb931399a8bfed1fd79d1717a61ce7f96e3" 2650 2616 dependencies = [ 2651 2617 "equivalent", 2652 2618 "hashbrown 0.16.1", ··· 3057 3023 3058 3024 [[package]] 3059 3025 name = "scc" 3060 - version = "3.6.9" 3026 + version = "3.6.10" 3061 3027 source = "registry+https://github.com/rust-lang/crates.io-index" 3062 - checksum = "45bb5ce9efd4a6e7b0f86c2697fe4c1d78d1f4e6d988c54b752d577cafe22fe8" 3028 + checksum = "ec3cea477634b93dee5619c6e437e2eae11c1e3730e59bad72af0e09dba9126f" 3063 3029 dependencies = [ 3064 3030 "saa", 3065 3031 "sdd", ··· 3067 3033 3068 3034 [[package]] 3069 3035 name = "schannel" 3070 - version = "0.1.28" 3036 + version = "0.1.29" 3071 3037 source = "registry+https://github.com/rust-lang/crates.io-index" 3072 - checksum = "891d81b926048e76efe18581bf793546b4c0eaf8448d72be8de2bbee5fd166e1" 3038 + checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" 3073 3039 dependencies = [ 3074 3040 "windows-sys 0.61.2", 3075 3041 ] ··· 3088 3054 3089 3055 [[package]] 3090 3056 name = "sdd" 3091 - version = "4.7.3" 3057 + version = "4.7.5" 3092 3058 source = "registry+https://github.com/rust-lang/crates.io-index" 3093 - checksum = "b21a75f5913ab130e4b369fb8693be25f29b983e2ecad4279df9bfa5dd8aaf3e" 3059 + checksum = "e6ca0e33fc1ae39e36b2d1fdfc8ee470b26397b642ff87572a59a36ff4f2340b" 3094 3060 3095 3061 [[package]] 3096 3062 name = "sec1" ··· 3260 3226 3261 3227 [[package]] 3262 3228 name = "serde_with" 3263 - version = "3.17.0" 3229 + version = "3.18.0" 3264 3230 source = "registry+https://github.com/rust-lang/crates.io-index" 3265 - checksum = "381b283ce7bc6b476d903296fb59d0d36633652b633b27f64db4fb46dcbfc3b9" 3231 + checksum = "dd5414fad8e6907dbdd5bc441a50ae8d6e26151a03b1de04d89a5576de61d01f" 3266 3232 dependencies = [ 3267 3233 "base64", 3268 3234 "chrono", ··· 3275 3241 3276 3242 [[package]] 3277 3243 name = "serde_with_macros" 3278 - version = "3.17.0" 3244 + version = "3.18.0" 3279 3245 source = "registry+https://github.com/rust-lang/crates.io-index" 3280 - checksum = "a6d4e30573c8cb306ed6ab1dca8423eec9a463ea0e155f45399455e0368b27e0" 3246 + checksum = "d3db8978e608f1fe7357e211969fd9abdcae80bac1ba7a3369bb7eb6b404eb65" 3281 3247 dependencies = [ 3282 - "darling 0.21.3", 3248 + "darling", 3283 3249 "proc-macro2", 3284 3250 "quote", 3285 3251 "syn", ··· 3531 3497 3532 3498 [[package]] 3533 3499 name = "tempfile" 3534 - version = "3.26.0" 3500 + version = "3.27.0" 3535 3501 source = "registry+https://github.com/rust-lang/crates.io-index" 3536 - checksum = "82a72c767771b47409d2345987fda8628641887d5466101319899796367354a0" 3502 + checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" 3537 3503 dependencies = [ 3538 3504 "fastrand", 3539 3505 "getrandom 0.4.2", ··· 3642 3608 3643 3609 [[package]] 3644 3610 name = "tinyvec" 3645 - version = "1.10.0" 3611 + version = "1.11.0" 3646 3612 source = "registry+https://github.com/rust-lang/crates.io-index" 3647 - checksum = "bfa5fdc3bce6191a1dbc8c02d5c8bffcf557bafa17c124c5264a458f1b0613fa" 3613 + checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" 3648 3614 dependencies = [ 3649 3615 "tinyvec_macros", 3650 3616 ] ··· 3821 3787 3822 3788 [[package]] 3823 3789 name = "tracing-subscriber" 3824 - version = "0.3.22" 3790 + version = "0.3.23" 3825 3791 source = "registry+https://github.com/rust-lang/crates.io-index" 3826 - checksum = "2f30143827ddab0d256fd843b7a66d164e9f271cfa0dde49142c5ca0ca291f1e" 3792 + checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" 3827 3793 dependencies = [ 3828 3794 "matchers", 3829 3795 "nu-ansi-term", ··· 4691 4657 4692 4658 [[package]] 4693 4659 name = "zerocopy" 4694 - version = "0.8.42" 4660 + version = "0.8.46" 4695 4661 source = "registry+https://github.com/rust-lang/crates.io-index" 4696 - checksum = "f2578b716f8a7a858b7f02d5bd870c14bf4ddbbcf3a4c05414ba6503640505e3" 4662 + checksum = "5c5030500cb2d66bdfbb4ebc9563be6ce7005a4b5d0f26be0c523870fe372ca6" 4697 4663 dependencies = [ 4698 4664 "zerocopy-derive", 4699 4665 ] 4700 4666 4701 4667 [[package]] 4702 4668 name = "zerocopy-derive" 4703 - version = "0.8.42" 4669 + version = "0.8.46" 4704 4670 source = "registry+https://github.com/rust-lang/crates.io-index" 4705 - checksum = "7e6cc098ea4d3bd6246687de65af3f920c430e236bee1e3bf2e441463f08a02f" 4671 + checksum = "a5f86989a046a79640b9d8867c823349a139367bda96549794fcc3313ce91f4e" 4706 4672 dependencies = [ 4707 4673 "proc-macro2", 4708 4674 "quote",
+6 -2
Cargo.toml
··· 18 18 serde_json = "1.0" 19 19 rmp-serde = { git = "https://github.com/90-008/msgpack-rust.git" } 20 20 21 - fjall = "3.1" 21 + fjall = { git = "https://github.com/90-008/fjall.git", features = ["lz4", "zstd"], version = "3.1.2" } 22 + lsm-tree = { git = "https://github.com/90-008/lsm-tree.git", features = ["lz4", "zstd"] } 22 23 serde_ipld_dagcbor = "0.6" 23 24 serde_urlencoded = "0.7" 24 25 ··· 49 50 rustls = { version = "0.23", features = ["aws-lc-rs"] } 50 51 tokio-tungstenite = { version = "0.28.0", features = ["rustls-tls-native-roots"] } 51 52 multibase = "0.9.2" 52 - lsm-tree = "3.1.0" 53 53 54 54 [dev-dependencies] 55 55 tempfile = "3.26.0" 56 56 57 57 [profile.dev] 58 58 opt-level = 1 59 + 60 + [profile.bench] 61 + opt-level = 3 62 + lto = "thin"
+23
src/api/mod.rs
··· 1 1 use crate::state::AppState; 2 + use axum::extract::State; 3 + use axum::routing::post; 2 4 use axum::{Router, routing::get}; 5 + use futures::FutureExt; 6 + use miette::IntoDiagnostic; 7 + use reqwest::StatusCode; 3 8 use std::{net::SocketAddr, sync::Arc}; 4 9 use tower_http::cors::CorsLayer; 5 10 use tower_http::trace::TraceLayer; ··· 15 20 let app = Router::new() 16 21 .route("/health", get(|| async { "OK" })) 17 22 .route("/stats", get(stats::get_stats)) 23 + .route("/_train_dict", post(handle_train_dict)) 18 24 .nest("/stream", stream::router()) 19 25 .merge(xrpc::router()) 20 26 .merge(filter::router()) ··· 37 43 .map_err(|e| miette::miette!("axum server error: {e}"))?; 38 44 39 45 Ok(()) 46 + } 47 + 48 + pub async fn handle_train_dict( 49 + State(state): State<Arc<AppState>>, 50 + ) -> Result<StatusCode, StatusCode> { 51 + let train = |name: &'static str| { 52 + let db = state.db.clone(); 53 + tokio::task::spawn_blocking(move || db.train_dict(name)) 54 + .map(|res| res.into_diagnostic().flatten()) 55 + }; 56 + let repos = train("repos"); 57 + let blocks = train("blocks"); 58 + let events = train("events"); 59 + 60 + tokio::try_join!(repos, blocks, events).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; 61 + 62 + Ok(StatusCode::OK) 40 63 } 41 64 42 65 pub async fn serve_debug(state: Arc<AppState>, port: u16) -> miette::Result<()> {
+3
src/config.rs
··· 8 8 #[derive(Debug, Clone, Copy, PartialEq, Eq)] 9 9 pub enum Compression { 10 10 Lz4, 11 + Zstd, 11 12 None, 12 13 } 13 14 ··· 16 17 fn from_str(s: &str) -> Result<Self> { 17 18 match s { 18 19 "lz4" => Ok(Self::Lz4), 20 + "zstd" => Ok(Self::Zstd), 19 21 "none" => Ok(Self::None), 20 22 _ => Err(miette::miette!("invalid compression type")), 21 23 } ··· 26 28 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 27 29 match self { 28 30 Self::Lz4 => write!(f, "lz4"), 31 + Self::Zstd => write!(f, "zstd"), 29 32 Self::None => write!(f, "none"), 30 33 } 31 34 }
+155 -7
src/db/mod.rs
··· 13 13 use scc::HashMap; 14 14 use smol_str::SmolStr; 15 15 16 + use std::cell::RefCell; 17 + use std::collections::HashSet; 16 18 use std::sync::Arc; 17 19 use std::sync::atomic::AtomicU64; 18 20 ··· 34 36 #[derive(Clone)] 35 37 pub struct Db { 36 38 pub inner: Arc<Database>, 39 + pub path: std::path::PathBuf, 37 40 pub repos: Keyspace, 38 41 pub records: Keyspace, 39 42 pub blocks: Keyspace, ··· 112 115 .manual_journal_persist(true) 113 116 .journal_compression(match cfg.journal_compression { 114 117 Compression::Lz4 => CompressionType::Lz4, 118 + Compression::Zstd => CompressionType::Zstd { level: 3 }, 115 119 Compression::None => CompressionType::None, 116 120 }) 117 121 .worker_threads(cfg.db_worker_threads) ··· 140 144 db.keyspace(name, move || opts).into_diagnostic() 141 145 }; 142 146 143 - let compression: CompressionType = match cfg.data_compression { 147 + let load_dict = |name: &str| -> Option<Arc<[u8]>> { 148 + let path = cfg.database_path.join(format!("dict_{name}.bin")); 149 + if path.exists() { 150 + if let Ok(bytes) = std::fs::read(&path) { 151 + tracing::info!( 152 + "loaded zstd dictionary for keyspace {name} ({} bytes)", 153 + bytes.len() 154 + ); 155 + return Some(bytes.into()); 156 + } 157 + } 158 + None 159 + }; 160 + let dicts = ["repos", "blocks", "events"].into_iter().fold( 161 + std::collections::HashMap::new(), 162 + |mut acc, name| { 163 + let Some(dict) = load_dict(name) else { 164 + return acc; 165 + }; 166 + acc.insert(name, dict); 167 + acc 168 + }, 169 + ); 170 + let get_compression = |name: &str, level: i32| match cfg.data_compression { 144 171 Compression::Lz4 => CompressionType::Lz4, 172 + Compression::Zstd => dicts 173 + .get(name) 174 + .map(|dict| CompressionType::ZstdDict { 175 + level, 176 + dict: dict.clone(), 177 + }) 178 + .unwrap_or_else(|| CompressionType::Zstd { level }), 145 179 Compression::None => CompressionType::None, 146 180 }; 147 181 ··· 156 190 // these block sizes work fine since we insert into repos constantly anyway 157 191 // whenever we update anything related to a repo, so the repos that arent 158 192 // being updated will be compacted away! 159 - .data_block_size_policy(BlockSizePolicy::new([kb(4), kb(4), kb(8), kb(64)])) 193 + .data_block_size_policy(BlockSizePolicy::new([kb(4), kb(4), kb(16), kb(64)])) 160 194 .data_block_compression_policy(CompressionPolicy::new([ 161 195 CompressionType::None, 162 196 CompressionType::None, 163 - compression, 197 + get_compression("repos", 3), 198 + get_compression("repos", 5), 164 199 ])) 165 200 // did plc are random so the interval wont rlly matter 166 201 .data_block_restart_interval_policy(RestartIntervalPolicy::new([2, 4])), ··· 197 232 // point reads are used a lot by stream, we know the blocks exist though 198 233 .expect_point_read_hits(true) 199 234 .max_memtable_size(mb(cfg.db_blocks_memtable_size_mb)) 200 - // 16 - 128 kb, as the newer blocks will be in the first level 235 + // 16 - 128 kb, as the newer blocks will be in the first level (or memtable) 201 236 // and any consumers will probably be streaming the newer events... 202 237 // and blocks are pretty big-ish like around 5kb... 203 238 // replaying will hit later levels so it will be slower but thats honestly ··· 208 243 // since we will be streaming them to consumers 209 244 .data_block_compression_policy(CompressionPolicy::new([ 210 245 CompressionType::None, 211 - compression, 246 + get_compression("blocks", 3), 247 + get_compression("blocks", 3), 248 + get_compression("blocks", 5), 212 249 ])) 213 250 .data_block_restart_interval_policy(RestartIntervalPolicy::new([8, 16, 32])), 214 251 )?; ··· 225 262 // cids arent compressable, most rkeys are TIDs so they will get compressed 226 263 // by prefix truncation anyway 227 264 .data_block_compression_policy(CompressionPolicy::disabled()) 228 - .data_block_restart_interval_policy(RestartIntervalPolicy::all(9)), 265 + .data_block_restart_interval_policy(RestartIntervalPolicy::new([9, 18])), 229 266 )?; 230 267 let cursors = open_ks( 231 268 "cursors", ··· 263 300 // we are streaming the new events to consumers so we dont want to compress them 264 301 .data_block_compression_policy(CompressionPolicy::new([ 265 302 CompressionType::None, 266 - compression, 303 + get_compression("events", 3), 304 + get_compression("events", 3), 305 + get_compression("events", 5), 267 306 ])) 268 307 // ids are int, we can prefix truncate a lot 269 308 .data_block_restart_interval_policy(RestartIntervalPolicy::new([64, 128])), ··· 349 388 350 389 Ok(Self { 351 390 inner: db, 391 + path: cfg.database_path.clone(), 352 392 repos, 353 393 records, 354 394 blocks, ··· 365 405 counts_map, 366 406 next_event_id: Arc::new(AtomicU64::new(last_id + 1)), 367 407 }) 408 + } 409 + 410 + pub fn train_dict(&self, ks_name: &str) -> Result<()> { 411 + let ks = match ks_name { 412 + "blocks" => &self.blocks, 413 + "events" => &self.events, 414 + "repos" => &self.repos, 415 + _ => miette::bail!("unknown keyspace for training: {ks_name}"), 416 + }; 417 + 418 + let dict_size = match ks_name { 419 + "blocks" => 32_768, 420 + "events" => 16_384, 421 + "repos" => 16_384, 422 + _ => 16_384, 423 + }; 424 + 425 + let samples = if ks_name == "blocks" { 426 + let mut collections = HashSet::new(); 427 + let mut iter = ks.iter(); 428 + while let Some(guard) = iter.next() { 429 + let key = guard.key().into_diagnostic()?; 430 + if let Some(sep_idx) = key.iter().position(|&b| b == keys::SEP) { 431 + if let Ok(col) = std::str::from_utf8(&key[..sep_idx]) { 432 + collections.insert(col.to_string()); 433 + let mut next_prefix = key[..sep_idx].to_vec(); 434 + next_prefix.push(keys::SEP + 1); 435 + iter = ks.range(next_prefix..); 436 + } else { 437 + break; 438 + } 439 + } else { 440 + break; 441 + } 442 + } 443 + 444 + let mut all_samples = Vec::new(); 445 + let mut seen_keys = HashSet::new(); 446 + let captured_keys = RefCell::new(Vec::new()); 447 + 448 + for t in collections { 449 + let prefix_str = format!("{t}|"); 450 + let prefix = prefix_str.as_bytes(); 451 + let mut end_prefix = prefix.to_vec(); 452 + if let Some(last) = end_prefix.last_mut() { 453 + *last = last.saturating_add(1); 454 + } 455 + 456 + let new = ks 457 + .sample_data_blocks(200, |first, last| { 458 + let passes = first < end_prefix.as_slice() 459 + && last >= prefix 460 + && !seen_keys.contains(&(first.to_vec(), last.to_vec())); 461 + if passes { 462 + captured_keys 463 + .borrow_mut() 464 + .push((first.to_vec(), last.to_vec())); 465 + } 466 + passes 467 + }) 468 + .into_diagnostic()?; 469 + 470 + for (s, keys) in new.into_iter().zip(captured_keys.borrow_mut().drain(..)) { 471 + if seen_keys.insert(keys) { 472 + all_samples.push(s.to_vec()); 473 + } 474 + } 475 + } 476 + all_samples 477 + } else { 478 + let mut seen_keys = HashSet::new(); 479 + let captured_keys = RefCell::new(Vec::new()); 480 + let new = ks 481 + .sample_data_blocks(1000, |first, last| { 482 + let passes = !seen_keys.contains(&(first.to_vec(), last.to_vec())); 483 + if passes { 484 + captured_keys 485 + .borrow_mut() 486 + .push((first.to_vec(), last.to_vec())); 487 + } 488 + passes 489 + }) 490 + .into_diagnostic()?; 491 + 492 + new.into_iter() 493 + .zip(captured_keys.into_inner()) 494 + .filter_map(|(s, keys)| seen_keys.insert(keys).then_some(s.to_vec())) 495 + .collect() 496 + }; 497 + 498 + if samples.is_empty() { 499 + miette::bail!("no samples found for keyspace {ks_name}, skipping training"); 500 + } 501 + 502 + tracing::info!( 503 + "training zstd dictionary for keyspace {ks_name} ({} samples, {dict_size} bytes limit)...", 504 + samples.len(), 505 + ); 506 + let dict_bytes = lsm_tree::train_zstd_dict(&samples, dict_size).into_diagnostic()?; 507 + let path = self.path.join(format!("dict_{ks_name}.bin")); 508 + std::fs::write(&path, &dict_bytes).into_diagnostic()?; 509 + tracing::info!( 510 + "saved zstd dictionary for keyspace {ks_name} ({} bytes) to {}. please restart the indexer to use the new dictionary.", 511 + dict_bytes.len(), 512 + path.display() 513 + ); 514 + 515 + Ok(()) 368 516 } 369 517 370 518 pub fn persist(&self) -> Result<()> {