Merge pull request #4 from eytans/master

Adding simple statistics collection
alecmocatta · Mar 26, 2023 · a635896 · a635896
2 parents 1ab159c + df8cf2c
commit a635896
Show file tree

Hide file tree

Showing 4 changed files with 113 additions and 14 deletions.
diff --git a/.cargo/config.toml b/.cargo/config.toml
@@ -0,0 +1,2 @@
+[env]
+RUST_TEST_THREADS = "1"
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "cap"
-version = "0.1.1"
+version = "0.1.2"
 license = "MIT OR Apache-2.0"
 authors = ["Alec Mocatta <alec@mocatta.net>"]
 categories = []
@@ -23,5 +23,6 @@ maintenance = { status = "passively-maintained" }
 
 [features]
 nightly = []
+stats = []
 
 [dependencies]
diff --git a/azure-pipelines.yml b/azure-pipelines.yml
@@ -14,9 +14,9 @@ jobs:
     endpoint: alecmocatta
     default:
       rust_toolchain: nightly
-      rust_lint_toolchain: nightly-2019-10-20
+      rust_lint_toolchain: nightly-2023-03-09
       rust_flags: ''
-      rust_features: ';nightly'
+      rust_features: ';stats'
       rust_target_check: ''
       rust_target_build: ''
       rust_target_run: ''
@@ -26,7 +26,7 @@ jobs:
         rust_target_run: 'x86_64-pc-windows-msvc i686-pc-windows-msvc' # currently broken building crate-type=lib: x86_64-pc-windows-gnu i686-pc-windows-gnu
       mac:
         imageName: 'macOS-10.15'
-        rust_target_run: 'x86_64-apple-darwin i686-apple-darwin'
+        rust_target_run: 'x86_64-apple-darwin'
       linux:
         imageName: 'ubuntu-18.04'
         rust_target_run: 'x86_64-unknown-linux-gnu i686-unknown-linux-gnu x86_64-unknown-linux-musl i686-unknown-linux-musl'
diff --git a/src/lib.rs b/src/lib.rs
@@ -40,7 +40,11 @@
 	unused_results,
 	clippy::pedantic
 )] // from https://github.com/rust-unofficial/patterns/blob/master/anti_patterns/deny-warnings.md
-#![allow()]
+#![allow(
+	clippy::result_unit_err,
+	clippy::let_underscore_untyped,
+	clippy::missing_errors_doc
+)]
 
 #[cfg(feature = "nightly")]
 use std::alloc::{Alloc, AllocErr, CannotReallocInPlace};
@@ -54,6 +58,10 @@ pub struct Cap<H> {
 	allocator: H,
 	remaining: AtomicUsize,
 	limit: AtomicUsize,
+	#[cfg(feature = "stats")]
+	total_allocated: AtomicUsize,
+	#[cfg(feature = "stats")]
+	max_allocated: AtomicUsize,
 }
 
 impl<H> Cap<H> {
@@ -65,6 +73,10 @@ impl<H> Cap<H> {
 			allocator,
 			remaining: AtomicUsize::new(limit),
 			limit: AtomicUsize::new(limit),
+			#[cfg(feature = "stats")]
+			total_allocated: AtomicUsize::new(0),
+			#[cfg(feature = "stats")]
+			max_allocated: AtomicUsize::new(0),
 		}
 	}
 
@@ -101,16 +113,16 @@ impl<H> Cap<H> {
 				}
 				if self
 					.limit
-					.compare_and_swap(limit_old, limit, Ordering::Relaxed)
-					!= limit_old
+					.compare_exchange(limit_old, limit, Ordering::Relaxed, Ordering::Relaxed)
+					.is_err()
 				{
 					continue;
 				}
 			} else {
 				if self
 					.limit
-					.compare_and_swap(limit_old, limit, Ordering::Relaxed)
-					!= limit_old
+					.compare_exchange(limit_old, limit, Ordering::Relaxed, Ordering::Relaxed)
+					.is_err()
 				{
 					continue;
 				}
@@ -134,6 +146,34 @@ impl<H> Cap<H> {
 			}
 		}
 	}
+
+	/// Get total amount of allocated memory. This includes already deallocated memory.
+	#[cfg(feature = "stats")]
+	pub fn total_allocated(&self) -> usize {
+		self.total_allocated.load(Ordering::Relaxed)
+	}
+
+	/// Get maximum amount of memory that was allocated at any point in time.
+	#[cfg(feature = "stats")]
+	pub fn max_allocated(&self) -> usize {
+		self.max_allocated.load(Ordering::Relaxed)
+	}
+
+	fn update_stats(&self, size: usize) {
+		#[cfg(feature = "stats")]
+		{
+			let _ = self.total_allocated.fetch_add(size, Ordering::Relaxed);
+			// If max_allocated is less than currently allocated, then it will be updated to limit - remaining.
+			// Otherwise, it will remain unchanged.
+			let _ = self
+				.max_allocated
+				.fetch_max(self.allocated(), Ordering::Relaxed);
+		}
+		#[cfg(not(feature = "stats"))]
+		{
+			let _ = (self, size);
+		}
+	}
 }
 
 unsafe impl<H> GlobalAlloc for Cap<H>
@@ -149,6 +189,8 @@ where
 		};
 		if res.is_null() {
 			let _ = self.remaining.fetch_add(size, Ordering::Release);
+		} else {
+			self.update_stats(size);
 		}
 		res
 	}
