DSM Full propagation mode for SQL Server

tracer/src/Datadog.Trace.Trimming/build/Datadog.Trace.Trimming.xml

@@ -202,9 +202,13 @@
       <type fullname="System.Data.CommandType" />
       <type fullname="System.Data.Common.DbCommand" />
       <type fullname="System.Data.Common.DbConnectionStringBuilder" />
+      <type fullname="System.Data.DbType" />
+      <type fullname="System.Data.IDataParameter" />
+      <type fullname="System.Data.IDataParameterCollection" />
       <type fullname="System.Data.IDataRecord" />
       <type fullname="System.Data.IDbCommand" />
       <type fullname="System.Data.IDbConnection" />
+      <type fullname="System.Data.IDbDataParameter" />
    </assembly>
    <assembly fullname="System.Data.SqlClient" />
    <assembly fullname="System.Data.SQLite" />

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AdoNet/DbScopeFactory.cs

@@ -8,14 +8,11 @@
 using System;
 using System.Collections.Generic;
 using System.Data;
-using System.Data.SqlTypes;
 using System.Diagnostics.CodeAnalysis;
 using System.Threading;
 using Datadog.Trace.AppSec;
-using Datadog.Trace.AppSec.Rasp;
 using Datadog.Trace.Configuration;
 using Datadog.Trace.DatabaseMonitoring;
-using Datadog.Trace.Iast;
 using Datadog.Trace.Logging;
 using Datadog.Trace.Tagging;
 using Datadog.Trace.Util;
