Correlate linked OTEL traces to tests

Agent-Logs-Url: https://github.com/thomhurst/TUnit/sessions/5339a2a6-32c0-4e8e-b65a-72c686b6873e

Co-authored-by: thomhurst <30480171+thomhurst@users.noreply.github.com>

TUnit.Core/TraceRegistry.cs

@@ -63,6 +63,36 @@ internal static bool IsRegistered(string traceId)
         return TraceToContextId.GetValueOrDefault(traceId);
     }
 
+    /// <summary>
+    /// Associates <paramref name="derivedTraceId"/> with the same test(s) as
+    /// <paramref name="sourceTraceId"/>. Useful for messaging/queue consumers that start
+    /// a new trace but keep a causal link to the original test trace via OTEL span links.
+    /// </summary>
+    internal static bool TryRegisterDerivedTrace(string derivedTraceId, string sourceTraceId)
+    {
+        if (string.Equals(derivedTraceId, sourceTraceId, StringComparison.OrdinalIgnoreCase))
+        {
+            return IsRegistered(sourceTraceId);
+        }
+
+        if (!TraceToTests.TryGetValue(sourceTraceId, out var testNodeUids))
+        {
+            return false;
+        }
+
+        foreach (var testNodeUid in testNodeUids.Keys)
+        {
+            Register(derivedTraceId, testNodeUid);
+        }
+
+        if (TraceToContextId.TryGetValue(sourceTraceId, out var contextId))
+        {
+            TraceToContextId[derivedTraceId] = contextId;
+        }
+
+        return true;
+    }
+
     /// <summary>
     /// Gets all trace IDs registered for the given test node UID.
     /// Used by HtmlReporter to populate additional trace IDs on test results.

TUnit.Engine/Reporters/Html/ActivityCollector.cs

@@ -186,7 +186,12 @@ internal void IngestExternalSpan(SpanData span)
     {
         if (!_knownTraceIds.ContainsKey(span.TraceId))
         {
-            return;
+            if (!TraceRegistry.IsRegistered(span.TraceId))
+            {
+                return;
+            }
+
+            _knownTraceIds.TryAdd(span.TraceId, 0);
         }
 
         // Prefer per-test cap when the span's direct parent is a known test case span.

TUnit.OpenTelemetry.Tests/OtlpReceiverIngestionTests.cs

@@ -47,4 +47,138 @@ public async Task Receiver_ParsedTrace_ReachesActivityCollector()
 
         await Assert.That(span).IsNotNull();
     }
