DSM Full propagation mode for SQL Server

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

@@ -202,8 +202,12 @@
       <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.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

@@ -10,11 +10,13 @@
 using System.Data;
 using System.Data.SqlTypes;
 using System.Diagnostics.CodeAnalysis;
+using System.Linq;
 using System.Threading;
 using Datadog.Trace.AppSec;
 using Datadog.Trace.AppSec.Rasp;
 using Datadog.Trace.Configuration;
 using Datadog.Trace.DatabaseMonitoring;
+using Datadog.Trace.ExtensionMethods;
 using Datadog.Trace.Iast;
 using Datadog.Trace.Logging;
 using Datadog.Trace.Tagging;
@@ -103,11 +105,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, tracer.Settings.DbmPropagationMode, integrationId, command.Connection, serviceName, scope, tags);
+                        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";
                             }

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

@@ -4,10 +4,12 @@
 // </copyright>
 
 using System;
+using System.Data;
 using Datadog.Trace.Configuration;
 using Datadog.Trace.Propagators;
 using Datadog.Trace.Tagging;
 using Datadog.Trace.Util;
+using Datadog.Trace.VendoredMicrosoftCode.System.Buffers.Binary;
 
 #nullable enable
 
@@ -23,7 +25,7 @@ internal static class DatabaseMonitoringPropagator
         private const string SqlCommentEnv = "dde";
         internal const string DbmPrefix = $"/*{SqlCommentSpanService}='";
 
-        internal static string PropagateSpanData(DbmPropagationLevel propagationStyle, string configuredServiceName, string? dbName, string? outhost, Span span, IntegrationId integrationId, out bool traceParentInjected)
+        internal static string PropagateDataViaComment(DbmPropagationLevel propagationStyle, string configuredServiceName, string? dbName, string? outhost, Span span, IntegrationId integrationId, out bool traceParentInjected)
         {
             traceParentInjected = false;
 
@@ -73,5 +75,79 @@ 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(Tracer tracer, DbmPropagationLevel propagationLevel, IntegrationId integrationId, IDbConnection? connection, string serviceName, Scope scope, SqlTags tags)
+        {
+            if (propagationLevel != DbmPropagationLevel.Full || integrationId != IntegrationId.SqlClient || connection == null)
+            {
+                return false;
+            }
+
+            // we want the instrumentation span to be a sibling of the actual query span
+            var instrumentationParent = scope.Parent?.Span?.Context;
+            // copy the tags so that modifications on one span don't impact the other
+            var copyProcessor = new ITags.CopyProcessor<SqlTags>();
+            tags.EnumerateTags(ref copyProcessor);
+            using (var instrumentationScope = tracer.StartActiveInternal("set context_info", instrumentationParent, tags: copyProcessor.TagsCopy, serviceName: serviceName))
+            {
+                instrumentationScope.Span.Type = SpanTypes.Sql;
+                // this tag serves as "documentation" for users to realize this is something done by the instrumentation
+                instrumentationScope.Span.Tags.SetTag("dd.instrumentation", "true");
+
+                byte version = 0; // version can have a maximum value of 15 in the current format
+                var sampled = SamplingPriorityValues.IsKeep(scope.Span.Context.GetOrMakeSamplingDecision() ?? SamplingPriorityValues.Default);
+                var contextValue = BuildContextValue(version, sampled, scope.Span.SpanId, scope.Span.TraceId128);
+                var injectionSql = "set context_info @context";
+                // important to set the resource name before running the command so that we don't re-instrument
+                instrumentationScope.Span.ResourceName = injectionSql;
+
+                using (var injectionCommand = connection.CreateCommand())
+                {
+                    injectionCommand.CommandText = injectionSql;
+
+                    var parameter = injectionCommand.CreateParameter();
+                    parameter.ParameterName = "@context";
+                    parameter.Value = contextValue;
+                    parameter.DbType = DbType.Binary;
+                    injectionCommand.Parameters.Add(parameter);
+
+                    injectionCommand.ExecuteNonQuery();
+                }
+            } // closing instrumentation span
+
+            // we don't want to measure the time spent in "set_context" in the actual query span
+            scope.Span.ResetStartTime();
+            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];
+
+            contextBytes[0] = versionAndSampling;
+            BinaryPrimitives.WriteUInt64BigEndian(
+                new VendoredMicrosoftCode.System.Span<byte>(contextBytes, start: 1, sizeof(ulong)),
+                spanId);
+            BinaryPrimitives.WriteUInt64BigEndian(
+                new VendoredMicrosoftCode.System.Span<byte>(contextBytes, 1 + sizeof(ulong), sizeof(ulong)),
+                traceId.Upper);
+            BinaryPrimitives.WriteUInt64BigEndian(
+                new VendoredMicrosoftCode.System.Span<byte>(contextBytes, 1 + sizeof(ulong) + sizeof(ulong), sizeof(ulong)),
+                traceId.Lower);
+
+            return contextBytes;
+        }
     }
 }

