SocialPub/PrivaPub/Infrastructure/Jobs/JobQueue.cs

using MongoDB.Bson;
using MongoDB.Driver;
using MongoDB.Entities;

using PrivaPub.Models.Jobs;

using System.Collections.Concurrent;
using System.Linq.Expressions;

namespace PrivaPub.Infrastructure.Jobs
{
	public sealed record JobOutcome(JobResult Result, string Error = default, DateTime? RetryAt = default)
	{
		public static readonly JobOutcome Done = new(JobResult.Done);
		public static JobOutcome Retry(string error) => new(JobResult.Retry, error);
		public static JobOutcome Dead(string error) => new(JobResult.Dead, error);
		public static JobOutcome Defer(DateTime until, string error) => new(JobResult.Defer, error, until);
	}

	public enum JobResult
	{
		Done,
		Retry,
		Dead,
		Defer
	}

	public interface IJobQueue
	{
		Task<bool> Enqueue(JobKind kind, string payload, string host, string dedupeKey, CancellationToken token);
		Task<string> Payload(string dedupeKey, CancellationToken token);
		Task<int> EnqueueMany(IEnumerable<Job> jobs, CancellationToken token);
		/// <summary>How long a lease lasts unless it is renewed.</summary>
		TimeSpan LeaseFor { get; }
		Task<Job> Lease(JobKind kind, IReadOnlyCollection<string> busyHosts, CancellationToken token);
		/// <summary>Extends a running job's lease; false once it is no longer this lease's (reaped and leased again).</summary>
		Task<bool> Renew(Job job, CancellationToken token);
		/// <summary>Records the outcome; false when the lease was lost meanwhile, and then nothing changes.</summary>
		Task<bool> Finish(Job job, JobOutcome outcome, int maxAttempts, CancellationToken token);
		Task<long> Reap(CancellationToken token);
		Task WaitForWork(JobKind kind, TimeSpan poll, CancellationToken token);
	}

	public class JobQueue : IJobQueue
	{
		public static readonly TimeSpan LeaseTime = TimeSpan.FromMinutes(2);

		// every lease is stamped with its own owner (the process and a fresh id): two workers of one process, the second
		// leasing a job the reaper took back from the first, must not renew nor finish each other's lease
		readonly string _owner = $"{Environment.MachineName}:{Environment.ProcessId}";
		readonly ConcurrentDictionary<JobKind, SemaphoreSlim> _signals = new();
		readonly FilterDefinition<Job> _scope = Builders<Job>.Filter.Empty;

		public JobQueue()
		{
		}

		public JobQueue(Expression<Func<Job, bool>> scope) => _scope = Builders<Job>.Filter.Where(scope);

		public JobQueue(TimeSpan leaseFor, Expression<Func<Job, bool>> scope = default)
		{
			LeaseFor = leaseFor;
			if (scope != default)
				_scope = Builders<Job>.Filter.Where(scope);
		}

		public TimeSpan LeaseFor { get; } = LeaseTime;

		public async Task<bool> Enqueue(JobKind kind, string payload, string host, string dedupeKey, CancellationToken token) =>
			await EnqueueMany(new[] { new Job { Kind = kind, Payload = payload, Host = host, DedupeKey = dedupeKey } }, token) == 1;

		// what the job a dedupe key names carries, or null when there is none
		public async Task<string> Payload(string dedupeKey, CancellationToken token) =>
			(await DB.Default.Find<Job>().Match(j => j.DedupeKey == dedupeKey).ExecuteFirstAsync(token))?.Payload;

		public async Task<int> EnqueueMany(IEnumerable<Job> jobs, CancellationToken token)
		{
			var inserted = 0;
			foreach (var job in jobs)
			{
				try
				{
					await DB.Default.SaveAsync(job, token);
					inserted++;
					Signal(job.Kind);
				}
				catch (MongoWriteException ex) when (ex.WriteError?.Category == ServerErrorCategory.DuplicateKey)
				{
				}
			}
			return inserted;
		}

