tracks lexicons and how many times they appeared on the jetstream
feat(server): zstd compression
+178
-51
+5
-1
server/Cargo.toml
+5
-1
server/Cargo.toml
···
3
3
version = "0.1.0"
4
4
edition = "2024"
5
5
6
+
[features]
7
+
default = ["compress"]
8
+
compress = ["dep:zstd"]
9
+
6
10
[dependencies]
7
11
anyhow = "1.0"
8
12
async-trait = "0.1"
···
30
34
rayon = "1.10.0"
31
35
parking_lot = { version = "0.12", features = ["send_guard", "hardware-lock-elision"] }
32
36
rclite = "0.2.7"
33
-
zstd = "0.13.3"
37
+
zstd = { version = "0.13.3", optional = true, features = ["experimental"] }
34
38
35
39
[target.'cfg(target_env = "msvc")'.dependencies]
36
40
snmalloc-rs = "0.3.8"
+113
-16
server/src/db/handle.rs
+113
-16
server/src/db/handle.rs
···
14
14
use rclite::Arc;
15
15
use smol_str::SmolStr;
16
16
17
+
#[cfg(feature = "compress")]
18
+
use zstd::bulk::{Compressor as ZstdCompressor, Decompressor as ZstdDecompressor};
19
+
17
20
use crate::{
18
21
db::{EventRecord, NsidHit, block},
19
22
error::AppResult,
20
23
utils::{CLOCK, DefaultRateTracker, RateTracker, ReadVariableExt, varints_unsigned_encoded},
21
24
};
22
25
23
-
pub type ItemDecoder = block::ItemDecoder<Cursor<Slice>, NsidHit>;
24
-
pub type ItemEncoder = block::ItemEncoder<Vec<u8>, NsidHit>;
26
+
#[cfg(feature = "compress")]
27
+
thread_local! {
28
+
static COMPRESSOR: std::cell::RefCell<Option<ZstdCompressor<'static>>> = std::cell::RefCell::new(None);
29
+
static DECOMPRESSOR: std::cell::RefCell<Option<ZstdDecompressor<'static>>> = std::cell::RefCell::new(None);
30
+
}
31
+
32
+
type ItemDecoder = block::ItemDecoder<Cursor<Vec<u8>>, NsidHit>;
33
+
type ItemEncoder = block::ItemEncoder<Vec<u8>, NsidHit>;
25
34
pub type Item = block::Item<NsidHit>;
26
35
36
+
#[derive(Clone)]
37
+
pub enum Compression {
38
+
None,
39
+
#[cfg(feature = "compress")]
40
+
Zstd(ByteView),
41
+
}
42
+
43
+
impl Compression {
44
+
#[cfg(feature = "compress")]
45
+
fn get_dict(&self) -> Option<&ByteView> {
46
+
match self {
47
+
Compression::None => None,
48
+
Compression::Zstd(dict) => Some(dict),
49
+
}
50
+
}
51
+
}
52
+
27
53
pub struct Block {
28
54
pub written: usize,
29
55
pub key: ByteView,
···
36
62
buf: Arc<Mutex<Vec<EventRecord>>>,
37
63
last_insert: AtomicU64, // relaxed
38
64
eps: DefaultRateTracker,
65
+
compress: Compression,
39
66
}
40
67
41
68
impl Debug for LexiconHandle {
···
55
82
}
56
83
57
84
impl LexiconHandle {
58
-
pub fn new(keyspace: &Keyspace, nsid: &str) -> Self {
85
+
pub fn new(keyspace: &Keyspace, nsid: &str, compress: Compression) -> Self {
59
86
let opts = PartitionCreateOptions::default()
60
87
.block_size(1024 * 128)
61
-
.compression(fjall::CompressionType::Miniz(9));
88
+
.compression(fjall::CompressionType::Lz4);
62
89
Self {
63
90
tree: keyspace.open_partition(nsid, opts).unwrap(),
64
91
nsid: nsid.into(),
65
92
buf: Default::default(),
66
93
last_insert: AtomicU64::new(0),
67
94
eps: RateTracker::new(Duration::from_secs(10)),
95
+
compress,
68
96
}
69
97
}
70
98
99
+
#[cfg(feature = "compress")]
100
+
fn with_compressor<T>(&self, mut f: impl FnMut(&mut ZstdCompressor<'static>) -> T) -> T {
101
+
COMPRESSOR.with_borrow_mut(|compressor| {
102
+
if compressor.is_none() {
103
+
*compressor = Some({
104
+
let mut c = ZstdCompressor::new(9).expect("cant construct zstd compressor");
105
+
c.include_checksum(false).unwrap();
106
+
if let Some(dict) = self.compress.get_dict() {
107
+
c.set_dictionary(9, dict).expect("cant set dict");
108
+
}
109
+
c
110
+
});
111
+
}
112
+
// SAFETY: this is safe because we just initialized the compressor
113
+
f(unsafe { compressor.as_mut().unwrap_unchecked() })
114
+
})
115
+
}
116
+
117
+
#[cfg(feature = "compress")]
118
+
pub fn compress(&self, data: impl AsRef<[u8]>) -> std::io::Result<Vec<u8>> {
119
+
self.with_compressor(|compressor| compressor.compress(data.as_ref()))
120
+
}
121
+
122
+
#[cfg(feature = "compress")]
123
+
fn with_decompressor<T>(&self, mut f: impl FnMut(&mut ZstdDecompressor<'static>) -> T) -> T {
124
+
DECOMPRESSOR.with_borrow_mut(|decompressor| {
125
+
if decompressor.is_none() {
126
+
*decompressor = Some({
127
+
let mut d = ZstdDecompressor::new().expect("cant construct zstd decompressor");
128
+
if let Some(dict) = self.compress.get_dict() {
129
+
d.set_dictionary(dict).expect("cant set dict");
130
+
}
131
+
d
132
+
});
133
+
}
134
+
// SAFETY: this is safe because we just initialized the decompressor
135
+
f(unsafe { decompressor.as_mut().unwrap_unchecked() })
136
+
})
137
+
}
138
+
139
+
#[cfg(feature = "compress")]
140
+
pub fn decompress(&self, data: impl AsRef<[u8]>) -> std::io::Result<Vec<u8>> {
141
+
self.with_decompressor(|decompressor| {
142
+
decompressor.decompress(data.as_ref(), 1024 * 1024 * 20)
143
+
})
144
+
}
145
+
71
146
pub fn nsid(&self) -> &SmolStr {
72
147
&self.nsid
73
148
}
···
123
198
}
124
199
125
200
let start_blocks_size = blocks_to_compact.len();
126
-
let keys_to_delete = blocks_to_compact.iter().map(|(key, _)| key);
201
+
let keys_to_delete = blocks_to_compact
202
+
.iter()
203
+
.map(|(key, _)| key)
204
+
.cloned()
205
+
.collect_vec();
127
206
let mut all_items =
128
207
blocks_to_compact
129
-
.iter()
208
+
.into_iter()
130
209
.try_fold(Vec::new(), |mut acc, (key, value)| {
131
-
let mut timestamps = Cursor::new(key);
132
-
let start_timestamp = timestamps.read_varint()?;
133
-
let decoder = block::ItemDecoder::new(Cursor::new(value), start_timestamp)?;
210
+
let decoder = self.get_decoder_for(key, value)?;
134
211
let mut items = decoder.collect::<Result<Vec<_>, _>>()?;
135
212
acc.append(&mut items);
136
213
AppResult::Ok(acc)
···
149
226
.into_par_iter()
150
227
.map(|chunk| {
151
228
let count = chunk.len();
152
-
Self::encode_block_from_items(chunk, count)
229
+
self.encode_block_from_items(chunk, count)
153
230
})
154
231
.collect::<Result<Vec<_>, _>>()?;
155
232
let end_blocks_size = new_blocks.len();
···
173
250
}
174
251
175
252
pub fn encode_block_from_items(
253
+
&self,
176
254
items: impl IntoIterator<Item = Item>,
177
255
count: usize,
178
256
) -> AppResult<Block> {
···
204
282
.into());
205
283
}
206
284
if let (Some(start_timestamp), Some(end_timestamp)) = (start_timestamp, end_timestamp) {
207
-
let value = writer.finish()?;
285
+
let data = self.put_raw_block(writer.finish()?)?;
208
286
let key = varints_unsigned_encoded([start_timestamp, end_timestamp]);
209
-
return Ok(Block {
210
-
written,
211
-
key,
212
-
data: value,
213
-
});
287
+
return Ok(Block { written, key, data });
214
288
}
215
289
Err(std::io::Error::new(std::io::ErrorKind::WriteZero, "no items are in queue").into())
216
290
}
···
228
302
)
229
303
})
230
304
.collect()
305
+
}
306
+
307
+
pub fn get_raw_block(&self, value: Slice) -> std::io::Result<Vec<u8>> {
308
+
match &self.compress {
309
+
Compression::None => Ok(value.as_ref().into()),
310
+
#[cfg(feature = "compress")]
311
+
Compression::Zstd(_) => self.decompress(value),
312
+
}
313
+
}
314
+
315
+
pub fn put_raw_block(&self, value: Vec<u8>) -> std::io::Result<Vec<u8>> {
316
+
match &self.compress {
317
+
Compression::None => Ok(value),
318
+
#[cfg(feature = "compress")]
319
+
Compression::Zstd(_) => self.compress(value),
320
+
}
321
+
}
322
+
323
+
pub fn get_decoder_for(&self, key: Slice, value: Slice) -> AppResult<ItemDecoder> {
324
+
let mut timestamps = Cursor::new(key);
325
+
let start_timestamp = timestamps.read_varint()?;
326
+
let decoder = ItemDecoder::new(Cursor::new(self.get_raw_block(value)?), start_timestamp)?;
327
+
Ok(decoder)
231
328
}
232
329
}
+48
-23
server/src/db/mod.rs
+48
-23
server/src/db/mod.rs
···
3
3
fmt::Debug,
4
4
io::Cursor,
5
5
ops::{Bound, Deref, RangeBounds},
6
-
path::{Path, PathBuf},
6
+
path::Path,
7
7
time::Duration,
8
8
};
9
9
10
10
use byteview::StrView;
11
-
use fjall::{Config, Keyspace, Partition, PartitionCreateOptions};
11
+
use fjall::{Keyspace, Partition, PartitionCreateOptions};
12
12
use itertools::{Either, Itertools};
13
-
use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
13
+
use rayon::iter::{IntoParallelIterator, ParallelIterator};
14
14
use rclite::Arc;
15
15
use rkyv::{Archive, Deserialize, Serialize, rancor::Error};
16
16
use smol_str::{SmolStr, ToSmolStr};
···
18
18
use tokio_util::sync::CancellationToken;
19
19
20
20
use crate::{
21
-
db::handle::{ItemDecoder, LexiconHandle},
21
+
db::handle::{Compression, LexiconHandle},
22
22
error::{AppError, AppResult},
23
23
jetstream::JetstreamEvent,
24
24
utils::{RateTracker, ReadVariableExt, varints_unsigned_encoded},
···
77
77
78
78
pub struct DbConfig {
79
79
pub ks_config: fjall::Config,
80
+
#[cfg(feature = "compress")]
81
+
pub dict_path: std::path::PathBuf,
80
82
pub min_block_size: usize,
81
83
pub max_block_size: usize,
82
84
pub max_last_activity: u64,
···
98
100
fn default() -> Self {
99
101
Self {
100
102
ks_config: fjall::Config::default(),
103
+
#[cfg(feature = "compress")]
104
+
dict_path: "zstd_dict".parse().unwrap(),
101
105
min_block_size: 512,
102
106
max_block_size: 500_000,
103
107
max_last_activity: Duration::from_secs(10).as_nanos() as u64,
···
116
120
event_broadcaster: broadcast::Sender<(SmolStr, NsidCounts)>,
117
121
eps: RateTracker<100>,
118
122
cancel_token: CancellationToken,
123
+
compression: Compression,
119
124
}
120
125
121
126
impl Db {
122
127
pub fn new(cfg: DbConfig, cancel_token: CancellationToken) -> AppResult<Self> {
123
128
tracing::info!("opening db...");
124
129
let ks = cfg.ks_config.clone().open()?;
130
+
let _compression = Compression::None;
131
+
#[cfg(feature = "compress")]
132
+
let dict = std::fs::File::open(&cfg.dict_path).ok().and_then(|mut f| {
133
+
let meta = f.metadata().ok()?;
134
+
byteview::ByteView::from_reader(&mut f, meta.len() as usize).ok()
135
+
});
136
+
#[cfg(feature = "compress")]
137
+
let _compression = match dict {
138
+
Some(dict) => {
139
+
tracing::info!(
140
+
"using zstd compression with dict from {}",
141
+
cfg.dict_path.to_string_lossy()
142
+
);
143
+
Compression::Zstd(dict)
144
+
}
145
+
None => Compression::None,
146
+
};
125
147
Ok(Self {
126
148
cfg,
127
149
hits: Default::default(),
···
136
158
event_broadcaster: broadcast::channel(1000).0,
137
159
eps: RateTracker::new(Duration::from_secs(1)),
138
160
cancel_token,
161
+
compression: _compression,
139
162
})
140
163
}
141
164
···
213
236
.into_par_iter()
214
237
.map(|(i, items, handle)| {
215
238
let count = items.len();
216
-
let block = LexiconHandle::encode_block_from_items(items, count)?;
239
+
let block = handle.encode_block_from_items(items, count)?;
217
240
tracing::info!(
218
241
"{}: encoded block with {} items",
219
242
handle.nsid(),
···
282
305
Some(handle) => handle.clone(),
283
306
None => {
284
307
if self.ks.partition_exists(nsid.as_ref()) {
285
-
let handle = Arc::new(LexiconHandle::new(&self.ks, nsid.as_ref()));
308
+
let handle = Arc::new(LexiconHandle::new(
309
+
&self.ks,
310
+
nsid.as_ref(),
311
+
self.compression.clone(),
312
+
));
286
313
let _ = self.hits.insert(SmolStr::new(nsid), handle.clone());
287
314
handle
288
315
} else {
···
295
322
296
323
#[inline(always)]
297
324
fn ensure_handle(&self, nsid: &SmolStr) -> impl Deref<Target = Arc<LexiconHandle>> + use<'_> {
298
-
self.hits
299
-
.entry(nsid.clone())
300
-
.or_insert_with(|| Arc::new(LexiconHandle::new(&self.ks, &nsid)))
325
+
self.hits.entry(nsid.clone()).or_insert_with(|| {
326
+
Arc::new(LexiconHandle::new(
327
+
&self.ks,
328
+
&nsid,
329
+
self.compression.clone(),
330
+
))
331
+
})
301
332
}
302
333
303
334
pub fn ingest_events(&self, events: impl Iterator<Item = EventRecord>) -> AppResult<()> {
···
366
397
};
367
398
let block_lens = handle.iter().rev().try_fold(Vec::new(), |mut acc, item| {
368
399
let (key, value) = item?;
369
-
let mut timestamps = Cursor::new(key);
370
-
let start_timestamp = timestamps.read_varint()?;
371
-
let decoder = ItemDecoder::new(Cursor::new(value), start_timestamp)?;
400
+
let decoder = handle.get_decoder_for(key, value)?;
372
401
acc.push(decoder.item_count());
373
402
AppResult::Ok(acc)
374
403
})?;
···
380
409
})
381
410
}
382
411
383
-
// train zstd dict with 100 blocks from every lexicon
412
+
// train zstd dict with 1000 blocks from every lexicon
413
+
#[cfg(feature = "compress")]
384
414
pub fn train_zstd_dict(&self) -> AppResult<Vec<u8>> {
385
415
let samples = self
386
416
.get_nsids()
···
388
418
.map(|handle| {
389
419
handle
390
420
.iter()
391
-
.rev()
392
-
.map(|res| {
421
+
.map(move |res| {
393
422
res.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
394
-
.map(|(_, value)| Cursor::new(value))
423
+
.and_then(|(_, value)| Ok(Cursor::new(handle.get_raw_block(value)?)))
395
424
})
396
425
.take(1000)
397
426
})
···
420
449
return Either::Right(std::iter::empty());
421
450
};
422
451
423
-
let map_block = move |(key, val)| {
424
-
let mut key_reader = Cursor::new(key);
425
-
let start_timestamp = key_reader.read_varint::<u64>()?;
426
-
if start_timestamp < start_limit {
427
-
return Ok(None);
428
-
}
429
-
let items = handle::ItemDecoder::new(Cursor::new(val), start_timestamp)?
452
+
let map_block = |(key, val)| {
453
+
let decoder = handle.get_decoder_for(key, val)?;
454
+
let items = decoder
430
455
.take_while(move |item| {
431
456
item.as_ref().map_or(true, |item| {
432
457
item.timestamp <= end_limit && item.timestamp >= start_limit
+11
-11
server/src/main.rs
+11
-11
server/src/main.rs
···
53
53
debug();
54
54
return;
55
55
}
56
-
Some("traindict") => {
57
-
train_zstd_dict();
58
-
return;
59
-
}
60
56
Some(x) => {
61
57
tracing::error!("unknown command: {}", x);
62
58
return;
···
211
207
db.sync(true).expect("cant sync db");
212
208
}
213
209
214
-
fn train_zstd_dict() {
215
-
let db = Db::new(DbConfig::default(), CancellationToken::new()).expect("couldnt create db");
216
-
let dict_data = db.train_zstd_dict().expect("cant train zstd dict");
217
-
std::fs::write("zstd_dict", dict_data).expect("cant save zstd dict")
218
-
}
219
-
220
210
fn debug() {
221
211
let db = Db::new(DbConfig::default(), CancellationToken::new()).expect("couldnt create db");
222
212
let info = db.info().expect("cant get db info");
···
246
236
DbConfig::default().ks(|c| {
247
237
c.max_journaling_size(u64::MAX)
248
238
.max_write_buffer_size(u64::MAX)
239
+
.compaction_workers(rayon::current_num_threads() * 4)
240
+
.flush_workers(rayon::current_num_threads() * 4)
249
241
}),
250
242
CancellationToken::new(),
251
243
)
···
269
261
270
262
fn migrate() {
271
263
let cancel_token = CancellationToken::new();
264
+
272
265
let from = Arc::new(
273
266
Db::new(
274
267
DbConfig::default().path(".fjall_data_from"),
···
276
269
)
277
270
.expect("couldnt create db"),
278
271
);
272
+
#[cfg(feature = "compress")]
273
+
std::fs::write(
274
+
"zstd_dict",
275
+
from.train_zstd_dict().expect("cant get zstd dict"),
276
+
)
277
+
.expect("cant write zstd dict");
278
+
279
279
let to = Arc::new(
280
280
Db::new(
281
281
DbConfig::default().path(".fjall_data_to").ks(|c| {
···
290
290
);
291
291
292
292
let nsids = from.get_nsids().collect::<Vec<_>>();
293
-
let eps_thread = std::thread::spawn({
293
+
let _eps_thread = std::thread::spawn({
294
294
let to = to.clone();
295
295
move || {
296
296
loop {