From 7a362c9cb9a72f7857c3ed5ac99a6b72f799ee7a Mon Sep 17 00:00:00 2001
From: Jared Goodwin <jared.goodwin@immy.bot>
Date: Mon, 31 Jul 2023 11:16:50 -0700
Subject: [PATCH] Add cancellation token to FileLogsManager methods.

---
 Agent/Services/AgentHubConnection.cs | 10 ++++++---
 Agent/Services/FileLogsManager.cs    | 33 ++++++++++++++++++++--------
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/Agent/Services/AgentHubConnection.cs b/Agent/Services/AgentHubConnection.cs
index 2e759d1d3..0aad376ca 100644
--- a/Agent/Services/AgentHubConnection.cs
+++ b/Agent/Services/AgentHubConnection.cs
@@ -1,6 +1,7 @@
 using Immense.RemoteControl.Desktop.Native.Windows;
 using Microsoft.AspNetCore.SignalR.Client;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Remotely.Agent.Extensions;
 using Remotely.Agent.Interfaces;
@@ -40,6 +41,7 @@ public class AgentHubConnection : IAgentHubConnection, IDisposable
     private readonly IWakeOnLanService _wakeOnLanService;
     private readonly ILogger<AgentHubConnection> _logger;
     private readonly IFileLogsManager _fileLogsManager;
+    private readonly IHostApplicationLifetime _appLifetime;
     private readonly IScriptExecutor _scriptExecutor;
     private readonly IUninstaller _uninstaller;
     private readonly IUpdater _updater;
@@ -60,6 +62,7 @@ public AgentHubConnection(
         IHttpClientFactory httpFactory,
         IWakeOnLanService wakeOnLanService,
         IFileLogsManager fileLogsManager,
+        IHostApplicationLifetime appLifetime,
         ILogger<AgentHubConnection> logger)
     {
         _configService = configService;
@@ -73,6 +76,7 @@ public AgentHubConnection(
         _wakeOnLanService = wakeOnLanService;
         _logger = logger;
         _fileLogsManager = fileLogsManager;
+        _appLifetime = appLifetime;
     }
 
     public bool IsConnected => _hubConnection?.State == HubConnectionState.Connected;
@@ -302,7 +306,7 @@ private void RegisterMessageHandlers()
 
         _hubConnection.On("DeleteLogs", () =>
        {
-           _fileLogsManager.DeleteLogs();
+           _fileLogsManager.DeleteLogs(_appLifetime.ApplicationStopping);
        });
 
 
@@ -367,14 +371,14 @@ private void RegisterMessageHandlers()
         {
             try
             {
-                if (!await _fileLogsManager.AnyLogsExist())
+                if (!await _fileLogsManager.AnyLogsExist(_appLifetime.ApplicationStopping))
                 {
                     var message = "There are no log entries written.";
                     await _hubConnection.InvokeAsync("SendLogs", message, senderConnectionId).ConfigureAwait(false);
                     return;
                 }
 
-                await foreach (var chunk in _fileLogsManager.ReadAllBytes())
+                await foreach (var chunk in _fileLogsManager.ReadAllBytes(_appLifetime.ApplicationStopping))
                 {
                     var lines = Encoding.UTF8.GetString(chunk);
                     await _hubConnection.InvokeAsync("SendLogs", lines, senderConnectionId).ConfigureAwait(false);
diff --git a/Agent/Services/FileLogsManager.cs b/Agent/Services/FileLogsManager.cs
index a9684b631..665ad7d88 100644
--- a/Agent/Services/FileLogsManager.cs
+++ b/Agent/Services/FileLogsManager.cs
@@ -4,23 +4,25 @@
 using System.IO;
 using System.Linq;
 using System.Reflection;
+using System.Runtime.CompilerServices;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Remotely.Shared.Services;
 
 public interface IFileLogsManager
 {
-    Task<bool> AnyLogsExist();
-    Task DeleteLogs();
-    IAsyncEnumerable<byte[]> ReadAllBytes();
+    Task<bool> AnyLogsExist(CancellationToken cancellationToken);
+    Task DeleteLogs(CancellationToken cancellationToken);
+    IAsyncEnumerable<byte[]> ReadAllBytes(CancellationToken cancellationToken);
 }
 
 public class FileLogsManager : IFileLogsManager
 {
-    public async Task<bool> AnyLogsExist()
+    public async Task<bool> AnyLogsExist(CancellationToken cancellationToken)
     {
-        using var logLock = await FileLoggerDefaults.AcquireLock();
+        using var logLock = await FileLoggerDefaults.AcquireLock(cancellationToken);
 
         var componentName = Assembly.GetExecutingAssembly().GetName().Name;
         var directory = Path.Combine(FileLoggerDefaults.LogsFolderPath, $"{componentName}");
@@ -29,6 +31,11 @@ public async Task<bool> AnyLogsExist()
         {
             foreach (var file in Directory.GetFiles(directory))
             {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    break;
+                }
+
                 if (new FileInfo(file).Length > 0)
                 {
                     return true;
@@ -38,9 +45,9 @@ public async Task<bool> AnyLogsExist()
         return false;
     }
 
-    public async Task DeleteLogs()
+    public async Task DeleteLogs(CancellationToken cancellationToken)
     {
-        using var logLock = await FileLoggerDefaults.AcquireLock();
+        using var logLock = await FileLoggerDefaults.AcquireLock(cancellationToken);
 
         var componentName = Assembly.GetExecutingAssembly().GetName().Name;
         var directory = Path.Combine(FileLoggerDefaults.LogsFolderPath, $"{componentName}");
@@ -49,6 +56,10 @@ public async Task DeleteLogs()
         {
             foreach (var file in Directory.GetFiles(directory))
             {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    break;
+                }
                 try
                 {
                     File.Delete(file);
@@ -58,9 +69,9 @@ public async Task DeleteLogs()
         }
     }
 
-    public async IAsyncEnumerable<byte[]> ReadAllBytes()
+    public async IAsyncEnumerable<byte[]> ReadAllBytes([EnumeratorCancellation] CancellationToken cancellationToken)
     {
-        using var logLock = await FileLoggerDefaults.AcquireLock();
+        using var logLock = await FileLoggerDefaults.AcquireLock(cancellationToken);
 
         var componentName = Assembly.GetExecutingAssembly().GetName().Name;
         var directory = Path.Combine(FileLoggerDefaults.LogsFolderPath, $"{componentName}");
@@ -78,6 +89,10 @@ public async IAsyncEnumerable<byte[]> ReadAllBytes()
         {
             foreach (var chunk in File.ReadAllBytes(file).Chunk(50_000))
             {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    yield break;
+                }
                 yield return File.ReadAllBytes(file);
             }
         }