sigma update

LightlessSync/Utils/TaskRegistry.cs

@@ -0,0 +1,81 @@
+using System.Collections.Concurrent;
+
+
+namespace LightlessSync.Utils;
+
+public sealed class TaskRegistry<HandleType> where HandleType : notnull
+{
+    private readonly ConcurrentDictionary<HandleType, ActiveTask> _activeTasks = new();
+
+    public Task GetOrStart(HandleType handle, Func<Task> taskFactory)
+    {
+        ActiveTask entry = _activeTasks.GetOrAdd(handle, i => new ActiveTask(() => ExecuteAndRemove(i, taskFactory)));
+        return entry.EnsureStarted();
+    }
+
+    public Task<T> GetOrStart<T>(HandleType handle, Func<Task<T>> taskFactory)
+    {
+        ActiveTask entry = _activeTasks.GetOrAdd(handle, i => new ActiveTask(() => ExecuteAndRemove(i, taskFactory)));
+        return (Task<T>)entry.EnsureStarted();
+    }
+
+    public bool TryGetExisting(HandleType handle, out Task task)
+    {
+        if (_activeTasks.TryGetValue(handle, out ActiveTask? entry))
+        {
+            task = entry.EnsureStarted();
+            return true;
+        }
+
+        task = Task.CompletedTask;
+        return false;
+    }
+
+    private async Task ExecuteAndRemove(HandleType handle, Func<Task> taskFactory)
+    {
+        try
+        {
+            await taskFactory().ConfigureAwait(false);
+        }
+        finally
+        {
+            _activeTasks.TryRemove(handle, out _);
+        }
+    }
+
+    private async Task<T> ExecuteAndRemove<T>(HandleType handle, Func<Task<T>> taskFactory)
+    {
+        try
+        {
+            return await taskFactory().ConfigureAwait(false);
+        }
+        finally
+        {
+            _activeTasks.TryRemove(handle, out _);
+        }
+    }
+
+    private sealed class ActiveTask
+    {
+        private readonly object _gate = new();
+        private readonly Func<Task> _starter;
+        private Task? _cached;
+
+        public ActiveTask(Func<Task> starter)
+        {
+            _starter = starter;
+        }
+
+        public Task EnsureStarted()
+        {
+            lock (_gate)
+            {
+                if (_cached == null || _cached.IsCompleted)
+                {
+                    _cached = _starter();
+                }
+                return _cached;
+            }
+        }
+    }
+}