Implement Thread.GetApartmentState/TrySetApartmentState (#6323)

Fixes #5776

10 files changed, 143 insertions, 40 deletions

diff --git a/src/Common/src/Interop/Windows/ole32/Interop.CoInitializeEx.cs b/src/Common/src/Interop/Windows/ole32/Interop.CoInitializeEx.cs
new file mode 100644
index 000000000..f953d5ec4
--- /dev/null
+++ b/src/Common/src/Interop/Windows/ole32/Interop.CoInitializeEx.cs
@@ -0,0 +1,19 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// 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;
+using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
+
+internal static partial class Interop
+{
+    internal static partial class Ole32
+    {
+        internal const uint COINIT_APARTMENTTHREADED = 2;
+        internal const uint COINIT_MULTITHREADED = 0;
+
+        [DllImport(Interop.Libraries.Ole32, ExactSpelling = true)]
+        internal extern static int CoInitializeEx(IntPtr reserved, uint dwCoInit);
+    }
+}
diff --git a/src/Common/src/Interop/Windows/ole32/Interop.CoUninitialize.cs b/src/Common/src/Interop/Windows/ole32/Interop.CoUninitialize.cs
new file mode 100644
index 000000000..ae0836688
--- /dev/null
+++ b/src/Common/src/Interop/Windows/ole32/Interop.CoUninitialize.cs
@@ -0,0 +1,16 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// 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;
+using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
+
+internal static partial class Interop
+{
+    internal static partial class Ole32
+    {
+        [DllImport(Interop.Libraries.Ole32, ExactSpelling = true)]
+        internal extern static int CoUninitialize();
+    }
+}
diff --git a/src/System.Private.CoreLib/shared/System/HResults.cs b/src/System.Private.CoreLib/shared/System/HResults.cs
index c242389db..4a5ec0d63 100644
--- a/src/System.Private.CoreLib/shared/System/HResults.cs
+++ b/src/System.Private.CoreLib/shared/System/HResults.cs
@@ -124,5 +124,6 @@ namespace System
         internal const int RO_E_CLOSED = unchecked((int)0x80000013);
         internal const int TYPE_E_TYPEMISMATCH = unchecked((int)0x80028CA0);
         internal const int CO_E_NOTINITIALIZED = unchecked((int)0x800401F0);
+        internal const int RPC_E_CHANGED_MODE = unchecked((int)0x80010106);
     }
 }
diff --git a/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.Unix.cs b/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.Unix.cs
index 4d4406027..bc3a042eb 100644
--- a/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.Unix.cs
+++ b/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.Unix.cs
@@ -144,6 +144,14 @@ namespace Internal.Runtime.Augments
             return IntPtr.Zero;
         }
 
+        private void InitializeComOnNewThread()
+        {
+        }
+
+        internal static void InitializeCom()
+        {
+        }
+
         public void Interrupt() => WaitSubsystem.Interrupt(this);
         internal static void UninterruptibleSleep0() => WaitSubsystem.UninterruptibleSleep0();
         private static void SleepInternal(int millisecondsTimeout) => WaitSubsystem.Sleep(millisecondsTimeout);
diff --git a/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.Windows.cs b/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.Windows.cs
index 4105dff16..e9fe5485b 100644
--- a/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.Windows.cs
+++ b/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.Windows.cs
@@ -21,8 +21,13 @@ namespace Internal.Runtime.Augments
         [ThreadStatic]
         private static ApartmentType t_apartmentType;
 
+        [ThreadStatic]
+        private static bool t_comInitializedByUs;
+
         private SafeWaitHandle _osHandle;
 
+        private ApartmentState _initialAppartmentState = ApartmentState.Unknown;
+
         /// <summary>
         /// Used by <see cref="WaitHandle"/>'s multi-wait functions
         /// </summary>
@@ -286,7 +291,11 @@ namespace Internal.Runtime.Augments
         public ApartmentState GetApartmentState()
         {
             if (this != CurrentThread)
-                throw new InvalidOperationException(SR.Thread_Operation_RequiresCurrentThread);
+            {
+                if (HasStarted())
+                    throw new ThreadStateException();
+                return _initialAppartmentState;
+            }
 
             switch (GetCurrentApartmentType())
             {
@@ -299,8 +308,81 @@ namespace Internal.Runtime.Augments
             }
         }
 
