SocialPub/PrivaPub.Tests/Infrastructure/JobQueueTests.cs

using Microsoft.Extensions.Logging.Abstractions;

using MongoDB.Entities;

using PrivaPub.Infrastructure.Jobs;
using PrivaPub.Models.Jobs;
using PrivaPub.Tests.Support;

namespace PrivaPub.Tests.Infrastructure
{
	public class BackoffTests
	{
		[Fact]
		public void Grows_like_mastodons()
		{
			Assert.InRange(Backoff.After(1).TotalSeconds, 16, 16 + 9 * 2);
			Assert.InRange(Backoff.After(10).TotalSeconds, 10015, 10015 + 9 * 11);
		}

		[Fact]
		public void Quarantines_a_host_only_past_the_threshold()
		{
			Assert.Equal(TimeSpan.Zero, Backoff.HostQuarantine(9, 10));
			Assert.Equal(TimeSpan.FromHours(1), Backoff.HostQuarantine(10, 10));
			Assert.Equal(TimeSpan.FromHours(4), Backoff.HostQuarantine(12, 10));
			Assert.Equal(TimeSpan.FromDays(7), Backoff.HostQuarantine(40, 10));
		}
	}

	[Trait("Category", "Integration")]
	public sealed class JobQueueTests : IAsyncLifetime
	{
		public ValueTask InitializeAsync()
		{
			Assert.SkipUnless(MongoFixture.Enabled, MongoFixture.Skip);
			return ValueTask.CompletedTask;
		}

		public ValueTask DisposeAsync() => ValueTask.CompletedTask;

		static Job NewJob(string host = default, string dedupe = default) => new()
		{
			Kind = (JobKind)99,
			Host = host,
			DedupeKey = dedupe,
			Payload = "{}"
		};

		[Fact]
		public async Task A_dedupe_key_is_queued_once()
		{
			var queue = new JobQueue();
			var key = Guid.NewGuid().ToString("N");

			var inserted = await queue.EnqueueMany(new[] { NewJob(dedupe: key), NewJob(dedupe: key) }, TestContext.Current.CancellationToken);

			Assert.Equal(1, inserted);
		}

		[Fact]
		public async Task A_busy_host_is_skipped_and_a_failed_job_backs_off()
		{
			var token = TestContext.Current.CancellationToken;
			var queue = new JobQueue();
			var kind = (JobKind)(100 + Random.Shared.Next(1000));
			var busy = $"busy{Guid.NewGuid():N}.example";
			var free = $"free{Guid.NewGuid():N}.example";
			await queue.EnqueueMany(new[] { new Job { Kind = kind, Host = busy }, new Job { Kind = kind, Host = free } }, token);

			var leased = await queue.Lease(kind, new[] { busy }, token);
			Assert.Equal(free, leased.Host);
			Assert.Equal(1, leased.Attempts);

			await queue.Finish(leased, JobOutcome.Retry("boom"), maxAttempts: 3, token);
			var after = await DB.Default.Find<Job>().OneAsync(leased.ID, token);
			Assert.Equal(JobState.Pending, after.State);
			Assert.True(after.RunAt > DateTime.UtcNow.AddSeconds(10));

			// leased again for its third attempt (a finish only counts for the lease it was given)
			await DB.Default.Update<Job>().MatchID(leased.ID).Modify(j => j.State, JobState.Running).Modify(j => j.LeaseOwner, leased.LeaseOwner).ExecuteAsync(token);
			leased.Attempts = 3;
			await queue.Finish(leased, JobOutcome.Retry("boom"), maxAttempts: 3, token);
			Assert.Equal(JobState.Dead, (await DB.Default.Find<Job>().OneAsync(leased.ID, token)).State);
		}

