Make _diag folder location customizable (#3797)

* Adding the option to override the _diag folder location via the enviroment variable * Fixing DiaLogManager to take agent logs if the path is overriden * Fixing some tests * Adding unit test for the change Co-authored-by: Flavia Saplacan <[email protected]> Co-authored-by: Anatoly Bolshakov <[email protected]> Co-authored-by: Martin Mrazik <[email protected]> Co-authored-by: Konstantin Tyukalov <[email protected]>
microsoft · May 24, 2022 · c0f5741 · c0f5741
1 parent dec2e81
commit c0f5741
Show file tree

Hide file tree

Showing 14 changed files with 108 additions and 40 deletions.
diff --git a/src/Agent.Listener/Program.cs b/src/Agent.Listener/Program.cs
@@ -25,7 +25,7 @@ public static int Main(string[] args)
                 AppContext.SetSwitch("System.Net.Http.UseSocketsHttpHandler", false);
             }
 
-            using (HostContext context = new HostContext("Agent"))
+            using (HostContext context = new HostContext(HostType.Agent))
             {
                 return MainAsync(context, args).GetAwaiter().GetResult();
             }

diff --git a/src/Agent.Listener/SelfUpdater.cs b/src/Agent.Listener/SelfUpdater.cs
@@ -556,7 +556,7 @@ private void DeletePreviousVersionAgentBackup(CancellationToken token)
         private string GenerateUpdateScript(bool restartInteractiveAgent)
         {
             int processId = Process.GetCurrentProcess().Id;
-            string updateLog = Path.Combine(HostContext.GetDirectory(WellKnownDirectory.Diag), $"SelfUpdate-{DateTime.UtcNow.ToString("yyyyMMdd-HHmmss")}.log");
+            string updateLog = Path.Combine(HostContext.GetDiagDirectory(), $"SelfUpdate-{DateTime.UtcNow.ToString("yyyyMMdd-HHmmss")}.log");
             string agentRoot = HostContext.GetDirectory(WellKnownDirectory.Root);
 
             string templateName = "update.sh.template";

diff --git a/src/Agent.Sdk/Knob/AgentKnobs.cs b/src/Agent.Sdk/Knob/AgentKnobs.cs
@@ -135,6 +135,19 @@ public class AgentKnobs
             new EnvironmentKnobSource("VSTSAGENT_DUMP_JOB_EVENT_LOGS"),
             new BuiltInDefaultKnobSource("false"));
 