+        public bool TrySetApartmentState(ApartmentState state)
+        {
+            if (this != CurrentThread)
+            {
+                using (LockHolder.Hold(_lock))
+                {
+                    if (HasStarted())
+                        throw new ThreadStateException();
+                    _initialAppartmentState = state;
+                    return true;
+                }
+            }
+
+            if (state != ApartmentState.Unknown)
+            {
+                InitializeCom(state);
+            }
+            else
+            {
+                UninitializeCom();
+            }
+
+            // Clear the cache and check whether new state matches the desired state
+            t_apartmentType = ApartmentType.Unknown;
+            return state == GetApartmentState();
+        }
+
+        private void InitializeComOnNewThread()
+        {
+            InitializeCom(_initialAppartmentState);
+        }
+
+        internal static void InitializeCom(ApartmentState state = ApartmentState.MTA)
+        {
+            if (t_comInitializedByUs)
+                return;
+
+#if ENABLE_WINRT
+            int hr = Interop.WinRT.RoInitialize(
+                (state == ApartmentState.STA) ? Interop.WinRT.RO_INIT_SINGLETHREADED
+                    : Interop.WinRT.RO_INIT_MULTITHREADED);
+#else
+            int hr = Interop.Ole32.CoInitializeEx(IntPtr.Zero,
+                (state == ApartmentState.STA) ? Interop.Ole32.COINIT_APARTMENTTHREADED
+                    : Interop.Ole32.COINIT_MULTITHREADED);
+#endif
+            // RPC_E_CHANGED_MODE indicates this thread has been already initialized with a different
+            // concurrency model. We stay away and let whoever else initialized the COM to be in control.
+            if (hr == HResults.RPC_E_CHANGED_MODE)
+                return;
+            if (hr < 0)
+                throw new OutOfMemoryException();
+
+            t_comInitializedByUs = true;
+
+            // If the thread has already been CoInitialized to the proper mode, then
+            // we don't want to leave an outstanding CoInit so we CoUninit.
+            if (hr > 0)
+                UninitializeCom();
+        }
+
+        private static void UninitializeCom()
+        {
+            if (!t_comInitializedByUs)
+                return;
+
+#if ENABLE_WINRT
+            Interop.WinRT.RoUninitialize();
+#else
+            Interop.Ole32.CoUninitialize();
+#endif
+            t_comInitializedByUs = false;
+        }
+
         // TODO: https://github.com/dotnet/corefx/issues/20766
-        public bool TrySetApartmentState(ApartmentState state) { throw new PlatformNotSupportedException(); }
         public void DisableComObjectEagerCleanup() { }
         public void Interrupt() { throw new PlatformNotSupportedException(); }
 
diff --git a/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.cs b/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.cs
index 4935750d5..b3e9a14fe 100644
--- a/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.cs
+++ b/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.cs
@@ -112,7 +112,7 @@ namespace Internal.Runtime.Augments
 
             if (threadPoolThread)
             {
-                RoInitialize();
+                InitializeCom();
             }
 
             return currentThread;
@@ -151,16 +151,6 @@ namespace Internal.Runtime.Augments
         }
 
         /// <summary>
-        /// Ensures the Windows Runtime is initialized on the current thread.
-        /// </summary>
-        internal static void RoInitialize()
-        {
-#if ENABLE_WINRT
-            Interop.WinRT.RoInitialize();
-#endif
-        }
-
-        /// <summary>
         /// Returns true if the underlying OS thread has been created and started execution of managed code.
         /// </summary>
         private bool HasStarted()
@@ -467,7 +457,7 @@ namespace Internal.Runtime.Augments
             {
                 t_currentThread = thread;
                 System.Threading.ManagedThreadId.SetForCurrentThread(thread._managedThreadId);
-                RoInitialize();
+                thread.InitializeComOnNewThread();
             }
             catch (OutOfMemoryException)
             {
diff --git a/src/System.Private.CoreLib/src/Interop/Interop.WinRT.cs b/src/System.Private.CoreLib/src/Interop/Interop.WinRT.cs
index 741ce4f21..0823dcffc 100644
--- a/src/System.Private.CoreLib/src/Interop/Interop.WinRT.cs
+++ b/src/System.Private.CoreLib/src/Interop/Interop.WinRT.cs
@@ -11,26 +11,13 @@ internal partial class Interop
     {
         private const string CORE_WINRT = "api-ms-win-core-winrt-l1-1-0.dll";
 
-        private const int RPC_E_CHANGED_MODE = unchecked((int)0x80010106);
+        internal const uint RO_INIT_SINGLETHREADED = 0;
+        internal const uint RO_INIT_MULTITHREADED = 1;
 
-        private enum RO_INIT_TYPE : uint
-        {
-            RO_INIT_MULTITHREADED = 1
-        }
-
-        internal static void RoInitialize()
-        {
-            int hr = RoInitialize((uint)RO_INIT_TYPE.RO_INIT_MULTITHREADED);
-
-            // RPC_E_CHANGED_MODE indicates this thread has been already initialized with a different
-            // concurrency model. That is fine; we just need to skip the RoUninitialize call on shutdown.
-            if ((hr < 0) && (hr != RPC_E_CHANGED_MODE))
-            {
-                throw new OutOfMemoryException();
-            }
-        }
+        [DllImport(CORE_WINRT, ExactSpelling = true)]
+        internal static extern int RoInitialize(uint initType);
 
         [DllImport(CORE_WINRT, ExactSpelling = true)]
-        private static extern int RoInitialize(uint initType);
+        internal static extern int RoUninitialize();
     }
 }
