feat: execute post-baseline jetstream parity plan

This commit is contained in:
Joseph Doherty
2026-02-23 12:11:19 -05:00
parent c3763e83d6
commit b41e6ff320
58 changed files with 1430 additions and 102 deletions

View File

@@ -2,22 +2,50 @@ namespace NATS.Server.JetStream.Consumers;
public sealed class AckProcessor
{
private readonly Dictionary<ulong, DateTime> _pending = new();
private readonly Dictionary<ulong, PendingState> _pending = new();
public void Register(ulong sequence, int ackWaitMs)
{
_pending[sequence] = DateTime.UtcNow.AddMilliseconds(Math.Max(ackWaitMs, 1));
if (_pending.ContainsKey(sequence))
return;
_pending[sequence] = new PendingState
{
DeadlineUtc = DateTime.UtcNow.AddMilliseconds(Math.Max(ackWaitMs, 1)),
Deliveries = 1,
};
}
public ulong? NextExpired()
public bool TryGetExpired(out ulong sequence, out int deliveries)
{
foreach (var (seq, deadline) in _pending)
foreach (var (seq, state) in _pending)
{
if (DateTime.UtcNow >= deadline)
return seq;
if (DateTime.UtcNow >= state.DeadlineUtc)
{
sequence = seq;
deliveries = state.Deliveries;
return true;
}
}
return null;
sequence = 0;
deliveries = 0;
return false;
}
public void ScheduleRedelivery(ulong sequence, int delayMs)
{
if (!_pending.TryGetValue(sequence, out var state))
return;
state.Deliveries++;
state.DeadlineUtc = DateTime.UtcNow.AddMilliseconds(Math.Max(delayMs, 1));
_pending[sequence] = state;
}
public void Drop(ulong sequence)
{
_pending.Remove(sequence);
}
public bool HasPending => _pending.Count > 0;
@@ -28,4 +56,10 @@ public sealed class AckProcessor
foreach (var key in _pending.Keys.Where(k => k <= sequence).ToArray())
_pending.Remove(key);
}
private sealed class PendingState
{
public DateTime DeadlineUtc { get; set; }
public int Deliveries { get; set; }
}
}

View File

@@ -16,11 +16,7 @@ public sealed class PullConsumerEngine
if (consumer.NextSequence == 1)
{
var state = await stream.Store.GetStateAsync(ct);
if (consumer.Config.DeliverPolicy == DeliverPolicy.Last && state.LastSeq > 0)
consumer.NextSequence = state.LastSeq;
else if (consumer.Config.DeliverPolicy == DeliverPolicy.New && state.LastSeq > 0)
consumer.NextSequence = state.LastSeq + 1;
consumer.NextSequence = await ResolveInitialSequenceAsync(stream, consumer.Config, ct);
}
if (request.NoWait)
@@ -32,9 +28,19 @@ public sealed class PullConsumerEngine
if (consumer.Config.AckPolicy == AckPolicy.Explicit)
{
var expired = consumer.AckProcessor.NextExpired();
if (expired is { } expiredSequence)
if (consumer.AckProcessor.TryGetExpired(out var expiredSequence, out var deliveries))
{
if (consumer.Config.MaxDeliver > 0 && deliveries > consumer.Config.MaxDeliver)
{
consumer.AckProcessor.Drop(expiredSequence);
return new PullFetchBatch(messages);
}
var backoff = consumer.Config.BackOffMs.Count >= deliveries
? consumer.Config.BackOffMs[deliveries - 1]
: consumer.Config.AckWaitMs;
consumer.AckProcessor.ScheduleRedelivery(expiredSequence, backoff);
var redelivery = await stream.Store.LoadAsync(expiredSequence, ct);
if (redelivery != null)
{
@@ -86,6 +92,27 @@ public sealed class PullConsumerEngine
return new PullFetchBatch(messages);
}
private static async ValueTask<ulong> ResolveInitialSequenceAsync(StreamHandle stream, ConsumerConfig config, CancellationToken ct)
{
var state = await stream.Store.GetStateAsync(ct);
return config.DeliverPolicy switch
{
DeliverPolicy.Last when state.LastSeq > 0 => state.LastSeq,
DeliverPolicy.New when state.LastSeq > 0 => state.LastSeq + 1,
DeliverPolicy.ByStartSequence when config.OptStartSeq > 0 => config.OptStartSeq,
DeliverPolicy.ByStartTime when config.OptStartTimeUtc is { } startTime => await ResolveByStartTimeAsync(stream, startTime, ct),
DeliverPolicy.LastPerSubject when state.LastSeq > 0 => state.LastSeq,
_ => 1,
};
}
private static async ValueTask<ulong> ResolveByStartTimeAsync(StreamHandle stream, DateTime startTimeUtc, CancellationToken ct)
{
var messages = await stream.Store.ListAsync(ct);
var match = messages.FirstOrDefault(m => m.TimestampUtc >= startTimeUtc);
return match?.Sequence ?? 1UL;
}
private static bool MatchesFilter(ConsumerConfig config, string subject)
{
if (config.FilterSubjects.Count > 0)

View File

@@ -7,20 +7,41 @@ public sealed class PushConsumerEngine
{
public void Enqueue(ConsumerHandle consumer, StoredMessage message)
{
var availableAtUtc = DateTime.UtcNow;
if (consumer.Config.RateLimitBps > 0)
{
if (consumer.NextPushDataAvailableAtUtc > availableAtUtc)
availableAtUtc = consumer.NextPushDataAvailableAtUtc;
var delayMs = (long)Math.Ceiling((double)message.Payload.Length * 1000 / consumer.Config.RateLimitBps);
consumer.NextPushDataAvailableAtUtc = availableAtUtc.AddMilliseconds(Math.Max(delayMs, 1));
}
consumer.PushFrames.Enqueue(new PushFrame
{
IsData = true,
Message = message,
AvailableAtUtc = availableAtUtc,
});
if (consumer.Config.AckPolicy is AckPolicy.Explicit or AckPolicy.All)
consumer.AckProcessor.Register(message.Sequence, consumer.Config.AckWaitMs);
if (consumer.Config.FlowControl)
{
consumer.PushFrames.Enqueue(new PushFrame
{
IsFlowControl = true,
AvailableAtUtc = availableAtUtc,
});
}
if (consumer.Config.HeartbeatMs > 0)
{
consumer.PushFrames.Enqueue(new PushFrame
{
IsHeartbeat = true,
AvailableAtUtc = availableAtUtc,
});
}
}
@@ -29,6 +50,8 @@ public sealed class PushConsumerEngine
public sealed class PushFrame
{
public bool IsData { get; init; }
public bool IsFlowControl { get; init; }
public bool IsHeartbeat { get; init; }
public StoredMessage? Message { get; init; }
public DateTime AvailableAtUtc { get; init; } = DateTime.UtcNow;
}