Fixes to address comments

Author: kr-igor

tracer/src/Datadog.Trace/DataStreamsMonitoring/DataStreamsWriter.cs

@@ -260,6 +260,11 @@ public async Task FlushAsync()
                 {
                     _aggregator.AddBacklog(in backlogPoint);
                 }
+
+                while (_transactionBuffer.TryDequeue(out var transactionPoint))
+                {
+                    _aggregator.AddTransaction(transactionPoint);
+                }
             }
             catch (Exception ex)
             {

tracer/src/Datadog.Trace/DataStreamsMonitoring/SpanExtensions.cs

@@ -18,7 +18,7 @@ internal static class SpanExtensions
     /// <param name="checkpointName">The checkpoint name at which the transaction is being tracked</param>
     internal static void TrackTransaction(this Span span, DataStreamsManager? manager, string transactionId, string checkpointName)
     {
-        if (manager is null)
+        if (manager is null || !manager.IsTransactionTrackingEnabled)
         {
             return;
         }

tracer/src/Datadog.Trace/DataStreamsMonitoring/TransactionTracking/DataStreamsExtractorRegistry.cs

@@ -4,13 +4,17 @@
 // </copyright>
 #nullable enable
 
+using System;
 using System.Collections.Generic;
+using Datadog.Trace.Logging;
 using Datadog.Trace.Util.Json;
 
 namespace Datadog.Trace.DataStreamsMonitoring.TransactionTracking;
 
 internal sealed class DataStreamsExtractorRegistry
 {
+    private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor<DataStreamsExtractorRegistry>();
+
     private readonly Dictionary<DataStreamsTransactionExtractor.Type, List<DataStreamsTransactionExtractor>> _extractors = new();
 
     internal DataStreamsExtractorRegistry(string extractorsJson)
@@ -20,7 +24,17 @@ internal DataStreamsExtractorRegistry(string extractorsJson)
             return;
         }
 
-        var deserialized = JsonHelper.DeserializeObject<List<DataStreamsTransactionExtractor>>(extractorsJson);
+        List<DataStreamsTransactionExtractor>? deserialized;
+        try
+        {
+            deserialized = JsonHelper.DeserializeObject<List<DataStreamsTransactionExtractor>>(extractorsJson);
+        }
+        catch (Exception ex)
+        {
+            Log.Warning(ex, "Failed to parse DD_DATA_STREAMS_TRANSACTION_EXTRACTORS value. Transaction tracking extractors will be disabled.");
+            return;
+        }
+
         if (deserialized == null)
         {
             return;

tracer/src/Datadog.Trace/DataStreamsMonitoring/TransactionTracking/DataStreamsTransactionInfo.cs

@@ -13,6 +13,8 @@ namespace Datadog.Trace.DataStreamsMonitoring.TransactionTracking;
 
 internal readonly struct DataStreamsTransactionInfo
 {
+    private const int MaxIdBytes = 255;
+
     private static readonly ConcurrentDictionary<string, int> Cache = new();
     private static int _counter;
 
@@ -22,22 +24,35 @@ internal readonly struct DataStreamsTransactionInfo
 
     internal DataStreamsTransactionInfo(string id, long timestamp, string checkpoint)
     {
-        _idBytes = Encoding.UTF8.GetBytes(id);
+        var encoded = Encoding.UTF8.GetBytes(id);
+        _idBytes = Truncate(encoded);
         _timestamp = timestamp;
-        _checkpointId = Cache.GetOrAdd(checkpoint, Interlocked.Increment(ref _counter));
+        _checkpointId = Cache.GetOrAdd(checkpoint, _ => Interlocked.Increment(ref _counter));
     }
 
     internal DataStreamsTransactionInfo(byte[] idBytes, long timestamp, string checkpoint)
     {
-        _idBytes = idBytes;
+        _idBytes = Truncate(idBytes);
         _timestamp = timestamp;
-        _checkpointId = Cache.GetOrAdd(checkpoint, Interlocked.Increment(ref _counter));
+        _checkpointId = Cache.GetOrAdd(checkpoint, _ => Interlocked.Increment(ref _counter));
     }
 
     internal long TimestampNs { get => _timestamp; }
 
     internal string TransactionId { get => Encoding.UTF8.GetString(_idBytes); }
 
+    private static byte[] Truncate(byte[] source)
+    {
+        if (source.Length <= MaxIdBytes)
+        {
+            return source;
+        }
+
+        var truncated = new byte[MaxIdBytes];
+        Array.Copy(source, truncated, MaxIdBytes);
+        return truncated;
+    }
+
     internal static byte[] GetCacheBytes()
     {
         var result = new byte[512];
@@ -92,7 +107,7 @@ internal void WriteTo(byte[] buffer, int offset)
         buffer[offset + 7] = (byte)(_timestamp >> 8);
         buffer[offset + 8] = (byte)_timestamp;
 
-        // id size, up to 256 bytes
+        // id size, up to 255 bytes
         buffer[offset + 9] = (byte)_idBytes.Length;
 
         // copy the ID

tracer/test/Datadog.Trace.Tests/DataStreamsMonitoring/DataStreamsExtractorRegistryTest.cs

@@ -18,4 +18,29 @@ public void DeserializeCorrectly()
         var registry = new DataStreamsExtractorRegistry("[{\"name\": \"transaction-origin\", \"type\": \"HTTP_OUT_HEADERS\", \"value\": \"transaction-id\"}]");
         registry.AsJson().Should().Be("{\"HttpOutHeaders\":[{\"name\":\"transaction-origin\",\"type\":\"HTTP_OUT_HEADERS\",\"value\":\"transaction-id\",\"ExtractorType\":1}]}");
     }
+
+    [Fact]
+    public void GetExtractorsByType_ReturnsAllExtractors_ForSameType()
+    {
+        var registry = new DataStreamsExtractorRegistry(
+            "[" +
+            "{\"name\": \"n1\", \"type\": \"HTTP_OUT_HEADERS\", \"value\": \"v1\"}," +
+            "{\"name\": \"n2\", \"type\": \"HTTP_OUT_HEADERS\", \"value\": \"v2\"}" +
+            "]");
+
+        var extractors = registry.GetExtractorsByType(DataStreamsTransactionExtractor.Type.HttpOutHeaders);
+
+        extractors.Should().HaveCount(2);
+        extractors.Should().Contain(e => e.Name == "n1" && e.Value == "v1");
+        extractors.Should().Contain(e => e.Name == "n2" && e.Value == "v2");
+    }
+
+    [Fact]
+    public void GetExtractorsByType_DoesNotReturnExtractors_ForOtherType()
+    {
+        var registry = new DataStreamsExtractorRegistry(
+            "[{\"name\": \"n1\", \"type\": \"HTTP_OUT_HEADERS\", \"value\": \"v1\"}]");
+
+        registry.GetExtractorsByType(DataStreamsTransactionExtractor.Type.HttpInHeaders).Should().BeNull();
+    }
 }

tracer/test/Datadog.Trace.Tests/DataStreamsMonitoring/DataStreamsManagerTests.cs

@@ -274,11 +274,14 @@ public void WhenEnabled_TrackTransaction_AddsTransactionAndTagsSpan()
     [Fact]
     public void WhenDisabled_TrackTransaction_DoesNothing()
     {
-        var dsm = GetDataStreamManager(false, out _);
+        var dsm = GetDataStreamManager(false, out var writer);
         var span = new Span(new SpanContext(traceId: 123, spanId: 456), DateTimeOffset.UtcNow);
 
-        var act = () => span.TrackTransaction(dsm, "tx-abc", "some-checkpoint");
-        act.Should().NotThrow();
+        span.TrackTransaction(dsm, "tx-abc", "some-checkpoint");
+
+        span.Tags.GetTag("dsm.transaction.id").Should().BeNull();
+        // writer is null when DSM is disabled, so nothing could have been enqueued
+        writer.Should().BeNull();
     }
 
     [Fact]

tracer/test/Datadog.Trace.Tests/DataStreamsMonitoring/DataStreamsTransactionInfoTest.cs

@@ -63,4 +63,57 @@ public void GetByteCount_MatchesGetBytesLength()
         var transaction = new DataStreamsTransactionInfo("hello-world", 12345L, "my-cp");
         transaction.GetByteCount().Should().Be(transaction.GetBytes().Length);
     }
+
+    [Fact]
+    public void CheckpointId_DoesNotIncrement_WhenCheckpointAlreadyCached()
+    {
+        DataStreamsTransactionInfo.ClearCache();
+
+        // First use of "cp-a" gets id 1, first use of "cp-b" gets id 2.
+        var a1 = new DataStreamsTransactionInfo("tx1", 1, "cp-a");
+        var b1 = new DataStreamsTransactionInfo("tx2", 1, "cp-b");
+
+        // Repeated uses of existing checkpoint names must not increment the counter.
+        for (var i = 0; i < 10; i++)
+        {
+            _ = new DataStreamsTransactionInfo("tx3", 1, "cp-a");
+            _ = new DataStreamsTransactionInfo("tx4", 1, "cp-b");
+        }
+
+        // A brand-new checkpoint must still get the next sequential id (3), not some higher value.
+        var c1 = new DataStreamsTransactionInfo("tx5", 1, "cp-c");
+
+        a1.GetBytes()[0].Should().Be(1);
+        b1.GetBytes()[0].Should().Be(2);
+        c1.GetBytes()[0].Should().Be(3, "counter must only advance for new checkpoint names");
+    }
+
+    [Theory]
+    [InlineData(255)]  // exact limit — no truncation
+    [InlineData(256)]  // one byte over
+    [InlineData(512)]  // well over
+    public void LongTransactionId_IsTruncatedTo255Bytes(int idByteLength)
+    {
+        var id = new string('a', idByteLength);
+        var transaction = new DataStreamsTransactionInfo(id, 1L, "cp");
+        var bytes = transaction.GetBytes();
+
+        // byte at offset 9 is the encoded length; must never exceed 255
+        bytes[9].Should().Be(255);
+        // total payload length must match: 10 header bytes + encoded id length
+        bytes.Length.Should().Be(10 + 255);
+        // GetByteCount must agree
+        transaction.GetByteCount().Should().Be(bytes.Length);
+    }
+
+    [Fact]
+    public void LongTransactionId_FromByteArray_IsTruncatedTo255Bytes()
+    {
+        var idBytes = new byte[300];
+        var transaction = new DataStreamsTransactionInfo(idBytes, 1L, "cp");
+        var bytes = transaction.GetBytes();
+
+        bytes[9].Should().Be(255);
+        bytes.Length.Should().Be(10 + 255);
+    }
 }

tracer/test/Datadog.Trace.Tests/DataStreamsMonitoring/DataStreamsWriterTests.cs

@@ -230,6 +230,41 @@ public async Task WhenNotSupported_DoesNotReportAnyBucketsOnClose(bool? isSuppor
         api.Sent.Should().BeEmpty();
     }
 
+    [Fact]
+    public async Task WhenSupported_WritesTransaction_OnClose()
+    {
+        // mirrors WhenSupported_WritesAStatsPoint_OnClose for transactions
+        var bucketDuration = 100_000_000;
+        var api = new StubApi();
+        var writer = CreateWriter(api, out var discovery, bucketDuration);
+        TriggerSupportUpdate(discovery, isSupported: true);
+
+        writer.AddTransaction(new DataStreamsTransactionInfo("tx-id", 1L, "cp"));
+
+        await writer.DisposeAsync();
+
+        api.Sent.Should().ContainSingle();
+    }
+
+    [Fact]
+    public async Task FlushAsync_DrainsPendingTransactions()
+    {
+        // int.MaxValue bucket: timer will never fire and ShouldFlushTransactions
+        // won't trigger for a small payload.  FlushAsync must drain _transactionBuffer
+        // directly — it is called immediately so ProcessQueueLoop has not yet woken
+        // from its initial 10 ms sleep.
+        var bucketDuration = int.MaxValue;
+        var api = new StubApi();
+        var writer = CreateWriter(api, out var discovery, bucketDuration);
+        TriggerSupportUpdate(discovery, isSupported: true);
+
+        writer.AddTransaction(new DataStreamsTransactionInfo("tx-id", 1L, "cp"));
+        await writer.FlushAsync();
+
+        api.Sent.Should().NotBeEmpty("FlushAsync must drain the transaction buffer");
+        await writer.DisposeAsync();
+    }
+
     [Fact]
     public async Task CanCreateWriterWithDefaultBucket()
     {