@@ -45,7 +42,7 @@ internal static class DbScopeFactory
 
                 if (parent is { Type: SpanTypes.Sql } &&
                     HasDbType(parent, dbType) &&
-                    (parent.ResourceName == commandText || commandText.StartsWith(DatabaseMonitoringPropagator.DbmPrefix)))
+                    (parent.ResourceName == commandText || commandText.StartsWith(DatabaseMonitoringPropagator.DbmPrefix) || commandText == DatabaseMonitoringPropagator.SetContextCommand))
                 {
                     // we are already instrumenting this,
                     // don't instrument nested methods that belong to the same stacktrace
@@ -103,11 +100,12 @@ internal static class DbScopeFactory
                     }
                     else
                     {
-                        var propagatedCommand = DatabaseMonitoringPropagator.PropagateSpanData(tracer.Settings.DbmPropagationMode, tracer.DefaultServiceName, tagsFromConnectionString.DbName, tagsFromConnectionString.OutHost, scope.Span, integrationId, out var traceParentInjected);
+                        var traceParentInjectedInContext = DatabaseMonitoringPropagator.PropagateDataViaContext(tracer.Settings.DbmPropagationMode, integrationId, command.Connection, scope.Span);
+                        var propagatedCommand = DatabaseMonitoringPropagator.PropagateDataViaComment(tracer.Settings.DbmPropagationMode, tracer.DefaultServiceName, tagsFromConnectionString.DbName, tagsFromConnectionString.OutHost, scope.Span, integrationId, out var traceParentInjectedInComment);
                         if (!string.IsNullOrEmpty(propagatedCommand))
                         {
                             command.CommandText = $"{propagatedCommand} {commandText}";
-                            if (traceParentInjected)
+                            if (traceParentInjectedInComment || traceParentInjectedInContext)
                             {
                                 tags.DbmTraceInjected = "true";
                             }
@@ -120,6 +118,10 @@ internal static class DbScopeFactory
                 Log.Error(ex, "Error propagating span data for DBM");
             }
 
+            // we have to start the span before doing the DBM propagation work (to have the span ID)
+            // but ultimately, we don't want to measure the time spent instrumenting.
+            scope.Span.ResetStartTime();
+
             return scope;
 
             static bool HasDbType(Span span, string dbType)

tracer/src/Datadog.Trace/DatabaseMonitoring/DatabaseMonitoringPropagator.cs

@@ -4,10 +4,14 @@
 // </copyright>
 
 using System;
+using System.Data;
 using Datadog.Trace.Configuration;
+using Datadog.Trace.Logging;
 using Datadog.Trace.Propagators;
 using Datadog.Trace.Tagging;
 using Datadog.Trace.Util;
+using Datadog.Trace.VendoredMicrosoftCode.System.Buffers.Binary;
+using Datadog.Trace.Vendors.Serilog.Events;
 
 #nullable enable
 
@@ -22,8 +26,12 @@ internal static class DatabaseMonitoringPropagator
         private const string SqlCommentVersion = "ddpv";
         private const string SqlCommentEnv = "dde";
         internal const string DbmPrefix = $"/*{SqlCommentSpanService}='";
+        private const string ContextInfoParameterName = "@dd_trace_context";
+        internal const string SetContextCommand = $"set context_info {ContextInfoParameterName}";
 
-        internal static string PropagateSpanData(DbmPropagationLevel propagationStyle, string configuredServiceName, string? dbName, string? outhost, Span span, IntegrationId integrationId, out bool traceParentInjected)
+        private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor(typeof(DatabaseMonitoringPropagator));
+
+        internal static string PropagateDataViaComment(DbmPropagationLevel propagationStyle, string configuredServiceName, string? dbName, string? outhost, Span span, IntegrationId integrationId, out bool traceParentInjected)
         {
             traceParentInjected = false;
 
@@ -73,5 +81,68 @@ internal static string PropagateSpanData(DbmPropagationLevel propagationStyle, s
 
             return string.Empty;
         }
+
+        /// <summary>
+        /// Uses a sql instruction to set a context for the current connection, bearing the span ID and trace ID.
+        /// This is meant to circumvent cache invalidation issues that occur when those values are injected in comment.
+        /// Currently only working for MSSQL (uses an instruction that is specific to it)
+        /// </summary>
+        /// <returns>True if the traceparent information was set</returns>
+        internal static bool PropagateDataViaContext(DbmPropagationLevel propagationLevel, IntegrationId integrationId, IDbConnection? connection, Span span)
+        {
+            if (propagationLevel != DbmPropagationLevel.Full || integrationId != IntegrationId.SqlClient || connection == null)
+            {
+                return false;
+            }
+
+            var stopwatch = System.Diagnostics.Stopwatch.StartNew();
+
+            const byte version = 0; // version can have a maximum value of 15 in the current format
+            var sampled = SamplingPriorityValues.IsKeep(span.Context.TraceContext.GetOrMakeSamplingDecision());
+            var contextValue = BuildContextValue(version, sampled, span.SpanId, span.TraceId128);
+
+            using (var injectionCommand = connection.CreateCommand())
+            {
+                injectionCommand.CommandText = SetContextCommand;
+
+                var parameter = injectionCommand.CreateParameter();
+                parameter.ParameterName = ContextInfoParameterName;
+                parameter.Value = contextValue;
+                parameter.DbType = DbType.Binary;
+                injectionCommand.Parameters.Add(parameter);
+
+                injectionCommand.ExecuteNonQuery();
+
+                if (Log.IsEnabled(LogEventLevel.Debug))
+                {
+                    // avoid building the string representation in the general case where debug is disabled
+                    Log.Debug("Span data for DBM propagated for {Integration} via context_info with value {ContextValue} (propagation level: {PropagationLevel}", integrationId, HexConverter.ToString(contextValue), propagationLevel);
+                }
+            }
+
+            // Since sending the query to the DB can be a bit long, we register the time it took for transparency.
+            // Not using _dd because we want the customers to be able to see that tag.
+            span.SetMetric("dd.instrumentation.time_ms", stopwatch.Elapsed.TotalMilliseconds);
+
+            return true;
+        }
+
+        /// <summary>
+        /// Writes the given info in a byte array with the following format:
+        /// 4 bits: protocol version, 3 bits: reserved, 1 bit: sampling decision, 64 bits: spanID, 128 bits: traceID
+        /// </summary>
+        private static byte[] BuildContextValue(byte version, bool isSampled, ulong spanId, TraceId traceId)
+        {
+            var sampled = isSampled ? 1 : 0;
+            var versionAndSampling = (byte)(((version << 4) & 0b1111_0000) | (sampled & 0b0000_0001));
+            var contextBytes = new byte[1 + sizeof(ulong) + TraceId.Size];
+
+            var span = new VendoredMicrosoftCode.System.Span<byte>(contextBytes) { [0] = versionAndSampling };
+            BinaryPrimitives.WriteUInt64BigEndian(span.Slice(1), spanId);
+            BinaryPrimitives.WriteUInt64BigEndian(span.Slice(1 + sizeof(ulong)), traceId.Upper);
+            BinaryPrimitives.WriteUInt64BigEndian(span.Slice(1 + sizeof(ulong) + sizeof(ulong)), traceId.Lower);
+
+            return contextBytes;
+        }
     }
 }

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/AdoNet/MicrosoftDataSqlClientTests.cs

@@ -26,15 +26,17 @@ public MicrosoftDataSqlClientTests(ITestOutputHelper output)
         public static IEnumerable<object[]> GetEnabledConfig()
             => from packageVersionArray in PackageVersions.MicrosoftDataSqlClient
                from metadataSchemaVersion in new[] { "v0", "v1" }
