Discord-History-Tracker/app/Utils/Tasks/RestartableTask.cs

using System;
using System.Threading;
using System.Threading.Tasks;

namespace DHT.Utils.Tasks;

public sealed class RestartableTask<T>(Action<T> resultProcessor, TaskScheduler resultScheduler) {
	private readonly object monitor = new ();

	private CancellationTokenSource? cancellationTokenSource;

	public void Restart(Func<CancellationToken, Task<T>> resultComputer) {
		lock (monitor) {
			Cancel();

			cancellationTokenSource = new CancellationTokenSource();
			
			var taskCancellationTokenSource = cancellationTokenSource;
			var taskCancellationToken = taskCancellationTokenSource.Token;

			Task.Run(() => resultComputer(taskCancellationToken), taskCancellationToken)
			    .ContinueWith(task => resultProcessor(task.Result), taskCancellationToken, TaskContinuationOptions.OnlyOnRanToCompletion, resultScheduler)
			    .ContinueWith(_ => OnTaskFinished(taskCancellationTokenSource), CancellationToken.None);
		}
	}

	public void Cancel() {
		lock (monitor) {
			if (cancellationTokenSource != null) {
				cancellationTokenSource.Cancel();
				cancellationTokenSource = null;
			}
		}
	}

	private void OnTaskFinished(CancellationTokenSource taskCancellationTokenSource) {
		lock (monitor) {
			taskCancellationTokenSource.Dispose();

			if (cancellationTokenSource == taskCancellationTokenSource) {
				cancellationTokenSource = null;
			}
		}
	}
}