testfx icon indicating copy to clipboard operation
testfx copied to clipboard

STA-Thread lost after async operation

Open SamuelMuellerKMS opened this issue 7 months ago • 3 comments

Describe the bug

When running a STATestMethod and calling a async Method, it's not syncing back to the STA Thread

Steps To Reproduce

Execute an async method in a STATestMethod. After the Execution you will not return on the STA-Thread.

[STATestMethod]
public async Task StaTestMethodIssueTest()
{
	Console.WriteLine($"Current Thread Apartment State: {Thread.CurrentThread.GetApartmentState()}");

	Console.WriteLine("Executing async Method...");
	await Task.Yield();
		
	Console.WriteLine($"Current Thread Apartment State: {Thread.CurrentThread.GetApartmentState()}");
	Assert.IsTrue(Thread.CurrentThread.GetApartmentState() == ApartmentState.STA, "Thread is not in Single Threaded Apartment (STA) mode");
}

Example Project: STATestIssue.zip

Expected behavior

After executing an async Method you should return to the calling thread (or not when .ConfigureAwait(false) is present)

Actual behavior

The STA Thread Context is lost and you get Exceptions when calling code that requires an STA-Thread

SamuelMuellerKMS avatar Jun 11 '25 07:06 SamuelMuellerKMS

Thanks for reporting this @SamuelMuellerKMS. Is this somehow a regression? Or that's how it always behaved?

The current behavior makes sense to me. MSTest doesn't (and shouldn't) plug any custom SynchronizationContext. Without a synchronization context, I don't see how you would expect the continuations to get back to an STA thread (whether it's the same thread or a different one).

If you need special needs, I think you should plug your own SynchronizationContext.

Youssef1313 avatar Jun 11 '25 07:06 Youssef1313

We may actually consider changing STAThreadAttribute behavior so that it either sets a single threaded synchronization context so that continuations are pushed back to the original thread (needs careful implementation, and may risk deadlocks), or alternatively, a synchronization context that always posts continuations to a new STA thread.

Youssef1313 avatar Jun 12 '25 17:06 Youssef1313

@SamuelMuellerKMS Meanwhile, you can try something like the following, which implements a new custom attribute, ExSTATestMethod that uses a single threaded synchronization context.

    internal sealed class ExSTATestMethodAttribute : TestMethodAttribute
    {
        private sealed class SingleThreadSynchronizationContext : SynchronizationContext
        {
            private readonly BlockingCollection<(SendOrPostCallback, object)> _queue = new();
            private readonly Thread _thread;

            public SingleThreadSynchronizationContext()
            {
                _thread = new Thread(() =>
                {
                    SetSynchronizationContext(this);
                    foreach (var (callback, state) in _queue.GetConsumingEnumerable())
                    {
                        callback(state);
                    }
                })
                {
                    IsBackground = true
                };
                _thread.SetApartmentState(ApartmentState.STA);
                _thread.Start();
            }

            public override void Post(SendOrPostCallback d, object state)
            {
                _queue.Add((d, state));
            }

            public override void Send(SendOrPostCallback d, object state)
            {
                if (Environment.CurrentManagedThreadId == _thread.ManagedThreadId)
                {
                    d(state);
                }
                else
                {
                    using var done = new ManualResetEventSlim();
                    _queue.Add((s =>
                    {
                        d(s);
                        done.Set();
                    }, state));
                    done.Wait();
                }
            }

            public void Complete() => _queue.CompleteAdding();
        }


        public override TestResult[] Execute(ITestMethod testMethod)
        {
            return ExecuteInner(testMethod).GetAwaiter().GetResult();
        }

        private async Task<TestResult[]> ExecuteInner(ITestMethod testMethod)
        {
            var originalContext = SynchronizationContext.Current;
            try
            {
                var syncContext = new SingleThreadSynchronizationContext();
                SynchronizationContext.SetSynchronizationContext(syncContext);
                await Task.Yield();
                var t = await testMethod.InvokeAsync(null);
                syncContext.Complete();
                return [t];
            }
            finally
            {
                SynchronizationContext.SetSynchronizationContext(originalContext);
            }
        }
    }

@nohwnd Does the above implementation sound reasonable to you?

Youssef1313 avatar Jun 12 '25 17:06 Youssef1313

@AArnott Can I have your thoughts please on modifying MSTest's STATestMethodAttribute per the above implementation? Today, what we do is very simple, we simply create a new STA thread, and run the test method (we do .GetAwaiter().GetResult()). This means that the test method is only STA until the first await is encountered.

I'm thinking we could (as a breaking change in MSTest v4) fix the implementation to set a custom synchronization context, as indicated above.

Youssef1313 avatar Jun 27 '25 10:06 Youssef1313

I've never used [STATestMethod] in MSTest, so I can't speak on compatibility for your proposed change. I can tell you that [StaFact] is probably one of the least popular attributes in Xunit.StaFact for the reason that it isn't 'sticky' across awaits though. It also ties tests to Windows.