tracer/src/Datadog.Trace/Tagging/ITags.cs

@@ -27,5 +27,24 @@ void EnumerateMetrics<TProcessor>(ref TProcessor processor)
 
         void EnumerateMetaStruct<TProcessor>(ref TProcessor processor)
             where TProcessor : struct, IItemProcessor<byte[]>;
+
+        /// <summary>
+        /// To be used in combination with <see cref="ITags.EnumerateTags"/> to create a copy of the tags.
+        /// </summary>
+        public readonly struct CopyProcessor<T> : IItemProcessor<string>
+            where T : ITags, new()
+        {
+            public readonly ITags TagsCopy;
+
+            public CopyProcessor()
+            {
+                TagsCopy = new T();
+            }
+
+            public void Process(TagItem<string> item)
+            {
+                TagsCopy.SetTag(item.Key, item.Value);
+            }
+        }
     }
 }

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)
@@ -65,20 +67,27 @@ public async Task SubmitsTraces(string packageVersion, string metadataSchemaVers
             }
 
             var expectedSpanCount = isVersion4 ? 91 : 147;
+            // there are as many spans for the instrumentation as regular spans, since we create one extra for each query.
+            var expectedInstrumentationSpanCount = propagation == "full" ? expectedSpanCount : 0;
             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
+            var instrumentationSpans = agent.WaitForSpans(expectedInstrumentationSpanCount, operationName: "set context_info");
 
             Assert.Equal(expectedSpanCount, actualSpanCount);
+            Assert.Equal(expectedInstrumentationSpanCount, instrumentationSpans.Count);
             ValidateIntegrationSpans(spans, metadataSchemaVersion, expectedServiceName: clientSpanServiceName, isExternalSpan);
             telemetry.AssertIntegrationEnabled(IntegrationId.SqlClient);
         }

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

@@ -646,6 +646,7 @@ 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")
                 .Matches("db.type", "sql-server")
                 .Matches("component", "SqlClient")
                 .Matches("span.kind", "client"));

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

@@ -888,6 +888,7 @@ 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")
                 .Matches("component", "SqlClient")
                 .Matches("span.kind", "client"));
 

tracer/test/Datadog.Trace.Tests/DatabaseMonitoring/DatabaseMonitoringPropagatorTests.cs

@@ -4,6 +4,8 @@
 // </copyright>
 
 using System;
+using System.Data;
+using System.Numerics;
 using Datadog.Trace.Agent;
 using Datadog.Trace.Configuration;
 using Datadog.Trace.Configuration.Telemetry;
