Add support for OperationContext/OperationContextScope flow through an await

.gitignore

@@ -276,3 +276,7 @@ Session.vim
 # Private test configuration and binaries.
 config.ps1
 **/IISApplications
+
+# Debug files
+*.etl
+*.dmp

Directory.Build.props

@@ -34,6 +34,12 @@
     <XUnitPublishTargetFramework>netcoreapp3.1</XUnitPublishTargetFramework>
   </PropertyGroup>
 
+  <PropertyGroup Condition="'$(IsTestProject)' == 'true'">
+    <!-- TODO: Remove this condition when VSTest is used in CI. -->
+    <EnableRunSettingsSupport Condition="'$(ContinuousIntegrationBuild)' != 'true'">true</EnableRunSettingsSupport>
+  </PropertyGroup>
+
+
   <ItemGroup>
     <!-- Include license and third party files to packages -->
     <Content Include="$(RepoRoot)THIRD-PARTY-NOTICES.TXT" Pack="true" PackagePath="\" />

Directory.Build.targets

@@ -4,4 +4,5 @@
   <Import Project="eng\FacadeAssemblies.targets" />
   <Import Project="eng\harvest.targets" Condition="'$(HarvestFromPackage)' == 'true'" />
   <Import Project="eng\uwp.targets" />
+  <Import Project="eng\testing\runsettings.targets" Condition="'$(EnableRunSettingsSupport)' == 'true'" />
 </Project>

eng/testing/.runsettings

@@ -0,0 +1,49 @@
+<?xml version="1.0" encoding="utf-8"?>
+<RunSettings>
+  <RunConfiguration>
+    <!-- Timeout in ms, 5 minutes -->
+    <TestSessionTimeout>300000</TestSessionTimeout>
+    <!-- Directory for test run reports. E.g. trx, coverage etc. -->
+    <ResultsDirectory>.\TestResults\</ResultsDirectory>
+    <!-- Working directory for test invocation. Results directory can be relative to this. Used by IDEs. -->
+    <SolutionDirectory>.\</SolutionDirectory>
+    <!-- Degree of parallelization, spawns n test hosts to run tests. -->
+    <MaxCpuCount>$$MAXCPUCOUNT$$</MaxCpuCount>
+    <!-- If true, an adapter should disable any test case parallelization. -->
+    <DisableParallelization>$$DISABLEPARALLELIZATION$$</DisableParallelization>
+    <!-- If true, an adapter shouldn't create appdomains to run tests. -->
+    <DisableAppDomain>$$DISABLEAPPDOMAIN$$</DisableAppDomain>
+    <!-- Filter out failing (wrong framwork, platform, runtime or activeissue) tests -->
+    <TestCaseFilter>$$TESTCASEFILTER$$</TestCaseFilter>
+    <DotNetHostPath>$$DOTNETHOSTPATH$$</DotNetHostPath>
+    <TargetFrameworkVersion>FrameworkCore10</TargetFrameworkVersion>
+  </RunConfiguration>
+  <LoggerRunSettings>
+    <Loggers>
+      <Logger friendlyName="trx" />
+      <Logger friendlyName="html" />
+      <Logger friendlyName="console">
+        <Configuration>
+          <Verbosity>Minimal</Verbosity>
+        </Configuration>
+      </Logger>
+    </Loggers>
+  </LoggerRunSettings>
+  <DataCollectionRunSettings>
+    <DataCollectors>
+	  <!-- Enabling the code coverage data collector as a later task -->
+      <!--<DataCollector friendlyName="XPlat code coverage" enabled="$$COVERAGE_ENABLED$$">
+        <Configuration>
+          <Include>$$COVERAGE_INCLUDE$$</Include>
+          <ExcludeByFile>$$COVERAGE_EXCLUDEBYFILE$$</ExcludeByFile>
+          <IncludeDirectory>$$COVERAGE_INCLUDEDIRECTORY$$</IncludeDirectory>
+          <Format>opencover</Format>
+          <SingleHit>false</SingleHit>
+          <UseSourceLink>true</UseSourceLink>
+          <IncludeTestAssembly>false</IncludeTestAssembly>
+        </Configuration>
+      </DataCollector>-->
+      <DataCollector friendlyName="blame" enabled="true" />
+    </DataCollectors>
+  </DataCollectionRunSettings>
+</RunSettings>

