From 8358fb8a6df7df00c9407c98280e392ac5ed0edf Mon Sep 17 00:00:00 2001
From: Arthur LE MOIGNE <arthur.lemoigne@gmail.com>
Date: Mon, 11 Sep 2023 13:35:01 +0200
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Add=20MmapStats?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/lib.rs          |  2 ++
 src/segment.rs      | 17 ++++++++++-------
 src/stats.rs        | 34 ++++++++++++++++++++++++++++++++++
 tests/test_stats.rs | 18 ++++++++++++++++++
 4 files changed, 64 insertions(+), 7 deletions(-)
 create mode 100644 src/stats.rs
 create mode 100644 tests/test_stats.rs

diff --git a/src/lib.rs b/src/lib.rs
index faffc68..e6044f7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -117,12 +117,14 @@ use std::{
 
 pub use segment::Segment;
 pub use segment_builder::{DefaultSegmentBuilder, SegmentBuilder};
+pub use stats::MmapStats;
 pub use vec_builder::MmapVecBuilder;
 
 use crate::utils::page_size;
 
 mod segment;
 mod segment_builder;
+mod stats;
 mod utils;
 mod vec_builder;
 
diff --git a/src/segment.rs b/src/segment.rs
index 104f85e..e588d48 100644
--- a/src/segment.rs
+++ b/src/segment.rs
@@ -5,8 +5,11 @@ use std::{
     os::{fd::AsRawFd, unix::prelude::FileExt},
     path::{Path, PathBuf},
     ptr, slice,
+    sync::atomic::Ordering,
 };
 
+use crate::stats::{COUNT_ACTIVE_SEGMENT, COUNT_MMAP_FAILED, COUNT_MUNMAP_FAILED};
+
 /// Segment is a constant slice of type T that is memory mapped to disk.
 ///
 /// It is the basic building block of memory mapped data structure.
@@ -66,8 +69,10 @@ impl<T> Segment<T> {
         };
 
         if addr == libc::MAP_FAILED {
+            COUNT_MMAP_FAILED.fetch_add(1, Ordering::Relaxed);
             Err(io::Error::last_os_error())
         } else {
+            COUNT_ACTIVE_SEGMENT.fetch_add(1, Ordering::Relaxed);
             Ok(Self {
                 addr: addr.cast(),
                 len: 0,
@@ -194,13 +199,11 @@ impl<T> Drop for Segment<T> {
             let unmap_code =
                 unsafe { libc::munmap(self.addr.cast(), self.capacity * mem::size_of::<T>()) };
 
-            // Assert that munmap has not failed to force use it's return value and do not optimize out
-            // above call
-            assert!(
-                unmap_code == 0,
-                "munmap failed: {}",
-                io::Error::last_os_error()
-            );
+            if unmap_code == 0 {
+                COUNT_ACTIVE_SEGMENT.fetch_sub(1, Ordering::Relaxed);
+            } else {
+                COUNT_MUNMAP_FAILED.fetch_add(1, Ordering::Relaxed);
+            }
         }
 
         if let Some(path) = &self.path {
diff --git a/src/stats.rs b/src/stats.rs
new file mode 100644
index 0000000..fae6522
--- /dev/null
+++ b/src/stats.rs
@@ -0,0 +1,34 @@
+use std::sync::atomic::{AtomicU64, Ordering};
+
+pub(crate) static COUNT_ACTIVE_SEGMENT: AtomicU64 = AtomicU64::new(0);
+pub(crate) static COUNT_MMAP_FAILED: AtomicU64 = AtomicU64::new(0);
+pub(crate) static COUNT_MUNMAP_FAILED: AtomicU64 = AtomicU64::new(0);
+
+/// Provides few statistics about low level segment allocation.
+///
+/// This stats can be useful to debug or to export in various monitoring
+/// systems.
+#[derive(Debug, Default)]
+pub struct MmapStats;
+
+impl MmapStats {
+    /// Get number of current segment mounted by this library.
+    ///
+    /// On linux there is a `systctl` limit you can access with:
+    /// ```shell
+    /// sysctl vm.max_map_count
+    /// ```
+    pub fn active_segment(&self) -> u64 {
+        COUNT_ACTIVE_SEGMENT.load(Ordering::Relaxed)
+    }
+
+    /// Get number of segment creation failed.
+    pub fn map_failed(&self) -> u64 {
+        COUNT_MMAP_FAILED.load(Ordering::Relaxed)
+    }
+
+    /// Get number of segment deletion failed.
+    pub fn unmap_failed(&self) -> u64 {
+        COUNT_MUNMAP_FAILED.load(Ordering::Relaxed)
+    }
+}
diff --git a/tests/test_stats.rs b/tests/test_stats.rs
new file mode 100644
index 0000000..2c8c4b6
--- /dev/null
+++ b/tests/test_stats.rs
@@ -0,0 +1,18 @@
+use mmap_vec::{MmapStats, MmapVec};
+
+#[test]
+fn test_stats() {
+    assert_eq!(MmapStats.active_segment(), 0);
+    assert_eq!(MmapStats.map_failed(), 0);
+    assert_eq!(MmapStats.unmap_failed(), 0);
+
+    let v = MmapVec::<u8>::with_capacity(500).unwrap();
+    assert_eq!(MmapStats.active_segment(), 1);
+    assert_eq!(MmapStats.map_failed(), 0);
+    assert_eq!(MmapStats.unmap_failed(), 0);
+
+    drop(v);
+    assert_eq!(MmapStats.active_segment(), 0);
+    assert_eq!(MmapStats.map_failed(), 0);
+    assert_eq!(MmapStats.unmap_failed(), 0);
+}