I generally recommend (when no specific GUI framework applies) that folks use [UIFact] in their tests when using Xunit with my xunit.stafact library, since that is xplat but creates an STA thread on Windows, and is 'sticky'.

You can see my implementation of SynchronizationContext that makes this sticky here:

https://github.com/AArnott/Xunit.StaFact/blob/main/src/Xunit.StaFact/Sdk/UISynchronizationContext.cs

TLDR: mine is significantly longer than yours. While some of that is due to xunit-specific interactions, you may still learn some things to add to your proposed version.

AArnott avatar Jun 27 '25 13:06 AArnott

I think i'm hitting the same issue.

I have some WPF testing, the Output Log shows (notice the ExecutionThreadApartmentState as STA)

[6/27/2025 3:05:12.920 PM]  Tests run settings for My.dll:
 <RunSettings>
  <!-- Configurations that affect the Test Framework -->
  <RunConfiguration>
    <!-- [x86] | x64  
      - You can also change it from menu Test, Test Settings, Default Processor Architecture -->
    <TargetPlatform>x64</TargetPlatform>
    <!-- Framework35 | [Framework40] | Framework45 -->
    <MaxCpuCount>0</MaxCpuCount>
    <DisableParallelization>true</DisableParallelization>
    <ExecutionThreadApartmentState>STA</ExecutionThreadApartmentState>
    <ResultsDirectory><bla>\TestResults</ResultsDirectory>
    <SolutionDirectory><bla></SolutionDirectory>
    <CollectSourceInformation>False</CollectSourceInformation>
  </RunConfiguration>
  <!-- MSTest adapter -->
  <MSTest>
    <DeleteDeploymentDirectoryAfterTestRunIsComplete>True</DeleteDeploymentDirectoryAfterTestRunIsComplete>
    <DeploymentEnabled>True</DeploymentEnabled>
  </MSTest>
  <Python>
    <TestCases />
  </Python>
</RunSettings>.

But some tests fails (in visual studio or on the build machine)

Image

i can reproduce with the following test class with a unit test configured to use WPF while the context is setup to STA :

[TestClass]
[STATestClass]
public sealed class Test1
{
    [TestMethod]
    public void TestMethod1()
    {
        _ = new Window();
    }

    [TestMethod]
    public async Task  TestMethod2()
    {
        await Task.Yield();
    }

    [TestMethod]
    public void TestMethod3()
    {
        _ = new Window();
    }
}

Image

@Youssef1313 this one is an impediment for us to execute our WPF test, note that the TestMethod2 may be completely unrelated to WPF and simply test a view model

fforjanAVEVA avatar Jun 27 '25 22:06 fforjanAVEVA

Thanks @fforjanAVEVA. As mentioned before, today's behavior of STA is that it just starts the test method in a STA thread and it doesn't set a synchronization context that schedules continuations on STA as well.

For now you can try out the ExSTATestMethodAttribute implementation that I pasted above. In MSTest v4, I'm considering having STATestMethodAttribute to, by default, set synchronization context.

Youssef1313 avatar Jun 28 '25 03:06 Youssef1313

I've never used [STATestMethod] in MSTest, so I can't speak on compatibility for your proposed change. I can tell you that [StaFact] is probably one of the least popular attributes in Xunit.StaFact for the reason that it isn't 'sticky' across awaits though. It also ties tests to Windows.

I generally recommend (when no specific GUI framework applies) that folks use [UIFact] in their tests when using Xunit with my xunit.stafact library, since that is xplat but creates an STA thread on Windows, and is 'sticky'.

You can see my implementation of SynchronizationContext that makes this sticky here:

https://github.com/AArnott/Xunit.StaFact/blob/main/src/Xunit.StaFact/Sdk/UISynchronizationContext.cs

TLDR: mine is significantly longer than yours. While some of that is due to xunit-specific interactions, you may still learn some things to add to your proposed version.

@AArnott It's so interesting to see the pumping mechanism you have. I'm wondering though what benefits/differences it has compared to using a simple BlockingCollection like the sample I have above. Performance? Something else?

Youssef1313 avatar Jun 28 '25 03:06 Youssef1313

@Youssef1313 in our case we are NOT using the STATestClass or any other attribtue we have configured the ExecutionThreadApartmentState into the runconfig, so any attribut will not help us. We are expecting the configuration to be enough

fforjanAVEVA avatar Jun 28 '25 14:06 fforjanAVEVA

@fforjanAVEVA You could still plug the synchronization context in assembly initialize, and hopefully it should work? My implementation above might need to be tweaked if it's to be used with assembly init, esp if you run in parallel, or if you want to ensure proper isolation between tests, for example, you may need to make sure that async locals won't get shared between tests if you plug a global synchronization context in assembly initialize - I understand that stuff like synchronization context can be complex to deal with, so I don't disagree that it's something MSTest itself should offer out of the box, maybe with some option, to be decided if it's opt-in or opt-out, to choose between new behavior and existing behavior.

Youssef1313 avatar Jun 28 '25 17:06 Youssef1313

