Refactoring to support DSM default state (#7082)

kr-igor · web-flow · commit ebf355011f8a · 2025-06-13T11:30:45.000-05:00
## Summary of changes This PR adds undefined state to `DATA_STREAMS_ENABLED` flag. The end-user behavior remains unchanged, as `undefined` defaults to `false` for now. This will change in followup PRs when DSM is enabled by default. ## Reason for change See details in [this](https://docs.google.com/document/d/111QQyLxz5q5G3wccj5Q2Zu7e3NemlXDSt6HcchA2LKU/edit?tab=t.0) doc. ## Implementation details DSM state is now represented by enum. The heaviest parts of DSM logic will be now disabled if DSM is in the default (undefined) state. ## Test coverage All existing tests were updated to support the change. ## Other details [Jira](https://datadoghq.atlassian.net/browse/DSMON-777)
diff --git a/tracer/src/Datadog.Trace/Activity/Handlers/AzureServiceBusActivityHandler.cs b/tracer/src/Datadog.Trace/Activity/Handlers/AzureServiceBusActivityHandler.cs
@@ -64,8 +64,8 @@ public void ActivityStopped<T>(string sourceName, T activity)
 
                     string namespaceString = span.Tags.GetTag("messaging.destination.name");
                     long? payloadSize = null;
-                    if (AzureServiceBusCommon.TryGetMessage(applicationProperties, out var message)
-                        && message.TryDuckCast<IServiceBusMessage>(out var serviceBusMessage))
+                    if (!dataStreamsManager.IsInDefaultState && (AzureServiceBusCommon.TryGetMessage(applicationProperties, out var message)
+                        && message.TryDuckCast<IServiceBusMessage>(out var serviceBusMessage)))
                     {
                         payloadSize = AzureServiceBusCommon.GetMessageSize(serviceBusMessage);
                     }
diff --git a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/ServiceBus/ProcessMessageIntegration.cs b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Azure/ServiceBus/ProcessMessageIntegration.cs
@@ -85,12 +85,12 @@ internal static CallTargetState OnMethodBegin<TTarget, TMessage>(TTarget instanc
                 var edgeTags = string.IsNullOrEmpty(namespaceString)
                                     ? new[] { "direction:in", "type:servicebus" }
                                     : new[] { "direction:in", $"topic:{namespaceString}", "type:servicebus" };
-
+                var msgSize = dataStreamsManager.IsInDefaultState ? 0 : AzureServiceBusCommon.GetMessageSize(message);
                 span.SetDataStreamsCheckpoint(
                     dataStreamsManager,
                     CheckpointKind.Consume,
                     edgeTags,
-                    AzureServiceBusCommon.GetMessageSize(message),
+                    msgSize,
                     (long)messageQueueTimeMs,
                     pathwayContext);
             }
diff --git a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Kafka/KafkaHelper.cs b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Kafka/KafkaHelper.cs
@@ -252,7 +252,7 @@ private static long GetMessageSize<T>(T message)
                         dataStreamsManager,
                         CheckpointKind.Consume,
                         edgeTags,
-                        message is null ? 0 : GetMessageSize(message),
+                        message is null || dataStreamsManager.IsInDefaultState ? 0 : GetMessageSize(message),
                         tags.MessageQueueTimeMs == null ? 0 : (long)tags.MessageQueueTimeMs,
                         pathwayContext);
 
@@ -341,8 +341,9 @@ internal static void TryInjectHeaders<TTopicPartitionMarker, TMessage>(
                     var edgeTags = string.IsNullOrEmpty(topic)
                         ? defaultProduceEdgeTags
                         : new[] { "direction:out", $"topic:{topic}", "type:kafka" };
+                    var msgSize = dataStreamsManager.IsInDefaultState ? 0 : GetMessageSize(message);
                     // produce is always the start of the edge, so defaultEdgeStartMs is always 0
-                    span.SetDataStreamsCheckpoint(dataStreamsManager, CheckpointKind.Produce, edgeTags, GetMessageSize(message), 0);
+                    span.SetDataStreamsCheckpoint(dataStreamsManager, CheckpointKind.Produce, edgeTags, msgSize, 0);
                     dataStreamsManager.InjectPathwayContext(span.Context.PathwayContext, adapter);
                 }
             }
diff --git a/tracer/src/Datadog.Trace/Configuration/TracerSettings.cs b/tracer/src/Datadog.Trace/Configuration/TracerSettings.cs
@@ -52,6 +52,7 @@ public record TracerSettings
         private readonly bool _traceEnabled;
         private readonly bool _apmTracingEnabled;
         private readonly bool _isDataStreamsMonitoringEnabled;
