feat: add message trace propagation across servers (Gap 10.4)

Add TraceContextPropagator and TraceContext to Internal/MessageTraceContext.cs for Nats-Trace-Parent header injection and extraction across server hops. Also add ConnzSortOption/ConnzSorter stubs to fix pre-existing build errors in Monitoring/Connz.cs. Covered by 10 new tests in TraceContextPropagationTests.cs.
2026-02-25 13:06:58 -05:00
parent eb801cd4cf
commit 9fb2ae205c
3 changed files with 391 additions and 1 deletions
--- a/tests/NATS.Server.Tests/Internal/TraceContextPropagationTests.cs
+++ b/tests/NATS.Server.Tests/Internal/TraceContextPropagationTests.cs
@@ -0,0 +1,136 @@
+using System.Text;
+using NATS.Server.Internal;
+
+namespace NATS.Server.Tests.Internal;
+
+/// <summary>
+/// Tests for TraceContextPropagator: trace creation, header injection/extraction,
+/// child span creation, round-trip fidelity, and ShouldTrace detection.
+/// Go reference: server/msgtrace.go — trace context embedding and extraction.
+/// </summary>
+public class TraceContextPropagationTests
+{
+    // Helper: build a minimal NATS/1.0 header block with the given headers.
+    private static byte[] BuildNatsHeaders(params (string key, string value)[] headers)
+    {
+        var sb = new StringBuilder("NATS/1.0\r\n");
+        foreach (var (key, value) in headers)
+            sb.Append($"{key}: {value}\r\n");
+        sb.Append("\r\n");
+        return Encoding.ASCII.GetBytes(sb.ToString());
+    }
+
+    [Fact]
+    public void CreateTrace_GeneratesValidContext()
+    {
+        var ctx = TraceContextPropagator.CreateTrace("abc123", "span456", destination: "trace.dest");
+
+        ctx.TraceId.ShouldBe("abc123");
+        ctx.SpanId.ShouldBe("span456");
+        ctx.Destination.ShouldBe("trace.dest");
+        ctx.TraceOnly.ShouldBeFalse();
+        ctx.CreatedAt.ShouldBeInRange(DateTime.UtcNow.AddSeconds(-5), DateTime.UtcNow.AddSeconds(1));
+    }
+
+    [Fact]
+    public void ExtractTrace_ValidHeaders_ReturnsContext()
+    {
+        var headers = BuildNatsHeaders((TraceContextPropagator.TraceParentHeader, "trace1-span1"));
+
+        var ctx = TraceContextPropagator.ExtractTrace(headers);
+
+        ctx.ShouldNotBeNull();
+        ctx.TraceId.ShouldBe("trace1");
+        ctx.SpanId.ShouldBe("span1");
+    }
+
+    [Fact]
+    public void ExtractTrace_NoTraceHeader_ReturnsNull()
+    {
+        var headers = BuildNatsHeaders(("Content-Type", "text/plain"));
+
+        var ctx = TraceContextPropagator.ExtractTrace(headers);
+
+        ctx.ShouldBeNull();
+    }
+
+    [Fact]
+    public void InjectTrace_AppendsToHeaders()
+    {
+        var existing = BuildNatsHeaders(("Content-Type", "text/plain"));
+        var ctx = TraceContextPropagator.CreateTrace("tid", "sid");
+
+        var result = TraceContextPropagator.InjectTrace(ctx, existing);
+
+        var text = Encoding.ASCII.GetString(result);
+        text.ShouldContain($"{TraceContextPropagator.TraceParentHeader}: tid-sid");
+        text.ShouldContain("Content-Type: text/plain");
+    }
+
+    [Fact]
+    public void InjectTrace_EmptyHeaders_CreatesNew()
+    {
+        var ctx = TraceContextPropagator.CreateTrace("newtrace", "newspan");
+
+        var result = TraceContextPropagator.InjectTrace(ctx, ReadOnlySpan<byte>.Empty);
+
+        var text = Encoding.ASCII.GetString(result);
+        text.ShouldStartWith("NATS/1.0\r\n");
+        text.ShouldContain($"{TraceContextPropagator.TraceParentHeader}: newtrace-newspan");
+    }
+
+    [Fact]
+    public void CreateChildSpan_PreservesTraceId()
+    {
+        var parent = TraceContextPropagator.CreateTrace("parentTrace", "parentSpan");
+
+        var child = TraceContextPropagator.CreateChildSpan(parent, "childSpan");
+
+        child.TraceId.ShouldBe("parentTrace");
+    }
+
+    [Fact]
+    public void CreateChildSpan_NewSpanId()
+    {
+        var parent = TraceContextPropagator.CreateTrace("parentTrace", "parentSpan");
+
+        var child = TraceContextPropagator.CreateChildSpan(parent, "childSpan");
+
+        child.SpanId.ShouldBe("childSpan");
+        child.SpanId.ShouldNotBe(parent.SpanId);
+    }
+
+    [Fact]
+    public void ShouldTrace_WithHeader_ReturnsTrue()
+    {
+        var headers = BuildNatsHeaders((TraceContextPropagator.TraceParentHeader, "trace1-span1"));
+
+        TraceContextPropagator.ShouldTrace(headers).ShouldBeTrue();
+    }
+
+    [Fact]
+    public void ShouldTrace_WithoutHeader_ReturnsFalse()
+    {
+        var headers = BuildNatsHeaders(("Content-Type", "text/plain"));
+
+        TraceContextPropagator.ShouldTrace(headers).ShouldBeFalse();
+    }
+
+    [Fact]
+    public void RoundTrip_CreateInjectExtract_Matches()
+    {
+        // Use hex-style IDs (no dashes) so the "{traceId}-{spanId}" wire format
+        // can be unambiguously split on the single separator dash.
+        var original = TraceContextPropagator.CreateTrace("0af7651916cd43dd8448eb211c80319c", "b7ad6b7169203331", destination: "trace.dest");
+
+        // Inject into empty headers
+        var injected = TraceContextPropagator.InjectTrace(original, ReadOnlySpan<byte>.Empty);
+
+        // Extract back from the injected headers
+        var extracted = TraceContextPropagator.ExtractTrace(injected);
+
+        extracted.ShouldNotBeNull();
+        extracted.TraceId.ShouldBe(original.TraceId);
+        extracted.SpanId.ShouldBe(original.SpanId);
+    }
+}