		[Fact]
		public async Task An_expired_lease_returns_to_the_queue()
		{
			var token = TestContext.Current.CancellationToken;
			var job = new Job { Kind = (JobKind)98, State = JobState.Running, LeasedUntil = DateTime.UtcNow.AddMinutes(-1) };
			await DB.Default.SaveAsync(job, token);

			await new JobQueue(j => j.ID == job.ID).Reap(token);

			Assert.Equal(JobState.Pending, (await DB.Default.Find<Job>().OneAsync(job.ID, token)).State);
		}

		// a handler running three times its lease keeps it: the reaper finds nothing to take back and it runs once
		[Fact]
		public async Task A_long_job_keeps_its_lease_and_runs_once()
		{
			var token = TestContext.Current.CancellationToken;
			var kind = (JobKind)(2000 + Random.Shared.Next(100000));
			var queue = new JobQueue(TimeSpan.FromSeconds(1.5), j => j.Kind == kind);
			var handler = new SlowHandler(kind, TimeSpan.FromSeconds(4.5));
			await queue.Enqueue(kind, "{}", default, default, token);

			using var worker = new JobWorker(queue, new IJobHandler[] { handler }, NullLogger<JobWorker>.Instance);
			await worker.StartAsync(token);
			try
			{
				var deadline = DateTime.UtcNow.AddSeconds(20);
				while (!await DB.Default.Find<Job>().Match(j => j.Kind == kind && j.State == JobState.Done).ExecuteAnyAsync(token))
				{
					Assert.True(DateTime.UtcNow < deadline, "the job never finished");
					Assert.Equal(0, await queue.Reap(token));
					await Task.Delay(250, token);
				}
			}
			finally
			{
				await worker.StopAsync(CancellationToken.None);
			}

			Assert.Equal(1, handler.Started);
		}

		// a lease taken by another worker (reaped, then leased again) stops the handler, and its outcome is not recorded
		[Fact]
		public async Task A_lost_lease_stops_the_handler_and_drops_its_outcome()
		{
			var token = TestContext.Current.CancellationToken;
			var kind = (JobKind)(2000 + Random.Shared.Next(100000));
			var queue = new JobQueue(TimeSpan.FromSeconds(1.5), j => j.Kind == kind);
			var handler = new SlowHandler(kind, TimeSpan.FromSeconds(30));
			await queue.Enqueue(kind, "{}", default, default, token);

			using var worker = new JobWorker(queue, new IJobHandler[] { handler }, NullLogger<JobWorker>.Instance);
			await worker.StartAsync(token);
			try
			{
				var deadline = DateTime.UtcNow.AddSeconds(20);
				while (handler.Started == 0)
				{
					Assert.True(DateTime.UtcNow < deadline, "the job never started");
					await Task.Delay(100, token);
				}
				await DB.Default.Update<Job>().Match(j => j.Kind == kind).Modify(j => j.LeaseOwner, "another worker").ExecuteAsync(token);
				while (!handler.Cancelled)
				{
					Assert.True(DateTime.UtcNow < deadline, "the handler was never stopped");
					await Task.Delay(100, token);
				}
				await Task.Delay(300, token);
			}
			finally
			{
				await worker.StopAsync(CancellationToken.None);
			}

			var job = await DB.Default.Find<Job>().Match(j => j.Kind == kind).ExecuteFirstAsync(token);
			Assert.Equal(JobState.Running, job.State);
			Assert.Equal("another worker", job.LeaseOwner);
		}

		sealed class SlowHandler(JobKind kind, TimeSpan takes) : IJobHandler
		{
			int _started;

			public int Started => _started;
			public bool Cancelled { get; private set; }
			public JobKind Kind => kind;
			public int Concurrency => 1;
			public int MaxAttempts => 3;
			public int PerHostLimit => 1;

			public async Task<JobOutcome> Handle(Job job, CancellationToken token)
			{
				Interlocked.Increment(ref _started);
				try
				{
					await Task.Delay(takes, token);
				}
				catch (OperationCanceledException)
				{
					Cancelled = true;
					throw;
				}
				return JobOutcome.Done;
			}
		}
	}
}