feat(consumer): add PriorityQueue-based RedeliveryTracker overloads

Add new constructor, Schedule(DateTimeOffset), GetDue(DateTimeOffset),
IncrementDeliveryCount, IsMaxDeliveries(), and GetBackoffDelay() to
RedeliveryTracker without breaking existing API. Uses PriorityQueue<ulong,
DateTimeOffset> for deadline-ordered dispatch mirroring Go consumer.go rdq.
This commit is contained in:
Joseph Doherty
2026-02-25 02:12:37 -05:00
parent 7e4a23a0b7
commit 55de052009
2 changed files with 219 additions and 1 deletions

View File

@@ -7,14 +7,37 @@ namespace NATS.Server.JetStream.Consumers;
public sealed class RedeliveryTracker
{
private readonly int[] _backoffMs;
private readonly long[]? _backoffMsLong;
// Go: consumer.go — pending maps sseq → (deadline, deliveries)
private readonly Dictionary<ulong, RedeliveryEntry> _entries = new();
// Go: consumer.go — rdc map tracks per-sequence delivery counts
private readonly Dictionary<ulong, int> _deliveryCounts = new();
// Go: consumer.go — rdq priority queue ordered by deadline for efficient dispatch
private readonly PriorityQueue<ulong, DateTimeOffset> _priorityQueue = new();
// Stored config for the new constructor overload
private readonly int _maxDeliveries;
private readonly long _ackWaitMs;
// Go: consumer.go:100 — BackOff []time.Duration in ConsumerConfig; empty falls back to ackWait
public RedeliveryTracker(int[] backoffMs)
{
_backoffMs = backoffMs;
_backoffMsLong = null;
_maxDeliveries = 0;
_ackWaitMs = 0;
}
// Go: consumer.go — ConsumerConfig maxDeliver + ackWait + backoff, new overload storing config fields
public RedeliveryTracker(int maxDeliveries, long ackWaitMs, long[]? backoffMs = null)
{
_backoffMs = [];
_backoffMsLong = backoffMs;
_maxDeliveries = maxDeliveries;
_ackWaitMs = ackWaitMs;
}
// Go: consumer.go:5540 — trackPending records delivery count and schedules deadline
@@ -34,6 +57,13 @@ public sealed class RedeliveryTracker
return deadline;
}
// Go: consumer.go — schedule with an explicit deadline into the priority queue
public void Schedule(ulong seq, DateTimeOffset deadline)
{
_deliveryCounts.TryAdd(seq, 0);
_priorityQueue.Enqueue(seq, deadline);
}
// Go: consumer.go — rdq entries are dispatched once their deadline has passed
public IReadOnlyList<ulong> GetDue()
{
@@ -52,8 +82,53 @@ public sealed class RedeliveryTracker
return due ?? (IReadOnlyList<ulong>)[];
}
// Go: consumer.go — drain the rdq priority queue of all entries whose deadline <= now,
// returning them in deadline order (earliest first).
public IEnumerable<ulong> GetDue(DateTimeOffset now)
{
List<(ulong seq, DateTimeOffset deadline)>? dequeued = null;
List<(ulong seq, DateTimeOffset deadline)>? future = null;
// Drain the entire queue, separating due from future
while (_priorityQueue.TryDequeue(out var seq, out var deadline))
{
// Skip sequences that were acknowledged
if (!_deliveryCounts.ContainsKey(seq))
continue;
if (deadline <= now)
{
dequeued ??= [];
dequeued.Add((seq, deadline));
}
else
{
future ??= [];
future.Add((seq, deadline));
}
}
// Re-enqueue future items
if (future is not null)
{
foreach (var (seq, deadline) in future)
_priorityQueue.Enqueue(seq, deadline);
}
if (dequeued is null)
return [];
// Already extracted in priority order since PriorityQueue dequeues min first
return dequeued.Select(x => x.seq);
}
// Go: consumer.go — acking a sequence removes it from the pending redelivery set
public void Acknowledge(ulong seq) => _entries.Remove(seq);
public void Acknowledge(ulong seq)
{
_entries.Remove(seq);
_deliveryCounts.Remove(seq);
// Priority queue entries are lazily skipped in GetDue when seq not in _deliveryCounts
}
// Go: consumer.go — maxdeliver check: drop sequence once delivery count exceeds max
public bool IsMaxDeliveries(ulong seq, int maxDeliver)
@@ -67,6 +142,36 @@ public sealed class RedeliveryTracker
return entry.DeliveryCount >= maxDeliver;
}
// Go: consumer.go — maxdeliver check using the stored _maxDeliveries from new constructor
public bool IsMaxDeliveries(ulong seq)
{
if (_maxDeliveries <= 0)
return false;
_deliveryCounts.TryGetValue(seq, out var count);
return count >= _maxDeliveries;
}
// Go: consumer.go — rdc map increment: track how many times a sequence has been delivered
public void IncrementDeliveryCount(ulong seq)
{
_deliveryCounts[seq] = _deliveryCounts.TryGetValue(seq, out var count) ? count + 1 : 1;
}
// Go: consumer.go — backoff delay lookup: index by deliveryCount, clamp to last entry,
// fall back to ackWait when no backoff array is configured.
public long GetBackoffDelay(int deliveryCount)
{
if (_backoffMsLong is { Length: > 0 })
{
var idx = Math.Min(deliveryCount - 1, _backoffMsLong.Length - 1);
if (idx < 0) idx = 0;
return _backoffMsLong[idx];
}
return _ackWaitMs;
}
public bool IsTracking(ulong seq) => _entries.ContainsKey(seq);
public int TrackedCount => _entries.Count;