@@ -57,7 +59,7 @@ public void ExpectedCommentInjected(string propagationMode, int? samplingPriorit
             var span = _v0Tracer.StartSpan(operationName: "mysql.query", parent: SpanContext.None, serviceName: dbServiceName, traceId: (TraceId)7021887840877922076, spanId: 407003698947780173);
             span.SetTraceSamplingPriority((SamplingPriority)samplingPriority.Value);
 
-            var returnedComment = DatabaseMonitoringPropagator.PropagateSpanData(dbmPropagationLevel, "Test.Service", "MyDatabase", "MyHost", span, integrationId, out var traceParentInjectedValue);
+            var returnedComment = DatabaseMonitoringPropagator.PropagateDataViaComment(dbmPropagationLevel, "Test.Service", "MyDatabase", "MyHost", span, integrationId, out var traceParentInjectedValue);
 
             traceParentInjectedValue.Should().Be(traceParentInjected);
             returnedComment.Should().Be(expectedComment);
@@ -75,7 +77,7 @@ public void ExpectedTagsInjected(string expectedComment, string env = null, stri
             span.Context.TraceContext.ServiceVersion = version;
             span.SetTraceSamplingPriority(SamplingPriority.AutoKeep);
 
-            var returnedComment = DatabaseMonitoringPropagator.PropagateSpanData(DbmPropagationLevel.Service, "Test.Service", "MyDatabase", "MyHost", span, IntegrationId.MySql, out var traceParentInjected);
+            var returnedComment = DatabaseMonitoringPropagator.PropagateDataViaComment(DbmPropagationLevel.Service, "Test.Service", "MyDatabase", "MyHost", span, IntegrationId.MySql, out var traceParentInjected);
 
             // Always false since this test never runs for full mode
             traceParentInjected.Should().Be(false);
@@ -94,7 +96,7 @@ public void ExpectedTagsEncoded(string expectedComment, string service, string d
             span.Context.TraceContext.ServiceVersion = version;
             span.SetTraceSamplingPriority(SamplingPriority.AutoKeep);
 
-            var returnedComment = DatabaseMonitoringPropagator.PropagateSpanData(DbmPropagationLevel.Service, service, dbName, host, span, IntegrationId.MySql, out var traceParentInjected);
+            var returnedComment = DatabaseMonitoringPropagator.PropagateDataViaComment(DbmPropagationLevel.Service, service, dbName, host, span, IntegrationId.MySql, out var traceParentInjected);
 
             // Always false since this test never runs for full mode
             traceParentInjected.Should().Be(false);
@@ -115,10 +117,55 @@ public void ExpectedCommentInjectedV1()
             var span = _v1Tracer.StartSpan(tags: new SqlV1Tags() { DbName = dbName }, operationName: "mysql.query", parent: SpanContext.None, serviceName: dbServiceName, traceId: (TraceId)7021887840877922076, spanId: 407003698947780173);
             span.SetTraceSamplingPriority(samplingPriority);
 
-            var returnedComment = DatabaseMonitoringPropagator.PropagateSpanData(dbmPropagationLevel, "Test.Service", "MyDatabase", "MyHost", span, integrationId, out var traceParentInjectedValue);
+            var returnedComment = DatabaseMonitoringPropagator.PropagateDataViaComment(dbmPropagationLevel, "Test.Service", "MyDatabase", "MyHost", span, integrationId, out var traceParentInjectedValue);
 
             traceParentInjectedValue.Should().Be(traceParentInjected);
             returnedComment.Should().Be(expectedComment);
         }
+
+        [Theory]
+        [InlineData("full", "sqlclient", SamplingPriorityValues.UserKeep, true, "01000000000000BEEF0000000000000000000000000000CAFE")]
+        [InlineData("full", "sqlclient", SamplingPriorityValues.UserReject, true, "00000000000000BEEF0000000000000000000000000000CAFE")]
+        [InlineData("nope", "sqlclient", SamplingPriorityValues.UserKeep, false, null)]
+        // disabled for all db types except mysql for now
+        [InlineData("full", "npgsql", SamplingPriorityValues.UserKeep, false, null)]
+        [InlineData("full", "sqlite", SamplingPriorityValues.UserKeep, false, null)]
+        [InlineData("full", "oracle", SamplingPriorityValues.UserKeep, false, null)]
+        [InlineData("full", "mysql", SamplingPriorityValues.UserKeep, false, null)]
+        public void ExpectedContextSet(string propagationMode, string integration, int? samplingPriority, bool shouldInject, string expectedContext)
+        {
+            Enum.TryParse(propagationMode, ignoreCase: true, out DbmPropagationLevel dbmPropagationLevel);
+            Enum.TryParse(integration, ignoreCase: true, out IntegrationId integrationId);
+
+            // capture command and parameter sent
+            string sql = null;
+            byte[] context = null;
+            var connectionMock = new Mock<IDbConnection>(MockBehavior.Strict);
+            var commandMock = new Mock<IDbCommand>();
+            var parameterMock = new Mock<IDbDataParameter>();
+            connectionMock.Setup(c => c.CreateCommand()).Returns(commandMock.Object);
+            commandMock.SetupSet(c => c.CommandText = It.IsAny<string>())
+                       .Callback<string>(value => sql = value);
+            commandMock.Setup(c => c.CreateParameter()).Returns(parameterMock.Object);
+            commandMock.SetupGet(c => c.Parameters).Returns(Mock.Of<IDataParameterCollection>());
+            parameterMock.SetupSet(p => p.Value = It.IsAny<byte[]>())
+                         .Callback<object>(value => context = (byte[])value);
+
+            var span = _v0Tracer.StartSpan("mysql.query", parent: SpanContext.None, serviceName: "pouet", traceId: (TraceId)0xCAFE, spanId: 0xBEEF);
+            span.SetTraceSamplingPriority((SamplingPriority)samplingPriority.Value);
+
+            DatabaseMonitoringPropagator.PropagateDataViaContext(_v0Tracer, dbmPropagationLevel, integrationId, connectionMock.Object, "pouet", new Scope(parent: null, span, scopeManager: null, finishOnClose: false), new SqlTags());
+
+            if (shouldInject)
+            {
+                sql.Should().StartWith("set context_info ");
+                BitConverter.ToString(context).Replace("-", string.Empty).Should().Be(expectedContext);
+            }
+            else
+            {
+                sql.Should().BeNull();
+                context.Should().BeNull();
+            }
+        }
     }
 }