+        private readonly bool _isDataStreamsMonitoringInDefaultState;
         private readonly ReadOnlyDictionary<string, string> _headerTags;
         private readonly ReadOnlyDictionary<string, string> _serviceNameMappings;
         private readonly ReadOnlyDictionary<string, string> _globalTags;
@@ -651,6 +652,9 @@ _ when x.ToBoolean() is { } boolean => boolean,
             _isDataStreamsMonitoringEnabled = config
                                             .WithKeys(ConfigurationKeys.DataStreamsMonitoring.Enabled)
                                             .AsBool(false);
+            _isDataStreamsMonitoringInDefaultState = config
+                                                    .WithKeys(ConfigurationKeys.DataStreamsMonitoring.Enabled)
+                                                    .AsBool() == null;
 
             IsDataStreamsLegacyHeadersEnabled = config
                                                .WithKeys(ConfigurationKeys.DataStreamsMonitoring.LegacyHeadersEnabled)
@@ -1184,6 +1188,11 @@ public bool DiagnosticSourceEnabled
         /// </summary>
         internal bool IsDataStreamsMonitoringEnabled => DynamicSettings.DataStreamsMonitoringEnabled ?? _isDataStreamsMonitoringEnabled;
 
+        /// <summary>
+        /// Gets a value indicating whether data streams configuration is present or not (set to true or false).
+        /// </summary>
+        internal bool IsDataStreamsMonitoringInDefaultState => DynamicSettings.DataStreamsMonitoringEnabled == null && _isDataStreamsMonitoringInDefaultState;
+
         /// <summary>
         /// Gets a value indicating whether to inject legacy binary headers for Data Streams.
         /// </summary>
diff --git a/tracer/src/Datadog.Trace/DataStreamsMonitoring/Aggregation/DataStreamsMessagePackFormatter.cs b/tracer/src/Datadog.Trace/DataStreamsMonitoring/Aggregation/DataStreamsMessagePackFormatter.cs
@@ -20,6 +20,7 @@ internal class DataStreamsMessagePackFormatter
         private readonly byte[] _environmentValueBytes;
         private readonly byte[] _serviceBytes = StringEncoding.UTF8.GetBytes("Service");
         private readonly long _productMask;
+        private readonly bool _isInDefaultState;
 
         private readonly byte[] _serviceValueBytes;
 
@@ -48,6 +49,7 @@ internal class DataStreamsMessagePackFormatter
         private readonly byte[] _backlogTagsBytes = StringEncoding.UTF8.GetBytes("Tags");
         private readonly byte[] _backlogValueBytes = StringEncoding.UTF8.GetBytes("Value");
         private readonly byte[] _productMaskBytes = StringEncoding.UTF8.GetBytes("ProductMask");