eng/testing/runsettings.targets

@@ -0,0 +1,55 @@
+<Project>
+  <PropertyGroup>
+    <RunSettingsInputFilePath>$(MSBuildThisFileDirectory).runsettings</RunSettingsInputFilePath>
+    <RunSettingsIntermediateOutputFilePath>$(ArtifactsObjDir)$(TargetOS)-$(Configuration)-$(TargetArchitecture).runsettings</RunSettingsIntermediateOutputFilePath>
+    <RunSettingsAppOutputFilePath>$(OutDir).runsettings</RunSettingsAppOutputFilePath>
+
+    <CreateIntermediateRunSettingsFile Condition="'$(CreateIntermediateRunSettingsFile)' == ''">false</CreateIntermediateRunSettingsFile>
+    <RunSettingsOutputFilePath Condition="'$(CreateIntermediateRunSettingsFile)' == 'true'">$(RunSettingsIntermediateOutputFilePath)</RunSettingsOutputFilePath>
+    <RunSettingsOutputFilePath Condition="'$(CreateIntermediateRunSettingsFile)' != 'true'">$(RunSettingsAppOutputFilePath)</RunSettingsOutputFilePath>
+
+    <!-- Set RunSettingsFilePath property which is read by VSTest. -->
+    <RunSettingsFilePath Condition="Exists('$(RunSettingsAppOutputFilePath)')">$(RunSettingsAppOutputFilePath)</RunSettingsFilePath>
+    <!-- Use an intermediate runsettings file if the app hasn't been built yet to enable VSTest discovery. -->
+    <RunSettingsFilePath Condition="'$(RunSettingsFilePath)' == '' and Exists('$(RunSettingsIntermediateOutputFilePath)')">$(RunSettingsIntermediateOutputFilePath)</RunSettingsFilePath>
+
+    <PrepareForRunDependsOn>GenerateRunSettingsFile;$(PrepareForRunDependsOn)</PrepareForRunDependsOn>
+  </PropertyGroup>
+
+  <PropertyGroup>
+    <_testFilter Condition="'$(_withCategories)' != ''">$(_withCategories.Replace(';', '&amp;amp;category='))</_testFilter>
+    <_testFilter Condition="'$(_withoutCategories)' != ''">$(_testFilter)$(_withoutCategories.Replace(';', '&amp;amp;category!='))</_testFilter>
+    <_testFilter>$(_testFilter.Trim('&amp;amp;'))</_testFilter>
+  </PropertyGroup>
+
+  <Target Name="GenerateRunSettingsFile" >
+    <PropertyGroup>
+      <RunSettingsFileContent>$([System.IO.File]::ReadAllText('$(RunSettingsInputFilePath)'))</RunSettingsFileContent>
+      <RunSettingsFileContent Condition="'$(TestDisableParallelization)' == 'true'">$(RunSettingsFileContent.Replace('$$MAXCPUCOUNT$$', '1'))</RunSettingsFileContent>
+      <RunSettingsFileContent Condition="'$(TestDisableParallelization)' != 'true'">$(RunSettingsFileContent.Replace('$$MAXCPUCOUNT$$', '0'))</RunSettingsFileContent>
+      <!-- Arm64 is currently not a known TargetPlatform value in VSTEST: https://github.com/microsoft/vstest/issues/2566 -->
+<!--
+      <RunSettingsFileContent Condition="'$(TargetArchitecture)' != 'arm64'">$(RunSettingsFileContent.Replace('$$TARGETPLATFORM$$', '<TargetPlatform>$(TargetArchitecture)</TargetPlatform>'))</RunSettingsFileContent>
+      <RunSettingsFileContent Condition="'$(TargetArchitecture)' == 'arm64'">$(RunSettingsFileContent.Replace('$$TARGETPLATFORM$$', ''))</RunSettingsFileContent>
+	  -->
+      <RunSettingsFileContent>$(RunSettingsFileContent.Replace('$$COVERAGE_INCLUDE$$', '$(CoverageIncludeFilter)')
+                                                      .Replace('$$COVERAGE_EXCLUDEBYFILE$$', '$(CoverageExcludeByFileFilter)')
+                                                      .Replace('$$COVERAGE_INCLUDEDIRECTORY$$', '$(CoverageIncludeDirectoryFilter)')
+                                                      .Replace('$$COVERAGE_ENABLED$$', '$([MSBuild]::ValueOrDefault('$(Coverage)', 'false'))')
+                                                      .Replace('$$DISABLEPARALLELIZATION$$', '$([MSBuild]::ValueOrDefault('$(TestDisableParallelization)', 'false'))')
+                                                      .Replace('$$DISABLEAPPDOMAIN$$', '$([MSBuild]::ValueOrDefault('$(TestDisableAppDomain)', 'false'))')
+                                                      .Replace('$$TESTCASEFILTER$$', '$(_testFilter)')
+                                                      .Replace('$$DOTNETHOSTPATH$$', '$(TestHostRootPath)$([System.IO.Path]::GetFileName('$(DotNetTool)'))'))</RunSettingsFileContent>
+    </PropertyGroup>
+
+    <WriteLinesToFile File="$(RunSettingsOutputFilePath)"
+                      Lines="$(RunSettingsFileContent)"
+                      WriteOnlyWhenDifferent="true"
+                      Overwrite="true" />
+
+    <!-- Set RunSettingsFilePath property which is read by VSTest. -->
+    <PropertyGroup>
+      <RunSettingsFilePath>$(RunSettingsOutputFilePath)</RunSettingsFilePath>
+    </PropertyGroup>
+  </Target>
+</Project>

