1use std::collections::HashSet;
16use std::sync::Arc;
17use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
18
19use crossbeam_skiplist::SkipMap;
20use dashmap::DashMap;
21use dashmap::mapref::entry::Entry;
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 && let Some(existing) = self.data.get(&key) {
254 let old_size = existing.value().estimated_size()
255 + std::mem::size_of::<MemtableKey>()
256 + NODE_OVERHEAD;
257 total_size_delta -= old_size.min(i64::MAX as usize) as i64;
258 }
259
260 prev_hash = cur_hash;
261 prev_ts = cur_ts;
262
263 self.data.insert(key, entry);
264 total_size_delta += entry_size.min(i64::MAX as usize) as i64;
266
267 new_index_entries.push((measurement, cur_hash));
268 }
269
270 if total_size_delta >= 0 {
272 self.estimated_size
273 .fetch_add(total_size_delta as usize, Ordering::Relaxed);
274 } else {
275 let _ =
276 self.estimated_size
277 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
278 Some(current.saturating_sub(total_size_delta.unsigned_abs() as usize))
279 });
280 }
281
282 let mut index_delta = 0usize;
284 for (measurement, hash) in new_index_entries {
285 index_delta += self.index_insert(measurement, hash);
286 }
287 if index_delta > 0 {
288 self.index_size.fetch_add(index_delta, Ordering::Relaxed);
289 }
290
291 metrics::gauge!("chronix_memtable_memory_bytes")
293 .set(self.estimated_size.load(Ordering::Relaxed) as f64);
294 metrics::gauge!("chronix_memtable_entries").set(self.data.len() as f64);
295
296 Ok(points.len())
297 }
298
299 pub fn insert_batch_with_wal_seq(&self, points: &[Point], wal_seq: u64) -> Result<usize> {
308 let count = self.insert_batch(points)?;
309 self.update_wal_seq(wal_seq);
310 Ok(count)
311 }
312
313 fn update_wal_seq(&self, seq: u64) {
315 let mut current = self.min_wal_seq.load(Ordering::Relaxed);
317 while seq < current {
318 match self.min_wal_seq.compare_exchange_weak(
319 current,
320 seq,
321 Ordering::Release,
322 Ordering::Relaxed,
323 ) {
324 Ok(_) => break,
325 Err(actual) => current = actual,
326 }
327 }
328
329 let mut current = self.max_wal_seq.load(Ordering::Relaxed);
331 while seq > current {
332 match self.max_wal_seq.compare_exchange_weak(
333 current,
334 seq,
335 Ordering::Release,
336 Ordering::Relaxed,
337 ) {
338 Ok(_) => break,
339 Err(actual) => current = actual,
340 }
341 }
342 }
343
344 #[must_use]
351 pub fn scan(&self, series_key: &SeriesKey, min_ts: Timestamp, max_ts: Timestamp) -> Vec<Point> {
352 let hash = series_key.hash_fnv();
353 let canonical = series_key.canonical_form();
354 let start = MemtableKey::new(hash, min_ts, Arc::from(canonical));
357 let end = MemtableKey::new(hash, max_ts, Arc::from(canonical));
358
359 self.data
360 .range(start..=end)
361 .filter(|entry| &*entry.key().series_canonical == canonical)
362 .filter_map(|entry| Self::entry_to_point(entry.key(), entry.value()))
363 .collect()
364 }
365
366 #[must_use]
375 pub fn scan_range(&self, min_ts: Timestamp, max_ts: Timestamp) -> Vec<Point> {
376 self.data
379 .iter()
380 .filter(|entry| {
381 let ts = entry.key().timestamp;
382 ts >= min_ts && ts <= max_ts
383 })
384 .filter_map(|entry| Self::entry_to_point(entry.key(), entry.value()))
385 .collect()
386 }
387
388 #[must_use]
391 pub fn scan_all(&self) -> Vec<Point> {
392 self.data
393 .iter()
394 .filter_map(|entry| Self::entry_to_point(entry.key(), entry.value()))
395 .collect()
396 }
397
398 #[must_use]
405 pub fn scan_measurement(
406 &self,
407 measurement: &str,
408 min_ts: Timestamp,
409 max_ts: Timestamp,
410 ) -> Vec<Point> {
411 let hashes = match self.measurement_index.get(measurement) {
412 Some(set) => set.iter().copied().collect::<Vec<_>>(),
413 None => return Vec::new(),
414 };
415
416 let mut results = Vec::new();
417 for hash in hashes {
418 let start = MemtableKey::lower_bound(hash, min_ts);
420 let end = MemtableKey::upper_bound(hash, max_ts);
421 results.extend(
422 self.data
423 .range(start..=end)
424 .filter(|entry| {
425 let ts = entry.key().timestamp;
426 ts >= min_ts && ts <= max_ts && &*entry.value().measurement == measurement
427 })
428 .filter_map(|entry| Self::entry_to_point(entry.key(), entry.value())),
429 );
430 }
431 results
432 }
433
434 #[must_use]
444 pub fn series_canonical_forms(&self) -> Vec<Arc<str>> {
445 self.series_tags
446 .iter()
447 .map(|e| Arc::clone(e.key()))
448 .collect()
449 }
450
451 pub fn visit_tags(&self, mut f: impl FnMut(&str, &str)) {
460 for entry in &self.series_tags {
461 for (key, value) in entry.value().iter() {
462 f(key, value);
463 }
464 }
465 }
466
467 #[must_use]
469 pub fn len(&self) -> usize {
470 self.data.len()
471 }
472
473 #[must_use]
475 pub fn is_empty(&self) -> bool {
476 self.data.is_empty()
477 }
478
479 #[must_use]
489 pub fn interner_bytes(&self) -> usize {
490 self.interner.memory_bytes()
491 }
492
493 #[must_use]
495 pub fn estimated_size(&self) -> usize {
496 self.estimated_size.load(Ordering::Relaxed) + self.index_size.load(Ordering::Relaxed)
497 }
498
499 fn series_tag_set(
505 &self,
506 canonical: &Arc<str>,
507 series_key: &SeriesKey,
508 ) -> Arc<[(Arc<str>, Arc<str>)]> {
509 if let Some(existing) = self.series_tags.get(canonical) {
510 return Arc::clone(existing.value());
511 }
512 let tags: Arc<[(Arc<str>, Arc<str>)]> = series_key
513 .tags()
514 .iter()
515 .map(|(k, v)| (self.interner.intern(k), self.interner.intern(v)))
516 .collect();
517 let mut charged = 0usize;
518 let stored = self
519 .series_tags
520 .entry(Arc::clone(canonical))
521 .or_insert_with(|| {
522 charged = canonical.len()
526 + tags.len() * std::mem::size_of::<(Arc<str>, Arc<str>)>()
527 + tags.iter().map(|(k, v)| k.len() + v.len()).sum::<usize>()
528 + 2 * std::mem::size_of::<usize>()
529 + 64;
530 tags
531 })
532 .clone();
533 if charged > 0 {
534 self.estimated_size.fetch_add(charged, Ordering::Relaxed);
535 }
536 stored
537 }
538
539 fn index_insert(&self, measurement: Arc<str>, hash: u64) -> usize {
548 const SET_OVERHEAD: usize = 56;
549 match self.measurement_index.entry(measurement) {
550 Entry::Occupied(mut occupied) => {
551 let before = occupied.get().capacity();
552 occupied.get_mut().insert(hash);
553 (occupied.get().capacity() - before) * std::mem::size_of::<u64>()
554 }
555 Entry::Vacant(vacant) => {
556 let key_bytes = std::mem::size_of::<usize>() * 2 + vacant.key().len();
557 let mut set = HashSet::new();
558 set.insert(hash);
559 let set_bytes = set.capacity() * std::mem::size_of::<u64>() + SET_OVERHEAD;
560 vacant.insert(set);
561 key_bytes + set_bytes
562 }
563 }
564 }
565
566 #[must_use]
570 pub fn min_wal_seq(&self) -> Option<u64> {
571 let val = self.min_wal_seq.load(Ordering::Acquire);
572 if val == u64::MAX { None } else { Some(val) }
573 }
574
575 #[must_use]
579 pub fn max_wal_seq(&self) -> Option<u64> {
580 let val = self.max_wal_seq.load(Ordering::Acquire);
581 if val == 0 && self.min_wal_seq.load(Ordering::Acquire) == u64::MAX {
582 None
583 } else {
584 Some(val)
585 }
586 }
587
588 #[must_use]
590 pub fn is_frozen(&self) -> bool {
591 self.frozen.load(Ordering::Acquire)
592 }
593
594 pub fn freeze(&self) {
599 self.frozen.store(true, Ordering::Release);
600 }
601
602 #[must_use]
604 pub fn to_points(&self) -> Vec<Point> {
605 self.scan_all()
606 }
607
608 fn entry_to_point(key: &MemtableKey, entry: &MemtableEntry) -> Option<Point> {
613 let tags = entry
614 .tags
615 .iter()
616 .map(|(k, v)| (k.to_string(), v.to_string()))
617 .collect();
618 let series_key = SeriesKey::new(entry.measurement.to_string(), tags)
619 .map_err(|e| {
620 tracing::error!(
623 hash = key.series_key_hash,
624 error = %e,
625 "DATA LOSS: failed to reconstruct series key during flush"
626 );
627 metrics::counter!("chronix_memtable_dropped_points_total").increment(1);
628 e
629 })
630 .ok()?;
631 Point::new(
632 series_key,
633 entry
634 .fields
635 .iter()
636 .map(|(k, v)| (k.to_string(), v.clone()))
637 .collect(),
638 key.timestamp,
639 )
640 .map_err(|e| {
641 tracing::error!(
642 hash = key.series_key_hash,
643 timestamp = key.timestamp,
644 error = %e,
645 "DATA LOSS: failed to reconstruct point during flush"
646 );
647 metrics::counter!("chronix_memtable_dropped_points_total").increment(1);
648 e
649 })
650 .ok()
651 }
652}
653
654impl Default for Memtable {
655 fn default() -> Self {
656 Self::new()
657 }
658}
659
660#[cfg(test)]
665mod tests {
666 use super::*;
667 use std::collections::BTreeMap;
668
669 use chronix_core::FieldValue;
670
671 fn make_point(measurement: &str, host: &str, ts: i64, value: f64) -> Point {
672 let tags: BTreeMap<String, String> = [("host".to_string(), host.to_string())]
673 .into_iter()
674 .collect();
675 let series_key = SeriesKey::new(measurement.to_string(), tags).unwrap();
676 let fields: BTreeMap<String, FieldValue> = [("value".to_string(), FieldValue::F64(value))]
677 .into_iter()
678 .collect();
679 Point::new(series_key, fields, ts).unwrap()
680 }
681
682 #[test]
683 fn insert_and_scan_all() {
684 let mt = Memtable::new();
685 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
686 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
687 mt.insert(&make_point("cpu", "a", 300, 3.0)).unwrap();
688
689 let points = mt.scan_all();
690 assert_eq!(points.len(), 3);
691 assert_eq!(points[0].timestamp(), 100);
692 assert_eq!(points[1].timestamp(), 200);
693 assert_eq!(points[2].timestamp(), 300);
694 }
695
696 #[test]
697 fn insert_and_scan_series() {
698 let mt = Memtable::new();
699 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
700 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
701 mt.insert(&make_point("cpu", "b", 150, 9.0)).unwrap();
702
703 let key_a = SeriesKey::new(
704 "cpu".to_string(),
705 [("host".to_string(), "a".to_string())]
706 .into_iter()
707 .collect(),
708 )
709 .unwrap();
710
711 let points = mt.scan(&key_a, 0, 300);
712 assert_eq!(points.len(), 2);
713 assert!(
714 points
715 .iter()
716 .all(|p| p.series_key().tag("host") == Some("a"))
717 );
718 }
719
720 #[test]
721 fn scan_time_range() {
722 let mt = Memtable::new();
723 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
724 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
725 mt.insert(&make_point("cpu", "a", 300, 3.0)).unwrap();
726
727 let key = SeriesKey::new(
728 "cpu".to_string(),
729 [("host".to_string(), "a".to_string())]
730 .into_iter()
731 .collect(),
732 )
733 .unwrap();
734
735 let points = mt.scan(&key, 150, 250);
736 assert_eq!(points.len(), 1);
737 assert_eq!(points[0].timestamp(), 200);
738 }
739
740 #[test]
741 fn last_write_wins() {
742 let mt = Memtable::new();
743 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
744 mt.insert(&make_point("cpu", "a", 100, 99.0)).unwrap();
745
746 let points = mt.scan_all();
747 assert_eq!(points.len(), 1);
748 assert_eq!(points[0].field("value"), Some(&FieldValue::F64(99.0)));
749 }
750
751 #[test]
752 fn empty_scan() {
753 let mt = Memtable::new();
754 let key = SeriesKey::new("cpu".to_string(), BTreeMap::new()).unwrap();
755 let points = mt.scan(&key, 0, 1000);
756 assert!(points.is_empty());
757 }
758
759 #[test]
762 fn index_accounting_matches_a_full_walk() {
763 let mt = Memtable::new();
764 for i in 0..500u64 {
765 let measurement = format!("m{}", i % 7);
766 let series = format!("s{}", i % 53);
767 mt.insert(&make_point(&measurement, &series, i as i64, 1.0))
768 .unwrap();
769 }
770
771 let batch: Vec<_> = (500..800u64)
773 .map(|i| {
774 make_point(
775 &format!("m{}", i % 11),
776 &format!("s{}", i % 97),
777 i as i64,
778 2.0,
779 )
780 })
781 .collect();
782 mt.insert_batch(&batch).unwrap();
783
784 let walked: usize = mt
785 .measurement_index
786 .iter()
787 .map(|entry| {
788 let key_bytes = std::mem::size_of::<usize>() * 2 + entry.key().len();
789 let set_bytes = entry.value().capacity() * std::mem::size_of::<u64>() + 56;
790 key_bytes + set_bytes
791 })
792 .sum();
793
794 assert_eq!(
795 mt.index_size.load(Ordering::Relaxed),
796 walked,
797 "incremental index accounting drifted from a full walk"
798 );
799 }
800
801 #[test]
802 fn estimated_size_increases() {
803 let mt = Memtable::new();
804 assert_eq!(mt.estimated_size(), 0);
805 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
806 assert!(mt.estimated_size() > 0);
807 let size_after_one = mt.estimated_size();
808 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
809 assert!(mt.estimated_size() > size_after_one);
810 }
811
812 #[test]
813 fn wal_seq_tracking() {
814 let mt = Memtable::new();
815 assert_eq!(mt.min_wal_seq(), None);
816 assert_eq!(mt.max_wal_seq(), None);
817
818 mt.insert_with_wal_seq(&make_point("cpu", "a", 100, 1.0), 10)
819 .unwrap();
820 assert_eq!(mt.min_wal_seq(), Some(10));
821 assert_eq!(mt.max_wal_seq(), Some(10));
822
823 mt.insert_with_wal_seq(&make_point("cpu", "a", 200, 2.0), 5)
824 .unwrap();
825 assert_eq!(mt.min_wal_seq(), Some(5));
826 assert_eq!(mt.max_wal_seq(), Some(10));
827
828 mt.insert_with_wal_seq(&make_point("cpu", "a", 300, 3.0), 20)
829 .unwrap();
830 assert_eq!(mt.min_wal_seq(), Some(5));
831 assert_eq!(mt.max_wal_seq(), Some(20));
832 }
833
834 #[test]
835 fn freeze_prevents_inserts() {
836 let mt = Memtable::new();
837 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
838 mt.freeze();
839 assert!(mt.is_frozen());
840
841 let result = mt.insert(&make_point("cpu", "a", 200, 2.0));
842 assert!(result.is_err());
843
844 let points = mt.scan_all();
846 assert_eq!(points.len(), 1);
847 }
848
849 #[test]
850 fn len_and_is_empty() {
851 let mt = Memtable::new();
852 assert!(mt.is_empty());
853 assert_eq!(mt.len(), 0);
854
855 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
856 assert!(!mt.is_empty());
857 assert_eq!(mt.len(), 1);
858 }
859
860 #[test]
861 fn multiple_series() {
862 let mt = Memtable::new();
863 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
864 mt.insert(&make_point("cpu", "b", 100, 2.0)).unwrap();
865 mt.insert(&make_point("mem", "a", 100, 3.0)).unwrap();
866
867 assert_eq!(mt.len(), 3);
868
869 let all = mt.scan_all();
870 assert_eq!(all.len(), 3);
871 }
872
873 #[test]
874 fn concurrent_inserts() {
875 use std::sync::Arc;
876 use std::thread;
877
878 let mt = Arc::new(Memtable::new());
879 let mut handles = Vec::new();
880
881 for thread_id in 0..4 {
882 let mt = Arc::clone(&mt);
883 handles.push(thread::spawn(move || {
884 for i in 0..100 {
885 let ts = i64::from(thread_id * 1000 + i);
886 let host = format!("host-{thread_id}");
887 mt.insert(&make_point("cpu", &host, ts, ts as f64)).unwrap();
888 }
889 }));
890 }
891
892 for h in handles {
893 h.join().unwrap();
894 }
895
896 assert_eq!(mt.len(), 400);
897 let all = mt.scan_all();
898 assert_eq!(all.len(), 400);
899
900 let mut prev_key: Option<MemtableKey> = None;
902 for p in &all {
903 let key = MemtableKey::new(
904 p.series_key().hash_fnv(),
905 p.timestamp(),
906 Arc::from(p.series_key().canonical_form()),
907 );
908 if let Some(prev) = prev_key {
909 assert!(key >= prev, "entries should be ordered");
910 }
911 prev_key = Some(key);
912 }
913 }
914
915 #[allow(deprecated)]
916 #[test]
917 fn to_points_is_scan_all() {
918 let mt = Memtable::new();
919 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
920 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
921
922 let scan = mt.scan_all();
923 let points = mt.to_points();
924 assert_eq!(scan.len(), points.len());
925 for (s, d) in scan.iter().zip(points.iter()) {
926 assert_eq!(s.timestamp(), d.timestamp());
927 assert_eq!(s.fields(), d.fields());
928 }
929 }
930
931 #[test]
932 fn scan_range_across_series() {
933 let mt = Memtable::new();
934 mt.insert(&make_point("cpu", "a", 50, 1.0)).unwrap();
935 mt.insert(&make_point("cpu", "a", 100, 2.0)).unwrap();
936 mt.insert(&make_point("cpu", "b", 75, 3.0)).unwrap();
937 mt.insert(&make_point("cpu", "b", 200, 4.0)).unwrap();
938
939 let points = mt.scan_range(60, 150);
940 assert_eq!(points.len(), 2);
942 for p in &points {
943 assert!(
944 p.timestamp() >= 60 && p.timestamp() <= 150,
945 "timestamp {} should be in [60, 150]",
946 p.timestamp()
947 );
948 }
949 }
950
951 #[test]
952 fn scan_measurement_filters_by_name() {
953 let mt = Memtable::new();
954 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
955 mt.insert(&make_point("cpu", "b", 200, 2.0)).unwrap();
956 mt.insert(&make_point("mem", "a", 150, 3.0)).unwrap();
957 mt.insert(&make_point("mem", "a", 300, 4.0)).unwrap();
958 mt.insert(&make_point("disk", "c", 250, 5.0)).unwrap();
959
960 let cpu = mt.scan_measurement("cpu", 0, i64::MAX);
962 assert_eq!(cpu.len(), 2);
963 assert!(cpu.iter().all(|p| p.series_key().measurement() == "cpu"));
964
965 let mem = mt.scan_measurement("mem", 0, i64::MAX);
967 assert_eq!(mem.len(), 2);
968 assert!(mem.iter().all(|p| p.series_key().measurement() == "mem"));
969
970 let disk = mt.scan_measurement("disk", 0, i64::MAX);
972 assert_eq!(disk.len(), 1);
973
974 let none = mt.scan_measurement("net", 0, i64::MAX);
976 assert!(none.is_empty());
977 }
978
979 #[test]
980 fn scan_measurement_respects_time_range() {
981 let mt = Memtable::new();
982 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
983 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
984 mt.insert(&make_point("cpu", "a", 300, 3.0)).unwrap();
985 mt.insert(&make_point("mem", "a", 200, 4.0)).unwrap();
986
987 let result = mt.scan_measurement("cpu", 150, 250);
988 assert_eq!(result.len(), 1);
989 assert_eq!(result[0].timestamp(), 200);
990 }
991
992 #[test]
996 fn hash_collision_preserves_both_series() {
997 let mt = Memtable::new();
998
999 let p1 = make_point("cpu", "host-a", 1000, 42.0);
1003 let p2 = make_point("cpu", "host-b", 1000, 99.0);
1004
1005 mt.insert(&p1).unwrap();
1006 mt.insert(&p2).unwrap();
1007
1008 assert_eq!(
1010 mt.len(),
1011 2,
1012 "two distinct series at same timestamp must both survive"
1013 );
1014
1015 let s1 = mt.scan(p1.series_key(), 0, i64::MAX);
1017 assert_eq!(s1.len(), 1);
1018 assert_eq!(
1019 *s1[0].field("value").unwrap(),
1020 chronix_core::FieldValue::F64(42.0)
1021 );
1022
1023 let s2 = mt.scan(p2.series_key(), 0, i64::MAX);
1024 assert_eq!(s2.len(), 1);
1025 assert_eq!(
1026 *s2[0].field("value").unwrap(),
1027 chronix_core::FieldValue::F64(99.0)
1028 );
1029 }
1030
1031 #[test]
1032 fn insert_batch_basic() {
1033 let mt = Memtable::new();
1034 let points = vec![
1035 make_point("cpu", "a", 100, 1.0),
1036 make_point("cpu", "a", 200, 2.0),
1037 make_point("cpu", "b", 300, 3.0),
1038 make_point("mem", "a", 400, 4.0),
1039 ];
1040 let count = mt.insert_batch(&points).unwrap();
1041 assert_eq!(count, 4);
1042
1043 let all = mt.scan_all();
1044 assert_eq!(all.len(), 4);
1045
1046 let cpu = mt.scan_measurement("cpu", 0, i64::MAX);
1048 assert_eq!(cpu.len(), 3);
1049 let mem = mt.scan_measurement("mem", 0, i64::MAX);
1050 assert_eq!(mem.len(), 1);
1051 }
1052
1053 #[test]
1054 fn insert_batch_overwrite() {
1055 let mt = Memtable::new();
1056 let initial = vec![
1057 make_point("cpu", "a", 100, 1.0),
1058 make_point("cpu", "a", 200, 2.0),
1059 ];
1060 mt.insert_batch(&initial).unwrap();
1061
1062 let overwrite = vec![
1064 make_point("cpu", "a", 100, 10.0),
1065 make_point("cpu", "a", 200, 20.0),
1066 ];
1067 mt.insert_batch(&overwrite).unwrap();
1068
1069 let all = mt.scan_all();
1070 assert_eq!(all.len(), 2);
1071 assert_eq!(
1072 *all[0].field("value").unwrap(),
1073 chronix_core::FieldValue::F64(10.0)
1074 );
1075 assert_eq!(
1076 *all[1].field("value").unwrap(),
1077 chronix_core::FieldValue::F64(20.0)
1078 );
1079 }
1080
1081 #[test]
1082 fn insert_batch_frozen_fails() {
1083 let mt = Memtable::new();
1084 mt.freeze();
1085 let points = vec![make_point("cpu", "a", 100, 1.0)];
1086 assert!(mt.insert_batch(&points).is_err());
1087 }
1088
1089 #[test]
1090 fn insert_batch_with_wal_seq_updates_range() {
1091 let mt = Memtable::new();
1092 let points = vec![
1093 make_point("cpu", "a", 100, 1.0),
1094 make_point("cpu", "a", 200, 2.0),
1095 ];
1096 mt.insert_batch_with_wal_seq(&points, 42).unwrap();
1097 assert_eq!(mt.min_wal_seq(), Some(42));
1098 assert_eq!(mt.max_wal_seq(), Some(42));
1099 }
1100}
1101
1102#[cfg(test)]
1108mod calibration {
1109 use std::alloc::{GlobalAlloc, Layout, System};
1110 use std::collections::BTreeMap;
1111 use std::sync::atomic::{AtomicUsize, Ordering};
1112
1113 use chronix_core::{FieldValue, Point, SeriesKey};
1114
1115 use super::Memtable;
1116
1117 struct Counting;
1118 static LIVE: AtomicUsize = AtomicUsize::new(0);
1119
1120 std::thread_local! {
1123 static THREAD_LIVE: std::cell::Cell<isize> = const { std::cell::Cell::new(0) };
1124 }
1125
1126 fn thread_live() -> isize {
1127 THREAD_LIVE.with(std::cell::Cell::get)
1128 }
1129
1130 #[allow(unsafe_code)]
1134 unsafe impl GlobalAlloc for Counting {
1135 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
1136 LIVE.fetch_add(layout.size(), Ordering::Relaxed);
1137 let _ = THREAD_LIVE.try_with(|c| c.set(c.get() + layout.size() as isize));
1138 unsafe { System.alloc(layout) }
1139 }
1140 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
1141 LIVE.fetch_sub(layout.size(), Ordering::Relaxed);
1142 let _ = THREAD_LIVE.try_with(|c| c.set(c.get() - layout.size() as isize));
1143 unsafe { System.dealloc(ptr, layout) }
1144 }
1145 }
1146
1147 #[global_allocator]
1148 static ALLOC: Counting = Counting;
1149
1150 #[test]
1154 fn memtable_estimate_tracks_real_allocation() {
1155 let keys: Vec<SeriesKey> = (0..50)
1156 .map(|m| {
1157 SeriesKey::new(
1158 "power",
1159 BTreeMap::from([
1160 ("meter".to_string(), format!("m{m:02}")),
1161 ("site".to_string(), "home".to_string()),
1162 ]),
1163 )
1164 .unwrap()
1165 })
1166 .collect();
1167 let points: Vec<Point> = (0..400i64)
1169 .flat_map(|s| {
1170 keys.iter().map(move |k| {
1171 Point::new(
1172 k.clone(),
1173 BTreeMap::from([
1174 ("w".to_string(), FieldValue::F64(230.0 + s as f64)),
1175 ("kwh".to_string(), FieldValue::F64(s as f64 / 1000.0)),
1176 ]),
1177 s * 1_000_000_000,
1178 )
1179 .unwrap()
1180 })
1181 })
1182 .collect();
1183
1184 let before = thread_live();
1185 let mt = Memtable::new();
1186 for p in &points {
1187 mt.insert(p).unwrap();
1188 }
1189 let measured = usize::try_from(thread_live() - before).unwrap_or(0);
1190 let estimated = mt.estimated_size();
1191 let ratio = estimated as f64 / measured as f64;
1192 assert!(
1193 (0.75..=1.25).contains(&ratio),
1194 "estimate {estimated} B vs allocator {measured} B for {} points: ratio {ratio:.2} \
1195 ({:.0} B/point measured, {:.0} B/point estimated)",
1196 points.len(),
1197 measured as f64 / points.len() as f64,
1198 estimated as f64 / points.len() as f64,
1199 );
1200 drop(mt);
1201 }
1202}
1203
1204#[derive(Debug)]
1206pub struct MeasurementBatch {
1207 pub measurement: String,
1209 pub batch: arrow::record_batch::RecordBatch,
1213 pub tag_columns: Vec<String>,
1215}
1216
1217enum FieldBuilder {
1219 F64(arrow::array::Float64Builder),
1220 I64(arrow::array::Int64Builder),
1221 U64(arrow::array::UInt64Builder),
1222 Bool(arrow::array::BooleanBuilder),
1223 Str(arrow::array::StringBuilder),
1224 Decimal(arrow::array::Decimal128Builder, u8),
1230 Histogram(arrow::array::BinaryBuilder),
1239}
1240
1241impl FieldBuilder {
1242 fn for_value(v: &chronix_core::FieldValue, rows_before: usize) -> Self {
1243 use chronix_core::FieldValue as V;
1244 let mut b = match v {
1245 V::F64(_) => Self::F64(arrow::array::Float64Builder::new()),
1246 V::I64(_) => Self::I64(arrow::array::Int64Builder::new()),
1247 V::U64(_) => Self::U64(arrow::array::UInt64Builder::new()),
1248 V::Bool(_) => Self::Bool(arrow::array::BooleanBuilder::new()),
1249 V::String(_) => Self::Str(arrow::array::StringBuilder::new()),
1250 V::Decimal(d) => Self::Decimal(
1256 arrow::array::Decimal128Builder::new().with_data_type(
1257 arrow::datatypes::DataType::Decimal128(
1258 chronix_core::DECIMAL_PRECISION,
1259 d.scale() as i8,
1260 ),
1261 ),
1262 d.scale(),
1263 ),
1264 V::Histogram(_) => Self::Histogram(arrow::array::BinaryBuilder::new()),
1265 };
1266 for _ in 0..rows_before {
1267 b.append_null();
1268 }
1269 b
1270 }
1271
1272 fn append_null(&mut self) {
1273 match self {
1274 Self::F64(b) => b.append_null(),
1275 Self::I64(b) => b.append_null(),
1276 Self::U64(b) => b.append_null(),
1277 Self::Bool(b) => b.append_null(),
1278 Self::Str(b) => b.append_null(),
1279 Self::Decimal(b, _) => b.append_null(),
1280 Self::Histogram(b) => b.append_null(),
1281 }
1282 }
1283
1284 fn append(&mut self, v: &chronix_core::FieldValue) {
1288 use chronix_core::FieldValue as V;
1289 match (self, v) {
1290 (Self::F64(b), V::F64(x)) => b.append_value(*x),
1291 (Self::I64(b), V::I64(x)) => b.append_value(*x),
1292 (Self::U64(b), V::U64(x)) => b.append_value(*x),
1293 (Self::Bool(b), V::Bool(x)) => b.append_value(*x),
1294 (Self::Str(b), V::String(x)) => b.append_value(x),
1295 (Self::Decimal(b, scale), V::Decimal(x)) => match x.rescale(*scale) {
1300 Ok(v) => b.append_value(v.mantissa()),
1301 Err(_) => b.append_null(),
1302 },
1303 (Self::Histogram(b), V::Histogram(h)) => match postcard::to_allocvec(h.as_ref()) {
1310 Ok(bytes) => b.append_value(bytes),
1311 Err(_) => b.append_null(),
1312 },
1313 (this, _) => this.append_null(),
1314 }
1315 }
1316
1317 fn finish(self) -> (arrow::datatypes::DataType, arrow::array::ArrayRef) {
1318 use arrow::datatypes::DataType;
1319 match self {
1320 Self::Histogram(mut b) => (DataType::Binary, Arc::new(b.finish())),
1321 Self::F64(mut b) => (DataType::Float64, Arc::new(b.finish())),
1322 Self::I64(mut b) => (DataType::Int64, Arc::new(b.finish())),
1323 Self::U64(mut b) => (DataType::UInt64, Arc::new(b.finish())),
1324 Self::Bool(mut b) => (DataType::Boolean, Arc::new(b.finish())),
1325 Self::Str(mut b) => (DataType::Utf8, Arc::new(b.finish())),
1326 Self::Decimal(mut b, scale) => {
1327 let dt = DataType::Decimal128(chronix_core::DECIMAL_PRECISION, scale as i8);
1328 (dt, Arc::new(b.finish()))
1329 }
1330 }
1331 }
1332}
1333
1334struct MeasurementBuilder {
1336 rows: usize,
1337 timestamps: arrow::array::Int64Builder,
1338 tags: std::collections::BTreeMap<Arc<str>, arrow::array::StringBuilder>,
1339 fields: std::collections::BTreeMap<Arc<str>, FieldBuilder>,
1340}
1341
1342impl MeasurementBuilder {
1343 fn new() -> Self {
1344 Self {
1345 rows: 0,
1346 timestamps: arrow::array::Int64Builder::new(),
1347 tags: std::collections::BTreeMap::new(),
1348 fields: std::collections::BTreeMap::new(),
1349 }
1350 }
1351
1352 fn push(&mut self, key: &MemtableKey, entry: &MemtableEntry) {
1353 self.timestamps.append_value(key.timestamp);
1354 for (k, v) in entry.tags.iter() {
1357 let rows = self.rows;
1358 self.tags.entry(Arc::clone(k)).or_insert_with(|| {
1359 let mut b = arrow::array::StringBuilder::new();
1360 for _ in 0..rows {
1361 b.append_null();
1362 }
1363 b
1364 });
1365 let _ = v;
1366 }
1367 for (name, b) in &mut self.tags {
1368 match entry.tags.iter().find(|(k, _)| k == name) {
1369 Some((_, v)) => b.append_value(v),
1370 None => b.append_null(),
1371 }
1372 }
1373 for (k, v) in entry.fields.iter() {
1374 let rows = self.rows;
1375 self.fields
1376 .entry(Arc::clone(k))
1377 .or_insert_with(|| FieldBuilder::for_value(v, rows));
1378 }
1379 for (name, b) in &mut self.fields {
1380 match entry.fields.iter().find(|(k, _)| k == name) {
1381 Some((_, v)) => b.append(v),
1382 None => b.append_null(),
1383 }
1384 }
1385 self.rows += 1;
1386 }
1387
1388 fn finish(mut self, measurement: &str) -> Option<MeasurementBatch> {
1389 use arrow::datatypes::{DataType, Field, Schema};
1390 if self.rows == 0 {
1391 return None;
1392 }
1393 use crate::segment::metadata::roles;
1394 let mut fields = vec![
1397 Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false)
1398 .with_metadata(roles::arrow_metadata(roles::TIMESTAMP)),
1399 ];
1400 let mut columns: Vec<arrow::array::ArrayRef> = vec![Arc::new(self.timestamps.finish())];
1401 let tag_columns: Vec<String> = self.tags.keys().map(ToString::to_string).collect();
1402 for (name, mut b) in self.tags {
1403 fields.push(
1404 Field::new(name.as_ref(), DataType::Utf8, true)
1405 .with_metadata(roles::arrow_metadata(roles::TAG)),
1406 );
1407 columns.push(Arc::new(b.finish()));
1408 }
1409 for (name, b) in self.fields {
1410 let (dt, arr) = b.finish();
1411 fields.push(
1412 Field::new(name.as_ref(), dt, true)
1413 .with_metadata(roles::arrow_metadata(roles::FIELD)),
1414 );
1415 columns.push(arr);
1416 }
1417 let batch =
1418 arrow::record_batch::RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
1419 .ok()?;
1420 Some(MeasurementBatch {
1421 measurement: measurement.to_string(),
1422 batch,
1423 tag_columns,
1424 })
1425 }
1426}
1427
1428impl Memtable {
1429 #[must_use]
1438 pub fn to_record_batches(&self) -> Vec<MeasurementBatch> {
1439 let mut builders: std::collections::BTreeMap<Arc<str>, MeasurementBuilder> =
1440 std::collections::BTreeMap::new();
1441 for entry in self.data.iter() {
1442 let value = entry.value();
1443 builders
1444 .entry(Arc::clone(&value.measurement))
1445 .or_insert_with(MeasurementBuilder::new)
1446 .push(entry.key(), value);
1447 }
1448 builders
1449 .into_iter()
1450 .filter_map(|(m, b)| b.finish(&m))
1451 .collect()
1452 }
1453}