LightlessClient/LightlessSync/Utils/TaskRegistry.cs

using System.Collections.Concurrent;

namespace LightlessSync.Utils;

public sealed class TaskRegistry<HandleType> where HandleType : notnull
{
    private readonly ConcurrentDictionary<HandleType, Lazy<Task>> _activeTasks = new();

    public Task GetOrStart(HandleType handle, Func<Task> taskFactory)
        => GetOrStartInternal(handle, taskFactory);

    public Task<T> GetOrStart<T>(HandleType handle, Func<Task<T>> taskFactory)
        => GetOrStartInternal(handle, taskFactory);

    public bool TryGetExisting(HandleType handle, out Task task)
    {
        if (_activeTasks.TryGetValue(handle, out Lazy<Task>? entry))
        {
            task = entry.Value;
            if (!task.IsCompleted)
            {
                return true;
            }

            _activeTasks.TryRemove(new KeyValuePair<HandleType, Lazy<Task>>(handle, entry));
        }

        task = Task.CompletedTask;
        return false;
    }

    private Task GetOrStartInternal(HandleType handle, Func<Task> taskFactory)
    {
        Lazy<Task> entry = _activeTasks.GetOrAdd(handle, _ => CreateEntry(handle, taskFactory));
        Task task = entry.Value;
        if (task.IsCompleted)
        {
            _activeTasks.TryRemove(new KeyValuePair<HandleType, Lazy<Task>>(handle, entry));
        }

        return task;
    }

    private Task<T> GetOrStartInternal<T>(HandleType handle, Func<Task<T>> taskFactory)
    {
        Lazy<Task> entry = _activeTasks.GetOrAdd(handle, _ => CreateEntry(handle, taskFactory));
        Task task = entry.Value;
        if (task.IsCompleted)
        {
            _activeTasks.TryRemove(new KeyValuePair<HandleType, Lazy<Task>>(handle, entry));
        }

        return (Task<T>)task;
    }

    private Lazy<Task> CreateEntry(HandleType handle, Func<Task> taskFactory)
    {
        Lazy<Task> entry = null!;
        entry = new Lazy<Task>(() => ExecuteAndRemove(handle, entry, taskFactory), LazyThreadSafetyMode.ExecutionAndPublication);
        return entry;
    }

    private Lazy<Task> CreateEntry<T>(HandleType handle, Func<Task<T>> taskFactory)
    {
        Lazy<Task> entry = null!;
        entry = new Lazy<Task>(() => ExecuteAndRemove(handle, entry, taskFactory), LazyThreadSafetyMode.ExecutionAndPublication);
        return entry;
    }

    private async Task ExecuteAndRemove(HandleType handle, Lazy<Task> entry, Func<Task> taskFactory)
    {
        try
        {
            await taskFactory().ConfigureAwait(false);
        }
        finally
        {
            _activeTasks.TryRemove(new KeyValuePair<HandleType, Lazy<Task>>(handle, entry));
        }
    }

    private async Task<T> ExecuteAndRemove<T>(HandleType handle, Lazy<Task> entry, Func<Task<T>> taskFactory)
    {
        try
        {
            return await taskFactory().ConfigureAwait(false);
        }
        finally
        {
            _activeTasks.TryRemove(new KeyValuePair<HandleType, Lazy<Task>>(handle, entry));
        }
    }
}