From 0936da11f5a153790315619f08e29c8a19099579 Mon Sep 17 00:00:00 2001
From: Jeff Kluge <jeffkl@microsoft.com>
Date: Tue, 7 Feb 2023 13:00:27 -0800
Subject: [PATCH] Embed NuGet files in binlog after restore

---
 .../MSBuildStaticGraphRestore.cs              | 22 +++++++++++++++++--
 .../ConsoleOutLogMessageType.cs               |  5 +++++
 .../NuGet.Build.Tasks/GlobalSuppressions.cs   |  2 ++
 .../NuGet.Build.Tasks/NuGet.RestoreEx.targets | 13 ++++++++---
 .../NuGet.Build.Tasks/NuGet.targets           | 10 ++++++++-
 .../NuGet.Build.Tasks/RestoreTask.cs          | 22 +++++++++++++++++++
 .../StaticGraphRestoreTaskBase.cs             | 22 ++++++++++++++++++-
 .../NuGet.Build.Tasks/TaskLoggingQueue.cs     | 12 +++++++++-
 .../TaskLoggingQueueTests.cs                  |  8 ++++---
 9 files changed, 105 insertions(+), 11 deletions(-)

diff --git a/src/NuGet.Core/NuGet.Build.Tasks.Console/MSBuildStaticGraphRestore.cs b/src/NuGet.Core/NuGet.Build.Tasks.Console/MSBuildStaticGraphRestore.cs
index 744150fafac..564d0e83ccc 100644
--- a/src/NuGet.Core/NuGet.Build.Tasks.Console/MSBuildStaticGraphRestore.cs
+++ b/src/NuGet.Core/NuGet.Build.Tasks.Console/MSBuildStaticGraphRestore.cs
@@ -114,7 +114,7 @@ public async Task<bool> RestoreAsync(string entryProjectFilePath, IDictionary<st
 
             try
             {
-                return await BuildTasksUtility.RestoreAsync(
+                bool result = await BuildTasksUtility.RestoreAsync(
                     dependencyGraphSpec: dependencyGraphSpec,
                     interactive: IsOptionTrue(nameof(RestoreTaskEx.Interactive), options),
                     recursive: IsOptionTrue(nameof(RestoreTaskEx.Recursive), options),
@@ -127,7 +127,25 @@ public async Task<bool> RestoreAsync(string entryProjectFilePath, IDictionary<st
                     restorePC: IsOptionTrue(nameof(RestoreTaskEx.RestorePackagesConfig), options),
                     cleanupAssetsForUnsupportedProjects: IsOptionTrue(nameof(RestoreTaskEx.CleanupAssetsForUnsupportedProjects), options),
                     log: MSBuildLogger,
-                    cancellationToken: CancellationToken.None);
+                cancellationToken: CancellationToken.None);
+
+                // If the MSBuildBinaryLoggerEnabled environment variable is set, return the paths to all of the projects used during restore so that the assets file and other NuGet files
+                // can be embedded in the binlog
+                if (string.Equals(Environment.GetEnvironmentVariable("MSBUILDBINARYLOGGERENABLED"), bool.TrueString, StringComparison.OrdinalIgnoreCase))
+                {
+                    IReadOnlyList<PackageSpec> projects = dependencyGraphSpec.Projects;
+
+                    for (int i = 0; i < projects.Count; i++)
+                    {
+                        LoggingQueue.Enqueue(new ConsoleOutLogMessage
+                        {
+                            Message = projects[i].RestoreMetadata.OutputPath,
+                            MessageType = ConsoleOutLogMessageType.RestoredProjectOutputPath
+                        });
+                    }
+                }
+
+                return result;
             }
             catch (Exception e)
             {
diff --git a/src/NuGet.Core/NuGet.Build.Tasks/ConsoleOutLogMessageType.cs b/src/NuGet.Core/NuGet.Build.Tasks/ConsoleOutLogMessageType.cs
index 9f0094454ba..9fdcd6ca032 100644
--- a/src/NuGet.Core/NuGet.Build.Tasks/ConsoleOutLogMessageType.cs
+++ b/src/NuGet.Core/NuGet.Build.Tasks/ConsoleOutLogMessageType.cs
@@ -27,5 +27,10 @@ public enum ConsoleOutLogMessageType
         /// The message is a message.
         /// </summary>
         Message,
+
+        /// <summary>
+        /// The message is a restored project's output path.
+        /// </summary>
+        RestoredProjectOutputPath
     }
 }
diff --git a/src/NuGet.Core/NuGet.Build.Tasks/GlobalSuppressions.cs b/src/NuGet.Core/NuGet.Build.Tasks/GlobalSuppressions.cs
index 92b22ee6598..c0da1eba205 100644
--- a/src/NuGet.Core/NuGet.Build.Tasks/GlobalSuppressions.cs
+++ b/src/NuGet.Core/NuGet.Build.Tasks/GlobalSuppressions.cs
@@ -68,3 +68,5 @@
 [assembly: SuppressMessage("Build", "CA1819:Properties should not return arrays", Justification = "<Pending>", Scope = "member", Target = "~P:NuGet.Build.Tasks.WriteRestoreGraphTask.RestoreGraphItems")]
 [assembly: SuppressMessage("Performance", "CA1819:Properties should not return arrays", Justification = "MSBuild tasks expect arrays.", Scope = "member", Target = "~P:NuGet.Build.Tasks.CheckForDuplicateNuGetItemsTask.Items")]
 [assembly: SuppressMessage("Performance", "CA1819:Properties should not return arrays", Justification = "MSBuild tasks expect arrays.", Scope = "member", Target = "~P:NuGet.Build.Tasks.CheckForDuplicateNuGetItemsTask.DeduplicatedItems")]
+[assembly: SuppressMessage("Performance", "CA1819:Properties should not return arrays", Justification = "MSBuild tasks expect arrays.", Scope = "member", Target = "~P:NuGet.Build.Tasks.RestoreTask.RestoredProjectOutputPaths")]
+[assembly: SuppressMessage("Performance", "CA1819:Properties should not return arrays", Justification = "MSBuild tasks expect arrays.", Scope = "member", Target = "~P:NuGet.Build.Tasks.StaticGraphRestoreTaskBase.RestoredProjectOutputPaths")]
diff --git a/src/NuGet.Core/NuGet.Build.Tasks/NuGet.RestoreEx.targets b/src/NuGet.Core/NuGet.Build.Tasks/NuGet.RestoreEx.targets
index db38689806e..26a62230f07 100644
--- a/src/NuGet.Core/NuGet.Build.Tasks/NuGet.RestoreEx.targets
+++ b/src/NuGet.Core/NuGet.Build.Tasks/NuGet.RestoreEx.targets
@@ -31,8 +31,15 @@ Copyright (c) .NET Foundation. All rights reserved.
         RestorePackagesConfig="$(RestorePackagesConfig)"
         SolutionPath="$(SolutionPath)"
         ProcessFileName="$(NuGetConsoleProcessFileName)"
-        MSBuildStartupDirectory="$(MSBuildStartupDirectory)"
-        />
+        MSBuildStartupDirectory="$(MSBuildStartupDirectory)">
+      <Output TaskParameter="RestoredProjectOutputPaths" ItemName="_RestoredProjectOutputPaths" />
+    </RestoreTaskEx>
+
+    <ItemGroup Condition="@(_RestoredProjectOutputPaths->Count()) > 0">
+      <EmbedInBinlog Include="%(_RestoredProjectOutputPaths.Identity)project.assets.json;
+                              %(_RestoredProjectOutputPaths.Identity)*.nuget.dgspec.json;
+                              %(_RestoredProjectOutputPaths.Identity)*.nuget.g.*" />
+    </ItemGroup>
   </Target>
 
  <!--
@@ -53,5 +60,5 @@ Copyright (c) .NET Foundation. All rights reserved.
         MSBuildStartupDirectory="$(MSBuildStartupDirectory)"
         />
   </Target>
-
+  
 </Project>
diff --git a/src/NuGet.Core/NuGet.Build.Tasks/NuGet.targets b/src/NuGet.Core/NuGet.Build.Tasks/NuGet.targets
index 62c5e106a91..f7d737927ee 100644
--- a/src/NuGet.Core/NuGet.Build.Tasks/NuGet.targets
+++ b/src/NuGet.Core/NuGet.Build.Tasks/NuGet.targets
@@ -139,7 +139,15 @@ Copyright (c) .NET Foundation. All rights reserved.
       HideWarningsAndErrors="$(HideWarningsAndErrors)"
       Interactive="$(NuGetInteractive)"
       RestoreForceEvaluate="$(RestoreForceEvaluate)"
-      RestorePackagesConfig="$(RestorePackagesConfig)"/>
+      RestorePackagesConfig="$(RestorePackagesConfig)">
+      <Output TaskParameter="RestoredProjectOutputPaths" ItemName="_RestoredProjectOutputPaths" />
+    </RestoreTask>
+
+    <ItemGroup Condition="@(_RestoredProjectOutputPaths->Count()) > 0">
+      <EmbedInBinlog Include="%(_RestoredProjectOutputPaths.Identity)project.assets.json;
+                              %(_RestoredProjectOutputPaths.Identity)*.nuget.dgspec.json;
+                              %(_RestoredProjectOutputPaths.Identity)*.nuget.g.*" />
+    </ItemGroup>
   </Target>
 
   <!--
diff --git a/src/NuGet.Core/NuGet.Build.Tasks/RestoreTask.cs b/src/NuGet.Core/NuGet.Build.Tasks/RestoreTask.cs
index b91c8c20cce..849c21de482 100644
--- a/src/NuGet.Core/NuGet.Build.Tasks/RestoreTask.cs
+++ b/src/NuGet.Core/NuGet.Build.Tasks/RestoreTask.cs
@@ -2,11 +2,13 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Build.Framework;
+using Microsoft.Build.Utilities;
 using NuGet.Commands;
 using NuGet.Common;
 
@@ -73,6 +75,12 @@ public class RestoreTask : Microsoft.Build.Utilities.Task, ICancelableTask, IDis
         /// <returns></returns>
         public bool RestorePackagesConfig { get; set; }
 
+        /// <summary>
+        /// Gets or sets the output path for all of the restored projects.
+        /// </summary>
+        [Output]
+        public ITaskItem[] RestoredProjectOutputPaths { get; set; }
+
         public override bool Execute()
         {
 #if DEBUG
@@ -127,6 +135,20 @@ private async Task<bool> ExecuteAsync(Common.ILogger log)
 
             var dgFile = MSBuildRestoreUtility.GetDependencySpec(wrappedItems);
 
+            // If the MSBuildBinaryLoggerEnabled environment variable is set, return the paths to all of the projects used during restore so that the assets file and other NuGet files
+            // can be embedded in the binlog
+            if (string.Equals(Environment.GetEnvironmentVariable("MSBUILDBINARYLOGGERENABLED"), bool.TrueString, StringComparison.OrdinalIgnoreCase))
+            {
+                IReadOnlyList<ProjectModel.PackageSpec> projects = dgFile.Projects;
+
+                RestoredProjectOutputPaths = new ITaskItem[projects.Count];
+
+                for (int i = 0; i < projects.Count; i++)
+                {
+                    RestoredProjectOutputPaths[i] = new TaskItem(projects[i].RestoreMetadata.OutputPath);
+                }
+            }
+
             return await BuildTasksUtility.RestoreAsync(
                 dependencyGraphSpec: dgFile,
                 interactive: Interactive,
diff --git a/src/NuGet.Core/NuGet.Build.Tasks/StaticGraphRestoreTaskBase.cs b/src/NuGet.Core/NuGet.Build.Tasks/StaticGraphRestoreTaskBase.cs
index bce5aa0b8d6..725c0ae7509 100644
--- a/src/NuGet.Core/NuGet.Build.Tasks/StaticGraphRestoreTaskBase.cs
+++ b/src/NuGet.Core/NuGet.Build.Tasks/StaticGraphRestoreTaskBase.cs
@@ -2,9 +2,11 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
+using System.Linq;
 
 #if !IS_CORECLR
 
@@ -14,6 +16,7 @@
 
 using System.Threading;
 using Microsoft.Build.Framework;
+using Microsoft.Build.Utilities;
 using NuGet.Common;
 
 namespace NuGet.Build.Tasks
@@ -62,6 +65,9 @@ public abstract class StaticGraphRestoreTaskBase : Microsoft.Build.Utilities.Tas
         /// </summary>
         public bool Recursive { get; set; }
 
+        [Output]
+        public ITaskItem[] RestoredProjectOutputPaths { get; set; }
+
         /// <summary>
         /// Gets or sets the full path to the solution file (if any) that is being built.
         /// </summary>
@@ -92,8 +98,10 @@ public override bool Execute()
 #endif
                 MSBuildLogger logger = new MSBuildLogger(Log);
 
+                List<string> restoredOutputPaths = new List<string>();
+
                 using (var semaphore = new SemaphoreSlim(initialCount: 0, maxCount: 1))
-                using (var loggingQueue = new TaskLoggingQueue(Log))
+                using (var loggingQueue = new TaskLoggingQueue(Log, restoredOutputPaths))
                 using (var process = new Process())
                 {
                     process.EnableRaisingEvents = true;
@@ -136,6 +144,18 @@ public override bool Execute()
                                 // The process may have exited, in this case ignore the exception
                             }
                         }
+
+                        if (restoredOutputPaths.Any())
+                        {
+                            RestoredProjectOutputPaths = new ITaskItem[restoredOutputPaths.Count];
+
+                            int i = 0;
+
+                            foreach (var item in restoredOutputPaths)
+                            {
+                                RestoredProjectOutputPaths[i++] = new TaskItem(item);
+                            }
+                        }
                     }
                     catch (Exception e) when (
                         e is OperationCanceledException
diff --git a/src/NuGet.Core/NuGet.Build.Tasks/TaskLoggingQueue.cs b/src/NuGet.Core/NuGet.Build.Tasks/TaskLoggingQueue.cs
index 09566b677e9..75763c3343f 100644
--- a/src/NuGet.Core/NuGet.Build.Tasks/TaskLoggingQueue.cs
+++ b/src/NuGet.Core/NuGet.Build.Tasks/TaskLoggingQueue.cs
@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Generic;
 using Microsoft.Build.Framework;
 using Microsoft.Build.Utilities;
 using Newtonsoft.Json;
@@ -19,13 +20,18 @@ internal class TaskLoggingQueue : LoggingQueue<string>
         /// </summary>
         private readonly TaskLoggingHelper _log;
 
+        private readonly ICollection<string> _restoredProjectOutputPaths;
+
         /// <summary>
         /// Initializes a new instance of the TaskLoggingHelperQueue class.
         /// </summary>
         /// <param name="taskLoggingHelper">The <see cref="Task" /> to create a logging queue for.</param>
-        public TaskLoggingQueue(TaskLoggingHelper taskLoggingHelper)
+        /// <param name="restoredProjectOutputPaths">Receives the list of output paths for the restored projects.</param>
+        public TaskLoggingQueue(TaskLoggingHelper taskLoggingHelper, ICollection<string> restoredProjectOutputPaths)
         {
             _log = taskLoggingHelper ?? throw new ArgumentNullException(nameof(taskLoggingHelper));
+
+            _restoredProjectOutputPaths = restoredProjectOutputPaths ?? throw new ArgumentNullException(nameof(restoredProjectOutputPaths));
         }
 
         /// <summary>
@@ -90,6 +96,10 @@ protected override void Process(string message)
                         _log.LogMessageFromText(consoleOutLogMessage.Message, consoleOutLogMessage.Importance);
                         return;
 
+                    case ConsoleOutLogMessageType.RestoredProjectOutputPath:
+                        _restoredProjectOutputPaths.Add(consoleOutLogMessage.Message);
+                        return;
+
                     default:
                         throw new ArgumentOutOfRangeException(
                             paramName: nameof(message),
diff --git a/test/NuGet.Core.Tests/NuGet.Build.Tasks.Test/TaskLoggingQueueTests.cs b/test/NuGet.Core.Tests/NuGet.Build.Tasks.Test/TaskLoggingQueueTests.cs
index b0a95e6d839..81cce6878c9 100644
--- a/test/NuGet.Core.Tests/NuGet.Build.Tasks.Test/TaskLoggingQueueTests.cs
+++ b/test/NuGet.Core.Tests/NuGet.Build.Tasks.Test/TaskLoggingQueueTests.cs
@@ -6,6 +6,8 @@
 #endif
 
 using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
 using FluentAssertions;
 using Microsoft.Build.Framework;
 #if IS_DESKTOP
@@ -27,7 +29,7 @@ public void TaskLoggingQueue_Process_HandleBadMessage()
 
             var buildEngine = new TestBuildEngine();
 
-            using (var loggingQueue = new TaskLoggingQueue(new TaskLoggingHelper(buildEngine, nameof(TaskLoggingQueueTests))))
+            using (var loggingQueue = new TaskLoggingQueue(new TaskLoggingHelper(buildEngine, nameof(TaskLoggingQueueTests)), new List<string>()))
             {
                 loggingQueue.Enqueue(badMessage);
             }
@@ -57,7 +59,7 @@ public void TaskLoggingQueue_Process_LogsMessages(string messageType, string mes
 
             var buildEngine = new TestBuildEngine();
 
-            using (var loggingQueue = new TaskLoggingQueue(new TaskLoggingHelper(buildEngine, nameof(TaskLoggingQueueTests))))
+            using (var loggingQueue = new TaskLoggingQueue(new TaskLoggingHelper(buildEngine, nameof(TaskLoggingQueueTests)), new List<string>()))
             {
                 loggingQueue.Enqueue(expected.ToJson());
             }
@@ -77,7 +79,7 @@ public void TaskLoggingQueue_Process_ThrowsIfInvalidMessageType()
 
             Action act = () =>
             {
-                using (var loggingQueue = new TaskLoggingQueue(new TaskLoggingHelper(buildEngine, nameof(TaskLoggingQueueTests))))
+                using (var loggingQueue = new TaskLoggingQueue(new TaskLoggingHelper(buildEngine, nameof(TaskLoggingQueueTests)), new List<string>()))
                 {
                     loggingQueue.Enqueue(badMessage);
                 }