Add JoinableTaskFactory.DisableProcessing()

.github/copilot-instructions.md

@@ -82,5 +82,6 @@ dotnet run --no-build -c Release --framework net9.0 -- --list-tests
 ## Coding style
 
 * Honor StyleCop rules and fix any reported build warnings *after* getting tests to pass.
-* In C# files, use namespace *statements* instead of namespace *blocks* for all new files.
+* In C# files, use namespace *statements* instead of namespace *blocks* for all new files that define namespaces.
+* Test files are *not* expected to declare namespaces.
 * Add API doc comments to all new public and internal members.

samples/DisableProcessing.cs

@@ -0,0 +1,50 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+#pragma warning disable VSTHRD103 // Call async methods when in an async method
+
+using System.IO;
+using Microsoft.VisualStudio.Threading;
+
+internal class DisableProcessing
+{
+    private readonly JoinableTaskFactory joinableTaskFactory = null!;
+
+    private void Simple()
+    {
+        #region Simple
+        this.joinableTaskFactory.Run(async delegate
+        {
+            this.joinableTaskFactory.DisableProcessing();
+
+            // Synchronous I/O and lock contentions will NOT result in any reentrancy within this JoinableTask.
+        });
+        #endregion
+    }
+
+    private void Exhaustive()
+    {
+        #region Exhaustive
+        this.joinableTaskFactory.Run(async delegate
+        {
+            // Async I/O isn't expected to synchronously block, and thus would never allow unwanted reentrancy.
+            string content = await File.ReadAllTextAsync(@"somefile.txt");
+
+            // Here, synchronous I/O and lock contentions MAY allow certain reentrancy (e.g. COM RPC messages).
+            content = File.ReadAllText(@"somefile.txt");
+
+            using (this.joinableTaskFactory.DisableProcessing())
+            {
+                // Within this block, synchronous I/O and lock contentions will NOT result in any reentrancy.
+                content = File.ReadAllText(@"somefile.txt");
+            }
+
+            // Just disable the synchronous wait message pump for the rest of this JoinableTask.
+            this.joinableTaskFactory.DisableProcessing();
+
+            // Sync I/O and lock contentions will NOT result in any reentrancy here.
+            content = File.ReadAllText(@"somefile.txt");
+        });
+        #endregion
+    }
+}

samples/Polyfill.cs

@@ -0,0 +1,18 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+#if NETFRAMEWORK
+
+using System;
+using System.IO;
+using System.Threading.Tasks;
+
+internal static class PolyfillExtensions
+{
+    extension(File)
+    {
+        internal static Task<string> ReadAllTextAsync(string path) => throw new NotImplementedException();
+    }
+}
+
+#endif

samples/ApiSamples.cs → samples/SuppressRelevance.cs

@@ -1,11 +1,10 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
-using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.VisualStudio.Threading;
 