@@ -166,13 +208,15 @@ where
 		};
 		if res.is_null() {
 			let _ = self.remaining.fetch_add(size, Ordering::Release);
+		} else {
+			self.update_stats(size);
 		}
 		res
 	}
 	unsafe fn realloc(&self, ptr: *mut u8, old_l: Layout, new_s: usize) -> *mut u8 {
 		let new_l = Layout::from_size_align_unchecked(new_s, old_l.align());
 		let (old_size, new_size) = (old_l.size(), new_l.size());
-		if new_size > old_size {
+		let res = if new_size > old_size {
 			let res = if self
 				.remaining
 				.fetch_sub(new_size - old_size, Ordering::Acquire)
@@ -195,8 +239,13 @@ where
 					.remaining
 					.fetch_add(old_size - new_size, Ordering::Release);
 			}
+			// Although this might just deaalocate, I will still update the stats as if it allocates to be on "the safe side"
 			res
+		};
+		if !res.is_null() {
+			self.update_stats(new_size);
 		}
+		res
 	}
 }
 
@@ -214,6 +263,8 @@ where
 		};
 		if res.is_err() {
 			let _ = self.remaining.fetch_add(size, Ordering::Release);
+		} else {
+			self.update_stats(size);
 		}
 		res
 	}
@@ -233,7 +284,7 @@ where
 			self.allocator.usable_size(&old_l).1,
 			self.allocator.usable_size(&new_l).1,
 		);
-		if new_size > old_size {
+		let res = if new_size > old_size {
 			let res = if self
 				.remaining
 				.fetch_sub(new_size - old_size, Ordering::Acquire)
@@ -257,7 +308,11 @@ where
 					.fetch_add(old_size - new_size, Ordering::Release);
 			}
 			res
+		};
+		if res.is_ok() {
+			self.update_stats(new_size);
 		}
+		res
 	}
 	unsafe fn alloc_zeroed(&mut self, l: Layout) -> Result<ptr::NonNull<u8>, AllocErr> {
 		let size = self.allocator.usable_size(&l).1;
@@ -268,6 +323,8 @@ where
 		};
 		if res.is_err() {
 			let _ = self.remaining.fetch_add(size, Ordering::Release);
+		} else {
+			self.update_stats(size);
 		}
 		res
 	}
@@ -292,6 +349,8 @@ where
 			let _ = self
 				.remaining
 				.fetch_add(new_size - old_size, Ordering::Release);
+		} else {
+			self.update_stats(new_size - old_size);
 		}
 		res
 	}
@@ -357,26 +416,63 @@ mod tests {
 				.into_iter()
 				.for_each(|thread| thread.join().unwrap());
 			let allocated2 = A.allocated();
+			#[cfg(feature = "stats")]
+			let total_allocated = A.total_allocated();
 			if cfg!(all(test, feature = "nightly")) {
 				assert_eq!(allocated, allocated2);
+				#[cfg(feature = "stats")]
+				assert!(total_allocated >= allocated);
 			}
 		}
+		#[cfg(feature = "stats")]
+		assert!(A.max_allocated() < A.total_allocated());
+	}
+
+	#[cfg(all(test, not(feature = "nightly")))]
+	#[test]
+	fn limit() {
+		#[cfg(feature = "stats")]
+		let initial = A.allocated();
+		let allocate_amount = 30 * 1024 * 1024;
+		A.set_limit(A.allocated() + allocate_amount).unwrap();
+		for _ in 0..10 {
+			let mut vec = Vec::<u8>::with_capacity(0);
+			if let Err(_e) = vec.try_reserve_exact(allocate_amount + 1) {
+			} else {
+				A.set_limit(usize::max_value()).unwrap();
+				panic!("{}", A.remaining());
+			};
+			assert_eq!(vec.try_reserve_exact(allocate_amount), Ok(()));
+			let mut vec2 = Vec::<u8>::with_capacity(0);
+			assert!(vec2.try_reserve_exact(1).is_err());
+		}
+		// Might have additional allocations of errors and what not along the way.
+		#[cfg(feature = "stats")]
+		{
+			assert!(A.total_allocated() >= initial + 10 * allocate_amount);
+			assert_eq!(A.max_allocated(), initial + allocate_amount);
+		}
 	}
 
 	#[cfg(all(test, feature = "nightly"))]
 	#[test]
 	fn limit() {
-		A.set_limit(A.allocated() + 30 * 1024 * 1024).unwrap();
+		let allocate_amount = 30 * 1024 * 1024;
+		A.set_limit(A.allocated() + allocate_amount).unwrap();
 		for _ in 0..10 {
 			let mut vec = Vec::<u8>::with_capacity(0);
 			if let Err(TryReserveError::AllocError { .. }) =
-				vec.try_reserve_exact(30 * 1024 * 1024 + 1)
+				vec.try_reserve_exact(allocate_amount + 1)
 			{
 			} else {
 				A.set_limit(usize::max_value()).unwrap();
 				panic!("{}", A.remaining())
 			};
-			assert_eq!(vec.try_reserve_exact(30 * 1024 * 1024), Ok(()));
+			assert_eq!(vec.try_reserve_exact(allocate_amount), Ok(()));
+			let mut vec2 = Vec::<u8>::with_capacity(0);
+			assert!(vec2.try_reserve_exact(1).is_err());
 		}
+		assert_eq!(A.total_allocated(), 10 * allocate_amount);
+		assert_eq!(A.max_allocated(), allocate_amount)
 	}
 }