JournaledJobShardState.cs

using System;
using System.Linq;
using Orleans.Journaling;

namespace Orleans.DurableJobs;

internal sealed class JournaledJobShardState : IJournaledState, IDurableValueCommandHandler<DurableJobShardJournalRecord>
{
    public const string StateName = "jobs";

    private readonly JobShardId _shardId;
    private readonly IDurableValueCommandCodec<DurableJobShardJournalRecord>? _codec;
    private readonly TimeProvider _timeProvider;
    private InMemoryJobQueue _jobQueue;
    private JournalStreamWriter _writer;

    public JournaledJobShardState(
        JobShardId shardId,
        DateTimeOffset startTime,
        DateTimeOffset endTime,
        IDurableValueCommandCodec<DurableJobShardJournalRecord> codec,
        TimeProvider? timeProvider = null)
        : this(shardId, startTime, endTime, codec, timeProvider, isAddingCompleted: false)
    {
        ArgumentNullException.ThrowIfNull(codec);
    }

    internal JournaledJobShardState(JobShardId shardId, DateTimeOffset startTime, DateTimeOffset endTime, TimeProvider? timeProvider = null)
        : this(shardId, startTime, endTime, codec: null, timeProvider: timeProvider, isAddingCompleted: false)
    {
    }

    private JournaledJobShardState(
        JobShardId shardId,
        DateTimeOffset startTime,
        DateTimeOffset endTime,
        IDurableValueCommandCodec<DurableJobShardJournalRecord>? codec,
        TimeProvider? timeProvider,
        bool isAddingCompleted)
    {
        if (endTime < startTime)
        {
            throw new ArgumentOutOfRangeException(nameof(endTime), "Shard end time must be greater than or equal to the start time.");
        }

        _shardId = shardId;
        _codec = codec;
        _timeProvider = timeProvider ?? TimeProvider.System;
        _jobQueue = new(_timeProvider);
        StartTime = startTime;
        EndTime = endTime;
        IsAddingCompleted = isAddingCompleted;
    }

    public string Id => _shardId.Value;

    public DateTimeOffset StartTime { get; }

    public DateTimeOffset EndTime { get; }

    public bool IsAddingCompleted { get; private set; }

    public int Count => _jobQueue.Count;

    public bool ContainsJob(string jobId) => _jobQueue.ContainsJob(jobId);

    public IAsyncEnumerable<IJobRunContext> ConsumeDurableJobsAsync() => _jobQueue;

    public DurableJob? TryScheduleJob(ScheduleJobRequest request)
    {
        if (IsAddingCompleted)
        {
            return null;
        }

        if (request.DueTime < StartTime || request.DueTime > EndTime)
        {
            throw new ArgumentOutOfRangeException(nameof(request), "Scheduled time is out of shard bounds.");
        }

        if (request.JobId is { } requestedJobId
            && _jobQueue.TryGetJob(requestedJobId, out var existingJob))
        {
            if (!request.Matches(existingJob!))
            {
                throw new InvalidOperationException(
                    $"Durable job ID '{requestedJobId}' is already scheduled with different properties.");
            }

            return existingJob;
        }

        var job = request.CreateJob(Id);

        Write(DurableJobShardJournalRecord.ForSchedule(job));
        ApplySchedule(job);
        return job;
    }

    public bool RemoveJob(string jobId)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(jobId);
        if (!_jobQueue.ContainsJob(jobId))
        {
            return false;
        }