+
+    [Test]
+    public async Task Receiver_ParsedLinkedTrace_RegistersAgainstOwningTest()
+    {
+        var collector = ActivityCollector.Current;
+        await Assert.That(collector).IsNotNull();
+
+        await using var receiver = new OtlpReceiver();
+        receiver.Start();
+
+        var linkedContext = Activity.Current!.Context;
+        var traceId = Guid.NewGuid().ToString("N");
+        var spanId = "0123456789abcdef";
+        var body = BuildLinkedTraceExportRequest(
+            traceId,
+            spanId,
+            "sut-linked-op",
+            linkedContext.TraceId.ToString(),
+            linkedContext.SpanId.ToString());
+
+        using var client = new HttpClient();
+        using var content = new ByteArrayContent(body);
+        content.Headers.ContentType = new("application/x-protobuf");
+        var response = await client.PostAsync($"http://127.0.0.1:{receiver.Port}/v1/traces", content);
+
+        await Assert.That(response.IsSuccessStatusCode).IsTrue();
+        await receiver.WhenIdle();
+
+        var span = collector!.GetAllSpans().FirstOrDefault(s =>
+            s.TraceId == traceId && s.Name == "sut-linked-op");
+
+        await Assert.That(span).IsNotNull();
+        await Assert.That(TraceRegistry.IsRegistered(traceId)).IsTrue();
+        await Assert.That(TraceRegistry.GetContextId(traceId)).IsEqualTo(TestContext.Current!.Id);
+    }
+
+    private static byte[] BuildLinkedTraceExportRequest(
+        string traceId,
+        string spanId,
+        string spanName,
+        string linkedTraceId,
+        string linkedSpanId)
+    {
+        using var exportStream = new MemoryStream();
+        var resourceSpans = BuildResourceSpans(
+            BuildScopeSpans(
+                BuildSpan(traceId, spanId, spanName, linkedTraceId, linkedSpanId)));
+        WriteField(exportStream, 1, resourceSpans);
+        return exportStream.ToArray();
+    }
+
+    private static byte[] BuildResourceSpans(byte[] scopeSpans)
+    {
+        using var stream = new MemoryStream();
+        WriteField(stream, 2, scopeSpans);
+        return stream.ToArray();
+    }
+
+    private static byte[] BuildScopeSpans(byte[] span)
+    {
+        using var stream = new MemoryStream();
+        WriteField(stream, 2, span);
+        return stream.ToArray();
+    }
+
+    private static byte[] BuildSpan(
+        string traceId,
+        string spanId,
+        string spanName,
+        string linkedTraceId,
+        string linkedSpanId)
+    {
+        using var stream = new MemoryStream();
+        WriteField(stream, 1, Convert.FromHexString(traceId));
+        WriteField(stream, 2, Convert.FromHexString(spanId));
+        WriteStringField(stream, 5, spanName);
+        WriteVarintField(stream, 6, 5); // CONSUMER
+        WriteFixed64Field(stream, 7, 1);
+        WriteFixed64Field(stream, 8, 2);
+        WriteField(stream, 13, BuildSpanLink(linkedTraceId, linkedSpanId));
+        return stream.ToArray();
+    }
+
+    private static byte[] BuildSpanLink(string traceId, string spanId)
+    {
+        using var stream = new MemoryStream();
+        WriteField(stream, 1, Convert.FromHexString(traceId));
+        WriteField(stream, 2, Convert.FromHexString(spanId));
+        return stream.ToArray();
+    }
+
+    private static void WriteStringField(MemoryStream stream, int fieldNumber, string value)
+    {
+        WriteField(stream, fieldNumber, System.Text.Encoding.UTF8.GetBytes(value));
+    }
+
+    private static void WriteVarintField(MemoryStream stream, int fieldNumber, ulong value)
+    {
+        WriteTag(stream, fieldNumber, 0);
+        WriteVarint(stream, value);
+    }
+
+    private static void WriteFixed64Field(MemoryStream stream, int fieldNumber, ulong value)
+    {
+        WriteTag(stream, fieldNumber, 1);
+        stream.Write(BitConverter.GetBytes(value));
+    }
+
+    private static void WriteField(MemoryStream stream, int fieldNumber, byte[] value)
+    {
+        WriteTag(stream, fieldNumber, 2);
+        WriteVarint(stream, (ulong)value.Length);
+        stream.Write(value);
+    }
+
+    private static void WriteTag(MemoryStream stream, int fieldNumber, int wireType)
+    {
+        WriteVarint(stream, (ulong)((fieldNumber << 3) | wireType));
+    }
+
+    private static void WriteVarint(MemoryStream stream, ulong value)
+    {
+        do
+        {
+            var current = (byte)(value & 0x7F);
+            value >>= 7;
+            if (value != 0)
+            {
+                current |= 0x80;
+            }
+
+            stream.WriteByte(current);
+        } while (value != 0);
+    }
 }

TUnit.OpenTelemetry/Receiver/OtlpReceiver.cs

@@ -230,10 +230,27 @@ private static void ProcessTraces(byte[] body)
 
         foreach (var span in spans)
         {
+            RegisterDerivedTrace(span);
             sink(ToSpanData(span));
         }
     }
 
+    private static void RegisterDerivedTrace(OtlpSpanRecord span)
+    {
+        if (TraceRegistry.IsRegistered(span.TraceId))
+        {
+            return;
+        }
+
+        foreach (var link in span.Links)
+        {
+            if (TraceRegistry.TryRegisterDerivedTrace(span.TraceId, link.TraceId))
+            {
+                return;
+            }
+        }
+    }
+
     private static SpanData ToSpanData(OtlpSpanRecord span)
     {
         ReportKeyValue[]? tags = null;