feat: implement jetstream pull consumer fetch

This commit is contained in:
Joseph Doherty
2026-02-23 06:07:02 -05:00
parent 40b940b1fd
commit 9a0de19c2d
4 changed files with 76 additions and 0 deletions

View File

@@ -1,5 +1,6 @@
using System.Collections.Concurrent;
using NATS.Server.JetStream.Api;
using NATS.Server.JetStream.Consumers;
using NATS.Server.JetStream.Models;
using NATS.Server.JetStream.Storage;
@@ -8,6 +9,7 @@ namespace NATS.Server.JetStream;
public sealed class ConsumerManager
{
private readonly ConcurrentDictionary<(string Stream, string Name), ConsumerHandle> _consumers = new();
private readonly PullConsumerEngine _pullConsumerEngine = new();
public int ConsumerCount => _consumers.Count;
@@ -48,6 +50,17 @@ public sealed class ConsumerManager
public bool TryGet(string stream, string durableName, out ConsumerHandle handle)
=> _consumers.TryGetValue((stream, durableName), out handle!);
public async ValueTask<PullFetchBatch> FetchAsync(string stream, string durableName, int batch, StreamManager streamManager, CancellationToken ct)
{
if (!_consumers.TryGetValue((stream, durableName), out var consumer))
return new PullFetchBatch([]);
if (!streamManager.TryGet(stream, out var streamHandle))
return new PullFetchBatch([]);
return await _pullConsumerEngine.FetchAsync(streamHandle, consumer, batch, ct);
}
}
public sealed record ConsumerHandle(string Stream, ConsumerConfig Config)

View File

@@ -0,0 +1,35 @@
using NATS.Server.JetStream.Storage;
namespace NATS.Server.JetStream.Consumers;
public sealed class PullConsumerEngine
{
public async ValueTask<PullFetchBatch> FetchAsync(StreamHandle stream, ConsumerHandle consumer, int batch, CancellationToken ct)
{
var messages = new List<StoredMessage>(batch);
var sequence = consumer.NextSequence;
for (var i = 0; i < batch; i++)
{
var message = await stream.Store.LoadAsync(sequence, ct);
if (message == null)
break;
messages.Add(message);
sequence++;
}
consumer.NextSequence = sequence;
return new PullFetchBatch(messages);
}
}
public sealed class PullFetchBatch
{
public IReadOnlyList<StoredMessage> Messages { get; }
public PullFetchBatch(IReadOnlyList<StoredMessage> messages)
{
Messages = messages;
}
}