Test to validate UnsafeRegister fixes holding a reference to Operatio…

…nContextScope

src/System.Private.ServiceModel/tests/Scenarios/Client/TypedClient/TypedProxyTests.4.0.0.cs

@@ -565,4 +565,77 @@ private static void ServiceContract_TypedProxy_AsyncTask_Call_TestImpl(Binding b
             ScenarioTestHelpers.CloseCommunicationObjects((ICommunicationObject)serviceProxy, factory);
         }
     }
+
+    [WcfFact]
+    [OuterLoop]
+    public static async Task ServiceContract_TypedProxy_Task_Call_AsyncLocal_NonCapture()
+    {
+        // This test verifies a task based call to a service operation doesn't capture any AsyncLocal's in an ExecutionContext
+        // This is indirectly checking that the OperationContext won't get captured by a registered CancellationToken callback
+        CustomBinding customBinding = null;
+        ChannelFactory<IWcfServiceGenerated> factory = null;
+        EndpointAddress endpointAddress = null;
+        IWcfServiceGenerated serviceProxy = null;
+        string requestString = "Hello";
+        string result = null;
+        WeakReference<OperationContextExtension> opContextReference = null;
+
+        try
+        {
+            // *** SETUP *** \\
+            customBinding = new CustomBinding();
+            customBinding.Elements.Add(new TextMessageEncodingBindingElement());
+            customBinding.Elements.Add(new HttpTransportBindingElement());
+            endpointAddress = new EndpointAddress(Endpoints.DefaultCustomHttp_Address);
+            factory = new ChannelFactory<IWcfServiceGenerated>(customBinding, endpointAddress);
+            serviceProxy = factory.CreateChannel();
+
+            // *** EXECUTE *** \\
+            var opExtension = new OperationContextExtension();
+            opContextReference = new WeakReference<OperationContextExtension>(opExtension);
+            using (new OperationContextScope((IContextChannel)serviceProxy))
+            {
+                OperationContext.Current.Extensions.Add(opExtension);
+                result = await serviceProxy.EchoAsync(requestString);
+            }
+
+            opExtension = null;
+
+            // The Task generated by the compiler for this method stores the current ExecutionContext when an await is hit.
+            // This means even after we've nulled out the OperationContext via disposing the OperationContextScope, it's still
+            // stored in the saved ExecutionContext which was captured by the Task. Forcing another async continuation via
+            // Task.Yield() causes the stored ExecutionContext to be replace with the current one which doesn't reference
+            // OperationContext any more. Without this, the OperationContext would be kept alive until the next await call,
+            // and the weak reference would still hold a reference to the OperationContextExtension.
+            await Task.Yield();
+
+            // Force a Gen2 GC so that the WeakReference will no longer have the OperationContextExtension as a target.
+            GC.Collect(2);
+
+            // *** VALIDATE *** \\
+            Assert.Equal(requestString, result);
+            Assert.False(opContextReference.TryGetTarget(out OperationContextExtension opContext), "OperationContextExtension should have been collected");
+            // *** CLEANUP *** \\
+            factory.Close();
+            ((ICommunicationObject)serviceProxy).Close();
+        }
+        finally
+        {
+            // *** ENSURE CLEANUP *** \\
+            ScenarioTestHelpers.CloseCommunicationObjects((ICommunicationObject)serviceProxy, factory);
+        }
+    }
+
+    public class OperationContextExtension : IExtension<OperationContext>
+    {
+        public void Attach(OperationContext owner)
+        {
+            // Do nothing
+        }
+        public void Detach(OperationContext owner)
+        {
+            // Do nothing
+        }
+    }
+
 }