From 73ae00f2bc8e3913d7de592e9bfee42d42b26d3f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Amaury=20Lev=C3=A9?= Date: Fri, 26 Jun 2026 16:39:10 +0200 Subject: [PATCH 1/3] perf: defer class-cleanup TestContextImplementation allocation to last test in class MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For every test except the last test in a class, the class-cleanup TestContextImplementation was allocated unconditionally at the start of RunSingleTestAsync (line 240), but then never used because isLastTestInClass was false. The context ctor copies the testContextProperties dictionary and registers a CancellationTokenRegistration, making it a non-trivial allocation. Move the allocation inside the `if (isLastTestInClass)` block. Safety invariant: the assembly-cleanup branch at line 279 reads executes only when ShouldRunEndOfAssemblyCleanup is true, which requires MarkClassComplete to have been called — and MarkClassComplete is called exclusively inside the isLastTestInClass block above, so testContextForClassCleanup is guaranteed non-null when line 279 is reached. Impact: for a test suite with C classes averaging K tests each, this saves C*(K-1) TestContextImplementation allocations — one dictionary copy + one CancellationTokenRegistration per non-last test per class. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../Execution/UnitTestRunner.cs | 11 +++-- .../Execution/UnitTestRunnerTests.cs | 46 +++++++++++++++++++ 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs b/src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs index fbf08aae43..66b81045e7 100644 --- a/src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs +++ b/src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs @@ -237,11 +237,13 @@ internal async Task RunSingleTestAsync(UnitTestElement unitTestEle } } - testContextForClassCleanup = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, testMethod.FullClassName, testContextProperties, messageLogger, testContextForTestExecution.Context.CurrentTestOutcome); - _classCleanupManager.MarkTestComplete(testMethod, out bool isLastTestInClass); if (isLastTestInClass) { + // Defer TestContextImplementation allocation to only the last test in each class, + // saving one dict-copy + CancellationTokenRegistration per non-last test. + testContextForClassCleanup = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, testMethod.FullClassName, testContextProperties, messageLogger, testContextForTestExecution.Context.CurrentTestOutcome); + if (testMethodInfo is not null) { // Flow properties set during AssemblyInitialize and ClassInitialize so the @@ -276,7 +278,10 @@ internal async Task RunSingleTestAsync(UnitTestElement unitTestEle if (testMethodInfo?.Parent.Parent.IsAssemblyInitializeExecuted == true && _classCleanupManager.ShouldRunEndOfAssemblyCleanup) { - testContextForAssemblyCleanup = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, null, testContextProperties, messageLogger, testContextForClassCleanup.Context.CurrentTestOutcome); + // testContextForClassCleanup is guaranteed non-null here: ShouldRunEndOfAssemblyCleanup + // becomes true only after MarkClassComplete, which is called exclusively inside the + // isLastTestInClass block above — where testContextForClassCleanup is allocated. + testContextForAssemblyCleanup = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, null, testContextProperties, messageLogger, testContextForClassCleanup!.Context.CurrentTestOutcome); // Flow properties set during AssemblyInitialize so the AssemblyCleanup method // observes them. Class-init properties are intentionally NOT flowed here because diff --git a/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs b/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs index 879021ff7e..dfc2582f9e 100644 --- a/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs +++ b/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs @@ -385,6 +385,46 @@ public async Task RunSingleTestWhenClassInitializeAlreadyFailedShouldTakeFastPat contextsAllocatedForSecondRun.Should().BeLessThan(contextsAllocatedForFirstRun); } + public async Task RunSingleTestShouldDeferClassCleanupContextAllocationToLastTestInClass() + { + Type type = typeof(DummyTestClassWithCleanupMethods); + TestMethod testMethod1 = CreateTestMethod("TestMethod", type.FullName!, "A", displayName: null); + TestMethod testMethod2 = CreateTestMethod("TestMethod2", type.FullName!, "A", displayName: null); + TestMethod testMethod3 = CreateTestMethod("TestMethod3", type.FullName!, "A", displayName: null); + var unitTestElement1 = new UnitTestElement(testMethod1); + var unitTestElement2 = new UnitTestElement(testMethod2); + var unitTestElement3 = new UnitTestElement(testMethod3); + var unitTestRunner = new UnitTestRunner(new MSTestSettings(), [unitTestElement1, unitTestElement2, unitTestElement3]); + + _testablePlatformServiceProvider.MockFileOperations.Setup(fo => fo.LoadAssembly("A")) + .Returns(Assembly.GetExecutingAssembly()); + + DummyTestClassWithCleanupMethods.ClassCleanupMethodBody = () => { }; + DummyTestClassWithCleanupMethods.AssemblyCleanupMethodBody = () => { }; + + // Run the first test to warm up assembly/class initialize (slow path, sets cached results). + TestResult[] firstResults = await unitTestRunner.RunSingleTestAsync(unitTestElement1, _testRunParameters, null!); + firstResults[0].Outcome.Should().Be(UnitTestOutcome.Passed); + + // Second test (non-last): both assembly-init and class-init take the fast path (no contexts), + // and the class-cleanup context is NOT allocated (deferred to the last test). + int countBeforeSecond = _testablePlatformServiceProvider.GetTestContextCallCount; + TestResult[] secondResults = await unitTestRunner.RunSingleTestAsync(unitTestElement2, _testRunParameters, null!); + secondResults[0].Outcome.Should().Be(UnitTestOutcome.Passed); + int allocationsForSecond = _testablePlatformServiceProvider.GetTestContextCallCount - countBeforeSecond; + + // Third test (last in class): takes fast paths for init, but allocates a class-cleanup context + // and an assembly-cleanup context. Allocation count must exceed the middle test's count. + int countBeforeThird = _testablePlatformServiceProvider.GetTestContextCallCount; + TestResult[] thirdResults = await unitTestRunner.RunSingleTestAsync(unitTestElement3, _testRunParameters, null!); + thirdResults[0].Outcome.Should().Be(UnitTestOutcome.Passed); + int allocationsForThird = _testablePlatformServiceProvider.GetTestContextCallCount - countBeforeThird; + + // The last-test run allocates more contexts because it creates the class/assembly cleanup contexts + // that are deferred from all non-last tests. + allocationsForThird.Should().BeGreaterThan(allocationsForSecond); + } + #endregion #region private helpers @@ -498,6 +538,12 @@ private class DummyTestClassWithCleanupMethods [TestMethod] public void TestMethod() => TestMethodBody?.Invoke(TestContext); + + [TestMethod] + public void TestMethod2() => TestMethodBody?.Invoke(TestContext); + + [TestMethod] + public void TestMethod3() => TestMethodBody?.Invoke(TestContext); } private class DummyTestClassAttribute : TestClassAttribute; From 0757a85d17aa17eab1640ad058d083d9760cd362 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 26 Jun 2026 16:57:59 +0200 Subject: [PATCH 2/3] Assert exact class-cleanup context allocation delta in deferred-cleanup test Address PR review feedback: replace the loose BeGreaterThan assertion with an exact delta check (allocationsForThird == allocationsForSecond + 2). This catches the regression the test guards against, where a non-last test mistakenly allocates the deferred class-cleanup context. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../Execution/UnitTestRunnerTests.cs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs b/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs index dfc2582f9e..7a50b5a48e 100644 --- a/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs +++ b/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs @@ -420,9 +420,11 @@ public async Task RunSingleTestShouldDeferClassCleanupContextAllocationToLastTes thirdResults[0].Outcome.Should().Be(UnitTestOutcome.Passed); int allocationsForThird = _testablePlatformServiceProvider.GetTestContextCallCount - countBeforeThird; - // The last-test run allocates more contexts because it creates the class/assembly cleanup contexts - // that are deferred from all non-last tests. - allocationsForThird.Should().BeGreaterThan(allocationsForSecond); + // The last-test run allocates exactly two more contexts than a non-last test: the class-cleanup + // context and the assembly-cleanup context, both of which are deferred from all non-last tests. + // Asserting the exact delta (rather than a loose "greater than") catches the regression this test + // guards against, where a non-last test mistakenly allocates the deferred class-cleanup context. + allocationsForThird.Should().Be(allocationsForSecond + 2); } #endregion From 5c615496e5732c633703f43b5d5a4a23975ed73b Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 26 Jun 2026 17:55:45 +0200 Subject: [PATCH 3/3] Address review nits: use DebugEx.Assert and CreateUnitTestRunner helper - Replace the null-forgiving operator on testContextForClassCleanup with a DebugEx.Assert, matching the codebase idiom for protocol invariants and keeping compiler nullability tracking intact. - Use the CreateUnitTestRunner helper in the new test instead of constructing UnitTestRunner directly, following the established pattern in sibling tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs | 3 ++- .../Execution/UnitTestRunnerTests.cs | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs b/src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs index 66b81045e7..d8fc5529d3 100644 --- a/src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs +++ b/src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs @@ -281,7 +281,8 @@ internal async Task RunSingleTestAsync(UnitTestElement unitTestEle // testContextForClassCleanup is guaranteed non-null here: ShouldRunEndOfAssemblyCleanup // becomes true only after MarkClassComplete, which is called exclusively inside the // isLastTestInClass block above — where testContextForClassCleanup is allocated. - testContextForAssemblyCleanup = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, null, testContextProperties, messageLogger, testContextForClassCleanup!.Context.CurrentTestOutcome); + DebugEx.Assert(testContextForClassCleanup is not null, "testContextForClassCleanup should not be null when running assembly cleanup."); + testContextForAssemblyCleanup = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, null, testContextProperties, messageLogger, testContextForClassCleanup.Context.CurrentTestOutcome); // Flow properties set during AssemblyInitialize so the AssemblyCleanup method // observes them. Class-init properties are intentionally NOT flowed here because diff --git a/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs b/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs index 7a50b5a48e..bfa71fc029 100644 --- a/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs +++ b/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs @@ -394,7 +394,7 @@ public async Task RunSingleTestShouldDeferClassCleanupContextAllocationToLastTes var unitTestElement1 = new UnitTestElement(testMethod1); var unitTestElement2 = new UnitTestElement(testMethod2); var unitTestElement3 = new UnitTestElement(testMethod3); - var unitTestRunner = new UnitTestRunner(new MSTestSettings(), [unitTestElement1, unitTestElement2, unitTestElement3]); + UnitTestRunner unitTestRunner = CreateUnitTestRunner([unitTestElement1, unitTestElement2, unitTestElement3]); _testablePlatformServiceProvider.MockFileOperations.Setup(fo => fo.LoadAssembly("A")) .Returns(Assembly.GetExecutingAssembly());