-               select new[] { packageVersionArray[0], metadataSchemaVersion };
+               from dbmEnabled in new[] { true, false }
+               from propagation in new[] { "disabled", "service", "full" }
+               select new[] { packageVersionArray[0], metadataSchemaVersion, dbmEnabled, propagation };
 
         public override Result ValidateIntegrationSpan(MockSpan span, string metadataSchemaVersion) => span.IsSqlClient(metadataSchemaVersion);
 
         [SkippableTheory]
         [MemberData(nameof(GetEnabledConfig))]
         [Trait("Category", "EndToEnd")]
         [Trait("RunOnWindows", "True")]
-        public async Task SubmitsTraces(string packageVersion, string metadataSchemaVersion)
+        public async Task SubmitsTraces(string packageVersion, string metadataSchemaVersion, bool dbmEnabled, string propagation)
         {
             // ALWAYS: 133 spans
             // - SqlCommand: 21 spans (3 groups * 7 spans)
@@ -68,15 +70,18 @@ public async Task SubmitsTraces(string packageVersion, string metadataSchemaVers
             const string dbType = "sql-server";
             const string expectedOperationName = dbType + ".query";
 
+            SetEnvironmentVariable(ConfigurationKeys.DataStreamsMonitoring.Enabled, dbmEnabled ? "1" : "0");
+            SetEnvironmentVariable("DD_DBM_PROPAGATION_MODE", propagation);
             SetEnvironmentVariable("DD_TRACE_SPAN_ATTRIBUTE_SCHEMA", metadataSchemaVersion);
             var isExternalSpan = metadataSchemaVersion == "v0";
             var clientSpanServiceName = isExternalSpan ? $"{EnvironmentHelper.FullSampleName}-{dbType}" : EnvironmentHelper.FullSampleName;
 
             using var telemetry = this.ConfigureTelemetry();
             using var agent = EnvironmentHelper.GetMockAgent();
             using var process = await RunSampleAndWaitForExit(agent, packageVersion: packageVersion);
+
             var spans = agent.WaitForSpans(expectedSpanCount, operationName: expectedOperationName);
-            int actualSpanCount = spans.Count(s => s.ParentId.HasValue); // Remove unexpected DB spans from the calculation
+            var actualSpanCount = spans.Count(s => s.ParentId.HasValue); // Remove unexpected DB spans from the calculation
 
             Assert.Equal(expectedSpanCount, actualSpanCount);
             ValidateIntegrationSpans(spans, metadataSchemaVersion, expectedServiceName: clientSpanServiceName, isExternalSpan);

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/AdoNet/SystemDataSqlClientTests.cs

@@ -89,6 +89,7 @@ public async Task SubmitsTraces(string packageVersion, string metadataSchemaVers
             settings.AddSimpleScrubber("peer.service: localhost", "peer.service: sqlserver");
             settings.AddSimpleScrubber("peer.service: (localdb)\\MSSQLLocalDB", "peer.service: sqlserver");
             settings.AddSimpleScrubber("peer.service: sqledge_arm64", "peer.service: sqlserver");
+            settings.AddRegexScrubber(new Regex("dd.instrumentation.time_ms: \\d+.\\d+"), "dd.instrumentation.time_ms: 123.456");
 
             var fileName = nameof(SystemDataSqlClientTests);
 

tracer/test/Datadog.Trace.TestHelpers/SpanMetadataV0Rules.cs

@@ -646,6 +646,8 @@ public static Result IsSqlClientV0(this MockSpan span) => Result.FromSpan(span)
                 .IsOptional("db.name")
                 .IsPresent("out.host")
                 .IsOptional("_dd.base_service")
+                .IsOptional("_dd.dbm_trace_injected")
+                .IsOptional("dd.instrumentation.time_ms")
                 .Matches("db.type", "sql-server")
                 .Matches("component", "SqlClient")
                 .Matches("span.kind", "client"));

tracer/test/Datadog.Trace.TestHelpers/SpanMetadataV1Rules.cs

@@ -888,6 +888,8 @@ public static Result IsSqlClientV1(this MockSpan span) => Result.FromSpan(span)
                 .IsOptional("peer.service.remapped_from")
                 .IsOptional("_dd.base_service")
                 .MatchesOneOf("_dd.peer.service.source", "db.name", "out.host", "peer.service")
+                .IsOptional("_dd.dbm_trace_injected")
+                .IsOptional("dd.instrumentation.time_ms")
                 .Matches("component", "SqlClient")
                 .Matches("span.kind", "client"));