        Write(DurableJobShardJournalRecord.ForRemove(jobId));
        return ApplyRemove(jobId);
    }

    public bool RetryJobLater(IJobRunContext jobContext, DateTimeOffset newDueTime)
    {
        ArgumentNullException.ThrowIfNull(jobContext);
        return RetryJobLater(jobContext.Job.Id, newDueTime, jobContext.DequeueCount, executionGeneration: null);
    }

    public bool RetryJobLater(string jobId, DateTimeOffset newDueTime, int dequeueCount)
        => RetryJobLater(jobId, newDueTime, dequeueCount, executionGeneration: null);

    public bool RetryJobLater(string jobId, DateTimeOffset newDueTime, int dequeueCount, long? executionGeneration)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(jobId);
        ValidateDequeueCount(dequeueCount);
        if (!_jobQueue.ContainsJob(jobId))
        {
            return false;
        }

        Write(DurableJobShardJournalRecord.ForRetry(jobId, newDueTime, dequeueCount, executionGeneration));
        return ApplyRetry(jobId, newDueTime, dequeueCount, executionGeneration);
    }

    public void MarkAsComplete()
    {
        IsAddingCompleted = true;
        _jobQueue.MarkAsComplete();
    }

    internal DurableJobShardSnapshot CaptureSnapshot()
    {
        var jobs = _jobQueue.GetSnapshot()
            .OrderBy(static item => item.Job.DueTime)
            .ThenBy(static item => item.Job.Id, StringComparer.Ordinal)
            .Select(static item => new DurableJobShardSnapshotEntry
            {
                Job = item.Job,
                DequeueCount = item.DequeueCount
            })
            .ToList();

        return new() { Jobs = jobs };
    }

    internal void Apply(DurableJobShardJournalRecord record)
    {
        ArgumentNullException.ThrowIfNull(record);

        switch (record.Kind)
        {
            case DurableJobShardJournalRecordKind.None:
                throw new InvalidOperationException(
                    "DurableJobs shard journal record has an uninitialized Kind (None). This usually " +
                    "indicates a malformed or default-constructed journal payload — for example, a JSON " +
                    "record whose 'kind' property was missing.");
            case DurableJobShardJournalRecordKind.Schedule:
                ApplySchedule(GetRequired(record.Schedule, nameof(record.Schedule)).Job);
                break;
            case DurableJobShardJournalRecordKind.Remove:
                ApplyRemove(GetRequired(record.Remove, nameof(record.Remove)).JobId);
                break;
            case DurableJobShardJournalRecordKind.Retry:
                var retry = GetRequired(record.Retry, nameof(record.Retry));
                ApplyRetry(retry.JobId, retry.DueTime, retry.DequeueCount, retry.ExecutionGeneration);
                break;
            case DurableJobShardJournalRecordKind.Snapshot:
                ApplySnapshot(GetRequired(record.Snapshot, nameof(record.Snapshot)));
                break;
            default:
                throw new NotSupportedException($"DurableJobs shard journal record kind '{record.Kind}' is not supported.");
        }
    }

    void IJournaledState.ReplayEntry(JournalEntry entry, JournalReplayContext context) =>
        context.GetRequiredCommandCodec(entry.FormatKey, GetCodec()).Apply(entry.Reader, this);

    void IDurableValueCommandHandler<DurableJobShardJournalRecord>.ApplySet(DurableJobShardJournalRecord value) => Apply(value);

    void IJournaledState.Reset(JournalStreamWriter writer)
    {
        _jobQueue = new(_timeProvider);
        IsAddingCompleted = false;
        _writer = writer;
    }

    void IJournaledState.AppendEntries(JournalStreamWriter writer)
    {
    }

    void IJournaledState.AppendSnapshot(JournalStreamWriter writer)
    {
        GetCodec().WriteSet(DurableJobShardJournalRecord.ForSnapshot(CaptureSnapshot()), writer);
    }

    IJournaledState IJournaledState.DeepCopy() => throw new NotSupportedException();

    private void Write(DurableJobShardJournalRecord record) => GetCodec().WriteSet(record, _writer);

    private void ApplySchedule(DurableJob job) => _jobQueue.Enqueue(job, dequeueCount: 0);

    private bool ApplyRemove(string jobId) => _jobQueue.RemoveJob(jobId);

    private bool ApplyRetry(string jobId, DateTimeOffset dueTime, int dequeueCount, long? executionGeneration)
    {
        ValidateDequeueCount(dequeueCount);
        return _jobQueue.RetryJobLater(jobId, dueTime, dequeueCount, executionGeneration);
    }

    private void ApplySnapshot(DurableJobShardSnapshot snapshot)
    {
        ArgumentNullException.ThrowIfNull(snapshot);

        _jobQueue.Clear();
        foreach (var entry in snapshot.Jobs)
        {
            ArgumentNullException.ThrowIfNull(entry.Job);
            ValidateDequeueCount(entry.DequeueCount);
            _jobQueue.Enqueue(entry.Job, entry.DequeueCount);
        }
    }

    private IDurableValueCommandCodec<DurableJobShardJournalRecord> GetCodec()
        => _codec ?? throw new InvalidOperationException("A DurableJobs shard journal operation codec is required before journal entries can be appended.");

    private static T GetRequired<T>(T? value, string propertyName) where T : class
        => value ?? throw new InvalidOperationException($"DurableJobs shard journal record is missing required '{propertyName}' payload.");

    private static void ValidateDequeueCount(int dequeueCount)
    {
        if (dequeueCount < 0)
        {
            throw new InvalidOperationException("DurableJobs shard journal dequeue count must not be negative.");
        }
    }
}

