1use std::collections::HashSet;
16use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
17use std::sync::Arc;
18
19use crossbeam_skiplist::SkipMap;
20use dashmap::mapref::entry::Entry;
21use dashmap::DashMap;
22
23use chronix_core::types::{Point, SeriesKey, Timestamp};
24
25use crate::memtable::error::{MemtableError, Result};
26use crate::memtable::interner::StringInterner;
27use crate::memtable::key::{MemtableEntry, MemtableKey, NODE_OVERHEAD};
28
29pub type TagSet = Arc<[(Arc<str>, Arc<str>)]>;
31
32pub struct Memtable {
43 data: SkipMap<MemtableKey, MemtableEntry>,
45 interner: StringInterner,
47 measurement_index: DashMap<Arc<str>, HashSet<u64>>,
56 series_tags: DashMap<Arc<str>, TagSet>,
59 estimated_size: AtomicUsize,
61 index_size: AtomicUsize,
67 min_wal_seq: AtomicU64,
69 max_wal_seq: AtomicU64,
71 frozen: AtomicBool,
73}
74
75impl std::fmt::Debug for Memtable {
76 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77 f.debug_struct("Memtable")
78 .field("len", &self.data.len())
79 .field(
80 "estimated_size",
81 &self.estimated_size.load(Ordering::Relaxed),
82 )
83 .field("interned_strings", &self.interner.len())
84 .field("min_wal_seq", &self.min_wal_seq())
85 .field("max_wal_seq", &self.max_wal_seq())
86 .field("frozen", &self.frozen.load(Ordering::Relaxed))
87 .finish()
88 }
89}
90
91impl Memtable {
92 #[must_use]
94 pub fn new() -> Self {
95 Self {
96 data: SkipMap::new(),
97 interner: StringInterner::new(),
98 measurement_index: DashMap::new(),
99 series_tags: DashMap::new(),
100 estimated_size: AtomicUsize::new(0),
101 index_size: AtomicUsize::new(0),
102 min_wal_seq: AtomicU64::new(u64::MAX),
103 max_wal_seq: AtomicU64::new(0),
104 frozen: AtomicBool::new(false),
105 }
106 }
107
108 pub fn insert(&self, point: &Point) -> Result<()> {
118 if self.frozen.load(Ordering::Acquire) {
119 return Err(MemtableError::Frozen);
120 }
121
122 let canonical: Arc<str> = self.interner.intern(point.series_key().canonical_form());
123 let key = MemtableKey::new(point.series_key().hash_fnv(), point.timestamp(), canonical);
124
125 let measurement = self.interner.intern(point.series_key().measurement());
126 let measurement_for_index = Arc::clone(&measurement);
127 let tags = self.series_tag_set(&key.series_canonical, point.series_key());
128
129 let entry = MemtableEntry {
130 measurement,
131 tags,
132 fields: point
133 .fields()
134 .iter()
135 .map(|(k, v)| (Arc::clone(k), v.clone()))
136 .collect(),
137 };
138
139 let entry_size =
140 entry.estimated_size() + std::mem::size_of::<MemtableKey>() + NODE_OVERHEAD;
141
142 if let Some(existing) = self.data.get(&key) {
150 let old_size = existing.value().estimated_size()
151 + std::mem::size_of::<MemtableKey>()
152 + NODE_OVERHEAD;
153 let _ =
154 self.estimated_size
155 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
156 Some(current.saturating_sub(old_size))
157 });
158 }
159
160 self.data.insert(key, entry);
161 let new_size = self.estimated_size.fetch_add(entry_size, Ordering::Relaxed) + entry_size;
162
163 metrics::gauge!("chronix_memtable_memory_bytes").set(new_size as f64);
166 metrics::gauge!("chronix_memtable_entries").set(self.data.len() as f64);
167
168 {
172 let hash = point.series_key().hash_fnv();
173 let delta = self.index_insert(measurement_for_index, hash);
174 if delta > 0 {
175 self.index_size.fetch_add(delta, Ordering::Relaxed);
176 }
177 }
178
179 Ok(())
180 }
181
182 pub fn insert_with_wal_seq(&self, point: &Point, wal_seq: u64) -> Result<()> {
191 self.insert(point)?;
192 self.update_wal_seq(wal_seq);
193 Ok(())
194 }
195
196 pub fn insert_batch(&self, points: &[Point]) -> Result<usize> {
211 if self.frozen.load(Ordering::Acquire) {
212 return Err(MemtableError::Frozen);
213 }
214
215 let mut total_size_delta: i64 = 0;
216 let mut new_index_entries: Vec<(Arc<str>, u64)> = Vec::new();
218
219 let mut prev_hash: u64 = 0;
223 let mut prev_ts: i64 = i64::MIN;
224
225 for point in points {
226 let canonical: Arc<str> = self.interner.intern(point.series_key().canonical_form());
227 let key = MemtableKey::new(point.series_key().hash_fnv(), point.timestamp(), canonical);
228
229 let measurement = self.interner.intern(point.series_key().measurement());
230 let tags = self.series_tag_set(&key.series_canonical, point.series_key());
231
232 let entry = MemtableEntry {
233 measurement: Arc::clone(&measurement),
234 tags,
235 fields: point
236 .fields()
237 .iter()
238 .map(|(k, v)| (Arc::clone(k), v.clone()))
239 .collect(),
240 };
241
242 let entry_size =
243 entry.estimated_size() + std::mem::size_of::<MemtableKey>() + NODE_OVERHEAD;
244
245 let cur_hash = point.series_key().hash_fnv();
250 let cur_ts = point.timestamp();
251 let may_overwrite = cur_hash == prev_hash && cur_ts == prev_ts;
252
253 if may_overwrite {
254 if let Some(existing) = self.data.get(&key) {
255 let old_size = existing.value().estimated_size()
256 + std::mem::size_of::<MemtableKey>()
257 + NODE_OVERHEAD;
258 total_size_delta -= old_size.min(i64::MAX as usize) as i64;
259 }
260 }
261
262 prev_hash = cur_hash;
263 prev_ts = cur_ts;
264
265 self.data.insert(key, entry);
266 total_size_delta += entry_size.min(i64::MAX as usize) as i64;
268
269 new_index_entries.push((measurement, cur_hash));
270 }
271
272 if total_size_delta >= 0 {
274 self.estimated_size
275 .fetch_add(total_size_delta as usize, Ordering::Relaxed);
276 } else {
277 let _ =
278 self.estimated_size
279 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
280 Some(current.saturating_sub(total_size_delta.unsigned_abs() as usize))
281 });
282 }
283
284 let mut index_delta = 0usize;
286 for (measurement, hash) in new_index_entries {
287 index_delta += self.index_insert(measurement, hash);
288 }
289 if index_delta > 0 {
290 self.index_size.fetch_add(index_delta, Ordering::Relaxed);
291 }
292
293 metrics::gauge!("chronix_memtable_memory_bytes")
295 .set(self.estimated_size.load(Ordering::Relaxed) as f64);
296 metrics::gauge!("chronix_memtable_entries").set(self.data.len() as f64);
297
298 Ok(points.len())
299 }
300
301 pub fn insert_batch_with_wal_seq(&self, points: &[Point], wal_seq: u64) -> Result<usize> {
310 let count = self.insert_batch(points)?;
311 self.update_wal_seq(wal_seq);
312 Ok(count)
313 }
314
315 fn update_wal_seq(&self, seq: u64) {
317 let mut current = self.min_wal_seq.load(Ordering::Relaxed);
319 while seq < current {
320 match self.min_wal_seq.compare_exchange_weak(
321 current,
322 seq,
323 Ordering::Release,
324 Ordering::Relaxed,
325 ) {
326 Ok(_) => break,
327 Err(actual) => current = actual,
328 }
329 }
330
331 let mut current = self.max_wal_seq.load(Ordering::Relaxed);
333 while seq > current {
334 match self.max_wal_seq.compare_exchange_weak(
335 current,
336 seq,
337 Ordering::Release,
338 Ordering::Relaxed,
339 ) {
340 Ok(_) => break,
341 Err(actual) => current = actual,
342 }
343 }
344 }
345
346 #[must_use]
353 pub fn scan(&self, series_key: &SeriesKey, min_ts: Timestamp, max_ts: Timestamp) -> Vec<Point> {
354 let hash = series_key.hash_fnv();
355 let canonical = series_key.canonical_form();
356 let start = MemtableKey::new(hash, min_ts, Arc::from(canonical));
359 let end = MemtableKey::new(hash, max_ts, Arc::from(canonical));
360
361 self.data
362 .range(start..=end)
363 .filter(|entry| &*entry.key().series_canonical == canonical)
364 .filter_map(|entry| Self::entry_to_point(entry.key(), entry.value()))
365 .collect()
366 }
367
368 #[must_use]
377 pub fn scan_range(&self, min_ts: Timestamp, max_ts: Timestamp) -> Vec<Point> {
378 self.data
381 .iter()
382 .filter(|entry| {
383 let ts = entry.key().timestamp;
384 ts >= min_ts && ts <= max_ts
385 })
386 .filter_map(|entry| Self::entry_to_point(entry.key(), entry.value()))
387 .collect()
388 }
389
390 #[must_use]
393 pub fn scan_all(&self) -> Vec<Point> {
394 self.data
395 .iter()
396 .filter_map(|entry| Self::entry_to_point(entry.key(), entry.value()))
397 .collect()
398 }
399
400 #[must_use]
407 pub fn scan_measurement(
408 &self,
409 measurement: &str,
410 min_ts: Timestamp,
411 max_ts: Timestamp,
412 ) -> Vec<Point> {
413 let hashes = match self.measurement_index.get(measurement) {
414 Some(set) => set.iter().copied().collect::<Vec<_>>(),
415 None => return Vec::new(),
416 };
417
418 let mut results = Vec::new();
419 for hash in hashes {
420 let start = MemtableKey::lower_bound(hash, min_ts);
422 let end = MemtableKey::upper_bound(hash, max_ts);
423 results.extend(
424 self.data
425 .range(start..=end)
426 .filter(|entry| {
427 let ts = entry.key().timestamp;
428 ts >= min_ts && ts <= max_ts && &*entry.value().measurement == measurement
429 })
430 .filter_map(|entry| Self::entry_to_point(entry.key(), entry.value())),
431 );
432 }
433 results
434 }
435
436 #[must_use]
438 pub fn len(&self) -> usize {
439 self.data.len()
440 }
441
442 #[must_use]
444 pub fn is_empty(&self) -> bool {
445 self.data.is_empty()
446 }
447
448 #[must_use]
458 pub fn interner_bytes(&self) -> usize {
459 self.interner.memory_bytes()
460 }
461
462 #[must_use]
464 pub fn estimated_size(&self) -> usize {
465 self.estimated_size.load(Ordering::Relaxed) + self.index_size.load(Ordering::Relaxed)
466 }
467
468 fn series_tag_set(
474 &self,
475 canonical: &Arc<str>,
476 series_key: &SeriesKey,
477 ) -> Arc<[(Arc<str>, Arc<str>)]> {
478 if let Some(existing) = self.series_tags.get(canonical) {
479 return Arc::clone(existing.value());
480 }
481 let tags: Arc<[(Arc<str>, Arc<str>)]> = series_key
482 .tags()
483 .iter()
484 .map(|(k, v)| (self.interner.intern(k), self.interner.intern(v)))
485 .collect();
486 let mut charged = 0usize;
487 let stored = self
488 .series_tags
489 .entry(Arc::clone(canonical))
490 .or_insert_with(|| {
491 charged = canonical.len()
495 + tags.len() * std::mem::size_of::<(Arc<str>, Arc<str>)>()
496 + tags.iter().map(|(k, v)| k.len() + v.len()).sum::<usize>()
497 + 2 * std::mem::size_of::<usize>()
498 + 64;
499 tags
500 })
501 .clone();
502 if charged > 0 {
503 self.estimated_size.fetch_add(charged, Ordering::Relaxed);
504 }
505 stored
506 }
507
508 fn index_insert(&self, measurement: Arc<str>, hash: u64) -> usize {
517 const SET_OVERHEAD: usize = 56;
518 match self.measurement_index.entry(measurement) {
519 Entry::Occupied(mut occupied) => {
520 let before = occupied.get().capacity();
521 occupied.get_mut().insert(hash);
522 (occupied.get().capacity() - before) * std::mem::size_of::<u64>()
523 }
524 Entry::Vacant(vacant) => {
525 let key_bytes = std::mem::size_of::<usize>() * 2 + vacant.key().len();
526 let mut set = HashSet::new();
527 set.insert(hash);
528 let set_bytes = set.capacity() * std::mem::size_of::<u64>() + SET_OVERHEAD;
529 vacant.insert(set);
530 key_bytes + set_bytes
531 }
532 }
533 }
534
535 #[must_use]
539 pub fn min_wal_seq(&self) -> Option<u64> {
540 let val = self.min_wal_seq.load(Ordering::Acquire);
541 if val == u64::MAX {
542 None
543 } else {
544 Some(val)
545 }
546 }
547
548 #[must_use]
552 pub fn max_wal_seq(&self) -> Option<u64> {
553 let val = self.max_wal_seq.load(Ordering::Acquire);
554 if val == 0 && self.min_wal_seq.load(Ordering::Acquire) == u64::MAX {
555 None
556 } else {
557 Some(val)
558 }
559 }
560
561 #[must_use]
563 pub fn is_frozen(&self) -> bool {
564 self.frozen.load(Ordering::Acquire)
565 }
566
567 pub fn freeze(&self) {
572 self.frozen.store(true, Ordering::Release);
573 }
574
575 #[must_use]
577 pub fn to_points(&self) -> Vec<Point> {
578 self.scan_all()
579 }
580
581 fn entry_to_point(key: &MemtableKey, entry: &MemtableEntry) -> Option<Point> {
586 let tags = entry
587 .tags
588 .iter()
589 .map(|(k, v)| (k.to_string(), v.to_string()))
590 .collect();
591 let series_key = SeriesKey::new(entry.measurement.to_string(), tags)
592 .map_err(|e| {
593 tracing::error!(
596 hash = key.series_key_hash,
597 error = %e,
598 "DATA LOSS: failed to reconstruct series key during flush"
599 );
600 metrics::counter!("chronix_memtable_dropped_points_total").increment(1);
601 e
602 })
603 .ok()?;
604 Point::new(
605 series_key,
606 entry
607 .fields
608 .iter()
609 .map(|(k, v)| (k.to_string(), v.clone()))
610 .collect(),
611 key.timestamp,
612 )
613 .map_err(|e| {
614 tracing::error!(
615 hash = key.series_key_hash,
616 timestamp = key.timestamp,
617 error = %e,
618 "DATA LOSS: failed to reconstruct point during flush"
619 );
620 metrics::counter!("chronix_memtable_dropped_points_total").increment(1);
621 e
622 })
623 .ok()
624 }
625}
626
627impl Default for Memtable {
628 fn default() -> Self {
629 Self::new()
630 }
631}
632
633#[cfg(test)]
638mod tests {
639 use super::*;
640 use std::collections::BTreeMap;
641
642 use chronix_core::FieldValue;
643
644 fn make_point(measurement: &str, host: &str, ts: i64, value: f64) -> Point {
645 let tags: BTreeMap<String, String> = [("host".to_string(), host.to_string())]
646 .into_iter()
647 .collect();
648 let series_key = SeriesKey::new(measurement.to_string(), tags).unwrap();
649 let fields: BTreeMap<String, FieldValue> = [("value".to_string(), FieldValue::F64(value))]
650 .into_iter()
651 .collect();
652 Point::new(series_key, fields, ts).unwrap()
653 }
654
655 #[test]
656 fn insert_and_scan_all() {
657 let mt = Memtable::new();
658 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
659 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
660 mt.insert(&make_point("cpu", "a", 300, 3.0)).unwrap();
661
662 let points = mt.scan_all();
663 assert_eq!(points.len(), 3);
664 assert_eq!(points[0].timestamp(), 100);
665 assert_eq!(points[1].timestamp(), 200);
666 assert_eq!(points[2].timestamp(), 300);
667 }
668
669 #[test]
670 fn insert_and_scan_series() {
671 let mt = Memtable::new();
672 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
673 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
674 mt.insert(&make_point("cpu", "b", 150, 9.0)).unwrap();
675
676 let key_a = SeriesKey::new(
677 "cpu".to_string(),
678 [("host".to_string(), "a".to_string())]
679 .into_iter()
680 .collect(),
681 )
682 .unwrap();
683
684 let points = mt.scan(&key_a, 0, 300);
685 assert_eq!(points.len(), 2);
686 assert!(points
687 .iter()
688 .all(|p| p.series_key().tag("host") == Some("a")));
689 }
690
691 #[test]
692 fn scan_time_range() {
693 let mt = Memtable::new();
694 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
695 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
696 mt.insert(&make_point("cpu", "a", 300, 3.0)).unwrap();
697
698 let key = SeriesKey::new(
699 "cpu".to_string(),
700 [("host".to_string(), "a".to_string())]
701 .into_iter()
702 .collect(),
703 )
704 .unwrap();
705
706 let points = mt.scan(&key, 150, 250);
707 assert_eq!(points.len(), 1);
708 assert_eq!(points[0].timestamp(), 200);
709 }
710
711 #[test]
712 fn last_write_wins() {
713 let mt = Memtable::new();
714 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
715 mt.insert(&make_point("cpu", "a", 100, 99.0)).unwrap();
716
717 let points = mt.scan_all();
718 assert_eq!(points.len(), 1);
719 assert_eq!(points[0].field("value"), Some(&FieldValue::F64(99.0)));
720 }
721
722 #[test]
723 fn empty_scan() {
724 let mt = Memtable::new();
725 let key = SeriesKey::new("cpu".to_string(), BTreeMap::new()).unwrap();
726 let points = mt.scan(&key, 0, 1000);
727 assert!(points.is_empty());
728 }
729
730 #[test]
733 fn index_accounting_matches_a_full_walk() {
734 let mt = Memtable::new();
735 for i in 0..500u64 {
736 let measurement = format!("m{}", i % 7);
737 let series = format!("s{}", i % 53);
738 mt.insert(&make_point(&measurement, &series, i as i64, 1.0))
739 .unwrap();
740 }
741
742 let batch: Vec<_> = (500..800u64)
744 .map(|i| {
745 make_point(
746 &format!("m{}", i % 11),
747 &format!("s{}", i % 97),
748 i as i64,
749 2.0,
750 )
751 })
752 .collect();
753 mt.insert_batch(&batch).unwrap();
754
755 let walked: usize = mt
756 .measurement_index
757 .iter()
758 .map(|entry| {
759 let key_bytes = std::mem::size_of::<usize>() * 2 + entry.key().len();
760 let set_bytes = entry.value().capacity() * std::mem::size_of::<u64>() + 56;
761 key_bytes + set_bytes
762 })
763 .sum();
764
765 assert_eq!(
766 mt.index_size.load(Ordering::Relaxed),
767 walked,
768 "incremental index accounting drifted from a full walk"
769 );
770 }
771
772 #[test]
773 fn estimated_size_increases() {
774 let mt = Memtable::new();
775 assert_eq!(mt.estimated_size(), 0);
776 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
777 assert!(mt.estimated_size() > 0);
778 let size_after_one = mt.estimated_size();
779 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
780 assert!(mt.estimated_size() > size_after_one);
781 }
782
783 #[test]
784 fn wal_seq_tracking() {
785 let mt = Memtable::new();
786 assert_eq!(mt.min_wal_seq(), None);
787 assert_eq!(mt.max_wal_seq(), None);
788
789 mt.insert_with_wal_seq(&make_point("cpu", "a", 100, 1.0), 10)
790 .unwrap();
791 assert_eq!(mt.min_wal_seq(), Some(10));
792 assert_eq!(mt.max_wal_seq(), Some(10));
793
794 mt.insert_with_wal_seq(&make_point("cpu", "a", 200, 2.0), 5)
795 .unwrap();
796 assert_eq!(mt.min_wal_seq(), Some(5));
797 assert_eq!(mt.max_wal_seq(), Some(10));
798
799 mt.insert_with_wal_seq(&make_point("cpu", "a", 300, 3.0), 20)
800 .unwrap();
801 assert_eq!(mt.min_wal_seq(), Some(5));
802 assert_eq!(mt.max_wal_seq(), Some(20));
803 }
804
805 #[test]
806 fn freeze_prevents_inserts() {
807 let mt = Memtable::new();
808 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
809 mt.freeze();
810 assert!(mt.is_frozen());
811
812 let result = mt.insert(&make_point("cpu", "a", 200, 2.0));
813 assert!(result.is_err());
814
815 let points = mt.scan_all();
817 assert_eq!(points.len(), 1);
818 }
819
820 #[test]
821 fn len_and_is_empty() {
822 let mt = Memtable::new();
823 assert!(mt.is_empty());
824 assert_eq!(mt.len(), 0);
825
826 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
827 assert!(!mt.is_empty());
828 assert_eq!(mt.len(), 1);
829 }
830
831 #[test]
832 fn multiple_series() {
833 let mt = Memtable::new();
834 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
835 mt.insert(&make_point("cpu", "b", 100, 2.0)).unwrap();
836 mt.insert(&make_point("mem", "a", 100, 3.0)).unwrap();
837
838 assert_eq!(mt.len(), 3);
839
840 let all = mt.scan_all();
841 assert_eq!(all.len(), 3);
842 }
843
844 #[test]
845 fn concurrent_inserts() {
846 use std::sync::Arc;
847 use std::thread;
848
849 let mt = Arc::new(Memtable::new());
850 let mut handles = Vec::new();
851
852 for thread_id in 0..4 {
853 let mt = Arc::clone(&mt);
854 handles.push(thread::spawn(move || {
855 for i in 0..100 {
856 let ts = i64::from(thread_id * 1000 + i);
857 let host = format!("host-{thread_id}");
858 mt.insert(&make_point("cpu", &host, ts, ts as f64)).unwrap();
859 }
860 }));
861 }
862
863 for h in handles {
864 h.join().unwrap();
865 }
866
867 assert_eq!(mt.len(), 400);
868 let all = mt.scan_all();
869 assert_eq!(all.len(), 400);
870
871 let mut prev_key: Option<MemtableKey> = None;
873 for p in &all {
874 let key = MemtableKey::new(
875 p.series_key().hash_fnv(),
876 p.timestamp(),
877 Arc::from(p.series_key().canonical_form()),
878 );
879 if let Some(prev) = prev_key {
880 assert!(key >= prev, "entries should be ordered");
881 }
882 prev_key = Some(key);
883 }
884 }
885
886 #[allow(deprecated)]
887 #[test]
888 fn to_points_is_scan_all() {
889 let mt = Memtable::new();
890 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
891 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
892
893 let scan = mt.scan_all();
894 let points = mt.to_points();
895 assert_eq!(scan.len(), points.len());
896 for (s, d) in scan.iter().zip(points.iter()) {
897 assert_eq!(s.timestamp(), d.timestamp());
898 assert_eq!(s.fields(), d.fields());
899 }
900 }
901
902 #[test]
903 fn scan_range_across_series() {
904 let mt = Memtable::new();
905 mt.insert(&make_point("cpu", "a", 50, 1.0)).unwrap();
906 mt.insert(&make_point("cpu", "a", 100, 2.0)).unwrap();
907 mt.insert(&make_point("cpu", "b", 75, 3.0)).unwrap();
908 mt.insert(&make_point("cpu", "b", 200, 4.0)).unwrap();
909
910 let points = mt.scan_range(60, 150);
911 assert_eq!(points.len(), 2);
913 for p in &points {
914 assert!(
915 p.timestamp() >= 60 && p.timestamp() <= 150,
916 "timestamp {} should be in [60, 150]",
917 p.timestamp()
918 );
919 }
920 }
921
922 #[test]
923 fn scan_measurement_filters_by_name() {
924 let mt = Memtable::new();
925 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
926 mt.insert(&make_point("cpu", "b", 200, 2.0)).unwrap();
927 mt.insert(&make_point("mem", "a", 150, 3.0)).unwrap();
928 mt.insert(&make_point("mem", "a", 300, 4.0)).unwrap();
929 mt.insert(&make_point("disk", "c", 250, 5.0)).unwrap();
930
931 let cpu = mt.scan_measurement("cpu", 0, i64::MAX);
933 assert_eq!(cpu.len(), 2);
934 assert!(cpu.iter().all(|p| p.series_key().measurement() == "cpu"));
935
936 let mem = mt.scan_measurement("mem", 0, i64::MAX);
938 assert_eq!(mem.len(), 2);
939 assert!(mem.iter().all(|p| p.series_key().measurement() == "mem"));
940
941 let disk = mt.scan_measurement("disk", 0, i64::MAX);
943 assert_eq!(disk.len(), 1);
944
945 let none = mt.scan_measurement("net", 0, i64::MAX);
947 assert!(none.is_empty());
948 }
949
950 #[test]
951 fn scan_measurement_respects_time_range() {
952 let mt = Memtable::new();
953 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
954 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
955 mt.insert(&make_point("cpu", "a", 300, 3.0)).unwrap();
956 mt.insert(&make_point("mem", "a", 200, 4.0)).unwrap();
957
958 let result = mt.scan_measurement("cpu", 150, 250);
959 assert_eq!(result.len(), 1);
960 assert_eq!(result[0].timestamp(), 200);
961 }
962
963 #[test]
967 fn hash_collision_preserves_both_series() {
968 let mt = Memtable::new();
969
970 let p1 = make_point("cpu", "host-a", 1000, 42.0);
974 let p2 = make_point("cpu", "host-b", 1000, 99.0);
975
976 mt.insert(&p1).unwrap();
977 mt.insert(&p2).unwrap();
978
979 assert_eq!(
981 mt.len(),
982 2,
983 "two distinct series at same timestamp must both survive"
984 );
985
986 let s1 = mt.scan(p1.series_key(), 0, i64::MAX);
988 assert_eq!(s1.len(), 1);
989 assert_eq!(
990 *s1[0].field("value").unwrap(),
991 chronix_core::FieldValue::F64(42.0)
992 );
993
994 let s2 = mt.scan(p2.series_key(), 0, i64::MAX);
995 assert_eq!(s2.len(), 1);
996 assert_eq!(
997 *s2[0].field("value").unwrap(),
998 chronix_core::FieldValue::F64(99.0)
999 );
1000 }
1001
1002 #[test]
1003 fn insert_batch_basic() {
1004 let mt = Memtable::new();
1005 let points = vec![
1006 make_point("cpu", "a", 100, 1.0),
1007 make_point("cpu", "a", 200, 2.0),
1008 make_point("cpu", "b", 300, 3.0),
1009 make_point("mem", "a", 400, 4.0),
1010 ];
1011 let count = mt.insert_batch(&points).unwrap();
1012 assert_eq!(count, 4);
1013
1014 let all = mt.scan_all();
1015 assert_eq!(all.len(), 4);
1016
1017 let cpu = mt.scan_measurement("cpu", 0, i64::MAX);
1019 assert_eq!(cpu.len(), 3);
1020 let mem = mt.scan_measurement("mem", 0, i64::MAX);
1021 assert_eq!(mem.len(), 1);
1022 }
1023
1024 #[test]
1025 fn insert_batch_overwrite() {
1026 let mt = Memtable::new();
1027 let initial = vec![
1028 make_point("cpu", "a", 100, 1.0),
1029 make_point("cpu", "a", 200, 2.0),
1030 ];
1031 mt.insert_batch(&initial).unwrap();
1032
1033 let overwrite = vec![
1035 make_point("cpu", "a", 100, 10.0),
1036 make_point("cpu", "a", 200, 20.0),
1037 ];
1038 mt.insert_batch(&overwrite).unwrap();
1039
1040 let all = mt.scan_all();
1041 assert_eq!(all.len(), 2);
1042 assert_eq!(
1043 *all[0].field("value").unwrap(),
1044 chronix_core::FieldValue::F64(10.0)
1045 );
1046 assert_eq!(
1047 *all[1].field("value").unwrap(),
1048 chronix_core::FieldValue::F64(20.0)
1049 );
1050 }
1051
1052 #[test]
1053 fn insert_batch_frozen_fails() {
1054 let mt = Memtable::new();
1055 mt.freeze();
1056 let points = vec![make_point("cpu", "a", 100, 1.0)];
1057 assert!(mt.insert_batch(&points).is_err());
1058 }
1059
1060 #[test]
1061 fn insert_batch_with_wal_seq_updates_range() {
1062 let mt = Memtable::new();
1063 let points = vec![
1064 make_point("cpu", "a", 100, 1.0),
1065 make_point("cpu", "a", 200, 2.0),
1066 ];
1067 mt.insert_batch_with_wal_seq(&points, 42).unwrap();
1068 assert_eq!(mt.min_wal_seq(), Some(42));
1069 assert_eq!(mt.max_wal_seq(), Some(42));
1070 }
1071}
1072
1073#[cfg(test)]
1079mod calibration {
1080 use std::alloc::{GlobalAlloc, Layout, System};
1081 use std::collections::BTreeMap;
1082 use std::sync::atomic::{AtomicUsize, Ordering};
1083
1084 use chronix_core::{FieldValue, Point, SeriesKey};
1085
1086 use super::Memtable;
1087
1088 struct Counting;
1089 static LIVE: AtomicUsize = AtomicUsize::new(0);
1090
1091 std::thread_local! {
1094 static THREAD_LIVE: std::cell::Cell<isize> = const { std::cell::Cell::new(0) };
1095 }
1096
1097 fn thread_live() -> isize {
1098 THREAD_LIVE.with(std::cell::Cell::get)
1099 }
1100
1101 #[allow(unsafe_code)]
1105 unsafe impl GlobalAlloc for Counting {
1106 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
1107 LIVE.fetch_add(layout.size(), Ordering::Relaxed);
1108 let _ = THREAD_LIVE.try_with(|c| c.set(c.get() + layout.size() as isize));
1109 unsafe { System.alloc(layout) }
1110 }
1111 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
1112 LIVE.fetch_sub(layout.size(), Ordering::Relaxed);
1113 let _ = THREAD_LIVE.try_with(|c| c.set(c.get() - layout.size() as isize));
1114 unsafe { System.dealloc(ptr, layout) }
1115 }
1116 }
1117
1118 #[global_allocator]
1119 static ALLOC: Counting = Counting;
1120
1121 #[test]
1125 fn memtable_estimate_tracks_real_allocation() {
1126 let keys: Vec<SeriesKey> = (0..50)
1127 .map(|m| {
1128 SeriesKey::new(
1129 "power",
1130 BTreeMap::from([
1131 ("meter".to_string(), format!("m{m:02}")),
1132 ("site".to_string(), "home".to_string()),
1133 ]),
1134 )
1135 .unwrap()
1136 })
1137 .collect();
1138 let points: Vec<Point> = (0..400i64)
1140 .flat_map(|s| {
1141 keys.iter().map(move |k| {
1142 Point::new(
1143 k.clone(),
1144 BTreeMap::from([
1145 ("w".to_string(), FieldValue::F64(230.0 + s as f64)),
1146 ("kwh".to_string(), FieldValue::F64(s as f64 / 1000.0)),
1147 ]),
1148 s * 1_000_000_000,
1149 )
1150 .unwrap()
1151 })
1152 })
1153 .collect();
1154
1155 let before = thread_live();
1156 let mt = Memtable::new();
1157 for p in &points {
1158 mt.insert(p).unwrap();
1159 }
1160 let measured = usize::try_from(thread_live() - before).unwrap_or(0);
1161 let estimated = mt.estimated_size();
1162 let ratio = estimated as f64 / measured as f64;
1163 assert!(
1164 (0.75..=1.25).contains(&ratio),
1165 "estimate {estimated} B vs allocator {measured} B for {} points: ratio {ratio:.2} \
1166 ({:.0} B/point measured, {:.0} B/point estimated)",
1167 points.len(),
1168 measured as f64 / points.len() as f64,
1169 estimated as f64 / points.len() as f64,
1170 );
1171 drop(mt);
1172 }
1173}
1174
1175#[derive(Debug)]
1177pub struct MeasurementBatch {
1178 pub measurement: String,
1180 pub batch: arrow::record_batch::RecordBatch,
1184 pub tag_columns: Vec<String>,
1186}
1187
1188enum FieldBuilder {
1190 F64(arrow::array::Float64Builder),
1191 I64(arrow::array::Int64Builder),
1192 U64(arrow::array::UInt64Builder),
1193 Bool(arrow::array::BooleanBuilder),
1194 Str(arrow::array::StringBuilder),
1195}
1196
1197impl FieldBuilder {
1198 fn for_value(v: &chronix_core::FieldValue, rows_before: usize) -> Self {
1199 use chronix_core::FieldValue as V;
1200 let mut b = match v {
1201 V::F64(_) => Self::F64(arrow::array::Float64Builder::new()),
1202 V::I64(_) => Self::I64(arrow::array::Int64Builder::new()),
1203 V::U64(_) => Self::U64(arrow::array::UInt64Builder::new()),
1204 V::Bool(_) => Self::Bool(arrow::array::BooleanBuilder::new()),
1205 V::String(_) => Self::Str(arrow::array::StringBuilder::new()),
1206 };
1207 for _ in 0..rows_before {
1208 b.append_null();
1209 }
1210 b
1211 }
1212
1213 fn append_null(&mut self) {
1214 match self {
1215 Self::F64(b) => b.append_null(),
1216 Self::I64(b) => b.append_null(),
1217 Self::U64(b) => b.append_null(),
1218 Self::Bool(b) => b.append_null(),
1219 Self::Str(b) => b.append_null(),
1220 }
1221 }
1222
1223 fn append(&mut self, v: &chronix_core::FieldValue) {
1227 use chronix_core::FieldValue as V;
1228 match (self, v) {
1229 (Self::F64(b), V::F64(x)) => b.append_value(*x),
1230 (Self::I64(b), V::I64(x)) => b.append_value(*x),
1231 (Self::U64(b), V::U64(x)) => b.append_value(*x),
1232 (Self::Bool(b), V::Bool(x)) => b.append_value(*x),
1233 (Self::Str(b), V::String(x)) => b.append_value(x),
1234 (this, _) => this.append_null(),
1235 }
1236 }
1237
1238 fn finish(self) -> (arrow::datatypes::DataType, arrow::array::ArrayRef) {
1239 use arrow::datatypes::DataType;
1240 match self {
1241 Self::F64(mut b) => (DataType::Float64, Arc::new(b.finish())),
1242 Self::I64(mut b) => (DataType::Int64, Arc::new(b.finish())),
1243 Self::U64(mut b) => (DataType::UInt64, Arc::new(b.finish())),
1244 Self::Bool(mut b) => (DataType::Boolean, Arc::new(b.finish())),
1245 Self::Str(mut b) => (DataType::Utf8, Arc::new(b.finish())),
1246 }
1247 }
1248}
1249
1250struct MeasurementBuilder {
1252 rows: usize,
1253 timestamps: arrow::array::Int64Builder,
1254 tags: std::collections::BTreeMap<Arc<str>, arrow::array::StringBuilder>,
1255 fields: std::collections::BTreeMap<Arc<str>, FieldBuilder>,
1256}
1257
1258impl MeasurementBuilder {
1259 fn new() -> Self {
1260 Self {
1261 rows: 0,
1262 timestamps: arrow::array::Int64Builder::new(),
1263 tags: std::collections::BTreeMap::new(),
1264 fields: std::collections::BTreeMap::new(),
1265 }
1266 }
1267
1268 fn push(&mut self, key: &MemtableKey, entry: &MemtableEntry) {
1269 self.timestamps.append_value(key.timestamp);
1270 for (k, v) in entry.tags.iter() {
1273 let rows = self.rows;
1274 self.tags.entry(Arc::clone(k)).or_insert_with(|| {
1275 let mut b = arrow::array::StringBuilder::new();
1276 for _ in 0..rows {
1277 b.append_null();
1278 }
1279 b
1280 });
1281 let _ = v;
1282 }
1283 for (name, b) in &mut self.tags {
1284 match entry.tags.iter().find(|(k, _)| k == name) {
1285 Some((_, v)) => b.append_value(v),
1286 None => b.append_null(),
1287 }
1288 }
1289 for (k, v) in entry.fields.iter() {
1290 let rows = self.rows;
1291 self.fields
1292 .entry(Arc::clone(k))
1293 .or_insert_with(|| FieldBuilder::for_value(v, rows));
1294 }
1295 for (name, b) in &mut self.fields {
1296 match entry.fields.iter().find(|(k, _)| k == name) {
1297 Some((_, v)) => b.append(v),
1298 None => b.append_null(),
1299 }
1300 }
1301 self.rows += 1;
1302 }
1303
1304 fn finish(mut self, measurement: &str) -> Option<MeasurementBatch> {
1305 use arrow::datatypes::{DataType, Field, Schema};
1306 if self.rows == 0 {
1307 return None;
1308 }
1309 use crate::segment::metadata::roles;
1310 let mut fields = vec![
1313 Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false)
1314 .with_metadata(roles::arrow_metadata(roles::TIMESTAMP)),
1315 ];
1316 let mut columns: Vec<arrow::array::ArrayRef> = vec![Arc::new(self.timestamps.finish())];
1317 let tag_columns: Vec<String> = self.tags.keys().map(ToString::to_string).collect();
1318 for (name, mut b) in self.tags {
1319 fields.push(
1320 Field::new(name.as_ref(), DataType::Utf8, true)
1321 .with_metadata(roles::arrow_metadata(roles::TAG)),
1322 );
1323 columns.push(Arc::new(b.finish()));
1324 }
1325 for (name, b) in self.fields {
1326 let (dt, arr) = b.finish();
1327 fields.push(
1328 Field::new(name.as_ref(), dt, true)
1329 .with_metadata(roles::arrow_metadata(roles::FIELD)),
1330 );
1331 columns.push(arr);
1332 }
1333 let batch =
1334 arrow::record_batch::RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
1335 .ok()?;
1336 Some(MeasurementBatch {
1337 measurement: measurement.to_string(),
1338 batch,
1339 tag_columns,
1340 })
1341 }
1342}
1343
1344impl Memtable {
1345 #[must_use]
1354 pub fn to_record_batches(&self) -> Vec<MeasurementBatch> {
1355 let mut builders: std::collections::BTreeMap<Arc<str>, MeasurementBuilder> =
1356 std::collections::BTreeMap::new();
1357 for entry in self.data.iter() {
1358 let value = entry.value();
1359 builders
1360 .entry(Arc::clone(&value.measurement))
1361 .or_insert_with(MeasurementBuilder::new)
1362 .push(entry.key(), value);
1363 }
1364 builders
1365 .into_iter()
1366 .filter_map(|(m, b)| b.finish(&m))
1367 .collect()
1368 }
1369}