Skip to content

perf(durable-jobs): reduce queue wake-ups #11103

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ReubenBond merged 1 commit into from
Sep 5, 2026
Merged

perf(durable-jobs): reduce queue wake-ups #11103

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 50 additions & 11 deletions src/Orleans.DurableJobs/InMemoryJobQueue.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ internal sealed class InMemoryJobQueue : IAsyncEnumerable<IJobRunContext>
private readonly PriorityQueue<JobBucket, DateTimeOffset> _queue = new();
private readonly Dictionary<string, JobBucket> _jobsIdToBucket = new();
private readonly Dictionary<DateTimeOffset, JobBucket> _buckets = new();
private TaskCompletionSource _queueChanged = CreateQueueChangedSource();
private TaskCompletionSource? _queueChangedWaiter;
private int _jobCount;
private bool _isComplete;
#if NET9_0_OR_GREATER
Expand Down Expand Up @@ -54,6 +54,8 @@ public void Enqueue(DurableJob job, int dequeueCount)
if (_isComplete)
throw new InvalidOperationException("Cannot enqueue job to a completed queue.");

var wakeCheckRequired = _queueChangedWaiter is not null;
var previousNextDueTime = wakeCheckRequired ? GetNextDueTime() : null;
var bucket = GetJobBucket(job.DueTime);
var isReplacement = _jobsIdToBucket.TryGetValue(job.Id, out var existingBucket);
if (existingBucket is not null && !ReferenceEquals(existingBucket, bucket))
Expand All @@ -68,7 +70,10 @@ public void Enqueue(DurableJob job, int dequeueCount)
Volatile.Write(ref _jobCount, _jobCount + 1);
}

SignalQueueChanged();
if (wakeCheckRequired)
{
SignalQueueChangedIfNextDueTimeChanged(previousNextDueTime);
}
}
}

Expand Down Expand Up @@ -98,14 +103,20 @@ public bool RemoveJob(string jobId)
ArgumentException.ThrowIfNullOrWhiteSpace(jobId);
lock (_syncLock)
{
var wakeCheckRequired = _queueChangedWaiter is not null;
var previousNextDueTime = wakeCheckRequired ? GetNextDueTime() : null;
if (_jobsIdToBucket.TryGetValue(jobId, out var bucket))
{
// Try to remove from bucket (may already be dequeued)
bucket.RemoveJob(jobId);
_jobsIdToBucket.Remove(jobId);
Volatile.Write(ref _jobCount, _jobCount - 1);
// Note: The bucket remains in the priority queue until processed
SignalQueueChanged();
if (wakeCheckRequired)
{
SignalQueueChangedIfNextDueTimeChanged(previousNextDueTime);
}

return true;
}

Expand Down Expand Up @@ -160,6 +171,8 @@ internal bool RetryJobLater(string jobId, DateTimeOffset newDueTime, int dequeue

lock (_syncLock)
{
var wakeCheckRequired = _queueChangedWaiter is not null;
var previousNextDueTime = wakeCheckRequired ? GetNextDueTime() : null;
if (!_jobsIdToBucket.TryGetValue(jobId, out var oldBucket) || !oldBucket.TryGetJob(jobId, out var existing))
{
return false;
Expand All @@ -183,7 +196,11 @@ internal bool RetryJobLater(string jobId, DateTimeOffset newDueTime, int dequeue
var newBucket = GetJobBucket(newDueTime);
newBucket.AddJob(newJob, dequeueCount);
_jobsIdToBucket[jobId] = newBucket;
SignalQueueChanged();
if (wakeCheckRequired)
{
SignalQueueChangedIfNextDueTimeChanged(previousNextDueTime);
}

return true;
}
}
Expand Down Expand Up @@ -216,12 +233,17 @@ public void Clear()
{
lock (_syncLock)
{
var wakeCheckRequired = _queueChangedWaiter is not null;
var previousNextDueTime = wakeCheckRequired ? GetNextDueTime() : null;
_queue.Clear();
_jobsIdToBucket.Clear();
_buckets.Clear();
Volatile.Write(ref _jobCount, 0);
_isComplete = false;
SignalQueueChanged();
if (wakeCheckRequired)
{
SignalQueueChangedIfNextDueTimeChanged(previousNextDueTime);
}
}
}

Expand Down Expand Up @@ -253,7 +275,7 @@ public async IAsyncEnumerator<IJobRunContext> GetAsyncEnumerator(CancellationTok
yield break; // Exit if the queue is frozen and empty
}

queueChanged = _queueChanged.Task;
queueChanged = GetQueueChangedTask();
}
else if (_queue.Count > 0)
{
Expand All @@ -275,13 +297,13 @@ public async IAsyncEnumerator<IJobRunContext> GetAsyncEnumerator(CancellationTok
}
else
{
queueChanged = _queueChanged.Task;
queueChanged = GetQueueChangedTask();
delay = nextBucket.DueTime - now;
}
}
else
{
queueChanged = _queueChanged.Task;
queueChanged = GetQueueChangedTask();
}
}

Expand Down Expand Up @@ -335,13 +357,30 @@ private void RemoveEmptyBuckets()
}
}