-public class SuppressRelevanceSample
+public class SuppressRelevance
 {
     private readonly ReentrantSemaphore semaphore = ReentrantSemaphore.Create(1, null, ReentrantSemaphore.ReentrancyMode.NotAllowed);
 
@@ -19,7 +18,7 @@ await this.semaphore.ExecuteAsync(async delegate
             await Task.Yield(); // represents some async work
 
             // Fire and forget code that uses the semaphore, but should *not*
-            // inherit our own posession of the semaphore.
+            // inherit our own possession of the semaphore.
             using (this.semaphore.SuppressRelevance())
             {
                 this.DoSomethingLaterAsync().Forget(); // Don't await this, or a deadlock will occur.

src/Microsoft.VisualStudio.Threading/DispatcherExtensions.cs

@@ -27,6 +27,13 @@ public static class DispatcherExtensions
     /// and for each asynchronous return to the main thread after an <see langword="await"/>.
     /// </param>
     /// <returns>A <see cref="JoinableTaskFactory"/> that may be used for scheduling async work with the specified priority.</returns>
+    /// <remarks>
+    /// In addition to scheduling work on the UI thread with the specified priority,
+    /// this <see cref="JoinableTaskFactory"/> also ensures that any synchronous waits
+    /// on the main thread within <see cref="JoinableTask" /> objects created with the returned factory
+    /// will honor calls to <see cref="Dispatcher.DisableProcessing()" />, producing similar behavior
+    /// to <see cref="JoinableTaskFactory.DisableProcessing()" />.
+    /// </remarks>
     public static JoinableTaskFactory WithPriority(this JoinableTaskFactory joinableTaskFactory, Dispatcher dispatcher, DispatcherPriority priority)
     {
         Requires.NotNull(joinableTaskFactory, nameof(joinableTaskFactory));
@@ -61,6 +68,9 @@ internal DispatcherJoinableTaskFactory(JoinableTaskFactory innerFactory, Dispatc
             : base(innerFactory)
         {
             this.dispatcher = dispatcher ?? throw new ArgumentNullException(nameof(dispatcher));
+#if NETFRAMEWORK // Avoid trim warnings on .NET, and only .NET Framework calls CoWait anyway.
+            this.DefaultWaitPolicy = new DispatcherSynchronizationContext(dispatcher);
+#endif
             this.priority = priority;
         }
 

src/Microsoft.VisualStudio.Threading/JoinableTask+JoinableTaskSynchronizationContext.cs

@@ -2,11 +2,11 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Text;
+using System.Runtime.InteropServices;
 using System.Threading;
 using System.Threading.Tasks;
+using Windows.Win32;
+using Windows.Win32.Foundation;
 
 namespace Microsoft.VisualStudio.Threading
 {
@@ -44,6 +44,11 @@ internal JoinableTaskSynchronizationContext(JoinableTaskFactory owner)
 
                 this.jobFactory = owner;
                 this.mainThreadAffinitized = true;
+
+                if (owner.DefaultWaitPolicy is not null)
+                {
+                    this.SetWaitNotificationRequired();
+                }
             }
 
             /// <summary>
@@ -56,6 +61,11 @@ internal JoinableTaskSynchronizationContext(JoinableTask joinableTask, bool main
             {
                 this.job = joinableTask;
                 this.mainThreadAffinitized = mainThreadAffinitized;
+
+                if (joinableTask.DisableProcessing > 0)
+                {
+                    this.SetWaitNotificationRequired();
+                }
             }
 
             /// <summary>
@@ -134,6 +144,50 @@ public override void Send(SendOrPostCallback d, object? state)
                 }
             }
 
+            /// <summary>
+            /// Synchronously blocks without a message pump.
+            /// </summary>
+            /// <param name="waitHandles">An array of type <see cref="IntPtr" /> that contains the native operating system handles.</param>
+            /// <param name="waitAll">true to wait for all handles; false to wait for any handle.</param>
+            /// <param name="millisecondsTimeout">The number of milliseconds to wait, or <see cref="Timeout.Infinite" /> (-1) to wait indefinitely.</param>
+            /// <returns>
+            /// The array index of the object that satisfied the wait.
+            /// </returns>
+            public override unsafe int Wait(IntPtr[] waitHandles, bool waitAll, int millisecondsTimeout)
+            {
+                Requires.NotNull(waitHandles, nameof(waitHandles));
+
+                if (this.job?.DisableProcessing > 0)
+                {
+                    // On .NET Framework we must take special care to NOT end up in a call to CoWait (which lets in RPC calls).
+                    // Off Windows, we can't p/invoke to kernel32, but it appears that .NET never calls CoWait, so we can rely on default behavior.
+                    // We're just going to use the OS as the switch instead of the runtime so that (one day) if we drop our .NET Framework specific target,
+                    // and if .NET ever adds CoWait support on Windows, we'll still behave properly.
+#if NET
+                    if (OperatingSystem.IsWindowsVersionAtLeast(5, 1, 2600))
+#else
+                    if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
+#endif
+                    {
+                        fixed (IntPtr* pHandles = waitHandles)
+                        {
+                            return (int)PInvoke.WaitForMultipleObjects((uint)waitHandles.Length, (HANDLE*)pHandles, waitAll, (uint)millisecondsTimeout);
+                        }
+                    }
+                }
+
+                // Use a surrogate default policy if provided.
+                if (this.jobFactory.DefaultWaitPolicy is { } waitPolicy)
+                {
+                    return waitPolicy.Wait(waitHandles, waitAll, millisecondsTimeout);
+                }
+
+                // Fallback to sync blocking such that CoWait might be called.
+                return WaitHelper(waitHandles, waitAll, millisecondsTimeout);
+            }
+
+            internal void ConsiderDisableProcessing() => this.SetWaitNotificationRequired();
+
             /// <summary>
             /// Called by the joinable task when it has completed.
             /// </summary>

src/Microsoft.VisualStudio.Threading/JoinableTask.cs

@@ -372,6 +372,23 @@ internal SynchronizationContext? ApplicableJobSyncContext
         }
     }
 
+    /// <summary>
+    /// Gets or sets a value indicating whether CoWait will be prohibited
+    /// during synchronously blocking waits from code actively running within this <see cref="JoinableTask"/>.
+    /// </summary>
+    internal int DisableProcessing
+    {
+        get => field;
+        set
+        {
+            field = value;
+            if (this.mainThreadJobSyncContext is { } syncContext)
+            {
+                syncContext.ConsiderDisableProcessing();
+            }
+        }
+    }
+
     /// <summary>
     /// Gets a weak reference to this object.
     /// </summary>

src/Microsoft.VisualStudio.Threading/JoinableTaskFactory.cs

@@ -98,6 +98,17 @@ internal JoinableTaskCollection? Collection
         get { return this.jobCollection; }
     }
 
+    /// <summary>
+    /// Gets a <see cref="SynchronizationContext"/> on which <see cref="SynchronizationContext.Wait(IntPtr[], bool, int)"/>
+    /// should be called from <see cref="JoinableTaskSynchronizationContext.Wait(IntPtr[], bool, int)"/>
+    /// when <see cref="DisableProcessing()"/> has not been called.
+    /// </summary>
+    /// <remarks>
+    /// This allows a WPF-aware <see cref="JoinableTaskFactory"/>-derived class within this assembly
+    /// to match <c>Dispatcher.DisableProcessing()</c> behavior.
+    /// </remarks>
+    internal SynchronizationContext? DefaultWaitPolicy { get; init; }
+
     /// <summary>
     /// Gets or sets the timeout after which no activity while synchronously blocking
     /// suggests a hang has occurred.
@@ -294,6 +305,55 @@ public JoinableTask<T> RunAsync<T>(Func<Task<T>> asyncMethod, string? parentToke
         return this.RunAsync(asyncMethod, synchronouslyBlocking: false, parentToken, creationOptions: creationOptions);
     }
 
+#pragma warning disable SA1629 // Documentation text should end with a period
+    /// <summary>
+    /// Prevents filtered message pumps from running during synchronous waits for the ambient <see cref="JoinableTask"/>.
+    /// </summary>
+    /// <returns>
+    /// A value that may be disposed of when the need to suppress synchronous wait message pumps is ended.
+    /// Alternatively it may be discarded if the rest of the <see cref="JoinableTask"/> is intended to have processing disabled.
+    /// </returns>
+    /// <exception cref="InvalidOperationException">Thrown when called outside the context of a <see cref="JoinableTask"/>.</exception>
+    /// <remarks>
+    /// <para>
+    /// During a yielding <see langword="await"/> within a <see cref="JoinableTask"/>, no message pump ever runs
+    /// regardless of whether this method is called, except for the internal one that lets in only relevant work.
+    /// When user code runs within the <see cref="JoinableTask"/> delegate or its callees that ends up requiring
+    /// a synchronous block of the main thread (e.g. synchronous I/O or lock contention), this wait is typically
+    /// implemented by calling <see cref="SynchronizationContext.Wait(IntPtr[], bool, int)"/> on <see cref="SynchronizationContext.Current"/>.
+    /// The default implementation of this method allows for certain interruptions (e.g. COM RPC calls), which
+    /// <em>may</em> avoid deadlocks in certain situations.
+    /// </para>
+    /// <para>
+    /// Calling this method will replace the default implementation of <see cref="SynchronizationContext.Wait(IntPtr[], bool, int)"/>
+    /// with one that will not allow such interruptions while that <see cref="JoinableTask"/> is active and in control of
+    /// <see cref="SynchronizationContext.Current"/>.
+    /// As this method may be called multiple times, this effect remains on the target <see cref="JoinableTask"/>
+    /// until all <see cref="DisableProcessing"/> invocations' return values are disposed (in any order).
+    /// The effect only applies to the direct <see cref="JoinableTask"/>. It does not affect any of its children or parents.
+    /// </para>
+    /// <para>
+    /// Disabling processing has no effect on non-Windows operating systems.
+    /// </para>
+    /// <para>
+    /// Disposing the resulting value will revert to the default behavior.
+    /// Callers need not ever dispose of this value if the intent is to disable processing for the remainder of that
+    /// <see cref="JoinableTask"/>'s execution.
+    /// </para>
+    /// </remarks>
+    /// <example>
+    /// <para>
+    /// Here is a simple, common usage of this method:
+    /// </para>
+    /// <code source="../../samples/DisableProcessing.cs" region="Simple" lang="C#" />
+    /// <para>
+    /// Following are more examples of how it might be used:
+    /// </para>
+    /// <code source="../../samples/DisableProcessing.cs" region="Exhaustive" lang="C#" />
+    /// </example>
+    public ProcessingDisabledOperation DisableProcessing() => new(this.Context.AmbientTask ?? throw new InvalidOperationException(Strings.NoAmbientTask));
+#pragma warning restore SA1629 // Documentation text should end with a period
+
     /// <summary>
     /// Responds to calls to <see cref="JoinableTaskFactory.MainThreadAwaiter.OnCompleted(Action)"/>
     /// by scheduling a continuation to execute on the Main thread.
@@ -682,6 +742,34 @@ static bool FailFast(Exception ex)
         }
     }
 