+        private readonly byte[] _isInDefaultStateBytes = StringEncoding.UTF8.GetBytes("IsInDefaultState");
 
         public DataStreamsMessagePackFormatter(TracerSettings tracerSettings, string defaultServiceName)
         {
@@ -59,8 +61,10 @@ public DataStreamsMessagePackFormatter(TracerSettings tracerSettings, string def
                                          : StringEncoding.UTF8.GetBytes(env);
             _serviceValueBytes = StringEncoding.UTF8.GetBytes(defaultServiceName);
             _productMask = GetProductsMask(tracerSettings);
+            _isInDefaultState = tracerSettings.IsDataStreamsMonitoringInDefaultState;
         }
 
+        // should be the same across all languages
         [Flags]
         private enum Products : long
         {
@@ -95,7 +99,7 @@ public int Serialize(Stream stream, long bucketDurationNs, List<SerializableStat
             // https://github.com/DataDog/dd-trace-java/blob/a4b7a7b177709e6bdfd9261904cb9a777e4febbe/dd-trace-core/src/main/java/datadog/trace/core/datastreams/MsgPackDatastreamsPayloadWriter.java#L35
             // -1 because we don't have a primary tag
             // -1 because service name override is not supported
-            bytesWritten += MessagePackBinary.WriteMapHeader(stream, 6);
+            bytesWritten += MessagePackBinary.WriteMapHeader(stream, 7);
 
             bytesWritten += MessagePackBinary.WriteStringBytes(stream, _environmentBytes);
             bytesWritten += MessagePackBinary.WriteStringBytes(stream, _environmentValueBytes);
@@ -189,6 +193,9 @@ public int Serialize(Stream stream, long bucketDurationNs, List<SerializableStat
             bytesWritten += MessagePackBinary.WriteStringBytes(stream, _productMaskBytes);
             bytesWritten += MessagePackBinary.WriteInt64(stream, _productMask);
 
+            bytesWritten += MessagePackBinary.WriteStringBytes(stream, _isInDefaultStateBytes);
+            bytesWritten += MessagePackBinary.WriteBoolean(stream, _isInDefaultState);
+
             return bytesWritten;
         }
 
diff --git a/tracer/src/Datadog.Trace/DataStreamsMonitoring/DataStreamsManager.cs b/tracer/src/Datadog.Trace/DataStreamsMonitoring/DataStreamsManager.cs
@@ -12,7 +12,6 @@
 using Datadog.Trace.Configuration;
 using Datadog.Trace.DataStreamsMonitoring.Aggregation;
 using Datadog.Trace.DataStreamsMonitoring.Hashes;
-using Datadog.Trace.ExtensionMethods;
 using Datadog.Trace.Headers;
 using Datadog.Trace.Logging;
 using Datadog.Trace.Vendors.Serilog.Events;
@@ -29,21 +28,26 @@ internal class DataStreamsManager
     private readonly ConcurrentDictionary<string, RateLimiter> _schemaRateLimiters = new();
     private readonly NodeHashBase _nodeHashBase;
     private bool _isEnabled;
+    private bool _isInDefaultState;
     private IDataStreamsWriter? _writer;
 
     public DataStreamsManager(
         string? env,
         string defaultServiceName,
-        IDataStreamsWriter? writer)
+        IDataStreamsWriter? writer,
+        bool isInDefaultState)
     {
         // We don't yet support primary tag in .NET yet
         _nodeHashBase = HashHelper.CalculateNodeHashBase(defaultServiceName, env, primaryTag: null);
         _isEnabled = writer is not null;
         _writer = writer;
+        _isInDefaultState = isInDefaultState;
     }
 
     public bool IsEnabled => Volatile.Read(ref _isEnabled);
 
+    public bool IsInDefaultState => Volatile.Read(ref _isInDefaultState);
+
     public static DataStreamsManager Create(
         TracerSettings settings,
         IDiscoveryService discoveryService,
@@ -53,7 +57,7 @@ public static DataStreamsManager Create(
                          ? DataStreamsWriter.Create(settings, discoveryService, defaultServiceName)
                          : null;
 
-        return new DataStreamsManager(settings.Environment, defaultServiceName, writer);
+        return new DataStreamsManager(settings.Environment, defaultServiceName, writer, settings.IsDataStreamsMonitoringInDefaultState);
     }
 
     public async Task DisposeAsync()
diff --git a/tracer/src/Datadog.Trace/DataStreamsMonitoring/SpanExtensions.cs b/tracer/src/Datadog.Trace/DataStreamsMonitoring/SpanExtensions.cs
@@ -3,6 +3,8 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
 // </copyright>
 #nullable enable
+using System.Globalization;
+
 namespace Datadog.Trace.DataStreamsMonitoring;
 
 internal static class SpanExtensions
@@ -25,10 +27,13 @@ internal static void SetDataStreamsCheckpoint(this Span span, DataStreamsManager
         }
 
         span.Context.SetCheckpoint(manager, checkpointKind, edgeTags, payloadSizeBytes, timeInQueueMs, parent);
-        var hash = span.Context.PathwayContext?.Hash.Value ?? 0;
-        if (hash != 0)
+        if (!manager.IsInDefaultState)
         {
-            span.SetTag("pathway.hash", hash.ToString());
+            var hash = span.Context.PathwayContext?.Hash.Value ?? 0;
+            if (hash != 0)
+            {
+                span.SetTag("pathway.hash", hash.ToString(CultureInfo.InvariantCulture));
+            }
         }
     }
 }
diff --git a/tracer/test/Datadog.Trace.Tests/DataStreamsMonitoring/DataStreamsManagerTests.cs b/tracer/test/Datadog.Trace.Tests/DataStreamsMonitoring/DataStreamsManagerTests.cs
@@ -303,7 +303,8 @@ private static DataStreamsManager GetDataStreamManager(bool enabled, out DataStr
         return new DataStreamsManager(
             env: "foo",
             defaultServiceName: "bar",
-            writer);
+            writer,
+            isInDefaultState: false);
     }
 
     internal class DataStreamsWriterMock : IDataStreamsWriter
diff --git a/tracer/test/Datadog.Trace.Tests/DataStreamsMonitoring/SpanContextDataStreamsManagerTests.cs b/tracer/test/Datadog.Trace.Tests/DataStreamsMonitoring/SpanContextDataStreamsManagerTests.cs
@@ -73,7 +73,8 @@ private static DataStreamsManager GetEnabledDataStreamManager()
         var dsm = new DataStreamsManager(
             env: "env",
             defaultServiceName: "service",
-            new Mock<IDataStreamsWriter>().Object);
+            new Mock<IDataStreamsWriter>().Object,
+            isInDefaultState: false);
         return dsm;
     }
 }
