natsdotnet/tests/NATS.Server.JetStream.Tests/JetStream/Api/ClusteredRequestTests.cs

// Go reference: jetstream_cluster.go:7620-7701 — jsClusteredStreamRequest lifecycle:
// propose to meta RAFT → wait for result → deliver or time out.
// ClusteredRequestProcessor tracks pending requests and delivers results when RAFT entries
// are applied, matching the Go server's callback-based completion mechanism.

using NATS.Server.JetStream.Api;

namespace NATS.Server.JetStream.Tests.JetStream.Api;

public class ClusteredRequestTests
{
    // ---------------------------------------------------------------
    // RegisterPending
    // ---------------------------------------------------------------

    /// <summary>
    /// Each call to RegisterPending returns a distinct, non-empty string identifier.
    /// Go reference: jetstream_cluster.go:7620 — each clustered request gets a unique ID
    /// used to correlate the RAFT apply callback with the waiting caller.
    /// </summary>
    [Fact]
    public void RegisterPending_returns_unique_id()
    {
        var processor = new ClusteredRequestProcessor();

        var id1 = processor.RegisterPending();
        var id2 = processor.RegisterPending();

        id1.ShouldNotBeNullOrWhiteSpace();
        id2.ShouldNotBeNullOrWhiteSpace();
        id1.ShouldNotBe(id2);
    }

    // ---------------------------------------------------------------
    // WaitForResult
    // ---------------------------------------------------------------

    /// <summary>
    /// When a result is delivered for a pending request, WaitForResultAsync returns that response.
    /// Go reference: jetstream_cluster.go:7620 — the waiting goroutine receives the result
    /// via channel once the RAFT leader applies the entry.
    /// </summary>
    [Fact]
    public async Task WaitForResult_returns_delivered_response()
    {
        var processor = new ClusteredRequestProcessor(timeout: TimeSpan.FromSeconds(5));
        var requestId = processor.RegisterPending();
        var expected = JetStreamApiResponse.SuccessResponse();

        // Use a semaphore so the wait starts before delivery occurs — no timing dependency.
        var waitStarted = new SemaphoreSlim(0, 1);
        var deliverTask = Task.Run(async () =>
        {
            // Wait until WaitForResultAsync has been entered before delivering.
            await waitStarted.WaitAsync();
            processor.DeliverResult(requestId, expected);
        });

        // Signal the deliver task once we begin waiting.
        waitStarted.Release();
        var result = await processor.WaitForResultAsync(requestId);

        await deliverTask;
        result.ShouldBeSameAs(expected);
    }

    /// <summary>
    /// When no result is delivered within the timeout, WaitForResultAsync returns a 408 error.
    /// Go reference: jetstream_cluster.go:7620 — if the RAFT group does not respond in time,
    /// the request is considered timed out and an error is returned to the client.
    /// </summary>
    [Fact]
    public async Task WaitForResult_times_out_after_timeout()
    {
        var processor = new ClusteredRequestProcessor(timeout: TimeSpan.FromMilliseconds(50));
        var requestId = processor.RegisterPending();

        var result = await processor.WaitForResultAsync(requestId);

        result.Error.ShouldNotBeNull();
        result.Error!.Code.ShouldBe(408);
        result.Error.Description.ShouldContain("timeout");
    }

    /// <summary>
    /// WaitForResultAsync returns a 500 error for an ID that was never registered.
    /// Go reference: jetstream_cluster.go — requesting a result for an unknown request ID
    /// is a programming error; return an internal server error.
    /// </summary>
    [Fact]
    public async Task WaitForResult_returns_error_for_unknown_id()
    {
        var processor = new ClusteredRequestProcessor();

        var result = await processor.WaitForResultAsync("nonexistent-id");

        result.Error.ShouldNotBeNull();
        result.Error!.Code.ShouldBe(500);
        result.Error.Description.ShouldContain("not found");
    }

    /// <summary>
    /// When the caller's CancellationToken is triggered, WaitForResultAsync returns a timeout error.
    /// Go reference: jetstream_cluster.go:7620 — callers can cancel waiting for a RAFT result
    /// if their own request context is cancelled.
    /// </summary>
    [Fact]
    public async Task WaitForResult_respects_cancellation_token()
    {
        var processor = new ClusteredRequestProcessor(timeout: TimeSpan.FromSeconds(30));
        var requestId = processor.RegisterPending();

        using var cts = new CancellationTokenSource(TimeSpan.FromMilliseconds(50));
        var result = await processor.WaitForResultAsync(requestId, cts.Token);

        result.Error.ShouldNotBeNull();
        result.Error!.Code.ShouldBe(408);
    }

    // ---------------------------------------------------------------
    // DeliverResult
    // ---------------------------------------------------------------

    /// <summary>
    /// DeliverResult returns true when the request ID is known and pending.
    /// Go reference: jetstream_cluster.go:7620 — the RAFT apply callback signals success
    /// by resolving the pending request.
    /// </summary>
    [Fact]
    public void DeliverResult_returns_true_for_pending_request()
    {
        var processor = new ClusteredRequestProcessor();
        var requestId = processor.RegisterPending();

        var delivered = processor.DeliverResult(requestId, JetStreamApiResponse.SuccessResponse());

        delivered.ShouldBeTrue();
    }