		public async Task<Job> Lease(JobKind kind, IReadOnlyCollection<string> busyHosts, CancellationToken token)
		{
			var now = DateTime.UtcNow;
			var busy = busyHosts?.ToList() ?? new List<string>();
			return await DB.Default.UpdateAndGet<Job>()
				.Match(f => f.Where(j => j.Kind == kind && j.State == JobState.Pending && j.RunAt <= now && !busy.Contains(j.Host)) & _scope)
				.Modify(j => j.State, JobState.Running)
				.Modify(j => j.LeasedUntil, now + LeaseFor)
				.Modify(j => j.LeaseOwner, $"{_owner}/{ObjectId.GenerateNewId()}")
				.Modify(b => b.Inc(j => j.Attempts, 1))
				.Option(o => o.Sort = Builders<Job>.Sort.Ascending(j => j.RunAt))
				.ExecuteAsync(token);
		}

		public async Task<bool> Renew(Job job, CancellationToken token)
		{
			var result = await DB.Default.Update<Job>()
				.Match(j => j.ID == job.ID && j.State == JobState.Running && j.LeaseOwner == job.LeaseOwner)
				.Modify(j => j.LeasedUntil, DateTime.UtcNow + LeaseFor)
				.ExecuteAsync(token);
			return result.MatchedCount == 1;
		}

		public async Task<bool> Finish(Job job, JobOutcome outcome, int maxAttempts, CancellationToken token)
		{
			var now = DateTime.UtcNow;
			var update = DB.Default.Update<Job>().Match(j => j.ID == job.ID && j.LeaseOwner == job.LeaseOwner)
				.Modify(j => j.LeasedUntil, null)
				.Modify(j => j.LeaseOwner, null)
				.Modify(j => j.LastError, outcome.Error);
			switch (outcome.Result)
			{
				case JobResult.Done:
					update.Modify(j => j.State, JobState.Done).Modify(j => j.FinishedAt, now);
					break;
				case JobResult.Retry when job.Attempts < maxAttempts:
					update.Modify(j => j.State, JobState.Pending).Modify(j => j.RunAt, now + Backoff.After(job.Attempts));
					break;
				case JobResult.Defer:
					update.Modify(j => j.State, JobState.Pending)
						.Modify(j => j.RunAt, outcome.RetryAt ?? now + Backoff.After(job.Attempts))
						.Modify(b => b.Inc(j => j.Attempts, -1));
					break;
				default:
					update.Modify(j => j.State, JobState.Dead).Modify(j => j.FinishedAt, now);
					break;
			}
			return (await update.ExecuteAsync(token)).MatchedCount == 1;
		}

		public async Task<long> Reap(CancellationToken token)
		{
			var now = DateTime.UtcNow;
			var result = await DB.Default.Update<Job>()
				.Match(f => f.Where(j => j.State == JobState.Running && j.LeasedUntil < now) & _scope)
				.Modify(j => j.State, JobState.Pending)
				.Modify(j => j.RunAt, now)
				.Modify(j => j.LeasedUntil, null)
				.Modify(j => j.LeaseOwner, null)
				.ExecuteAsync(token);
			return result.ModifiedCount;
		}

		public async Task WaitForWork(JobKind kind, TimeSpan poll, CancellationToken token)
		{
			try
			{
				await SignalFor(kind).WaitAsync(poll, token);
			}
			catch (OperationCanceledException) when (token.IsCancellationRequested)
			{
			}
		}

		void Signal(JobKind kind)
		{
			try
			{
				SignalFor(kind).Release();
			}
			catch (SemaphoreFullException)
			{
			}
		}

		SemaphoreSlim SignalFor(JobKind kind) => _signals.GetOrAdd(kind, _ => new SemaphoreSlim(0, 64));
	}
}