diff --git a/tracer/test/Datadog.Trace.Tests/TracerManagerFactoryTests.cs b/tracer/test/Datadog.Trace.Tests/TracerManagerFactoryTests.cs
@@ -83,7 +83,7 @@ private static TracerManager CreateTracerManager(TracerSettings settings)
             BuildLogSubmissionManager(),
             Mock.Of<ITelemetryController>(),
             Mock.Of<IDiscoveryService>(),
-            new DataStreamsManager("env", "service", Mock.Of<IDataStreamsWriter>()),
+            new DataStreamsManager("env", "service", Mock.Of<IDataStreamsWriter>(), isInDefaultState: false),
             remoteConfigurationManager: null,
             dynamicConfigurationManager: null,
             tracerFlareManager: null,

Original file line number	Diff line number	Diff line change
`@@ -64,8 +64,8 @@ public void ActivityStopped<T>(string sourceName, T activity)`
`64`	`64`
`65`	`65`	`string namespaceString = span.Tags.GetTag("messaging.destination.name");`
`66`	`66`	`long? payloadSize = null;`
`67`		`- if (AzureServiceBusCommon.TryGetMessage(applicationProperties, out var message)`
`68`		`- && message.TryDuckCast<IServiceBusMessage>(out var serviceBusMessage))`
	`67`	`+ if (!dataStreamsManager.IsInDefaultState && (AzureServiceBusCommon.TryGetMessage(applicationProperties, out var message)`
	`68`	`+ && message.TryDuckCast<IServiceBusMessage>(out var serviceBusMessage)))`
`69`	`69`	`{`
`70`	`70`	`payloadSize = AzureServiceBusCommon.GetMessageSize(serviceBusMessage);`
`71`	`71`	`}`
Original file line number	Diff line number	Diff line change
`@@ -3,6 +3,8 @@`
`3`	`3`	`// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.`
`4`	`4`	`// </copyright>`
`5`	`5`	`#nullable enable`
	`6`	`+using System.Globalization;`
	`7`	`+`
`6`	`8`	`namespace Datadog.Trace.DataStreamsMonitoring;`
`7`	`9`
`8`	`10`	`internal static class SpanExtensions`
`@@ -25,10 +27,13 @@ internal static void SetDataStreamsCheckpoint(this Span span, DataStreamsManager`
`25`	`27`	`}`
`26`	`28`
`27`	`29`	`span.Context.SetCheckpoint(manager, checkpointKind, edgeTags, payloadSizeBytes, timeInQueueMs, parent);`
`28`		`- var hash = span.Context.PathwayContext?.Hash.Value ?? 0;`
`29`		`- if (hash != 0)`
	`30`	`+ if (!manager.IsInDefaultState)`
`30`	`31`	`{`
`31`		`- span.SetTag("pathway.hash", hash.ToString());`
	`32`	`+ var hash = span.Context.PathwayContext?.Hash.Value ?? 0;`
	`33`	`+ if (hash != 0)`
	`34`	`+ {`
	`35`	`+ span.SetTag("pathway.hash", hash.ToString(CultureInfo.InvariantCulture));`
	`36`	`+ }`
`32`	`37`	`}`
`33`	`38`	`}`
`34`	`39`	`}`
Original file line number	Diff line number	Diff line change
`@@ -303,7 +303,8 @@ private static DataStreamsManager GetDataStreamManager(bool enabled, out DataStr`
`303`	`303`	`return new DataStreamsManager(`
`304`	`304`	`env: "foo",`
`305`	`305`	`defaultServiceName: "bar",`
`306`		`- writer);`
	`306`	`+ writer,`
	`307`	`+ isInDefaultState: false);`
`307`	`308`	`}`
`308`	`309`
`309`	`310`	`internal class DataStreamsWriterMock : IDataStreamsWriter`
Original file line number	Diff line number	Diff line change
`@@ -73,7 +73,8 @@ private static DataStreamsManager GetEnabledDataStreamManager()`
`73`	`73`	`var dsm = new DataStreamsManager(`
`74`	`74`	`env: "env",`
`75`	`75`	`defaultServiceName: "service",`
`76`		`- new Mock<IDataStreamsWriter>().Object);`
	`76`	`+ new Mock<IDataStreamsWriter>().Object,`
	`77`	`+ isInDefaultState: false);`
`77`	`78`	`return dsm;`
`78`	`79`	`}`
`79`	`80`	`}`