Refactor and fix ObjectPool<T>

Source.Common.Tests/ObjectPoolTests.cs

@@ -0,0 +1,101 @@
+namespace Source.Common.Tests;
+
+public class ObjectPoolTests
+{
+	sealed class TestPoolable : IPoolableObject
+	{
+		public int InitCalls;
+		public int ResetCalls;
+		public int InUse;
+		public bool HandedOutTwice;
+
+		public void Init() {
+			InitCalls++;
+			if (Interlocked.Exchange(ref InUse, 1) != 0)
+				HandedOutTwice = true;
+		}
+
+		public void Reset() {
+			ResetCalls++;
+			Interlocked.Exchange(ref InUse, 0);
+		}
+	}
+
+	[Fact]
+	public void AllocInitsAndFreeResets() {
+		ObjectPool<TestPoolable> pool = new();
+
+		TestPoolable obj = pool.Alloc();
+		Assert.Equal(1, obj.InitCalls);
+		Assert.Equal(0, obj.ResetCalls);
+
+		pool.Free(obj);
+		Assert.Equal(1, obj.ResetCalls);
+	}
+
+	[Fact]
+	public void FreedInstanceIsReused() {
+		ObjectPool<TestPoolable> pool = new();
+
+		TestPoolable first = pool.Alloc();
+		pool.Free(first);
+		TestPoolable second = pool.Alloc();
+
+		Assert.Same(first, second);
+		Assert.Equal(2, second.InitCalls);
+	}
+
+	[Fact]
+	public void LiveInstancesAreDistinct() {
+		ObjectPool<TestPoolable> pool = new();
+
+		TestPoolable a = pool.Alloc();
+		TestPoolable b = pool.Alloc();
+
+		Assert.NotSame(a, b);
+	}
+
+	[Fact]
+	public void CountTracksOutstandingInstances() {
+		ObjectPool<TestPoolable> pool = new();
+
+		TestPoolable a = pool.Alloc();
+		TestPoolable b = pool.Alloc();
+		Assert.Equal(2, pool.Count());
+
+		pool.Free(a);
+		Assert.Equal(1, pool.Count());
+
+		pool.Free(b);
+		Assert.Equal(0, pool.Count());
+	}
+
+	[Fact]
+	public void FreeNullIsIgnored() {
+		ObjectPool<TestPoolable> pool = new();
+
+		pool.Free(null);
+
+		Assert.Equal(0, pool.Count());
+	}
+
+	[Fact]
+	public void ConcurrentUseNeverHandsOutAnInstanceTwice() {
+		ObjectPool<TestPoolable> pool = new();
+		const int threads = 8;
+		const int iterations = 20000;
+		bool handedOutTwice = false;
+
+		Parallel.For(0, threads, _ => {
+			for (int i = 0; i < iterations; i++) {
+				TestPoolable obj = pool.Alloc();
+				if (obj.HandedOutTwice)
+					Volatile.Write(ref handedOutTwice, true);
+				pool.Free(obj);
+			}
+		});
+
+		Assert.False(handedOutTwice);
+		Assert.Equal(0, pool.Count());
+	}
+}

Source.Common/Exts.cs

@@ -136,53 +136,6 @@ public void Purge() {
 	public V this[ulong index] => dict[index];
 }
 
