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]
446 pub fn series_canonical_forms(&self) -> Vec<Arc<str>> {
447 self.series_tags
448 .iter()
449 .map(|e| Arc::clone(e.key()))
450 .collect()
451 }
452
453 pub fn visit_tags(&self, mut f: impl FnMut(&str, &str)) {
462 for entry in &self.series_tags {
463 for (key, value) in entry.value().iter() {
464 f(key, value);
465 }
466 }
467 }
468
469 #[must_use]
471 pub fn len(&self) -> usize {
472 self.data.len()
473 }
474
475 #[must_use]
477 pub fn is_empty(&self) -> bool {
478 self.data.is_empty()
479 }
480
481 #[must_use]
491 pub fn interner_bytes(&self) -> usize {
492 self.interner.memory_bytes()
493 }
494
495 #[must_use]
497 pub fn estimated_size(&self) -> usize {
498 self.estimated_size.load(Ordering::Relaxed) + self.index_size.load(Ordering::Relaxed)
499 }
500
501 fn series_tag_set(
507 &self,
508 canonical: &Arc<str>,
509 series_key: &SeriesKey,
510 ) -> Arc<[(Arc<str>, Arc<str>)]> {
511 if let Some(existing) = self.series_tags.get(canonical) {
512 return Arc::clone(existing.value());
513 }
514 let tags: Arc<[(Arc<str>, Arc<str>)]> = series_key
515 .tags()
516 .iter()
517 .map(|(k, v)| (self.interner.intern(k), self.interner.intern(v)))
518 .collect();
519 let mut charged = 0usize;
520 let stored = self
521 .series_tags
522 .entry(Arc::clone(canonical))
523 .or_insert_with(|| {
524 charged = canonical.len()
528 + tags.len() * std::mem::size_of::<(Arc<str>, Arc<str>)>()
529 + tags.iter().map(|(k, v)| k.len() + v.len()).sum::<usize>()
530 + 2 * std::mem::size_of::<usize>()
531 + 64;
532 tags
533 })
534 .clone();
535 if charged > 0 {
536 self.estimated_size.fetch_add(charged, Ordering::Relaxed);
537 }
538 stored
539 }
540
541 fn index_insert(&self, measurement: Arc<str>, hash: u64) -> usize {
550 const SET_OVERHEAD: usize = 56;
551 match self.measurement_index.entry(measurement) {
552 Entry::Occupied(mut occupied) => {
553 let before = occupied.get().capacity();
554 occupied.get_mut().insert(hash);
555 (occupied.get().capacity() - before) * std::mem::size_of::<u64>()
556 }
557 Entry::Vacant(vacant) => {
558 let key_bytes = std::mem::size_of::<usize>() * 2 + vacant.key().len();
559 let mut set = HashSet::new();
560 set.insert(hash);
561 let set_bytes = set.capacity() * std::mem::size_of::<u64>() + SET_OVERHEAD;
562 vacant.insert(set);
563 key_bytes + set_bytes
564 }
565 }
566 }
567
568 #[must_use]
572 pub fn min_wal_seq(&self) -> Option<u64> {
573 let val = self.min_wal_seq.load(Ordering::Acquire);
574 if val == u64::MAX {
575 None
576 } else {
577 Some(val)
578 }
579 }
580
581 #[must_use]
585 pub fn max_wal_seq(&self) -> Option<u64> {
586 let val = self.max_wal_seq.load(Ordering::Acquire);
587 if val == 0 && self.min_wal_seq.load(Ordering::Acquire) == u64::MAX {
588 None
589 } else {
590 Some(val)
591 }
592 }
593
594 #[must_use]
596 pub fn is_frozen(&self) -> bool {
597 self.frozen.load(Ordering::Acquire)
598 }
599
600 pub fn freeze(&self) {
605 self.frozen.store(true, Ordering::Release);
606 }
607
608 #[must_use]
610 pub fn to_points(&self) -> Vec<Point> {
611 self.scan_all()
612 }
613
614 fn entry_to_point(key: &MemtableKey, entry: &MemtableEntry) -> Option<Point> {
619 let tags = entry
620 .tags
621 .iter()
622 .map(|(k, v)| (k.to_string(), v.to_string()))
623 .collect();
624 let series_key = SeriesKey::new(entry.measurement.to_string(), tags)
625 .map_err(|e| {
626 tracing::error!(
629 hash = key.series_key_hash,
630 error = %e,
631 "DATA LOSS: failed to reconstruct series key during flush"
632 );
633 metrics::counter!("chronix_memtable_dropped_points_total").increment(1);
634 e
635 })
636 .ok()?;
637 Point::new(
638 series_key,
639 entry
640 .fields
641 .iter()
642 .map(|(k, v)| (k.to_string(), v.clone()))
643 .collect(),
644 key.timestamp,
645 )
646 .map_err(|e| {
647 tracing::error!(
648 hash = key.series_key_hash,
649 timestamp = key.timestamp,
650 error = %e,
651 "DATA LOSS: failed to reconstruct point during flush"
652 );
653 metrics::counter!("chronix_memtable_dropped_points_total").increment(1);
654 e
655 })
656 .ok()
657 }
658}
659
660impl Default for Memtable {
661 fn default() -> Self {
662 Self::new()
663 }
664}
665
666#[cfg(test)]
671mod tests {
672 use super::*;
673 use std::collections::BTreeMap;
674
675 use chronix_core::FieldValue;
676
677 fn make_point(measurement: &str, host: &str, ts: i64, value: f64) -> Point {
678 let tags: BTreeMap<String, String> = [("host".to_string(), host.to_string())]
679 .into_iter()
680 .collect();
681 let series_key = SeriesKey::new(measurement.to_string(), tags).unwrap();
682 let fields: BTreeMap<String, FieldValue> = [("value".to_string(), FieldValue::F64(value))]
683 .into_iter()
684 .collect();
685 Point::new(series_key, fields, ts).unwrap()
686 }
687
688 #[test]
689 fn insert_and_scan_all() {
690 let mt = Memtable::new();
691 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
692 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
693 mt.insert(&make_point("cpu", "a", 300, 3.0)).unwrap();
694
695 let points = mt.scan_all();
696 assert_eq!(points.len(), 3);
697 assert_eq!(points[0].timestamp(), 100);
698 assert_eq!(points[1].timestamp(), 200);
699 assert_eq!(points[2].timestamp(), 300);
700 }
701
702 #[test]
703 fn insert_and_scan_series() {
704 let mt = Memtable::new();
705 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
706 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
707 mt.insert(&make_point("cpu", "b", 150, 9.0)).unwrap();
708
709 let key_a = SeriesKey::new(
710 "cpu".to_string(),
711 [("host".to_string(), "a".to_string())]
712 .into_iter()
713 .collect(),
714 )
715 .unwrap();
716
717 let points = mt.scan(&key_a, 0, 300);
718 assert_eq!(points.len(), 2);
719 assert!(points
720 .iter()
721 .all(|p| p.series_key().tag("host") == Some("a")));
722 }
723
724 #[test]
725 fn scan_time_range() {
726 let mt = Memtable::new();
727 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
728 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
729 mt.insert(&make_point("cpu", "a", 300, 3.0)).unwrap();
730
731 let key = SeriesKey::new(
732 "cpu".to_string(),
733 [("host".to_string(), "a".to_string())]
734 .into_iter()
735 .collect(),
736 )
737 .unwrap();
738
739 let points = mt.scan(&key, 150, 250);
740 assert_eq!(points.len(), 1);
741 assert_eq!(points[0].timestamp(), 200);
742 }
743
744 #[test]
745 fn last_write_wins() {
746 let mt = Memtable::new();
747 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
748 mt.insert(&make_point("cpu", "a", 100, 99.0)).unwrap();
749
750 let points = mt.scan_all();
751 assert_eq!(points.len(), 1);
752 assert_eq!(points[0].field("value"), Some(&FieldValue::F64(99.0)));
753 }
754
755 #[test]
756 fn empty_scan() {
757 let mt = Memtable::new();
758 let key = SeriesKey::new("cpu".to_string(), BTreeMap::new()).unwrap();
759 let points = mt.scan(&key, 0, 1000);
760 assert!(points.is_empty());
761 }
762
763 #[test]
766 fn index_accounting_matches_a_full_walk() {
767 let mt = Memtable::new();
768 for i in 0..500u64 {
769 let measurement = format!("m{}", i % 7);
770 let series = format!("s{}", i % 53);
771 mt.insert(&make_point(&measurement, &series, i as i64, 1.0))
772 .unwrap();
773 }
774
775 let batch: Vec<_> = (500..800u64)
777 .map(|i| {
778 make_point(
779 &format!("m{}", i % 11),
780 &format!("s{}", i % 97),
781 i as i64,
782 2.0,
783 )
784 })
785 .collect();
786 mt.insert_batch(&batch).unwrap();
787
788 let walked: usize = mt
789 .measurement_index
790 .iter()
791 .map(|entry| {
792 let key_bytes = std::mem::size_of::<usize>() * 2 + entry.key().len();
793 let set_bytes = entry.value().capacity() * std::mem::size_of::<u64>() + 56;
794 key_bytes + set_bytes
795 })
796 .sum();
797
798 assert_eq!(
799 mt.index_size.load(Ordering::Relaxed),
800 walked,
801 "incremental index accounting drifted from a full walk"
802 );
803 }
804
805 #[test]
806 fn estimated_size_increases() {
807 let mt = Memtable::new();
808 assert_eq!(mt.estimated_size(), 0);
809 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
810 assert!(mt.estimated_size() > 0);
811 let size_after_one = mt.estimated_size();
812 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
813 assert!(mt.estimated_size() > size_after_one);
814 }
815
816 #[test]
817 fn wal_seq_tracking() {
818 let mt = Memtable::new();
819 assert_eq!(mt.min_wal_seq(), None);
820 assert_eq!(mt.max_wal_seq(), None);
821
822 mt.insert_with_wal_seq(&make_point("cpu", "a", 100, 1.0), 10)
823 .unwrap();
824 assert_eq!(mt.min_wal_seq(), Some(10));
825 assert_eq!(mt.max_wal_seq(), Some(10));
826
827 mt.insert_with_wal_seq(&make_point("cpu", "a", 200, 2.0), 5)
828 .unwrap();
829 assert_eq!(mt.min_wal_seq(), Some(5));
830 assert_eq!(mt.max_wal_seq(), Some(10));
831
832 mt.insert_with_wal_seq(&make_point("cpu", "a", 300, 3.0), 20)
833 .unwrap();
834 assert_eq!(mt.min_wal_seq(), Some(5));
835 assert_eq!(mt.max_wal_seq(), Some(20));
836 }
837
838 #[test]
839 fn freeze_prevents_inserts() {
840 let mt = Memtable::new();
841 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
842 mt.freeze();
843 assert!(mt.is_frozen());
844
845 let result = mt.insert(&make_point("cpu", "a", 200, 2.0));
846 assert!(result.is_err());
847
848 let points = mt.scan_all();
850 assert_eq!(points.len(), 1);
851 }
852
853 #[test]
854 fn len_and_is_empty() {
855 let mt = Memtable::new();
856 assert!(mt.is_empty());
857 assert_eq!(mt.len(), 0);
858
859 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
860 assert!(!mt.is_empty());
861 assert_eq!(mt.len(), 1);
862 }
863
864 #[test]
865 fn multiple_series() {
866 let mt = Memtable::new();
867 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
868 mt.insert(&make_point("cpu", "b", 100, 2.0)).unwrap();
869 mt.insert(&make_point("mem", "a", 100, 3.0)).unwrap();
870
871 assert_eq!(mt.len(), 3);
872
873 let all = mt.scan_all();
874 assert_eq!(all.len(), 3);
875 }
876
877 #[test]
878 fn concurrent_inserts() {
879 use std::sync::Arc;
880 use std::thread;
881
882 let mt = Arc::new(Memtable::new());
883 let mut handles = Vec::new();
884
885 for thread_id in 0..4 {
886 let mt = Arc::clone(&mt);
887 handles.push(thread::spawn(move || {
888 for i in 0..100 {
889 let ts = i64::from(thread_id * 1000 + i);
890 let host = format!("host-{thread_id}");
891 mt.insert(&make_point("cpu", &host, ts, ts as f64)).unwrap();
892 }
893 }));
894 }
895
896 for h in handles {
897 h.join().unwrap();
898 }
899
900 assert_eq!(mt.len(), 400);
901 let all = mt.scan_all();
902 assert_eq!(all.len(), 400);
903
904 let mut prev_key: Option<MemtableKey> = None;
906 for p in &all {
907 let key = MemtableKey::new(
908 p.series_key().hash_fnv(),
909 p.timestamp(),
910 Arc::from(p.series_key().canonical_form()),
911 );
912 if let Some(prev) = prev_key {
913 assert!(key >= prev, "entries should be ordered");
914 }
915 prev_key = Some(key);
916 }
917 }
918
919 #[allow(deprecated)]
920 #[test]
921 fn to_points_is_scan_all() {
922 let mt = Memtable::new();
923 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
924 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
925
926 let scan = mt.scan_all();
927 let points = mt.to_points();
928 assert_eq!(scan.len(), points.len());
929 for (s, d) in scan.iter().zip(points.iter()) {
930 assert_eq!(s.timestamp(), d.timestamp());
931 assert_eq!(s.fields(), d.fields());
932 }
933 }
934
935 #[test]
936 fn scan_range_across_series() {
937 let mt = Memtable::new();
938 mt.insert(&make_point("cpu", "a", 50, 1.0)).unwrap();
939 mt.insert(&make_point("cpu", "a", 100, 2.0)).unwrap();
940 mt.insert(&make_point("cpu", "b", 75, 3.0)).unwrap();
941 mt.insert(&make_point("cpu", "b", 200, 4.0)).unwrap();
942
943 let points = mt.scan_range(60, 150);
944 assert_eq!(points.len(), 2);
946 for p in &points {
947 assert!(
948 p.timestamp() >= 60 && p.timestamp() <= 150,
949 "timestamp {} should be in [60, 150]",
950 p.timestamp()
951 );
952 }
953 }
954
955 #[test]
956 fn scan_measurement_filters_by_name() {
957 let mt = Memtable::new();
958 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
959 mt.insert(&make_point("cpu", "b", 200, 2.0)).unwrap();
960 mt.insert(&make_point("mem", "a", 150, 3.0)).unwrap();
961 mt.insert(&make_point("mem", "a", 300, 4.0)).unwrap();
962 mt.insert(&make_point("disk", "c", 250, 5.0)).unwrap();
963
964 let cpu = mt.scan_measurement("cpu", 0, i64::MAX);
966 assert_eq!(cpu.len(), 2);
967 assert!(cpu.iter().all(|p| p.series_key().measurement() == "cpu"));
968
969 let mem = mt.scan_measurement("mem", 0, i64::MAX);
971 assert_eq!(mem.len(), 2);
972 assert!(mem.iter().all(|p| p.series_key().measurement() == "mem"));
973
974 let disk = mt.scan_measurement("disk", 0, i64::MAX);
976 assert_eq!(disk.len(), 1);
977
978 let none = mt.scan_measurement("net", 0, i64::MAX);
980 assert!(none.is_empty());
981 }
982
983 #[test]
984 fn scan_measurement_respects_time_range() {
985 let mt = Memtable::new();
986 mt.insert(&make_point("cpu", "a", 100, 1.0)).unwrap();
987 mt.insert(&make_point("cpu", "a", 200, 2.0)).unwrap();
988 mt.insert(&make_point("cpu", "a", 300, 3.0)).unwrap();
989 mt.insert(&make_point("mem", "a", 200, 4.0)).unwrap();
990
991 let result = mt.scan_measurement("cpu", 150, 250);
992 assert_eq!(result.len(), 1);
993 assert_eq!(result[0].timestamp(), 200);
994 }
995
996 #[test]
1000 fn hash_collision_preserves_both_series() {
1001 let mt = Memtable::new();
1002
1003 let p1 = make_point("cpu", "host-a", 1000, 42.0);
1007 let p2 = make_point("cpu", "host-b", 1000, 99.0);
1008
1009 mt.insert(&p1).unwrap();
1010 mt.insert(&p2).unwrap();
1011
1012 assert_eq!(
1014 mt.len(),
1015 2,
1016 "two distinct series at same timestamp must both survive"
1017 );
1018
1019 let s1 = mt.scan(p1.series_key(), 0, i64::MAX);
1021 assert_eq!(s1.len(), 1);
1022 assert_eq!(
1023 *s1[0].field("value").unwrap(),
1024 chronix_core::FieldValue::F64(42.0)
1025 );
1026
1027 let s2 = mt.scan(p2.series_key(), 0, i64::MAX);
1028 assert_eq!(s2.len(), 1);
1029 assert_eq!(
1030 *s2[0].field("value").unwrap(),
1031 chronix_core::FieldValue::F64(99.0)
1032 );
1033 }
1034
1035 #[test]
1036 fn insert_batch_basic() {
1037 let mt = Memtable::new();
1038 let points = vec![
1039 make_point("cpu", "a", 100, 1.0),
1040 make_point("cpu", "a", 200, 2.0),
1041 make_point("cpu", "b", 300, 3.0),
1042 make_point("mem", "a", 400, 4.0),
1043 ];
1044 let count = mt.insert_batch(&points).unwrap();
1045 assert_eq!(count, 4);
1046
1047 let all = mt.scan_all();
1048 assert_eq!(all.len(), 4);
1049
1050 let cpu = mt.scan_measurement("cpu", 0, i64::MAX);
1052 assert_eq!(cpu.len(), 3);
1053 let mem = mt.scan_measurement("mem", 0, i64::MAX);
1054 assert_eq!(mem.len(), 1);
1055 }
1056
1057 #[test]
1058 fn insert_batch_overwrite() {
1059 let mt = Memtable::new();
1060 let initial = vec![
1061 make_point("cpu", "a", 100, 1.0),
1062 make_point("cpu", "a", 200, 2.0),
1063 ];
1064 mt.insert_batch(&initial).unwrap();
1065
1066 let overwrite = vec![
1068 make_point("cpu", "a", 100, 10.0),
1069 make_point("cpu", "a", 200, 20.0),
1070 ];
1071 mt.insert_batch(&overwrite).unwrap();
1072
1073 let all = mt.scan_all();
1074 assert_eq!(all.len(), 2);
1075 assert_eq!(
1076 *all[0].field("value").unwrap(),
1077 chronix_core::FieldValue::F64(10.0)
1078 );
1079 assert_eq!(
1080 *all[1].field("value").unwrap(),
1081 chronix_core::FieldValue::F64(20.0)
1082 );
1083 }
1084
1085 #[test]
1086 fn insert_batch_frozen_fails() {
1087 let mt = Memtable::new();
1088 mt.freeze();
1089 let points = vec![make_point("cpu", "a", 100, 1.0)];
1090 assert!(mt.insert_batch(&points).is_err());
1091 }
1092
1093 #[test]
1094 fn insert_batch_with_wal_seq_updates_range() {
1095 let mt = Memtable::new();
1096 let points = vec![
1097 make_point("cpu", "a", 100, 1.0),
1098 make_point("cpu", "a", 200, 2.0),
1099 ];
1100 mt.insert_batch_with_wal_seq(&points, 42).unwrap();
1101 assert_eq!(mt.min_wal_seq(), Some(42));
1102 assert_eq!(mt.max_wal_seq(), Some(42));
1103 }
1104}
1105
1106#[cfg(test)]
1112mod calibration {
1113 use std::alloc::{GlobalAlloc, Layout, System};
1114 use std::collections::BTreeMap;
1115 use std::sync::atomic::{AtomicUsize, Ordering};
1116
1117 use chronix_core::{FieldValue, Point, SeriesKey};
1118
1119 use super::Memtable;
1120
1121 struct Counting;
1122 static LIVE: AtomicUsize = AtomicUsize::new(0);
1123
1124 std::thread_local! {
1127 static THREAD_LIVE: std::cell::Cell<isize> = const { std::cell::Cell::new(0) };
1128 }
1129
1130 fn thread_live() -> isize {
1131 THREAD_LIVE.with(std::cell::Cell::get)
1132 }
1133
1134 #[allow(unsafe_code)]
1138 unsafe impl GlobalAlloc for Counting {
1139 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
1140 LIVE.fetch_add(layout.size(), Ordering::Relaxed);
1141 let _ = THREAD_LIVE.try_with(|c| c.set(c.get() + layout.size() as isize));
1142 unsafe { System.alloc(layout) }
1143 }
1144 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
1145 LIVE.fetch_sub(layout.size(), Ordering::Relaxed);
1146 let _ = THREAD_LIVE.try_with(|c| c.set(c.get() - layout.size() as isize));
1147 unsafe { System.dealloc(ptr, layout) }
1148 }
1149 }
1150
1151 #[global_allocator]
1152 static ALLOC: Counting = Counting;
1153
1154 #[test]
1158 fn memtable_estimate_tracks_real_allocation() {
1159 let keys: Vec<SeriesKey> = (0..50)
1160 .map(|m| {
1161 SeriesKey::new(
1162 "power",
1163 BTreeMap::from([
1164 ("meter".to_string(), format!("m{m:02}")),
1165 ("site".to_string(), "home".to_string()),
1166 ]),
1167 )
1168 .unwrap()
1169 })
1170 .collect();
1171 let points: Vec<Point> = (0..400i64)
1173 .flat_map(|s| {
1174 keys.iter().map(move |k| {
1175 Point::new(
1176 k.clone(),
1177 BTreeMap::from([
1178 ("w".to_string(), FieldValue::F64(230.0 + s as f64)),
1179 ("kwh".to_string(), FieldValue::F64(s as f64 / 1000.0)),
1180 ]),
1181 s * 1_000_000_000,
1182 )
1183 .unwrap()
1184 })
1185 })
1186 .collect();
1187
1188 let before = thread_live();
1189 let mt = Memtable::new();
1190 for p in &points {
1191 mt.insert(p).unwrap();
1192 }
1193 let measured = usize::try_from(thread_live() - before).unwrap_or(0);
1194 let estimated = mt.estimated_size();
1195 let ratio = estimated as f64 / measured as f64;
1196 assert!(
1197 (0.75..=1.25).contains(&ratio),
1198 "estimate {estimated} B vs allocator {measured} B for {} points: ratio {ratio:.2} \
1199 ({:.0} B/point measured, {:.0} B/point estimated)",
1200 points.len(),
1201 measured as f64 / points.len() as f64,
1202 estimated as f64 / points.len() as f64,
1203 );
1204 drop(mt);
1205 }
1206}
1207
1208#[derive(Debug)]
1210pub struct MeasurementBatch {
1211 pub measurement: String,
1213 pub batch: arrow::record_batch::RecordBatch,
1217 pub tag_columns: Vec<String>,
1219}
1220
1221enum FieldBuilder {
1223 F64(arrow::array::Float64Builder),
1224 I64(arrow::array::Int64Builder),
1225 U64(arrow::array::UInt64Builder),
1226 Bool(arrow::array::BooleanBuilder),
1227 Str(arrow::array::StringBuilder),
1228 Decimal(arrow::array::Decimal128Builder, u8),
1234}
1235
1236impl FieldBuilder {
1237 fn for_value(v: &chronix_core::FieldValue, rows_before: usize) -> Self {
1238 use chronix_core::FieldValue as V;
1239 let mut b = match v {
1240 V::F64(_) => Self::F64(arrow::array::Float64Builder::new()),
1241 V::I64(_) => Self::I64(arrow::array::Int64Builder::new()),
1242 V::U64(_) => Self::U64(arrow::array::UInt64Builder::new()),
1243 V::Bool(_) => Self::Bool(arrow::array::BooleanBuilder::new()),
1244 V::String(_) => Self::Str(arrow::array::StringBuilder::new()),
1245 V::Decimal(d) => Self::Decimal(
1251 arrow::array::Decimal128Builder::new().with_data_type(
1252 arrow::datatypes::DataType::Decimal128(
1253 chronix_core::DECIMAL_PRECISION,
1254 d.scale() as i8,
1255 ),
1256 ),
1257 d.scale(),
1258 ),
1259 };
1260 for _ in 0..rows_before {
1261 b.append_null();
1262 }
1263 b
1264 }
1265
1266 fn append_null(&mut self) {
1267 match self {
1268 Self::F64(b) => b.append_null(),
1269 Self::I64(b) => b.append_null(),
1270 Self::U64(b) => b.append_null(),
1271 Self::Bool(b) => b.append_null(),
1272 Self::Str(b) => b.append_null(),
1273 Self::Decimal(b, _) => b.append_null(),
1274 }
1275 }
1276
1277 fn append(&mut self, v: &chronix_core::FieldValue) {
1281 use chronix_core::FieldValue as V;
1282 match (self, v) {
1283 (Self::F64(b), V::F64(x)) => b.append_value(*x),
1284 (Self::I64(b), V::I64(x)) => b.append_value(*x),
1285 (Self::U64(b), V::U64(x)) => b.append_value(*x),
1286 (Self::Bool(b), V::Bool(x)) => b.append_value(*x),
1287 (Self::Str(b), V::String(x)) => b.append_value(x),
1288 (Self::Decimal(b, scale), V::Decimal(x)) => match x.rescale(*scale) {
1293 Ok(v) => b.append_value(v.mantissa()),
1294 Err(_) => b.append_null(),
1295 },
1296 (this, _) => this.append_null(),
1297 }
1298 }
1299
1300 fn finish(self) -> (arrow::datatypes::DataType, arrow::array::ArrayRef) {
1301 use arrow::datatypes::DataType;
1302 match self {
1303 Self::F64(mut b) => (DataType::Float64, Arc::new(b.finish())),
1304 Self::I64(mut b) => (DataType::Int64, Arc::new(b.finish())),
1305 Self::U64(mut b) => (DataType::UInt64, Arc::new(b.finish())),
1306 Self::Bool(mut b) => (DataType::Boolean, Arc::new(b.finish())),
1307 Self::Str(mut b) => (DataType::Utf8, Arc::new(b.finish())),
1308 Self::Decimal(mut b, scale) => {
1309 let dt = DataType::Decimal128(chronix_core::DECIMAL_PRECISION, scale as i8);
1310 (dt, Arc::new(b.finish()))
1311 }
1312 }
1313 }
1314}
1315
1316struct MeasurementBuilder {
1318 rows: usize,
1319 timestamps: arrow::array::Int64Builder,
1320 tags: std::collections::BTreeMap<Arc<str>, arrow::array::StringBuilder>,
1321 fields: std::collections::BTreeMap<Arc<str>, FieldBuilder>,
1322}
1323
1324impl MeasurementBuilder {
1325 fn new() -> Self {
1326 Self {
1327 rows: 0,
1328 timestamps: arrow::array::Int64Builder::new(),
1329 tags: std::collections::BTreeMap::new(),
1330 fields: std::collections::BTreeMap::new(),
1331 }
1332 }
1333
1334 fn push(&mut self, key: &MemtableKey, entry: &MemtableEntry) {
1335 self.timestamps.append_value(key.timestamp);
1336 for (k, v) in entry.tags.iter() {
1339 let rows = self.rows;
1340 self.tags.entry(Arc::clone(k)).or_insert_with(|| {
1341 let mut b = arrow::array::StringBuilder::new();
1342 for _ in 0..rows {
1343 b.append_null();
1344 }
1345 b
1346 });
1347 let _ = v;
1348 }
1349 for (name, b) in &mut self.tags {
1350 match entry.tags.iter().find(|(k, _)| k == name) {
1351 Some((_, v)) => b.append_value(v),
1352 None => b.append_null(),
1353 }
1354 }
1355 for (k, v) in entry.fields.iter() {
1356 let rows = self.rows;
1357 self.fields
1358 .entry(Arc::clone(k))
1359 .or_insert_with(|| FieldBuilder::for_value(v, rows));
1360 }
1361 for (name, b) in &mut self.fields {
1362 match entry.fields.iter().find(|(k, _)| k == name) {
1363 Some((_, v)) => b.append(v),
1364 None => b.append_null(),
1365 }
1366 }
1367 self.rows += 1;
1368 }
1369
1370 fn finish(mut self, measurement: &str) -> Option<MeasurementBatch> {
1371 use arrow::datatypes::{DataType, Field, Schema};
1372 if self.rows == 0 {
1373 return None;
1374 }
1375 use crate::segment::metadata::roles;
1376 let mut fields = vec![
1379 Field::new(chronix_core::TIME_COLUMN, DataType::Int64, false)
1380 .with_metadata(roles::arrow_metadata(roles::TIMESTAMP)),
1381 ];
1382 let mut columns: Vec<arrow::array::ArrayRef> = vec![Arc::new(self.timestamps.finish())];
1383 let tag_columns: Vec<String> = self.tags.keys().map(ToString::to_string).collect();
1384 for (name, mut b) in self.tags {
1385 fields.push(
1386 Field::new(name.as_ref(), DataType::Utf8, true)
1387 .with_metadata(roles::arrow_metadata(roles::TAG)),
1388 );
1389 columns.push(Arc::new(b.finish()));
1390 }
1391 for (name, b) in self.fields {
1392 let (dt, arr) = b.finish();
1393 fields.push(
1394 Field::new(name.as_ref(), dt, true)
1395 .with_metadata(roles::arrow_metadata(roles::FIELD)),
1396 );
1397 columns.push(arr);
1398 }
1399 let batch =
1400 arrow::record_batch::RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
1401 .ok()?;
1402 Some(MeasurementBatch {
1403 measurement: measurement.to_string(),
1404 batch,
1405 tag_columns,
1406 })
1407 }
1408}
1409
1410impl Memtable {
1411 #[must_use]
1420 pub fn to_record_batches(&self) -> Vec<MeasurementBatch> {
1421 let mut builders: std::collections::BTreeMap<Arc<str>, MeasurementBuilder> =
1422 std::collections::BTreeMap::new();
1423 for entry in self.data.iter() {
1424 let value = entry.value();
1425 builders
1426 .entry(Arc::clone(&value.measurement))
1427 .or_insert_with(MeasurementBuilder::new)
1428 .push(entry.key(), value);
1429 }
1430 builders
1431 .into_iter()
1432 .filter_map(|(m, b)| b.finish(&m))
1433 .collect()
1434 }
1435}