[backend/queue] Improved AsyncAutoResetEvent implementation

This implementation improves the memory behavior of AsyncAutoResetEvents that fire in rapid succession.

Iceshrimp.Backend/Core/Extensions/TaskExtensions.cs

@@ -13,7 +13,7 @@ public static class TaskExtensions
 			// ignored
 		}
 	}
-	
+
 	public static async Task SafeWaitAsync(this Task task, CancellationToken token)
 	{
 		try
@@ -25,7 +25,7 @@ public static class TaskExtensions
 			// ignored
 		}
 	}
-	
+
 	public static async Task SafeWaitAsync(this Task task, CancellationToken token, Action action)
 	{
 		try
@@ -37,7 +37,7 @@ public static class TaskExtensions
 			action();
 		}
 	}
-	
+
 	public static async Task SafeWaitAsync(this Task task, CancellationToken token, Func<Task> action)
 	{
 		try

Iceshrimp.Backend/Core/Helpers/EventHelpers.cs

@@ -1,57 +1,50 @@
 namespace Iceshrimp.Backend.Core.Helpers;
 
-public sealed class AsyncAutoResetEvent(bool signaled)
+public sealed class AsyncAutoResetEvent(bool signaled = false)
 {
-	private readonly Queue<TaskCompletionSource> _queue = new();
+	private readonly List<TaskCompletionSource<bool>> _taskCompletionSources = [];
 
-	public Task WaitAsync(CancellationToken cancellationToken = default)
+	public Task<bool> WaitAsync(CancellationToken cancellationToken = default)
 	{
-		lock (_queue)
+		lock (_taskCompletionSources)
 		{
 			if (signaled)
 			{
 				signaled = false;
-				return Task.CompletedTask;
+				return Task.FromResult(true);
 			}
 
-			var tcs = new TaskCompletionSource();
+			var tcs = new TaskCompletionSource<bool>();
-			if (cancellationToken.CanBeCanceled)
+			cancellationToken.Register(Callback, (this, tcs));
-			{
+			_taskCompletionSources.Add(tcs);
-				// If the token is cancelled, cancel the waiter.
-				var registration =
-					cancellationToken.Register(() => tcs.TrySetCanceled(), false);
-
-				// If the waiter completes or faults, unregister our interest in cancellation.
-				tcs.Task.ContinueWith(
-				                      _ => registration.Unregister(),
-				                      cancellationToken,
-				                      TaskContinuationOptions.OnlyOnRanToCompletion |
-				                      TaskContinuationOptions.NotOnFaulted,
-				                      TaskScheduler.Default);
-			}
-
-			_queue.Enqueue(tcs);
 			return tcs.Task;
 		}
 	}
 
 	public void Set()
 	{
-		TaskCompletionSource? toRelease = null;
+		lock (_taskCompletionSources)
-
-		lock (_queue)
 		{
-			if (_queue.Count > 0)
+			if (_taskCompletionSources.Count > 0)
 			{
-				toRelease = _queue.Dequeue();
+				var tcs = _taskCompletionSources[0];
+				_taskCompletionSources.RemoveAt(0);
+				tcs.TrySetResult(true);
+				return;
 			}
-			else if (!signaled)
-			{
-				signaled = true;
-			}
-		}
 
-		// It's possible that the TCS has already been cancelled.
+			signaled = true;
-		toRelease?.TrySetResult();
+		}
+	}
+
+	private static void Callback(object? state)
+	{
+		var (ev, tcs) = ((AsyncAutoResetEvent, TaskCompletionSource<bool>))state!;
+		lock (ev._taskCompletionSources)
+		{
+			if (tcs.Task.IsCompleted) return;
+			tcs.TrySetCanceled();
+			ev._taskCompletionSources.Remove(tcs);
+		}
 	}
 }

Iceshrimp.Backend/Core/Services/QueueService.cs

@@ -203,8 +203,8 @@ public class PostgresJobQueue<T>(
 	int parallelism
 ) : IPostgresJobQueue where T : class
 {
-	private readonly AsyncAutoResetEvent  _delayedChannel = new(false);
+	private readonly AsyncAutoResetEvent  _delayedChannel = new();
-	private readonly AsyncAutoResetEvent  _queuedChannel  = new(false);
+	private readonly AsyncAutoResetEvent  _queuedChannel  = new();
 	private          IServiceScopeFactory _scopeFactory   = null!;
 	public           string               Name => name;
 	private          string?              _workerId;