-public class ObjectPool<T> where T : IPoolableObject, new()
-{
-	public static readonly ObjectPool<T> Shared = new();
-
-
-	readonly ConcurrentDictionary<T, bool> valueStates = [];
-
-	public T Alloc() {
-		foreach (var kvp in valueStates) {
-			if (kvp.Value == false) { // We found something free
-				valueStates[kvp.Key] = true;
-				kvp.Key.Init();
-				return kvp.Key;
-			}
-		}
-
-		// Make an new instance of the class
-		var instance = new T();
-		valueStates[instance] = true;
-		instance.Init();
-		return instance;
-	}
-
-	public int Count() {
-		int count = 0;
-		foreach (KeyValuePair<T, bool> kvp in valueStates)
-			if (kvp.Value)
-				count++;
-
-		return count;
-	}
-
-	public bool IsMemoryPoolAllocated(T value) => valueStates.TryGetValue(value, out _);
-	public void Free(T value) {
-		if (value == null)
-			return;
-		if (!valueStates.TryGetValue(value, out bool state))
-			AssertMsg(false, $"Passed an instance of {typeof(T).Name} to {nameof(Free)}(T value) that was not allocated by {nameof(Alloc)}()");
-		else if (state == false)
-			AssertMsg(false, $"Attempted to free {typeof(T).Name} instance twice in ClassPool<T>, please verify\n");
-		else {
-			value.Reset();
-			valueStates[value] = false;
-		}
-	}
-}
-
 public sealed class PooledLinkedList<T> where T : struct
 {
 	struct Node
@@ -282,61 +235,6 @@ void Grow(int num = 0) {
 	}
 }
 
-public class StructMemoryPool<T> where T : struct
-{
-	public static readonly StructMemoryPool<T> Shared = new();
-
-	readonly RefStack<T> instances = new();
-	readonly ConcurrentDictionary<int, bool> valueStates = [];
-
-	public ref T Alloc() {
-		foreach (var kvp in valueStates) {
-			ref T existing = ref instances[kvp.Key];
-			if (kvp.Value == false) {
-				valueStates[kvp.Key] = true;
-				return ref existing;
-			}
-		}
-
-		lock (instances) {
-			ref T instance = ref instances.Push();
-			valueStates[instances.Count - 1] = true;
-			return ref instance;
-		}
-	}
-
-	public unsafe bool IsMemoryPoolAllocated(ref T value) {
-		lock (instances) {
-			for (int i = 0; i < instances.Count; i++) {
-				ref T instance = ref instances[i];
-				if (Unsafe.AreSame(ref value, ref instance))
-					return true;
-			}
-		}
-
-		return false;
-	}
-
-
-	public void Free(ref T value) {
-		lock (instances) {
-			for (int i = 0; i < instances.Count; i++) {
-				ref T instance = ref instances[i];
-				if (Unsafe.AreSame(ref value, ref instance)) {
-					if (!valueStates.TryGetValue(i, out bool state))
-						AssertMsg(false, $"Passed an instance of {typeof(T).Name} to {nameof(Free)}(T value) that was not allocated by {nameof(Alloc)}()");
-					else if (state == false)
-						AssertMsg(false, $"Attempted to free {typeof(T).Name} instance twice in StructPool<T>, please verify\n");
-					else {
-						valueStates[i] = false;
-						instance = default; // Zero out the instance
-					}
-				}
-			}
-		}
-	}
-}
-
 public static class StrTools
 {
 	// We're going to use / everywhere instead of \ on Windows, since C#'s API

Source.Common/ObjectPool.cs

@@ -0,0 +1,31 @@
+using System.Collections.Concurrent;
+
+namespace Source;
+
+public class ObjectPool<T> where T : IPoolableObject, new()
+{
+	public static readonly ObjectPool<T> Shared = new();
+
+	readonly ConcurrentBag<T> free = [];
+	int allocated;
+
+	public T Alloc() {
+		if (!free.TryTake(out T? instance))
+			instance = new T();
+
+		Interlocked.Increment(ref allocated);
+		instance.Init();
+		return instance;
+	}
+
+	public int Count() => Volatile.Read(ref allocated);
+
+	public void Free(T? value) {
+		if (value == null)
+			return;
+
+		value.Reset();
+		Interlocked.Decrement(ref allocated);
+		free.Add(value);
+	}
+}

Source.Engine/ClientFrame.cs

@@ -191,8 +191,5 @@ public void DeleteClientFrames(int tick) {
 
 	public ClientFrame AllocateFrame() => ClientFramePool.Alloc();
 
-	public void FreeFrame(ClientFrame frame) {
-		if (ClientFramePool.IsMemoryPoolAllocated(frame))
-			ClientFramePool.Free(frame);
-	}
+	public void FreeFrame(ClientFrame frame) => ClientFramePool.Free(frame);
 }