Added support for some statistics collection

alecmocatta · Mar 24, 2023 · 33ceda9 · 33ceda9
1 parent 1ab159c
commit 33ceda9
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 5 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -23,5  23,6 @@ maintenance = { status = "passively-maintained" }
 
 [features]
 nightly = []
 stats = []
 
 [dependencies]
diff --git a/src/lib.rs b/src/lib.rs
@@ -54,6  54,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  69,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),
 		}
 	}
 
@@ -134,6  142,27 @@ 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.
 		#[cfg(feature = "stats")]
 		let _ = self.max_allocated.fetch_max(self.allocated(), Ordering::Relaxed);
 	}
 }
 
 unsafe impl<H> GlobalAlloc for Cap<H>
@@ -149,6  178,8 @@ where
 		};
 		if res.is_null() {
 			let _ = self.remaining.fetch_add(size, Ordering::Release);
 		} else {
 			self.update_stats(size);
 		}
 		res
 	}
@@ -166,13  197,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  228,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  252,8 @@ where
 		};
 		if res.is_err() {
 			let _ = self.remaining.fetch_add(size, Ordering::Release);
 		} else {
 			self.update_stats(size);
 		}
 		res
 	}
@@ -233,7  273,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  297,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  312,8 @@ where
 		};
 		if res.is_err() {
 			let _ = self.remaining.fetch_add(size, Ordering::Release);
 		} else {
 			self.update_stats(size);
 		}
 		res
 	}
@@ -292,6  338,8 @@ where
 			let _ = self
 				.remaining
 				.fetch_add(new_size - old_size, Ordering::Release);
 		} else {
 			self.update_stats(new_size - old_size);
 		}
 		res
 	}
@@ -357,26  405,58 @@ mod tests {
 				.into_iter()
 				.for_each(|thread| thread.join().unwrap());
 			let allocated2 = A.allocated();
 			let total_allocated = A.total_allocated();
 			if cfg!(all(test, feature = "nightly")) {
 				assert_eq!(allocated, allocated2);
 				assert!(total_allocated >= allocated);
 			}
 		}
 		assert!(A.max_allocated() < A.total_allocated());
 	}
 
 	#[cfg(all(test, not(feature = "nightly")))]
 	#[test]
 	fn limit() {
 		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.
 		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)
 	}
 }