From 7b9afd9de21119fec6788f19aa851023b3198063 Mon Sep 17 00:00:00 2001
From: Mihai Codoban <micodoba@microsoft.com>
Date: Tue, 2 Feb 2021 17:42:18 -0800
Subject: [PATCH] Fix bad plugin EndBuild exception handling during graph
 builds

---
 .../ProjectCache/ProjectCacheTests.cs         | 68 ++++++++++++++++---
 .../BackEnd/BuildManager/BuildManager.cs      | 10 ++-
 src/Shared/UnitTests/ObjectModelHelpers.cs    | 10 +--
 3 files changed, 71 insertions(+), 17 deletions(-)

diff --git a/src/Build.UnitTests/ProjectCache/ProjectCacheTests.cs b/src/Build.UnitTests/ProjectCache/ProjectCacheTests.cs
index f7d7eca1c2c..df36e0dfd2e 100644
--- a/src/Build.UnitTests/ProjectCache/ProjectCacheTests.cs
+++ b/src/Build.UnitTests/ProjectCache/ProjectCacheTests.cs
@@ -818,16 +818,7 @@ public void EngineShouldHandleExceptionsFromCachePlugin(ExceptionLocations excep
                         </Target>
                     </Project>".Cleanup());
 
-            foreach (var enumValue in Enum.GetValues(typeof(ExceptionLocations)))
-            {
-                var typedValue = (ExceptionLocations) enumValue;
-                if (exceptionLocations.HasFlag(typedValue))
-                {
-                    var exceptionLocation = typedValue.ToString();
-                    _env.SetEnvironmentVariable(exceptionLocation, "1");
-                    _output.WriteLine($"Set exception location: {exceptionLocation}");
-                }
-            }
+            SetEnvironmentForExceptionLocations(exceptionLocations);
 
             using var buildSession = new Helpers.BuildManagerSession(
                 _env,
@@ -880,5 +871,62 @@ public void EngineShouldHandleExceptionsFromCachePlugin(ExceptionLocations excep
                 logger.FullLog.ShouldContain($"{AssemblyMockCache}: EndBuildAsync");
             }
         }
+
+        [Fact]
+        public void EndBuildShouldGetCalledOnceWhenItThrowsExceptionsFromGraphBuilds()
+        {
+            _env.DoNotLaunchDebugger();
+
+            var project = _env.CreateFile(
+                "1.proj",
+                @$"
+                    <Project>
+                        <ItemGroup>
+                            <{ItemTypeNames.ProjectCachePlugin} Include=`{SamplePluginAssemblyPath.Value}` />
+                        </ItemGroup>
+                        <Target Name=`Build`>
+                            <Message Text=`Hello EngineShouldHandleExceptionsFromCachePlugin` Importance=`High` />
+                        </Target>
+                    </Project>".Cleanup());
+
+            SetEnvironmentForExceptionLocations(ExceptionLocations.EndBuildAsync);
+
+            using var buildSession = new Helpers.BuildManagerSession(
+                _env,
+                new BuildParameters
+                {
+                    UseSynchronousLogging = true
+                });
+
+            var logger = buildSession.Logger;
+
+            GraphBuildResult? buildResult = null;
+            Should.NotThrow(
+                () =>
+                {
+                    buildResult = buildSession.BuildGraph(new ProjectGraph(project.Path));
+                });
+
+            buildResult!.OverallResult.ShouldBe(BuildResultCode.Failure);
+            buildResult.Exception.Message.ShouldContain("Cache plugin exception from EndBuildAsync");
+
+            buildSession.Dispose();
+
+            Regex.Matches(logger.FullLog, $"{nameof(AssemblyMockCache)}: EndBuildAsync").Count.ShouldBe(1);
+        }
+
+        private void SetEnvironmentForExceptionLocations(ExceptionLocations exceptionLocations)
+        {
+            foreach (var enumValue in Enum.GetValues(typeof(ExceptionLocations)))
+            {
+                var typedValue = (ExceptionLocations) enumValue;
+                if (exceptionLocations.HasFlag(typedValue))
+                {
+                    var exceptionLocation = typedValue.ToString();
+                    _env.SetEnvironmentVariable(exceptionLocation, "1");
+                    _output.WriteLine($"Set exception location: {exceptionLocation}");
+                }
+            }
+        }
     }
 }
diff --git a/src/Build/BackEnd/BuildManager/BuildManager.cs b/src/Build/BackEnd/BuildManager/BuildManager.cs
index a81d091cf14..f8264107f8b 100644
--- a/src/Build/BackEnd/BuildManager/BuildManager.cs
+++ b/src/Build/BackEnd/BuildManager/BuildManager.cs
@@ -1943,8 +1943,14 @@ public void Dispose()
 
                 lock (_buildManager._syncLock)
                 {
-                    _buildManager._projectCacheService?.Result.ShutDown().GetAwaiter().GetResult();
-                    _buildManager._projectCacheService = null;
+                    try
+                    {
+                        _buildManager._projectCacheService?.Result.ShutDown().GetAwaiter().GetResult();
+                    }
+                    finally
+                    {
+                        _buildManager._projectCacheService = null;
+                    }
                 }
             }
         }
diff --git a/src/Shared/UnitTests/ObjectModelHelpers.cs b/src/Shared/UnitTests/ObjectModelHelpers.cs
index 5c721b50c7e..60a2c0451b9 100644
--- a/src/Shared/UnitTests/ObjectModelHelpers.cs
+++ b/src/Shared/UnitTests/ObjectModelHelpers.cs
@@ -1950,6 +1950,11 @@ public GraphBuildResult BuildGraphSubmission(GraphBuildRequestData requestData)
                 return _buildManager.BuildRequest(requestData);
             }
 
+            public GraphBuildResult BuildGraph(ProjectGraph graph, string[] entryTargets = null)
+            {
+                return _buildManager.BuildRequest(new GraphBuildRequestData(graph, entryTargets ?? new string[0]));
+            }
+
             public void Dispose()
             {
                 if (_disposed)
@@ -1962,11 +1967,6 @@ public void Dispose()
                 _buildManager.EndBuild();
                 _buildManager.Dispose();
             }
-
-            public GraphBuildResult BuildGraph(ProjectGraph graph, string[] entryTargets = null)
-            {
-                return _buildManager.BuildRequest(new GraphBuildRequestData(graph, entryTargets ?? new string[0]));
-            }
         }
 
         internal class LoggingFileSystem : MSBuildFileSystemBase