+    /// <summary>
+    /// A struct whose disposal will revert the effect of an earlier call to <see cref="DisableProcessing"/>.
+    /// </summary>
+    public struct ProcessingDisabledOperation : IDisposable
+    {
+        private JoinableTask? owner;
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="ProcessingDisabledOperation"/> struct.
+        /// </summary>
+        /// <param name="owner">The owner of this struct.</param>
+        internal ProcessingDisabledOperation(JoinableTask owner)
+        {
+            owner.DisableProcessing++;
+            this.owner = owner;
+        }
+
+        /// <inheritdoc/>
+        public void Dispose()
+        {
+            if (this.owner is { } owner)
+            {
+                owner.DisableProcessing--;
+                this.owner = null;
+            }
+        }
+    }
+
     /// <summary>
     /// An awaitable struct that facilitates an asynchronous transition to the Main thread.
     /// </summary>

src/Microsoft.VisualStudio.Threading/NoMessagePumpSyncContext.cs

@@ -2,7 +2,6 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System;
-using System.Buffers;
 using System.Runtime.InteropServices;
 using System.Threading;
 using global::Windows.Win32;
@@ -52,10 +51,10 @@ public override unsafe int Wait(IntPtr[] waitHandles, bool waitAll, int millisec
         Requires.NotNull(waitHandles, nameof(waitHandles));
 
         // On .NET Framework we must take special care to NOT end up in a call to CoWait (which lets in RPC calls).
-        // Off Windows, we can't p/invoke to kernel32, but it appears that .NET Core never calls CoWait, so we can rely on default behavior.
-        // We're just going to use the OS as the switch instead of the framework so that (one day) if we drop our .NET Framework specific target,
-        // and if .NET Core ever adds CoWait support on Windows, we'll still behave properly.
-#if NET5_0_OR_GREATER
+        // Off Windows, we can't p/invoke to kernel32, but it appears that .NET never calls CoWait, so we can rely on default behavior.
+        // We're just going to use the OS as the switch instead of the runtime so that (one day) if we drop our .NET Framework specific target,
+        // and if .NET ever adds CoWait support on Windows, we'll still behave properly.
+#if NET
         if (OperatingSystem.IsWindowsVersionAtLeast(5, 1, 2600))
 #else
         if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))

src/Microsoft.VisualStudio.Threading/ReentrantSemaphore.cs

@@ -191,7 +191,7 @@ public static ReentrantSemaphore Create(int initialCount = 1, JoinableTaskContex
     /// <para>
     /// The following snippet demonstrates a way to use this method.
     /// </para>
-    /// <code source="../../samples/ApiSamples.cs" region="SuppressRelevance" lang="C#" />
+    /// <code source="../../samples/SuppressRelevance.cs" region="SuppressRelevance" lang="C#" />
     /// </example>
     public virtual RevertRelevance SuppressRelevance() => default;
 

src/Microsoft.VisualStudio.Threading/Strings.resx

@@ -193,4 +193,7 @@
   <data name="SyncContextNotSet" xml:space="preserve">
     <value>No SynchronizationContext to reach the main thread has been set.</value>
   </data>
+  <data name="NoAmbientTask" xml:space="preserve">
+    <value>No JoinableTask is active.</value>
+  </data>
 </root>