+        // Diag logging
+        public static readonly Knob AgentDiagLogPath = new Knob(
+            nameof(AgentDiagLogPath),
+            "If set to anything, the folder containing the agent diag log will be created here.",
+            new EnvironmentKnobSource("AGENT_DIAGLOGPATH"),
+            new BuiltInDefaultKnobSource(string.Empty));
+
+        public static readonly Knob WorkerDiagLogPath = new Knob(
+            nameof(WorkerDiagLogPath),
+            "If set to anything, the folder containing the agent worker diag log will be created here.",
+            new EnvironmentKnobSource("WORKER_DIAGLOGPATH"),
+            new BuiltInDefaultKnobSource(string.Empty));
+
         // Timeouts
         public static readonly Knob AgentChannelTimeout = new Knob(
             nameof(AgentChannelTimeout),

diff --git a/src/Agent.Worker/DiagnosticLogManager.cs b/src/Agent.Worker/DiagnosticLogManager.cs
@@ -83,8 +83,8 @@ public async Task UploadDiagnosticLogsAsync(IExecutionContext executionContext,
             File.WriteAllText(capabilitiesFile, capabilitiesContent);
 
             // Copy worker diag log files
-            List<string> workerDiagLogFiles = GetWorkerDiagLogFiles(HostContext.GetDirectory(WellKnownDirectory.Diag), jobStartTimeUtc);
-            executionContext.Debug($"Copying {workerDiagLogFiles.Count()} worker diag logs.");
+            List<string> workerDiagLogFiles = GetWorkerDiagLogFiles(HostContext.GetDiagDirectory(), jobStartTimeUtc);
+            executionContext.Debug($"Copying {workerDiagLogFiles.Count()} worker diag logs from {HostContext.GetDiagDirectory()}.");
 
             foreach (string workerLogFile in workerDiagLogFiles)
             {
@@ -94,9 +94,9 @@ public async Task UploadDiagnosticLogsAsync(IExecutionContext executionContext,
                 File.Copy(workerLogFile, destination);
             }
 
-            // Copy agent diag log files
-            List<string> agentDiagLogFiles = GetAgentDiagLogFiles(HostContext.GetDirectory(WellKnownDirectory.Diag), jobStartTimeUtc);
-            executionContext.Debug($"Copying {agentDiagLogFiles.Count()} agent diag logs.");
+            // Copy agent diag log files - we are using the worker Host Context and we need the diag folder form the Agent.
+            List<string> agentDiagLogFiles = GetAgentDiagLogFiles(HostContext.GetDiagDirectory(HostType.Agent), jobStartTimeUtc);
+            executionContext.Debug($"Copying {agentDiagLogFiles.Count()} agent diag logs from {HostContext.GetDiagDirectory(HostType.Agent)}.");
 
             foreach (string agentLogFile in agentDiagLogFiles)
             {
@@ -143,7 +143,7 @@ public async Task UploadDiagnosticLogsAsync(IExecutionContext executionContext,
             {
                 executionContext.Debug("Dumping cloud-init logs.");
 
-                string logsFilePath = $"{HostContext.GetDirectory(WellKnownDirectory.Diag)}/cloudinit-{jobStartTimeUtc.ToString("yyyyMMdd-HHmmss")}-logs.tar.gz";
+                string logsFilePath = $"{HostContext.GetDiagDirectory()}/cloudinit-{jobStartTimeUtc.ToString("yyyyMMdd-HHmmss")}-logs.tar.gz";
                 string resultLogs = await DumpCloudInitLogs(logsFilePath);
                 executionContext.Debug(resultLogs);
 
@@ -169,7 +169,7 @@ public async Task UploadDiagnosticLogsAsync(IExecutionContext executionContext,
 
                 try
                 {
-                    string eventLogsFile = $"{HostContext.GetDirectory(WellKnownDirectory.Diag)}/EventViewer-{ jobStartTimeUtc.ToString("yyyyMMdd-HHmmss") }.csv";
+                    string eventLogsFile = $"{HostContext.GetDiagDirectory()}/EventViewer-{ jobStartTimeUtc.ToString("yyyyMMdd-HHmmss") }.csv";
                     await DumpCurrentJobEventLogs(executionContext, eventLogsFile, jobStartTimeUtc);
 
                     string destination = Path.Combine(supportFilesFolder, Path.GetFileName(eventLogsFile));
@@ -197,7 +197,7 @@ public async Task UploadDiagnosticLogsAsync(IExecutionContext executionContext,
                             .Split("\n")
                             .Where((line) => !String.IsNullOrEmpty(line) && !line.EndsWith("OK"));
 
-                        string brokenPackagesLogsPath = $"{HostContext.GetDirectory(WellKnownDirectory.Diag)}/BrokenPackages-{ jobStartTimeUtc.ToString("yyyyMMdd-HHmmss") }.log";
+                        string brokenPackagesLogsPath = $"{HostContext.GetDiagDirectory()}/BrokenPackages-{ jobStartTimeUtc.ToString("yyyyMMdd-HHmmss") }.log";
                         File.AppendAllLines(brokenPackagesLogsPath, brokenPackagesInfo);
 
                         string destination = Path.Combine(supportFilesFolder, Path.GetFileName(brokenPackagesLogsPath));
@@ -313,7 +313,7 @@ private bool DumpAgentExtensionLogs(IExecutionContext executionContext, string s
 
             executionContext.Debug($"Path to agent extension logs: {pathToLogs}");
 
-            string archivePath = Path.Combine(HostContext.GetDirectory(WellKnownDirectory.Diag), archiveName);
+            string archivePath = Path.Combine(HostContext.GetDiagDirectory(), archiveName);
             executionContext.Debug($"Archiving agent extension logs to: {archivePath}");
             ZipFile.CreateFromDirectory(pathToLogs, archivePath);
 
@@ -416,7 +416,7 @@ private List<string> GetWorkerDiagLogFiles(string diagFolder, DateTime jobStartT
             int bufferInSeconds = -30;
             DateTime searchTimeUtc = jobStartTimeUtc.AddSeconds(bufferInSeconds);
 
-            foreach (FileInfo file in directoryInfo.GetFiles().Where(f => f.Name.StartsWith("Worker_")))
+            foreach (FileInfo file in directoryInfo.GetFiles().Where(f => f.Name.StartsWith($"{HostType.Worker}_")))
             {
                 // The format of the logs is:
                 // Worker_20171003-143110-utc.log
@@ -444,7 +444,7 @@ private List<string> GetAgentDiagLogFiles(string diagFolder, DateTime jobStartTi
             String recentLog = null;
             DateTime recentTimeUtc = DateTime.MinValue;
 
-            foreach (FileInfo file in directoryInfo.GetFiles().Where(f => f.Name.StartsWith("Agent_")))
+            foreach (FileInfo file in directoryInfo.GetFiles().Where(f => f.Name.StartsWith($"{HostType.Agent}_")))
             {
                 // The format of the logs is:
                 // Agent_20171003-143110-utc.log

diff --git a/src/Agent.Worker/ExecutionContext.cs b/src/Agent.Worker/ExecutionContext.cs
@@ -177,7 +177,7 @@ public override void Initialize(IHostContext hostContext)
 
             if (_disableLogUploads)
             {
-                _buildLogsFolderPath = Path.Combine(hostContext.GetDirectory(WellKnownDirectory.Diag), _buildLogsFolderName);
+                _buildLogsFolderPath = Path.Combine(hostContext.GetDiagDirectory(), _buildLogsFolderName);
                 Directory.CreateDirectory(_buildLogsFolderPath);
             }
 

diff --git a/src/Agent.Worker/Program.cs b/src/Agent.Worker/Program.cs
@@ -19,7 +19,7 @@ public static int Main(string[] args)
                 AppContext.SetSwitch("System.Net.Http.UseSocketsHttpHandler", false);
             }
 
-            using (HostContext context = new HostContext("Worker"))
+            using (HostContext context = new HostContext(HostType.Worker))
             {
                 return MainAsync(context, args).GetAwaiter().GetResult();
             }

diff --git a/src/Microsoft.VisualStudio.Services.Agent/Constants.cs b/src/Microsoft.VisualStudio.Services.Agent/Constants.cs
@@ -10,7 +10,6 @@ namespace Microsoft.VisualStudio.Services.Agent
     public enum WellKnownDirectory
     {
         Bin,
-        Diag,
         Externals,
         LegacyPSHost,
         Root,

diff --git a/src/Microsoft.VisualStudio.Services.Agent/HostContext.cs b/src/Microsoft.VisualStudio.Services.Agent/HostContext.cs
@@ -32,6 +32,7 @@ public interface IHostContext : IDisposable, IKnobValueContext
         ILoggedSecretMasker SecretMasker { get; }
         ProductInfoHeaderValue UserAgent { get; }
         string GetDirectory(WellKnownDirectory directory);
+        string GetDiagDirectory(HostType hostType = HostType.Undefined);
         string GetConfigFile(WellKnownConfigFile configFile);
         Tracing GetTrace(string name);
         Task Delay(TimeSpan delay, CancellationToken cancellationToken);
@@ -51,6 +52,13 @@ public enum StartupType
         AutoStartup
     }
 
+    public enum HostType
+    {
+        Undefined, // Default value, used when getting the current hostContext type
+        Worker,
+        Agent
+    }
+
     public class HostContext : EventListener, IObserver<DiagnosticListener>, IObserver<KeyValuePair<string, object>>, IHostContext
     {
         private const int _defaultLogPageSize = 8;  //MB
@@ -73,15 +81,20 @@ public class HostContext : EventListener, IObserver<DiagnosticListener>, IObserv
         private IDisposable _diagListenerSubscription;
         private StartupType _startupType;
         private string _perfFile;
+        private HostType _hostType;
         public event EventHandler Unloading;
         public CancellationToken AgentShutdownToken => _agentShutdownTokenSource.Token;
         public ShutdownReason AgentShutdownReason { get; private set; }
         public ILoggedSecretMasker SecretMasker => _secretMasker;
         public ProductInfoHeaderValue UserAgent => _userAgent;
-        public HostContext(string hostType, string logFile = null)
+        public HostContext(HostType hostType, string logFile = null)
         {
+
             // Validate args.
-            ArgUtil.NotNullOrEmpty(hostType, nameof(hostType));
+            if (hostType == HostType.Undefined) {
+                throw new ArgumentException(message: $"HostType cannot be {HostType.Undefined}");
+            }
+            _hostType = hostType;
 
             _loadContext = AssemblyLoadContext.GetLoadContext(typeof(HostContext).GetTypeInfo().Assembly);
             _loadContext.Unloading += LoadContext_Unloading;
@@ -102,22 +115,23 @@ public HostContext(string hostType, string logFile = null)
             if (string.IsNullOrEmpty(logFile))
             {
                 int logPageSize;
-                string logSizeEnv = Environment.GetEnvironmentVariable($"{hostType.ToUpperInvariant()}_LOGSIZE");
+                string logSizeEnv = Environment.GetEnvironmentVariable($"{_hostType.ToString().ToUpperInvariant()}_LOGSIZE");
                 if (!string.IsNullOrEmpty(logSizeEnv) || !int.TryParse(logSizeEnv, out logPageSize))
                 {
                     logPageSize = _defaultLogPageSize;
                 }
 
                 int logRetentionDays;
-                string logRetentionDaysEnv = Environment.GetEnvironmentVariable($"{hostType.ToUpperInvariant()}_LOGRETENTION");
+                string logRetentionDaysEnv = Environment.GetEnvironmentVariable($"{_hostType.ToString().ToUpperInvariant()}_LOGRETENTION");
                 if (!string.IsNullOrEmpty(logRetentionDaysEnv) || !int.TryParse(logRetentionDaysEnv, out logRetentionDays))
                 {
                     logRetentionDays = _defaultLogRetentionDays;
                 }
 
-                // this should give us _diag folder under agent root directory
-                string diagLogDirectory = Path.Combine(new DirectoryInfo(Path.GetDirectoryName(Assembly.GetEntryAssembly().Location)).Parent.FullName, Constants.Path.DiagDirectory);
-                _traceManager = new TraceManager(new HostTraceListener(diagLogDirectory, hostType, logPageSize, logRetentionDays), this.SecretMasker);
+                // this should give us _diag folder under agent root directory as default value for diagLogDirctory
+                string diagLogPath = GetDiagDirectory(_hostType);
+                _traceManager = new TraceManager(new HostTraceListener(diagLogPath, hostType.ToString(), logPageSize, logRetentionDays), this.SecretMasker);
+
             }
             else
             {
@@ -168,12 +182,6 @@ public virtual string GetDirectory(WellKnownDirectory directory)
                     path = Path.GetDirectoryName(Assembly.GetEntryAssembly().Location);
                     break;
 
-                case WellKnownDirectory.Diag:
-                    path = Path.Combine(
-                        GetDirectory(WellKnownDirectory.Root),
-                        Constants.Path.DiagDirectory);
-                    break;
-
                 case WellKnownDirectory.Externals:
                     path = Path.Combine(
                         GetDirectory(WellKnownDirectory.Root),
@@ -260,6 +268,29 @@ public virtual string GetDirectory(WellKnownDirectory directory)
             return path;
         }
 
+        public string GetDiagDirectory(HostType hostType = HostType.Undefined)
+        {
+            return hostType switch
+            {
+                HostType.Undefined => GetDiagDirectory(_hostType),
+                HostType.Agent => GetDiagOrDefault(AgentKnobs.AgentDiagLogPath.GetValue(this).AsString()),
+                HostType.Worker => GetDiagOrDefault(AgentKnobs.WorkerDiagLogPath.GetValue(this).AsString()),
+                _ => throw new NotSupportedException($"Unexpected host type: '{hostType}'"),
+            };
+        }
+
+        private string GetDiagOrDefault(string diagFolder)
+        {
+            if (!string.IsNullOrEmpty(diagFolder))
+            {
+                return diagFolder;
+            }
+
+            return Path.Combine(
+                new DirectoryInfo(Path.GetDirectoryName(Assembly.GetEntryAssembly().Location)).Parent.FullName,
+                Constants.Path.DiagDirectory);          
+        }
+
         public string GetConfigFile(WellKnownConfigFile configFile)
         {
             string path;

diff --git a/src/Microsoft.VisualStudio.Services.Agent/Logging.cs b/src/Microsoft.VisualStudio.Services.Agent/Logging.cs
@@ -46,7 +46,7 @@ public override void Initialize(IHostContext hostContext)
             base.Initialize(hostContext);
             _totalLines = 0;
             _pageId = Guid.NewGuid().ToString();
-            _pagesFolder = Path.Combine(hostContext.GetDirectory(WellKnownDirectory.Diag), PagingFolder);
+            _pagesFolder = Path.Combine(hostContext.GetDiagDirectory(), PagingFolder);
             _jobServerQueue = HostContext.GetService<IJobServerQueue>();
             Directory.CreateDirectory(_pagesFolder);
         }

diff --git a/src/Test/L0/HostContextL0.cs b/src/Test/L0/HostContextL0.cs
@@ -118,10 +118,34 @@ public void OtherSecretsAreMasked(string input, string expected)
             }
         }
 
+        [Fact]
+        public void LogFileChangedAccordingToEnvVariable()
+        {
+            try
+            {
+                var newPath = Path.Combine(Path.GetDirectoryName(Assembly.GetEntryAssembly().Location), "logs");
+                Environment.SetEnvironmentVariable("AGENT_DIAGLOGPATH", newPath);
+
+                using (var _hc = new HostContext(HostType.Agent))
+                {
+                    // Act.
+                    var diagFolder = _hc.GetDiagDirectory();
+
+                    // Assert
+                    Assert.Equal(Path.Combine(newPath, Constants.Path.DiagDirectory), diagFolder);
+                    Directory.Exists(diagFolder);
+                }
+            }
+            finally
+            {
+                Environment.SetEnvironmentVariable("AGENT_DIAGLOGPATH", null);
+            }
+        }
+
         public HostContext Setup([CallerMemberName] string testName = "")
         {
             return new HostContext(
-                hostType: "L0Test",
+                hostType: HostType.Agent,
                 logFile: Path.Combine(Path.GetDirectoryName(Assembly.GetEntryAssembly().Location), $"trace_{nameof(HostContextL0)}_{testName}.log"));
         }
     }

diff --git a/src/Test/L0/PagingLoggerL0.cs b/src/Test/L0/PagingLoggerL0.cs
@@ -25,7 +25,7 @@ private void CleanLogFolder()
             using (TestHostContext hc = new TestHostContext(this))
             {
                 //clean test data if any old test forgot
-                string pagesFolder = Path.Combine(hc.GetDirectory(WellKnownDirectory.Diag), PagingLogger.PagingFolder);
+                string pagesFolder = Path.Combine(hc.GetDiagDirectory(), PagingLogger.PagingFolder);
                 if (Directory.Exists(pagesFolder))
                 {
                     Directory.Delete(pagesFolder, true);

diff --git a/src/Test/L0/TestHostContext.cs b/src/Test/L0/TestHostContext.cs
@@ -179,12 +179,6 @@ public string GetDirectory(WellKnownDirectory directory)
                     path = Path.GetDirectoryName(Assembly.GetEntryAssembly().Location);
                     break;
 
-                case WellKnownDirectory.Diag:
-                    path = Path.Combine(
-                        GetDirectory(WellKnownDirectory.Root),
-                        Constants.Path.DiagDirectory);
-                    break;
-
                 case WellKnownDirectory.Externals:
                     path = Path.Combine(
                         GetDirectory(WellKnownDirectory.Root),
@@ -267,6 +261,13 @@ public string GetDirectory(WellKnownDirectory directory)
             return path;
         }
 
+        public string GetDiagDirectory(HostType hostType = HostType.Undefined)
+        {
+            return Path.Combine(
+                        GetDirectory(WellKnownDirectory.Root),
+                        Constants.Path.DiagDirectory);
+        }
+
         public string GetConfigFile(WellKnownConfigFile configFile)
         {
             string path;

diff --git a/src/Test/L1/L1HostContext.cs b/src/Test/L1/L1HostContext.cs
@@ -9,7 +9,7 @@ namespace Microsoft.VisualStudio.Services.Agent.Tests.L1.Worker
 {
     public class L1HostContext : HostContext
     {
-        public L1HostContext(string hostType, string logFile = null)
+        public L1HostContext(HostType hostType, string logFile = null)
             : base(hostType, logFile)
         {
         }

diff --git a/src/Test/L1/Worker/L1TestBase.cs b/src/Test/L1/Worker/L1TestBase.cs
@@ -194,7 +194,7 @@ public void SetupL1([CallerMemberName] string testName = "")
             var stringFile = Path.Combine(assemblyLocation, "en-US", "strings.json");
             StringUtil.LoadExternalLocalization(stringFile);
 
-            _l1HostContext = new L1HostContext("Agent", GetLogFile(this, testName));
+            _l1HostContext = new L1HostContext(HostType.Agent, GetLogFile(this, testName));
             SetupMocks(_l1HostContext);
 
             // Use different working directories for each test
-Original file line number
+Diff line change
@@ Expand Up / @@ -10,7 +10,6 @@ namespace Microsoft.VisualStudio.Services.Agent @@
         public enum WellKnownDirectory
         {
             Bin,
-            Diag,
             Externals,
             LegacyPSHost,
             Root,
@@ Expand Down @@