[Memtrie] (2/n) Construction, refcounting, and lookup of memtrie. (ne…

…ar#9646)

Construction: efficient bottom-up algorithm for constructing in-memory
trie from flat storage.
Refcounting: maintenance of multiple roots, and GC based on heights.
Lookup: looking up a FlatStateValue from a root of the memtrie.

Also includes a CLI tool to load the trie, to easily measure timing and
memory usage.

core/store/src/trie/mem/arena/alloc.rs

@@ -1,11 +1,26 @@
+use near_o11y::metrics::IntGauge;
+
+use super::metrics::MEM_TRIE_ARENA_ACTIVE_ALLOCS_COUNT;
 use super::{ArenaMemory, ArenaSliceMut};
+use crate::trie::mem::arena::metrics::{
+    MEM_TRIE_ARENA_ACTIVE_ALLOCS_BYTES, MEM_TRIE_ARENA_MEMORY_USAGE_BYTES,
+};
 
 /// Simple bump allocator with freelists. Allocations are rounded up to its
 /// allocation class, so that deallocated memory can be reused by a similarly
 /// sized allocation.
 pub struct Allocator {
     freelists: [usize; NUM_ALLOCATION_CLASSES],
     next_ptr: usize,
+
+    // Stats. Note that keep the bytes and count locally too because the
+    // gauges are process-wide, so stats-keeping directly with those may not be
+    // accurate in case multiple instances of the allocator share the same name.
+    active_allocs_bytes: usize,
+    active_allocs_count: usize,
+    active_allocs_bytes_gauge: IntGauge,
+    active_allocs_count_gauge: IntGauge,
+    memory_usage_gauge: IntGauge,
 }
 
 const MAX_ALLOC_SIZE: usize = 16 * 1024;
@@ -39,13 +54,27 @@ const fn allocation_size(size_class: usize) -> usize {
 const NUM_ALLOCATION_CLASSES: usize = allocation_class(MAX_ALLOC_SIZE) + 1;
 
 impl Allocator {
-    pub fn new() -> Self {
-        Self { freelists: [usize::MAX; NUM_ALLOCATION_CLASSES], next_ptr: 0 }
+    pub fn new(name: String) -> Self {
+        Self {
+            freelists: [usize::MAX; NUM_ALLOCATION_CLASSES],
+            next_ptr: 0,
+            active_allocs_bytes: 0,
+            active_allocs_count: 0,
+            active_allocs_bytes_gauge: MEM_TRIE_ARENA_ACTIVE_ALLOCS_BYTES
+                .with_label_values(&[&name]),
+            active_allocs_count_gauge: MEM_TRIE_ARENA_ACTIVE_ALLOCS_COUNT
+                .with_label_values(&[&name]),
+            memory_usage_gauge: MEM_TRIE_ARENA_MEMORY_USAGE_BYTES.with_label_values(&[&name]),
+        }
     }
 
     /// Allocates a slice of the given size in the arena.
     pub fn allocate<'a>(&mut self, arena: &'a mut ArenaMemory, size: usize) -> ArenaSliceMut<'a> {
         assert!(size <= MAX_ALLOC_SIZE, "Cannot allocate {} bytes", size);
+        self.active_allocs_bytes += size;
+        self.active_allocs_count += 1;
+        self.active_allocs_bytes_gauge.set(self.active_allocs_bytes as i64);
+        self.active_allocs_count_gauge.set(self.active_allocs_count as i64);
         let size_class = allocation_class(size);
         let allocation_size = allocation_size(size_class);
         if self.freelists[size_class] == usize::MAX {
@@ -60,6 +89,7 @@ impl Allocator {
             }
             let ptr = self.next_ptr;
             self.next_ptr += allocation_size;
+            self.memory_usage_gauge.set(self.next_ptr as i64);
             arena.slice_mut(ptr, size)
         } else {
             let pos = self.freelists[size_class];
@@ -71,25 +101,32 @@ impl Allocator {
     /// Deallocates the given slice from the arena; the slice's `pos` and `len`
     /// must be the same as an allocation that was returned earlier.
     pub fn deallocate(&mut self, arena: &mut ArenaMemory, pos: usize, len: usize) {
+        self.active_allocs_bytes -= len;
+        self.active_allocs_count -= 1;
+        self.active_allocs_bytes_gauge.set(self.active_allocs_bytes as i64);
+        self.active_allocs_count_gauge.set(self.active_allocs_count as i64);
         let size_class = allocation_class(len);
         arena
             .slice_mut(pos, allocation_size(size_class))
             .write_usize_at(0, self.freelists[size_class]);
         self.freelists[size_class] = pos;
     }
+
+    #[cfg(test)]
+    pub fn num_active_allocs(&self) -> usize {
+        self.active_allocs_count
+    }
 }
 
 #[cfg(test)]
 mod test {
-    use std::mem::size_of;
-
-    use crate::trie::mem::arena::Arena;
-
     use super::MAX_ALLOC_SIZE;
+    use crate::trie::mem::arena::Arena;
+    use std::mem::size_of;
 
     #[test]
     fn test_allocate_deallocate() {
-        let mut arena = Arena::new(10000);
+        let mut arena = Arena::new(10000, "".to_owned());
         // Repeatedly allocate and deallocate; we should not run out of memory.
         for i in 0..1000 {
             let mut slices = Vec::new();

core/store/src/trie/mem/arena/metrics.rs

@@ -0,0 +1,29 @@
+use near_o11y::metrics::{try_create_int_gauge_vec, IntGaugeVec};
+use once_cell::sync::Lazy;
+
+pub static MEM_TRIE_ARENA_ACTIVE_ALLOCS_BYTES: Lazy<IntGaugeVec> = Lazy::new(|| {
+    try_create_int_gauge_vec(
+        "near_mem_trie_arena_active_allocs_bytes",
+        "Total size of active allocations on the in-memory trie arena",
+        &["shard_uid"],
+    )
+    .unwrap()
+});
+
+pub static MEM_TRIE_ARENA_ACTIVE_ALLOCS_COUNT: Lazy<IntGaugeVec> = Lazy::new(|| {
+    try_create_int_gauge_vec(
+        "near_mem_trie_arena_active_allocs_count",
+        "Total number of active allocations on the in-memory trie arena",
+        &["shard_uid"],
+    )
+    .unwrap()
+});
+
+pub static MEM_TRIE_ARENA_MEMORY_USAGE_BYTES: Lazy<IntGaugeVec> = Lazy::new(|| {
+    try_create_int_gauge_vec(
+        "near_mem_trie_arena_memory_usage_bytes",
+        "Memory usage of the in-memory trie arena",
+        &["shard_uid"],
+    )
+    .unwrap()
+});

core/store/src/trie/mem/arena/mod.rs

@@ -1,4 +1,5 @@
 mod alloc;
+mod metrics;
 use self::alloc::Allocator;
 use borsh::{BorshDeserialize, BorshSerialize};
 use memmap2::{MmapMut, MmapOptions};
@@ -67,8 +68,8 @@ impl Arena {
     /// The `max_size_in_bytes` can be conservatively large as long as it
     /// can fit into virtual memory (which there are terabytes of). The actual
     /// memory usage will only be as much as is needed.
-    pub fn new(max_size_in_bytes: usize) -> Self {
-        Self { memory: ArenaMemory::new(max_size_in_bytes), allocator: Allocator::new() }
+    pub fn new(max_size_in_bytes: usize, name: String) -> Self {
+        Self { memory: ArenaMemory::new(max_size_in_bytes), allocator: Allocator::new(name) }
     }
 
     /// Allocates a slice of the given size in the arena.
@@ -82,6 +83,12 @@ impl Arena {
         self.allocator.deallocate(&mut self.memory, pos, len);
     }
 
+    /// Number of active allocations (alloc calls minus dealloc calls).
+    #[cfg(test)]
+    pub fn num_active_allocs(&self) -> usize {
+        self.allocator.num_active_allocs()
+    }
+
     pub fn memory(&self) -> &ArenaMemory {
         &self.memory
     }