From a4362336fa1fc6f0d73f7a653df243ca6262fed3 Mon Sep 17 00:00:00 2001
From: Jordan Dominion <Cyberboss@users.noreply.github.com>
Date: Mon, 10 Jul 2023 14:35:29 -0400
Subject: [PATCH] Add FifoSemaphore and test

---
 .../Utils/FifoSemaphore.cs                    | 111 ++++++++++++++++++
 .../Utils/TestFifoSemaphore.cs                |  60 ++++++++++
 2 files changed, 171 insertions(+)
 create mode 100644 src/Tgstation.Server.Host/Utils/FifoSemaphore.cs
 create mode 100644 tests/Tgstation.Server.Host.Tests/Utils/TestFifoSemaphore.cs

diff --git a/src/Tgstation.Server.Host/Utils/FifoSemaphore.cs b/src/Tgstation.Server.Host/Utils/FifoSemaphore.cs
new file mode 100644
index 0000000000..7a4e200c9c
--- /dev/null
+++ b/src/Tgstation.Server.Host/Utils/FifoSemaphore.cs
@@ -0,0 +1,111 @@
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Tgstation.Server.Host.Utils
+{
+	/// <summary>
+	/// A first-in first-out async semaphore.
+	/// </summary>
+	/// <remarks>This is contentious and could be re-written using <see cref="TaskCompletionSource"/>s to make it promise based. However, it has a lower memory footprint without them and is fine for our uses.</remarks>
+	sealed class FifoSemaphore : IDisposable
+	{
+		/// <summary>
+		/// <see langword="class"/> to represent a ticket in the <see cref="ticketQueue"/> and whether or not it is <see cref="Abandoned"/>.
+		/// </summary>
+		class FifoSemaphoreTicket
+		{
+			/// <summary>
+			/// Set if the wait operation on a <see cref="FifoSemaphore"/> was cancelled to avoid clogging the queue.
+			/// </summary>
+			public bool Abandoned { get; set; }
+		}
+
+		/// <summary>
+		/// The backing <see cref="SemaphoreSlim"/>.
+		/// </summary>
+		readonly SemaphoreSlim semaphore;
+
+		/// <summary>
+		/// The <see cref="Queue{T}"/> of ticket <see cref="FifoSemaphoreTicket"/>s.
+		/// </summary>
+		readonly Queue<FifoSemaphoreTicket> ticketQueue;
+
+		/// <summary>
+		/// Initializes a new instance of the <see cref="FifoSemaphore"/> class.
+		/// </summary>
+		public FifoSemaphore()
+		{
+			ticketQueue = new Queue<FifoSemaphoreTicket>();
+			semaphore = new SemaphoreSlim(1);
+		}
+
+		/// <inheritdoc />
+		public void Dispose() => semaphore.Dispose();
+
+		/// <summary>
+		/// Locks the <see cref="FifoSemaphore"/>.
+		/// </summary>
+		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
+		/// <returns>A <see cref="ValueTask{TResult}"/> resulting in the locked <see cref="SemaphoreSlimContext"/>.</returns>
+		public async ValueTask<SemaphoreSlimContext> Lock(CancellationToken cancellationToken)
+		{
+			FifoSemaphoreTicket ticket = null;
+			using (cancellationToken.Register(
+				() =>
+				{
+					if (ticket != null)
+						ticket.Abandoned = true;
+				}))
+				while (true)
+				{
+					var context = await SemaphoreSlimContext.Lock(semaphore, cancellationToken);
+					try
+					{
+						FifoSemaphoreTicket peekedTicket = null;
+						while (ticketQueue.Count > 0)
+						{
+							peekedTicket = ticketQueue.Peek();
+							if (peekedTicket.Abandoned)
+								ticketQueue.Dequeue();
+							else
+								break;
+						}
+
+						cancellationToken.ThrowIfCancellationRequested();
+
+						bool goTime;
+						if (ticketQueue.Count == 0)
+							goTime = true;
+						else if (ticket == null)
+						{
+							ticket = new FifoSemaphoreTicket();
+							cancellationToken.ThrowIfCancellationRequested();
+							ticketQueue.Enqueue(ticket);
+							goTime = false;
+						}
+						else
+						{
+							goTime = peekedTicket == ticket;
+							if (goTime)
+								ticketQueue.Dequeue();
+						}
+
+						if (goTime)
+						{
+							var localContext = context;
+							context = null;
+							return localContext;
+						}
+					}
+					finally
+					{
+						context?.Dispose();
+					}
+
+					await Task.Yield();
+				}
+		}
+	}
+}
diff --git a/tests/Tgstation.Server.Host.Tests/Utils/TestFifoSemaphore.cs b/tests/Tgstation.Server.Host.Tests/Utils/TestFifoSemaphore.cs
new file mode 100644
index 0000000000..d09716610c
--- /dev/null
+++ b/tests/Tgstation.Server.Host.Tests/Utils/TestFifoSemaphore.cs
@@ -0,0 +1,60 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+namespace Tgstation.Server.Host.Utils.Tests
+{
+	[TestClass]
+	public sealed class TestFifoSemaphore
+	{
+		[TestMethod]
+		public async Task TestContention()
+		{
+			const int Count = 1000000;
+
+			using var cts = new CancellationTokenSource();
+			cts.Cancel();
+
+			using var semaphore = new FifoSemaphore();
+			var tcs = new TaskCompletionSource();
+
+			var orderTotal = 0;
+			var orderActual = 0;
+
+			async Task LockAndUnlock(int? expectedOrder)
+			{
+				try
+				{
+					using (await semaphore.Lock(expectedOrder.HasValue ? CancellationToken.None : cts.Token))
+					{
+						await tcs.Task;
+						Assert.AreEqual(expectedOrder.Value, ++orderActual);
+					}
+				}
+				catch (OperationCanceledException)
+				{
+					Assert.IsFalse(expectedOrder.HasValue);
+				}
+			}
+
+			var tasks = new List<Task>(Count);
+			for (var i = 0; i < Count; ++i)
+				tasks.Add(
+					LockAndUnlock(
+						(i % 3 == 0)
+							? null
+							: ++orderTotal));
+
+			var totalCancelled = tasks.Count(x => x.IsCompleted);
+			Assert.AreEqual((Count / 3) + 1, totalCancelled);
+
+			tcs.SetResult();
+
+			await Task.WhenAll(tasks).WaitAsync(TimeSpan.FromMinutes(2));
+		}
+	}
+}