@Youssef1313 is the option ExecutionThreadApartmentState from the config file ? if not what is the use of this option ? It got introduce to ensure tests where running on the same thread recently ?

fforjanAVEVA avatar Jun 30 '25 13:06 fforjanAVEVA

Today, ExecutionThreadApartmentState won't handle async properly. It only ensures that the tests are run in STA, without setting a SynchronizationContext that ensures you will get back to STA as well. This is something we can look into changing in v4, but you might be able to workaround it via a custom synchronization context that you plug in assembly initialize.

Youssef1313 avatar Jun 30 '25 14:06 Youssef1313

@Youssef1313 the funny part, is we are currently using mstest 3.8.3 and we do not have this issue, either locally or on the build machine, but fails immediately when using mstest 3.9.3

fforjanAVEVA avatar Jun 30 '25 14:06 fforjanAVEVA

This should have never worked, even in 3.8.3. Can you show an example of how it works in 3.8.3?

Following is the behavior I see in 3.8.3, which makes sense per the current design of the feature.

Image

After the await, the continuation happens on a threadpool thread (because that's where the task completes)

Youssef1313 avatar Jun 30 '25 15:06 Youssef1313

@Youssef1313 Sorry if I wasn't clear I was refering to my sample with the 3 tests methods, not from the same method:

[TestClass]
[STATestClass]
public sealed class Test1
{
    [TestMethod]
    public void TestMethod1()
    {
        _ = new Window();
    }

    [TestMethod]
    public async Task  TestMethod2()
    {
        await Task.Yield();
    }

    [TestMethod]
    public void TestMethod3()
    {
        _ = new Window();
    }
}

fforjanAVEVA avatar Jun 30 '25 15:06 fforjanAVEVA

Ah okay. I misinterpreted that. In that case, yes there could have been slight change in behavior. I need to double check that. I assume you are not enabling test parallelism right?

Youssef1313 avatar Jun 30 '25 15:06 Youssef1313

@Youssef1313 we have <DisableParallelization>true</DisableParallelization> in our runsettings

fforjanAVEVA avatar Jun 30 '25 15:06 fforjanAVEVA

In a framework where an attribute does not exist that keeps async test methods on the STA thread, the test can always apply its own SynchronizationContext that achieves the same thing.

This class for example would do the trick: https://microsoft.github.io/vs-threading/api/Microsoft.VisualStudio.Threading.SingleThreadedSynchronizationContext.html

@Youssef1313 Your use of BlockingCollection looks like it simplifies things nicely. I haven't thoroughly compared our two implementations to make a list of differences. But one I notice is that if a delegate is given to your Send method that throws, your caller will hang, because you don't set the ManualResetEvent in a finally block.

AArnott avatar Jun 30 '25 16:06 AArnott

I followed the pattern of using the SingleThreadedSynchronizationContext but this does not seem to work with the [TestInitialize] method being async (in 3.09 at least). If I await a task it can deadlock.

seanhalliday avatar Aug 15 '25 03:08 seanhalliday

@seanhalliday Can you please upload a complete repro project? Also did you try MSTest v4 preview?

Youssef1313 avatar Aug 15 '25 03:08 Youssef1313

Sorry, my mistake, it is when we use the WPF dispatcher. It will be some effort to extract a test project but I will try.

seanhalliday avatar Aug 15 '25 03:08 seanhalliday

using System.Windows;

namespace TestProject1
{
	[TestClass]
	public sealed class Test1
	{
		[AssemblyInitialize]
		public static void AssemblyInit(TestContext context)
		{
			System.Windows.Application app = new System.Windows.Application()
			{
				ShutdownMode = ShutdownMode.OnExplicitShutdown
			};
		}

		[TestInitialize]
		public async Task TestInit()
		{
			// Hangs here in 3.10.2
			await Task.Delay(10);
		}

		[UITestMethod]
		public async Task TestMethod1()
		{
			await Task.Delay(10);
		}
	}
}
------------------------------------------------------------------------------------
using System.Windows;
using System.Windows.Threading;

namespace TestProject1
{

	public class UITestMethodAttribute : TestMethodAttribute
	{
		public override TestResult?[] Execute(ITestMethod testMethod)
		{
			if (typeof(Task).IsAssignableFrom(testMethod.MethodInfo.ReturnType))
			{
				TestResult? result = null;

				Application.Current.Dispatcher.Invoke(() =>
				{
					var task = testMethod.InvokeAsync(null);
					var dispatcherFrame = new DispatcherFrame();
					task.ContinueWith(_ => dispatcherFrame.Continue = false);
					Dispatcher.PushFrame(dispatcherFrame);
					result = task.Result;
				});

				return new[] { result };
			}

			return new[] { testMethod.Invoke(null) };
		}
	}
}

seanhalliday avatar Aug 15 '25 16:08 seanhalliday

We can work around it by having our TestInit be void and use a helper to run a async function like the above with the dispatcher frame code.

		[TestInitialize]
		public void TestInit()
		{
			WPFTestHelper.WaitForTask(InitAsync());
		}

seanhalliday avatar Aug 15 '25 16:08 seanhalliday