From 126505391e45a5f2260c0aee3022aa69f1a52f25 Mon Sep 17 00:00:00 2001
From: Kuinox <github@kuinox.io>
Date: Sat, 27 Jan 2024 17:00:45 +0100
Subject: [PATCH] Refactor IoWorker of debugger. Fixes #336. (#377)

---
 src/Draco.Debugger/Debugger.cs    |  8 ++---
 src/Draco.Debugger/IO/IoWorker.cs | 55 +++++++++++++++++--------------
 2 files changed, 34 insertions(+), 29 deletions(-)

diff --git a/src/Draco.Debugger/Debugger.cs b/src/Draco.Debugger/Debugger.cs
index 4e74ab5b2..22144d110 100644
--- a/src/Draco.Debugger/Debugger.cs
+++ b/src/Draco.Debugger/Debugger.cs
@@ -17,7 +17,7 @@ public sealed class Debugger
     /// <summary>
     /// The task that is completed, when the process has terminated.
     /// </summary>
-    public Task Terminated => this.terminatedCompletionSource.Task;
+    public Task Terminated => this.ioWorker.WorkLoopTask;
 
     /// <summary>
     /// The main user-module.
@@ -94,7 +94,6 @@ public event EventHandler<string> OnStandardError
     private readonly IoWorker<CorDebugProcess> ioWorker;
 
     private readonly TaskCompletionSource terminatedCompletionSource = new();
-    private readonly CancellationTokenSource terminateTokenSource = new();
 
     private readonly SessionCache sessionCache = new();
 
@@ -111,7 +110,7 @@ internal Debugger(
         this.ioWorker = ioWorker;
 
         this.InitializeEventHandler(cb);
-        ioWorker.Run(this.terminateTokenSource.Token);
+        ioWorker.Start();
     }
 
     private void ClearCache()
@@ -322,8 +321,7 @@ private void OnUnloadModuleHandler(object? sender, UnloadModuleCorDebugManagedCa
 
     private void OnExitProcessHandler(object? sender, ExitProcessCorDebugManagedCallbackEventArgs args)
     {
-        this.terminateTokenSource.Cancel();
-        this.terminatedCompletionSource.SetResult();
+        this.ioWorker.SignalStop();
         // TODO: Get exit code properly
         this.OnExited?.Invoke(sender, 0);
         this.Continue();
diff --git a/src/Draco.Debugger/IO/IoWorker.cs b/src/Draco.Debugger/IO/IoWorker.cs
index 8fa431f48..7a74109bf 100644
--- a/src/Draco.Debugger/IO/IoWorker.cs
+++ b/src/Draco.Debugger/IO/IoWorker.cs
@@ -28,8 +28,13 @@ internal sealed class IoWorker<TProcess>
     /// </summary>
     public StreamWriter StandardInput { get; }
 
+    /// <summary>
+    /// The task of the worker loop. Is always completed when the loop is not running.
+    /// </summary>
+    public Task WorkLoopTask { get; private set; } = Task.CompletedTask;
     private readonly TProcess process;
     private readonly RemoteIoHandles handles;
+    private readonly CancellationTokenSource _readCTS = new();
 
     public IoWorker(TProcess process, RemoteIoHandles handles)
     {
@@ -43,36 +48,38 @@ public IoWorker(TProcess process, RemoteIoHandles handles)
     /// </summary>
     /// <param name="cancellationToken">Can be used to cancel the worker.</param>
     /// <returns>The task of the worker loop.</returns>
-    public Task Run(CancellationToken cancellationToken) => Task.Run(async () =>
+    public void Start() =>
+        this.WorkLoopTask = Task.WhenAll(
+            this.ReadStdout(),
+            this.ReadStderr()
+        );
+
+    public void SignalStop() => this._readCTS.Cancel();
+
+    private async Task ReadStdout()
     {
         var stdoutReader = new StreamReader(this.handles.StandardOutputReader);
-        var stderrReader = new StreamReader(this.handles.StandardErrorReader);
-
         var stdoutBuffer = new char[BufferSize];
-        var stderrBuffer = new char[BufferSize];
-
-        var stdoutTask = null as Task<int>;
-        var stderrTask = null as Task<int>;
-
+        var cancellationToken = this._readCTS.Token;
         while (!cancellationToken.IsCancellationRequested)
         {
-            stdoutTask ??= stdoutReader.ReadAsync(stdoutBuffer, cancellationToken).AsTask();
-            stderrTask ??= stderrReader.ReadAsync(stderrBuffer, cancellationToken).AsTask();
+            var result = await stdoutReader.ReadAsync(stdoutBuffer, cancellationToken);
+            var str = new string(stdoutBuffer, 0, result);
+            this.OnStandardOut?.Invoke(this.process, str);
+        }
+    }
 
-            await Task.WhenAny(stdoutTask, stderrTask);
+    private async Task ReadStderr()
+    {
+        var stderrReader = new StreamReader(this.handles.StandardErrorReader);
+        var stderrBuffer = new char[BufferSize];
+        var cancellationToken = this._readCTS.Token;
 
-            if (stdoutTask.IsCompleted)
-            {
-                var str = new string(stdoutBuffer, 0, stdoutTask.Result);
-                this.OnStandardOut?.Invoke(this.process, str);
-                stdoutTask = null;
-            }
-            if (stderrTask.IsCompleted)
-            {
-                var str = new string(stderrBuffer, 0, stderrTask.Result);
-                this.OnStandardError?.Invoke(this.process, str);
-                stderrTask = null;
-            }
+        while (!cancellationToken.IsCancellationRequested)
+        {
+            var result = await stderrReader.ReadAsync(stderrBuffer, cancellationToken);
+            var str = new string(stderrBuffer, 0, result);
+            this.OnStandardError?.Invoke(this.process, str);
         }
-    }, cancellationToken);
+    }
 }