src/System.Private.ServiceModel/src/System/ServiceModel/Dispatcher/MessageRpc.cs

@@ -409,15 +409,12 @@ internal bool Process(bool isOperationContextSet)
                     NextProcessor = null;
 
                     OperationContext originalContext;
-                    OperationContext.Holder contextHolder;
                     if (!isOperationContextSet)
                     {
-                        contextHolder = OperationContext.CurrentHolder;
-                        originalContext = contextHolder.Context;
+                        originalContext = OperationContext.Current;
                     }
                     else
                     {
-                        contextHolder = null;
                         originalContext = null;
                     }
                     IncrementBusyCount();
@@ -426,7 +423,7 @@ internal bool Process(bool isOperationContextSet)
                     {
                         if (!isOperationContextSet)
                         {
-                            contextHolder.Context = OperationContext;
+                            OperationContext.Current = OperationContext;
                         }
 
                         processor(ref this);
@@ -455,7 +452,7 @@ internal bool Process(bool isOperationContextSet)
 
                             if (!isOperationContextSet)
                             {
-                                contextHolder.Context = originalContext;
+                                OperationContext.Current = originalContext;
                             }
 
                             completed = !IsPaused;

src/System.Private.ServiceModel/src/System/ServiceModel/OperationContext.cs

@@ -8,13 +8,24 @@
 using System.Security.Principal;
 using System.ServiceModel.Channels;
 using System.ServiceModel.Dispatcher;
+using System.Threading;
 
 namespace System.ServiceModel
 {
     public sealed class OperationContext : IExtensibleObject<OperationContext>
     {
+        static OperationContext()
+        {
+            DisableAsyncFlow = AppContext.TryGetSwitch("System.ServiceModel.OperationContext.DisableAsyncFlow", out var enabled) && enabled;
+            if (!DisableAsyncFlow)
+            {
+                s_asyncContext = new AsyncLocal<OperationContext>();
+            }
+        }
+
         [ThreadStatic]
         private static Holder s_currentContext;
+        private static AsyncLocal<OperationContext> s_asyncContext;
         private Message _clientReply;
         private bool _closeClientReply;
         private ExtensionCollection<OperationContext> _extensions;
@@ -80,29 +91,47 @@ public static OperationContext Current
         {
             get
             {
-                return CurrentHolder.Context;
+                if (DisableAsyncFlow)
+                {
+                    return CurrentHolder.Context;
+                }
+                else
+                {
+                    return s_asyncContext.Value;
+                }
             }
 
             set
             {
-                CurrentHolder.Context = value;
+                if (DisableAsyncFlow)
+                {
+                    CurrentHolder.Context = value;
+                }
+                else
+                {
+                    s_asyncContext.Value = value;
+                }
             }
         }
 
         internal static Holder CurrentHolder
         {
             get
             {
-                Holder holder = OperationContext.s_currentContext;
+                Holder holder;
+                holder = s_currentContext;
                 if (holder == null)
                 {
                     holder = new Holder();
-                    OperationContext.s_currentContext = holder;
+                    s_currentContext = holder;
                 }
+
                 return holder;
             }
         }
 
+        internal static bool DisableAsyncFlow { get; }
+
         public EndpointDispatcher EndpointDispatcher
         {
             get

src/System.Private.ServiceModel/src/System/ServiceModel/OperationContextScope.cs

@@ -2,19 +2,27 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-
+using System.Threading;
 
 namespace System.ServiceModel
 {
     public sealed class OperationContextScope : IDisposable
     {
+        static OperationContextScope()
+        {
+            if (!OperationContext.DisableAsyncFlow)
+            {
+                s_asyncCurrentScope = new AsyncLocal<OperationContextScope>();
+            }
+        }
+
         [ThreadStatic]
         private static OperationContextScope s_currentScope;
-
+        private static AsyncLocal<OperationContextScope> s_asyncCurrentScope;
         private OperationContext _currentContext;
         private bool _disposed;
         private readonly OperationContext _originalContext = OperationContext.Current;
-        private readonly OperationContextScope _originalScope = OperationContextScope.s_currentScope;
+        private readonly OperationContextScope _originalScope = OperationContext.DisableAsyncFlow ? s_currentScope : s_asyncCurrentScope.Value;
 
         public OperationContextScope(IContextChannel channel)
         {
@@ -38,13 +46,21 @@ public void Dispose()
         private void PushContext(OperationContext context)
         {
             _currentContext = context;
-            OperationContextScope.s_currentScope = this;
+            if (OperationContext.DisableAsyncFlow)
+            {
+                s_currentScope = this;
+            }
+            else
+            {
+                s_asyncCurrentScope.Value = this;
+            }
+
             OperationContext.Current = _currentContext;
         }
 
         private void PopContext()
         {
-            if (OperationContextScope.s_currentScope != this)
+            if ((OperationContext.DisableAsyncFlow ? s_currentScope : s_asyncCurrentScope.Value) != this)
             {
                 throw DiagnosticUtility.ExceptionUtility.ThrowHelperError(new InvalidOperationException(SR.SFxInterleavedContextScopes0));
             }
@@ -54,7 +70,15 @@ private void PopContext()
                 throw DiagnosticUtility.ExceptionUtility.ThrowHelperError(new InvalidOperationException(SR.SFxContextModifiedInsideScope0));
             }
 
-            OperationContextScope.s_currentScope = _originalScope;
+            if (OperationContext.DisableAsyncFlow)
+            {
+                s_currentScope = _originalScope;
+            }
+            else
+            {
+                s_asyncCurrentScope.Value = _originalScope;
+            }
+
             OperationContext.Current = _originalContext;
 
             if (_currentContext != null)

src/System.Private.ServiceModel/tests/Common/Infrastructure/ServiceUtilHelper.cs

@@ -43,6 +43,9 @@ public static X509Certificate2 RootCertificate
     {
         get
         {
+            // When running under VS, the Condition checks are run in a different process than the test run
+            // which means the test running process won't have the certificate installed by the condition check code.
+            EnsureRootCertificateInstalled();
             ThrowIfRootCertificateInstallationError();
             return s_rootCertificate;
         }
@@ -52,6 +55,9 @@ public static X509Certificate2 ClientCertificate
     {
         get
         {
+            // When running under VS, the Condition checks are run in a different process than the test run
+            // which means the test running process won't have the certificate installed by the condition check code.
+            EnsureClientCertificateInstalled();
             ThrowIfClientCertificateInstallationError();
             return s_clientCertificate;
         }
@@ -61,6 +67,9 @@ public static X509Certificate2 PeerCertificate
     {
         get
         {
+            // When running under VS, the Condition checks are run in a different process than the test run
+            // which means the test running process won't have the certificate installed by the condition check code.
+            EnsurePeerCertificateInstalled();
             ThrowIfPeerCertificateInstallationError();
             return s_peerCertificate;
         }