private DateTimeOffset? GetNextDueTime()
{
RemoveEmptyBuckets();
return _queue.Count == 0 ? null : _queue.Peek().DueTime;
}

private void SignalQueueChangedIfNextDueTimeChanged(DateTimeOffset? previousNextDueTime)
{
if ((_isComplete && _jobCount == 0) || previousNextDueTime != GetNextDueTime())
{
SignalQueueChanged();
}
}

private void SignalQueueChanged()
{
var previous = _queueChanged;
_queueChanged = CreateQueueChangedSource();
previous.TrySetResult();
var waiter = _queueChangedWaiter;
_queueChangedWaiter = null;
waiter?.TrySetResult();
}

private Task GetQueueChangedTask()
=> (_queueChangedWaiter ??= CreateQueueChangedSource()).Task;

private async Task WaitForQueueChangeOrDelayAsync(Task queueChanged, TimeSpan? delay, CancellationToken cancellationToken)
{
if (delay is null)
Expand Down
40 changes: 40 additions & 0 deletions test/Orleans.Core.Tests/DurableJobs/InMemoryJobQueueTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,6 +202,25 @@ public async Task GetAsyncEnumerator_WhenDueJobIsEnqueued_WakesWithoutAdvancingT
Assert.Equal(job.Id, enumerator.Current.Job.Id);
}

[Fact]
public async Task RetryJobLater_WhenNextDueTimeMovesEarlier_WakesWaitingEnumerator()
{
var timeProvider = new FakeTimeProvider(new DateTimeOffset(2026, 1, 1, 0, 0, 0, TimeSpan.Zero));
var queue = new InMemoryJobQueue(timeProvider);
var job = CreateJob("job1", timeProvider.GetUtcNow().AddHours(1));
queue.Enqueue(job, 0);

await using var enumerator = queue.GetAsyncEnumerator(TestContext.Current.CancellationToken);
var moveNextTask = enumerator.MoveNextAsync().AsTask();
Assert.False(moveNextTask.IsCompleted);

Assert.True(queue.RetryJobLater(job.Id, timeProvider.GetUtcNow(), 3));

Assert.True(await moveNextTask.WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken));
Assert.Equal(job.Id, enumerator.Current.Job.Id);
Assert.Equal(4, enumerator.Current.DequeueCount);
}

[Fact]
public async Task GetAsyncEnumerator_CompletesWhenQueueIsMarkedComplete()
{
Expand Down Expand Up @@ -544,6 +563,27 @@ public async Task Enqueue_SameJobIdMovedDuringBucketDrain_YieldsOnlyReplacement(
Assert.Equal(0, queue.Count);
}

[Fact]
public async Task RemoveLastDequeuedJob_AfterCompletion_WakesWaitingEnumerator()
{
var timeProvider = new FakeTimeProvider(new DateTimeOffset(2026, 8, 16, 0, 0, 0, TimeSpan.Zero));
var queue = new InMemoryJobQueue(timeProvider);
var job = CreateJob("in-flight", timeProvider.GetUtcNow().AddMilliseconds(-100));
queue.Enqueue(job, 0);

await using var enumerator = queue.GetAsyncEnumerator(TestContext.Current.CancellationToken);
Assert.True(await enumerator.MoveNextAsync());
Assert.Equal(1, queue.Count);

queue.MarkAsComplete();
var completion = enumerator.MoveNextAsync().AsTask();
Assert.False(completion.IsCompleted);

Assert.True(queue.RemoveJob(job.Id));
Assert.False(await completion.WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken));
Assert.Equal(0, queue.Count);
}

private static DurableJob CreateJob(string id, DateTimeOffset dueTime, string? traceParent = null, string? traceState = null)
{
return new DurableJob
Expand Down
Loading