diff --git a/src/System.Private.CoreLib/src/Interop/Interop.manual.cs b/src/System.Private.CoreLib/src/Interop/Interop.manual.cs
index 261aad609..b54b15522 100644
--- a/src/System.Private.CoreLib/src/Interop/Interop.manual.cs
+++ b/src/System.Private.CoreLib/src/Interop/Interop.manual.cs
@@ -12,11 +12,6 @@ internal partial class Interop
         WaitObject0 = 0x0u,
         FailFastGenerateExceptionAddress = 0x1u,
         ExceptionNonContinuable = 0x1u,
-        CreateMutexInitialOwner = 0x1u,
-        CreateEventManualReset = 0x1u,
-        CreateEventInitialSet = 0x2u,
-        SemaphoreModifyState = 0x2u,
-        EventModifyState = 0x2u,
         DuplicateSameAccess = 0x2u,
         CreateSuspended = 0x4u,
         WaitAbandoned0 = 0x80u,
@@ -25,7 +20,6 @@ internal partial class Interop
         WaitFailed = 0xFFFFFFFFu,
     }
 
-    // MCG doesn't currently support constants that are not uint.
     internal static IntPtr InvalidHandleValue => new IntPtr(-1);
 
 #pragma warning disable 649
diff --git a/src/System.Private.CoreLib/src/System.Private.CoreLib.csproj b/src/System.Private.CoreLib/src/System.Private.CoreLib.csproj
index 923a7bfd4..5bacc1219 100644
--- a/src/System.Private.CoreLib/src/System.Private.CoreLib.csproj
+++ b/src/System.Private.CoreLib/src/System.Private.CoreLib.csproj
@@ -411,6 +411,12 @@
     <Compile Include="..\..\Common\src\Interop\Windows\ole32\Interop.CoTaskMemAllocFree.cs">
       <Link>Interop\Windows\ole32\Interop.CoTaskMemAllocFree.cs</Link>
     </Compile>
+    <Compile Include="..\..\Common\src\Interop\Windows\ole32\Interop.CoInitializeEx.cs">
+      <Link>Interop\Windows\ole32\Interop.CoInitializeEx.cs</Link>
+    </Compile>
+    <Compile Include="..\..\Common\src\Interop\Windows\ole32\Interop.CoUninitialize.cs">
+      <Link>Interop\Windows\ole32\Interop.CoUninitialize.cs</Link>
+    </Compile>
     <Compile Include="..\..\Common\src\Interop\Windows\ole32\Interop.CoGetApartmentType.cs">
       <Link>Interop\Windows\ole32\Interop.CoGetApartmentType.cs</Link>
     </Compile>
diff --git a/src/System.Private.CoreLib/src/System/Runtime/InitializeFinalizerThread.cs b/src/System.Private.CoreLib/src/System/Runtime/InitializeFinalizerThread.cs
index 2b0e67bc8..cc52fda24 100644
--- a/src/System.Private.CoreLib/src/System/Runtime/InitializeFinalizerThread.cs
+++ b/src/System.Private.CoreLib/src/System/Runtime/InitializeFinalizerThread.cs
@@ -13,9 +13,9 @@ namespace System.Runtime
         [RuntimeExport("InitializeFinalizerThread")]
         public static void DoInitialize()
         {
-            // Make sure that the finalizer thread is RoInitialized before any objects are finalized.  If this
+            // Make sure that the finalizer thread is CoInitialized before any objects are finalized.  If this
             // fails, it will throw an exception and that will go unhandled, triggering a FailFast.
-            RuntimeThread.RoInitialize();
+            RuntimeThread.InitializeCom();
         }
     }
 }