Fix orphaned spans in WCF traces (#7842)

## Summary of changes

Fixes orphan traces in WCF

## Reason for change

Many customers have reported that in v3, with "delayed WCF"
instrumentation enabled, they get orphan spans. This can be easily
reproduced, so this PR fixes the issue.

## Background

There are two "modes" WCF can operate in: delayed instrumentation, and
"simple" instrumentation. Delayed instrumentation was introduced in
#1992 to avoid a
problematic error scenario. Instead of instrumenting early across the
whole "channel", we delayed instrumentation. A subsequent PR
#5206 updated the
instrumentation point to ensure the wcf span was active while
`IDispatchMessageInspector` were running.

This was somewhat complex, due to the convoluted and legacy way wcf uses
async and threading. Basically, we use a weak table to hold the current
span, and us the `MessageRpc` object as a key. Later on, we retrieve the
scope and dispose it.

## Implementation details

The missing part was that we were not "restoring" the wcf scope for when
the actual Invoker implementations executed. The process was:

- Update the tests and snapshots to:
- Record _all_ spans and traces (it makes it easier to [understand the
snapshots
using](https://andrewlock.github.io/DatadogSpanSnapshotViewer/))
  - Create a span inside the "service" calls.
- These _should_ have been children of the wcf server spans, but they
were not. They were orphaned
- Update the "Invoker" implementations to grab the wcf spans, make it
active, and then deactivate it afterwards
- Regenerate the snapshots, and confirm the custom spans are no longer
orphaned 🎉

Additionally, do a small amount of refactoring:
- "Fix" nullability in some places
- Extract the "restore wcf span" functionality to `WcfCommon` as it's
used in multiple places

## Test coverage

Covered by existing, after the tweaks above. The snapshots change like
this:

Before the fix, after adding the additional custom spans,
`WSHttpBinding` traces:

<img width="1468" height="622" alt="image"
src="https://github.com/user-attachments/assets/7973514c-660e-4871-8c4d-7495473e9944"
/>

After the fix, `WSHttpBinding` traces:

<img width="1459" height="96" alt="image"
src="https://github.com/user-attachments/assets/0698e5d3-cb06-4ea6-ab00-5a3680f7e3ee"
/>

Before the fix, after adding the additional custom spans, WebHttp
traces:

<img width="1458" height="268" alt="image"
src="https://github.com/user-attachments/assets/b0a9f4c9-8584-4fa2-9441-74268d6fd279"
/>

After the fix, Web Http traces

<img width="1465" height="126" alt="image"
src="https://github.com/user-attachments/assets/b06f0255-0818-4516-a6a1-b1807c9781ad"
/>

Unfortunately, this ends up being somewhat unstable in regards to
sorting, so I did some more refactoring to separate the root span for
each separate httpcall. This should make sure we have more stable
ordering of the spans using the built-in sort order

The result is we get separate traces for each API call instead:

<img width="1458" height="713" alt="image"
src="https://github.com/user-attachments/assets/989a0cc8-befa-429f-9a9e-ede09c2abc48"
/>



## Other details

I wonder if we really need a weak table here? 🤔 seeing as we have an
instrumentation point that explicitly disposes the span, shouldn't we
just remove the entry at the same time?

Author: andrewlock

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Wcf/AfterReceiveRequestIntegration.cs

@@ -70,7 +70,7 @@ internal static CallTargetState OnMethodBegin<TTarget, TMessageRpc>(TTarget inst
         /// <param name="exception">Exception instance in case the original code threw an exception.</param>
         /// <param name="state">Calltarget state value</param>
         /// <returns>A response value, in an async scenario will be T of Task of T</returns>
-        internal static CallTargetReturn OnMethodEnd<TTarget>(TTarget instance, Exception exception, in CallTargetState state)
+        internal static CallTargetReturn OnMethodEnd<TTarget>(TTarget instance, Exception? exception, in CallTargetState state)
         {
             // Add an exception and close the span only if there was an uncaught exception,
             // which should only happen if one of the IClientMessageInspector's has thrown an exception.
@@ -80,17 +80,10 @@ internal static CallTargetReturn OnMethodEnd<TTarget>(TTarget instance, Exceptio
                 state.Scope.DisposeWithException(exception);
             }
 
-            if (state.Scope != null)
-            {
-                // OnMethodBegin started an active span that can be accessed by IDispatchMessageInspector's and the actual WCF endpoint
-                // Before returning, we must reset the scope to the previous active scope, so that callers of this method do not see this scope
-                // Don't worry, this will be accessed and closed by the BeforeSendReplyIntegration
-                if (Tracer.Instance.ScopeManager is IScopeRawAccess rawAccess)
-                {
-                    rawAccess.Active = state.PreviousScope;
-                    DistributedTracer.Instance.SetSpanContext(state.PreviousDistributedSpanContext);
-                }
-            }
+            // OnMethodBegin started an active span that can be accessed by IDispatchMessageInspector's and the actual WCF endpoint
+            // Before returning, we must reset the scope to the previous active scope, so that callers of this method do not see this scope
+            // Don't worry, this will be accessed and closed by the BeforeSendReplyIntegration
+            WcfCommon.RestorePreviousScope(in state);
 
             return CallTargetReturn.GetDefault();
         }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Wcf/AsyncMethodInvoker_InvokeBegin_Integration.cs

@@ -9,7 +9,6 @@
 using System;
 using System.ComponentModel;
 using Datadog.Trace.ClrProfiler.CallTarget;
-using Datadog.Trace.DuckTyping;
 
 namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Wcf
 {
@@ -29,27 +28,27 @@ namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Wcf
     [EditorBrowsable(EditorBrowsableState.Never)]
     public class AsyncMethodInvoker_InvokeBegin_Integration
     {
-        internal static CallTargetReturn<TReturn> OnMethodEnd<TTarget, TReturn>(TTarget instance, TReturn returnValue, Exception exception, in CallTargetState state)
+        internal static CallTargetState OnMethodBegin<TTarget>(TTarget instance, object? instanceArg, object[]? inputs, ref AsyncCallback? callback, object? state)
         {
-            if (exception is not null)
+            var tracer = Tracer.Instance;
+            if (!tracer.CurrentTraceSettings.Settings.IsIntegrationEnabled(WcfCommon.IntegrationId) || !tracer.Settings.DelayWcfInstrumentationEnabled || WcfCommon.GetCurrentOperationContext is null)
             {
-                var operationContext = WcfCommon.GetCurrentOperationContext?.Invoke();
+                return CallTargetState.GetDefault();
+            }
 
-                if (operationContext != null && operationContext.TryDuckCast<IOperationContextStruct>(out var operationContextProxy))
-                {
-                    var requestContext = operationContextProxy.RequestContext;
+            return WcfCommon.ActivateScopeFromContext();
+        }
 
-                    // Retrieve the scope that we saved during InvokeBegin
-                    if (((IDuckType?)requestContext)?.Instance is object requestContextInstance
-                        && WcfCommon.Scopes.TryGetValue(requestContextInstance, out var scope))
-                    {
-                        // Add the exception but do not dispose the span.
-                        // BeforeSendReplyIntegration is responsible for closing the span.
-                        scope.Span?.SetException(exception);
-                    }
-                }
+        internal static CallTargetReturn<TReturn> OnMethodEnd<TTarget, TReturn>(TTarget instance, TReturn returnValue, Exception? exception, in CallTargetState state)
+        {
+            if (state.Scope is not null && exception is not null)
+            {
+                // Add the exception but do not dispose the scope.
+                // BeforeSendReplyIntegration is responsible for closing the span.
+                state.Scope.Span.SetException(exception);
             }
 
+            WcfCommon.RestorePreviousScope(in state);
             return new CallTargetReturn<TReturn>(returnValue);
         }
     }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Wcf/AsyncMethodInvoker_InvokeEnd_Integration.cs

@@ -9,7 +9,6 @@
 using System;
 using System.ComponentModel;
 using Datadog.Trace.ClrProfiler.CallTarget;
-using Datadog.Trace.DuckTyping;
 
 namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Wcf
 {
@@ -29,6 +28,17 @@ namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Wcf
     [EditorBrowsable(EditorBrowsableState.Never)]
     public class AsyncMethodInvoker_InvokeEnd_Integration
     {
+        internal static CallTargetState OnMethodBegin<TTarget>(TTarget instance, object? instanceArg, object[]? outputs, IAsyncResult? result)
+        {
+            var tracer = Tracer.Instance;
+            if (!tracer.CurrentTraceSettings.Settings.IsIntegrationEnabled(WcfCommon.IntegrationId) || !tracer.Settings.DelayWcfInstrumentationEnabled || WcfCommon.GetCurrentOperationContext is null)
+            {
+                return CallTargetState.GetDefault();
+            }
+
+            return WcfCommon.ActivateScopeFromContext();
+        }
+
         /// <summary>
         /// OnMethodEnd callback
         /// </summary>
@@ -39,27 +49,16 @@ public class AsyncMethodInvoker_InvokeEnd_Integration
         /// <param name="exception">Exception instance in case the original code threw an exception.</param>
         /// <param name="state">Calltarget state value</param>
         /// <returns>A response value, in an async scenario will be T of Task of T</returns>
-        internal static CallTargetReturn<TReturn> OnMethodEnd<TTarget, TReturn>(TTarget instance, TReturn returnValue, Exception exception, in CallTargetState state)
+        internal static CallTargetReturn<TReturn> OnMethodEnd<TTarget, TReturn>(TTarget instance, TReturn returnValue, Exception? exception, in CallTargetState state)
         {
-            if (exception is not null)
+            if (state.Scope is not null && exception is not null)
             {
-                var operationContext = WcfCommon.GetCurrentOperationContext?.Invoke();
-
-                if (operationContext != null && operationContext.TryDuckCast<IOperationContextStruct>(out var operationContextProxy))
-                {
-                    var requestContext = operationContextProxy.RequestContext;
-
-                    // Retrieve the scope that we saved during InvokeBegin
-                    if (((IDuckType?)requestContext)?.Instance is object requestContextInstance
-                        && WcfCommon.Scopes.TryGetValue(requestContextInstance, out var scope))
-                    {
-                        // Add the exception but do not dispose the span.
-                        // BeforeSendReplyIntegration is responsible for closing the span.
-                        scope.Span?.SetException(exception);
-                    }
-                }
+                // Add the exception but do not dispose the scope.
+                // BeforeSendReplyIntegration is responsible for closing the span.
+                state.Scope.Span.SetException(exception);
             }
 
+            WcfCommon.RestorePreviousScope(in state);
             return new CallTargetReturn<TReturn>(returnValue);
         }
     }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Wcf/BeforeSendReplyIntegration.cs

@@ -48,7 +48,7 @@ internal static CallTargetState OnMethodBegin<TTarget, TMessageRpc>(TTarget inst
             }
 
             var rpcProxy = rpc.DuckCast<MessageRpcStruct>();
-            if (((IDuckType?)rpcProxy.OperationContext.RequestContext)?.Instance is object requestContextInstance
+            if (rpcProxy.OperationContext.RequestContext?.Instance is { } requestContextInstance
                 && WcfCommon.Scopes.TryGetValue(requestContextInstance, out var scope))
             {
                 return new CallTargetState(scope);
@@ -65,7 +65,7 @@ internal static CallTargetState OnMethodBegin<TTarget, TMessageRpc>(TTarget inst
         /// <param name="exception">Exception instance in case the original code threw an exception.</param>
         /// <param name="state">Calltarget state value</param>
         /// <returns>A response value, in an async scenario will be T of Task of T</returns>
-        internal static CallTargetReturn OnMethodEnd<TTarget>(TTarget instance, Exception exception, in CallTargetState state)
+        internal static CallTargetReturn OnMethodEnd<TTarget>(TTarget instance, Exception? exception, in CallTargetState state)
         {
             state.Scope.DisposeWithException(exception);
             return CallTargetReturn.GetDefault();

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Wcf/ChannelHandlerIntegration.cs

@@ -63,7 +63,7 @@ internal static CallTargetState OnMethodBegin<TTarget, TRequestContext, TOperati
         /// <param name="exception">Exception instance in case the original code threw an exception.</param>
         /// <param name="state">Calltarget state value</param>
         /// <returns>A response value, in an async scenario will be T of Task of T</returns>
-        internal static CallTargetReturn<TReturn> OnMethodEnd<TTarget, TReturn>(TTarget instance, TReturn returnValue, Exception exception, in CallTargetState state)
+        internal static CallTargetReturn<TReturn> OnMethodEnd<TTarget, TReturn>(TTarget instance, TReturn returnValue, Exception? exception, in CallTargetState state)
         {
             state.Scope.DisposeWithException(exception);
             return new CallTargetReturn<TReturn>(returnValue);

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Wcf/IRequestContext.cs

@@ -5,12 +5,14 @@
 
 #nullable enable
 
+using Datadog.Trace.DuckTyping;
+
 namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Wcf
 {
     /// <summary>
     /// System.ServiceModel.Channels.RequestContext interface for duck-typing
     /// </summary>
-    internal interface IRequestContext
+    internal interface IRequestContext : IDuckType
     {
         /// <summary>
         /// Gets the request message

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Wcf/SyncMethodInvokerIntegration.cs

@@ -9,7 +9,6 @@
 using System;
 using System.ComponentModel;
 using Datadog.Trace.ClrProfiler.CallTarget;
-using Datadog.Trace.DuckTyping;
 
 namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Wcf
 {
@@ -40,12 +39,13 @@ public class SyncMethodInvokerIntegration
         /// <returns>Calltarget state value</returns>
         internal static CallTargetState OnMethodBegin<TTarget>(TTarget instance, object instanceArg, object[] inputs, ref object[] outputs)
         {
-            if (!Tracer.Instance.CurrentTraceSettings.Settings.IsIntegrationEnabled(WcfCommon.IntegrationId) || !Tracer.Instance.Settings.DelayWcfInstrumentationEnabled || WcfCommon.GetCurrentOperationContext is null)
+            var tracer = Tracer.Instance;
+            if (!tracer.CurrentTraceSettings.Settings.IsIntegrationEnabled(WcfCommon.IntegrationId) || !tracer.Settings.DelayWcfInstrumentationEnabled || WcfCommon.GetCurrentOperationContext is null)
             {
                 return CallTargetState.GetDefault();
             }
 
-            return new CallTargetState(Tracer.Instance.ActiveScope as Scope);
+            return WcfCommon.ActivateScopeFromContext();
         }
 
         /// <summary>
@@ -58,15 +58,16 @@ internal static CallTargetState OnMethodBegin<TTarget>(TTarget instance, object
         /// <param name="exception">Exception instance in case the original code threw an exception.</param>
         /// <param name="state">Calltarget state value</param>
         /// <returns>A response value, in an async scenario will be T of Task of T</returns>
-        internal static CallTargetReturn<TReturn> OnMethodEnd<TTarget, TReturn>(TTarget instance, TReturn returnValue, Exception exception, in CallTargetState state)
+        internal static CallTargetReturn<TReturn> OnMethodEnd<TTarget, TReturn>(TTarget instance, TReturn returnValue, Exception? exception, in CallTargetState state)
         {
             if (state.Scope is not null && exception is not null)
             {
                 // Add the exception but do not dispose the scope.
                 // BeforeSendReplyIntegration is responsible for closing the span.
-                state.Scope.Span?.SetException(exception);
+                state.Scope.Span.SetException(exception);
             }
 
+            WcfCommon.RestorePreviousScope(in state);
             return new CallTargetReturn<TReturn>(returnValue);
         }
     }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Wcf/TaskMethodInvokerIntegration.cs

@@ -9,8 +9,6 @@
 using System;
 using System.ComponentModel;
 using Datadog.Trace.ClrProfiler.CallTarget;
-using Datadog.Trace.Configuration;
-using Datadog.Trace.DuckTyping;
 
 namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Wcf
 {
@@ -40,12 +38,13 @@ public class TaskMethodInvokerIntegration
         /// <returns>Calltarget state value</returns>
         internal static CallTargetState OnMethodBegin<TTarget>(TTarget instance, object instanceArg, object[] inputs)
         {
-            if (!Tracer.Instance.CurrentTraceSettings.Settings.IsIntegrationEnabled(WcfCommon.IntegrationId) || !Tracer.Instance.Settings.DelayWcfInstrumentationEnabled || WcfCommon.GetCurrentOperationContext is null)
+            var tracer = Tracer.Instance;
+            if (!tracer.CurrentTraceSettings.Settings.IsIntegrationEnabled(WcfCommon.IntegrationId) || !tracer.Settings.DelayWcfInstrumentationEnabled || WcfCommon.GetCurrentOperationContext is null)
             {
                 return CallTargetState.GetDefault();
             }
 
-            return new CallTargetState(Tracer.Instance.ActiveScope as Scope);
+            return WcfCommon.ActivateScopeFromContext();
         }
 
         /// <summary>
@@ -58,15 +57,16 @@ internal static CallTargetState OnMethodBegin<TTarget>(TTarget instance, object
         /// <param name="exception">Exception instance in case the original code threw an exception.</param>
         /// <param name="state">Calltarget state value</param>
         /// <returns>A response value, in an async scenario will be T of Task of T</returns>
-        internal static TResponse OnAsyncMethodEnd<TTarget, TResponse>(TTarget instance, TResponse returnValue, Exception exception, in CallTargetState state)
+        internal static TResponse OnAsyncMethodEnd<TTarget, TResponse>(TTarget instance, TResponse returnValue, Exception? exception, in CallTargetState state)
         {
             if (state.Scope is not null && exception is not null)
             {
                 // Add the exception but do not dispose the scope.
                 // BeforeSendReplyIntegration is responsible for closing the span.
-                state.Scope.Span?.SetException(exception);
+                state.Scope.Span.SetException(exception);
             }
 
+            WcfCommon.RestorePreviousScope(in state);
             return returnValue;
         }
     }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Wcf/WcfCommon.cs

@@ -12,6 +12,7 @@
 using System.Net;
 using System.Reflection;
 using System.Runtime.CompilerServices;
+using Datadog.Trace.ClrProfiler.CallTarget;
 using Datadog.Trace.Configuration;
 using Datadog.Trace.DuckTyping;
 using Datadog.Trace.ExtensionMethods;
@@ -77,7 +78,11 @@ internal class WcfCommon
                     httpMethod = httpRequestPropertyProxy.Method?.ToUpperInvariant();
 
                     // try to extract propagated context values from http headers
-                    if (tracer.ActiveScope == null)
+                    if (tracer.ActiveScope is { } activeScope)
+                    {
+                        Log.Warning("Skipped extracting headers due to existing scope: {ActiveScope}", activeScope.Span);
+                    }
+                    else
                     {
                         try
                         {
@@ -216,6 +221,49 @@ internal class WcfCommon
 
             return null;
         }
+
+        public static CallTargetState ActivateScopeFromContext()
+        {
+            // First, capture the active scope
+            var tracer = Tracer.Instance;
+            var activeScope = tracer.InternalActiveScope;
+
+            var operationContext = WcfCommon.GetCurrentOperationContext?.Invoke();
+            if (operationContext != null && operationContext.TryDuckCast<IOperationContextStruct>(out var operationContextProxy))
+            {
+                var requestContext = operationContextProxy.RequestContext;
+
+                // Retrieve the scope that we saved during InvokeBegin
+                if (requestContext?.Instance is { } requestContextInstance
+                 && Scopes.TryGetValue(requestContextInstance, out var scope)
+                    && scope is not null
+                    && scope.Span.SpanId != activeScope?.Span.SpanId)
+                {
+                    // capture the raw context for later activation
+                    var spanContextRaw = DistributedTracer.Instance.GetSpanContextRaw() ?? activeScope?.Span.Context;
+                    Log.Debug("Activating scope from operation context {ActivatedSpan}", scope.Span);
+
+                    tracer.ActivateSpan(scope.Span);
+                    // Add the exception but do not dispose the span.
+                    // BeforeSendReplyIntegration is responsible for closing the span.
+                    return new CallTargetState(scope, activeScope, spanContextRaw);
+                }
+            }
+
+            return CallTargetState.GetDefault();
+        }
+
+        public static void RestorePreviousScope(in CallTargetState state)
+        {
+            // Reset the scope to the previous active scope, so that callers of this method do not see this scope
+            // Don't worry, this will be accessed and closed by the BeforeSendReplyIntegration
+            if (state.Scope is not null && Tracer.Instance.ScopeManager is IScopeRawAccess rawAccess)
+            {
+                Log.Debug("Restoring previous scope from {Previous} to {Restored}", state.Scope.Span, state.PreviousScope?.Span);
+                rawAccess.Active = state.PreviousScope;
+                DistributedTracer.Instance.SetSpanContext(state.PreviousDistributedSpanContext);
+            }
+        }
     }
 }
 #endif