    /// <summary>
    /// DeliverResult returns false when the request ID is not found.
    /// Go reference: jetstream_cluster.go — delivering a result for an unknown or already-completed
    /// request is a no-op; return false so the caller knows the result was not consumed.
    /// </summary>
    [Fact]
    public void DeliverResult_returns_false_for_unknown_request()
    {
        var processor = new ClusteredRequestProcessor();

        var delivered = processor.DeliverResult("unknown-id", JetStreamApiResponse.SuccessResponse());

        delivered.ShouldBeFalse();
    }

    // ---------------------------------------------------------------
    // PendingCount
    // ---------------------------------------------------------------

    /// <summary>
    /// PendingCount increases with each RegisterPending call and decreases when a result is
    /// delivered or the request times out.
    /// Go reference: jetstream_cluster.go — the server tracks pending RAFT proposals for
    /// observability and to detect stuck requests.
    /// </summary>
    [Fact]
    public async Task PendingCount_tracks_active_requests()
    {
        var processor = new ClusteredRequestProcessor(timeout: TimeSpan.FromMilliseconds(50));

        processor.PendingCount.ShouldBe(0);

        var id1 = processor.RegisterPending();
        processor.PendingCount.ShouldBe(1);

        var id2 = processor.RegisterPending();
        processor.PendingCount.ShouldBe(2);

        // Deliver one request.
        processor.DeliverResult(id1, JetStreamApiResponse.SuccessResponse());
        processor.PendingCount.ShouldBe(1);

        // Let id2 time out.
        await processor.WaitForResultAsync(id2);
        processor.PendingCount.ShouldBe(0);
    }

    // ---------------------------------------------------------------
    // CancelAll
    // ---------------------------------------------------------------

    /// <summary>
    /// CancelAll completes all pending requests with a 503 error response.
    /// Go reference: jetstream_cluster.go — when this node loses RAFT leadership, all
    /// in-flight proposals must be failed so callers do not hang indefinitely.
    /// </summary>
    [Fact]
    public async Task CancelAll_completes_all_pending_with_error()
    {
        var processor = new ClusteredRequestProcessor(timeout: TimeSpan.FromSeconds(30));

        var id1 = processor.RegisterPending();
        var id2 = processor.RegisterPending();

        var task1 = processor.WaitForResultAsync(id1);
        var task2 = processor.WaitForResultAsync(id2);

        processor.CancelAll("leadership changed");

        var result1 = await task1;
        var result2 = await task2;

        result1.Error.ShouldNotBeNull();
        result1.Error!.Code.ShouldBe(503);
        result1.Error.Description.ShouldContain("leadership changed");

        result2.Error.ShouldNotBeNull();
        result2.Error!.Code.ShouldBe(503);
        result2.Error.Description.ShouldContain("leadership changed");
    }

    /// <summary>
    /// After CancelAll, PendingCount drops to zero.
    /// Go reference: jetstream_cluster.go — a leadership change clears all pending state.
    /// </summary>
    [Fact]
    public void CancelAll_clears_pending_count()
    {
        var processor = new ClusteredRequestProcessor(timeout: TimeSpan.FromSeconds(30));

        processor.RegisterPending();
        processor.RegisterPending();
        processor.RegisterPending();

        processor.PendingCount.ShouldBe(3);

        processor.CancelAll();

        processor.PendingCount.ShouldBe(0);
    }

    /// <summary>
    /// CancelAll uses a default reason of "leadership changed" when no reason is provided.
    /// Go reference: jetstream_cluster.go — default cancellation reason matches NATS cluster semantics.
    /// </summary>
    [Fact]
    public async Task CancelAll_uses_default_reason()
    {
        var processor = new ClusteredRequestProcessor(timeout: TimeSpan.FromSeconds(30));

        var id = processor.RegisterPending();
        var task = processor.WaitForResultAsync(id);

        processor.CancelAll(); // no reason argument

        var result = await task;

        result.Error.ShouldNotBeNull();
        result.Error!.Description.ShouldContain("leadership changed");
    }

    // ---------------------------------------------------------------
    // Concurrency
    // ---------------------------------------------------------------

    /// <summary>
    /// Concurrent registrations and deliveries all receive the correct response.
    /// Go reference: jetstream_cluster.go — in a cluster, many API requests may be in-flight
    /// simultaneously, each waiting for its own RAFT entry to be applied.
    /// </summary>
    [Fact]
    public async Task Concurrent_register_and_deliver()
    {
        const int count = 50;
        var processor = new ClusteredRequestProcessor(timeout: TimeSpan.FromSeconds(10));

        var requestIds = new string[count];
        for (var i = 0; i < count; i++)
            requestIds[i] = processor.RegisterPending();

        // Start all waits concurrently.
        var waitTasks = requestIds.Select(id => processor.WaitForResultAsync(id)).ToArray();

        // Deliver all results concurrently — no delay needed; the ThreadPool provides
        // sufficient interleaving to exercise concurrent access patterns.
        var deliverTasks = requestIds.Select((id, i) => Task.Run(() =>
        {
            processor.DeliverResult(id, JetStreamApiResponse.ErrorResponse(200 + i, $"response-{i}"));
        })).ToArray();

        await Task.WhenAll(deliverTasks);
        var results = await Task.WhenAll(waitTasks);

        // Every result should be a valid response (no null errors from "not found").
        results.Length.ShouldBe(count);
        foreach (var result in results)
        {
            // Each result was an explicitly delivered response with a known code.
            result.Error.ShouldNotBeNull();
            result.Error!.Code.ShouldBeGreaterThanOrEqualTo(200);
            result.Error.Code.ShouldBeLessThan(200 + count);
        }
    }
}