[GenerateSerializer]
[Alias("Orleans.DurableJobs.DurableJobShardJournalRecordKind")]
internal enum DurableJobShardJournalRecordKind : byte
{
    /// <summary>
    /// Reserved sentinel. A record whose <see cref="DurableJobShardJournalRecord.Kind"/>
    /// is <see cref="None"/> is invalid — it indicates a default-constructed instance or a
    /// malformed journal payload (for example, a JSON record where <c>kind</c> was omitted
    /// because the serializer drops default-valued properties). Apply paths must reject it.
    /// Keep this at value 0 so that the default surfaces as an explicit failure rather than
    /// silently aliasing to another kind.
    /// </summary>
    None = 0,
    Schedule = 1,
    Remove = 2,
    Retry = 3,
    Snapshot = 4,
}

[GenerateSerializer]
[Alias("Orleans.DurableJobs.DurableJobShardJournalRecord")]
internal sealed class DurableJobShardJournalRecord
{
    [Id(0)]
    public DurableJobShardJournalRecordKind Kind { get; init; }

    [Id(1)]
    public DurableJobShardScheduleOperation? Schedule { get; init; }

    [Id(2)]
    public DurableJobShardRemoveOperation? Remove { get; init; }

    [Id(3)]
    public DurableJobShardRetryOperation? Retry { get; init; }

    [Id(4)]
    public DurableJobShardSnapshot? Snapshot { get; init; }

    public static DurableJobShardJournalRecord ForSchedule(DurableJob job)
    {
        ArgumentNullException.ThrowIfNull(job);

        return new()
        {
            Kind = DurableJobShardJournalRecordKind.Schedule,
            Schedule = new() { Job = job }
        };
    }

    public static DurableJobShardJournalRecord ForRemove(string jobId)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(jobId);

        return new()
        {
            Kind = DurableJobShardJournalRecordKind.Remove,
            Remove = new() { JobId = jobId }
        };
    }

    public static DurableJobShardJournalRecord ForRetry(
        string jobId,
        DateTimeOffset dueTime,
        int dequeueCount,
        long? executionGeneration = null)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(jobId);

        return new()
        {
            Kind = DurableJobShardJournalRecordKind.Retry,
            Retry = new()
            {
                JobId = jobId,
                DueTime = dueTime,
                DequeueCount = dequeueCount,
                ExecutionGeneration = executionGeneration
            }
        };
    }

    public static DurableJobShardJournalRecord ForSnapshot(DurableJobShardSnapshot snapshot)
    {
        ArgumentNullException.ThrowIfNull(snapshot);

        return new()
        {
            Kind = DurableJobShardJournalRecordKind.Snapshot,
            Snapshot = snapshot
        };
    }
}

[GenerateSerializer]
[Alias("Orleans.DurableJobs.DurableJobShardScheduleOperation")]
internal sealed class DurableJobShardScheduleOperation
{
    [Id(0)]
    public DurableJob Job { get; init; } = null!;
}

[GenerateSerializer]
[Alias("Orleans.DurableJobs.DurableJobShardRemoveOperation")]
internal sealed class DurableJobShardRemoveOperation
{
    [Id(0)]
    public string JobId { get; init; } = string.Empty;
}

[GenerateSerializer]
[Alias("Orleans.DurableJobs.DurableJobShardRetryOperation")]
internal sealed class DurableJobShardRetryOperation
{
    [Id(0)]
    public string JobId { get; init; } = string.Empty;

    [Id(1)]
    public DateTimeOffset DueTime { get; init; }

    [Id(2)]
    public int DequeueCount { get; init; }

    [Id(3)]
    public long? ExecutionGeneration { get; init; }
}

[GenerateSerializer]
[Alias("Orleans.DurableJobs.DurableJobShardSnapshot")]
internal sealed class DurableJobShardSnapshot
{
    [Id(0)]
    public List<DurableJobShardSnapshotEntry> Jobs { get; init; } = [];
}

[GenerateSerializer]
[Alias("Orleans.DurableJobs.DurableJobShardSnapshotEntry")]
internal sealed class DurableJobShardSnapshotEntry
{
    [Id(0)]
    public DurableJob Job { get; init; } = null!;

    [Id(1)]
    public int DequeueCount { get; init; }
}