From b5d79d0a4f30485809a51c05bc44384bf01b9816 Mon Sep 17 00:00:00 2001
From: Dongbo Wang <dongbow@microsoft.com>
Date: Tue, 19 Apr 2022 10:08:30 -0700
Subject: [PATCH] Update the cancellation related code

---
 PSReadLine/ReadLine.cs | 35 ++++++++++++++++++-----------------
 1 file changed, 18 insertions(+), 17 deletions(-)

diff --git a/PSReadLine/ReadLine.cs b/PSReadLine/ReadLine.cs
index 2e8f21a2..e76c78d6 100644
--- a/PSReadLine/ReadLine.cs
+++ b/PSReadLine/ReadLine.cs
@@ -91,7 +91,7 @@ bool IPSConsoleReadLineMockableMethods.RunspaceIsRemote(Runspace runspace)
             return runspace?.ConnectionInfo != null;
         }
 
-        private void ReadOneOrMoreKeys(CancellationToken cancellationToken)
+        private void ReadOneOrMoreKeys()
         {
             _readkeyStopwatch.Restart();
             while (_console.KeyAvailable)
@@ -143,8 +143,8 @@ private void ReadOneOrMoreKeys(CancellationToken cancellationToken)
                 }
                 while (_charMap.KeyAvailable)
                 {
-                    var key = PSKeyInfo.FromConsoleKeyInfo(_charMap.ReadKey());
-                    if (cancellationToken.IsCancellationRequested)
+                    ConsoleKeyInfo keyInfo = _charMap.ReadKey();
+                    if (_cancelReadCancellationToken.IsCancellationRequested)
                     {
                         // If PSReadLine is running under a host that can cancel it, the
                         // cancellation will come at a time when ReadKey is stuck waiting for input.
@@ -152,6 +152,8 @@ private void ReadOneOrMoreKeys(CancellationToken cancellationToken)
                         // discard this key since we were already canceled.
                         continue;
                     }
+
+                    var key = PSKeyInfo.FromConsoleKeyInfo(keyInfo);
                     _lastNKeys.Enqueue(key);
                     _queuedKeys.Enqueue(key);
                 }
@@ -163,13 +165,12 @@ private void ReadKeyThreadProc()
             while (true)
             {
                 // Wait until ReadKey tells us to read a key (or it's time to exit).
-                int handleId = WaitHandle.WaitAny(_singleton._threadProcWaitHandles);
+                int handleId = WaitHandle.WaitAny(_threadProcWaitHandles);
                 if (handleId == 1) // It was the _closingWaitHandle that was signaled.
                     break;
 
-                var localCancellationToken = _singleton._cancelReadCancellationToken;
-                ReadOneOrMoreKeys(localCancellationToken);
-                if (localCancellationToken.IsCancellationRequested)
+                ReadOneOrMoreKeys();
+                if (_cancelReadCancellationToken.IsCancellationRequested)
                 {
                     continue;
                 }
@@ -395,7 +396,7 @@ public static string ReadLine(
                     }
 
                     _singleton._cancelReadCancellationToken = cancellationToken;
-                    _singleton._requestKeyWaitHandles[2] = _singleton._cancelReadCancellationToken.WaitHandle;
+                    _singleton._requestKeyWaitHandles[2] = cancellationToken.WaitHandle;
                     return _singleton.InputLoop();
                 }
                 catch (OperationCanceledException)
@@ -873,11 +874,11 @@ private void DelayedOneTimeInitialize()
             _killIndex = -1; // So first add indexes 0.
             _killRing = new List<string>(Options.MaximumKillRingCount);
 
-            _singleton._readKeyWaitHandle = new AutoResetEvent(false);
-            _singleton._keyReadWaitHandle = new AutoResetEvent(false);
-            _singleton._closingWaitHandle = new ManualResetEvent(false);
-            _singleton._requestKeyWaitHandles = new WaitHandle[] {_singleton._keyReadWaitHandle, _singleton._closingWaitHandle, _defaultCancellationToken.WaitHandle};
-            _singleton._threadProcWaitHandles = new WaitHandle[] {_singleton._readKeyWaitHandle, _singleton._closingWaitHandle};
+            _readKeyWaitHandle = new AutoResetEvent(false);
+            _keyReadWaitHandle = new AutoResetEvent(false);
+            _closingWaitHandle = new ManualResetEvent(false);
+            _requestKeyWaitHandles = new WaitHandle[] {_keyReadWaitHandle, _closingWaitHandle, null};
+            _threadProcWaitHandles = new WaitHandle[] {_readKeyWaitHandle, _closingWaitHandle};
 
             // This is for a "being hosted in an alternate appdomain scenario" (the
             // DomainUnload event is not raised for the default appdomain). It allows us
@@ -887,13 +888,13 @@ private void DelayedOneTimeInitialize()
             {
                 AppDomain.CurrentDomain.DomainUnload += (x, y) =>
                 {
-                    _singleton._closingWaitHandle.Set();
-                    _singleton._readKeyThread.Join(); // may need to wait for history to be written
+                    _closingWaitHandle.Set();
+                    _readKeyThread.Join(); // may need to wait for history to be written
                 };
             }
 
-            _singleton._readKeyThread = new Thread(_singleton.ReadKeyThreadProc) {IsBackground = true, Name = "PSReadLine ReadKey Thread"};
-            _singleton._readKeyThread.Start();
+            _readKeyThread = new Thread(ReadKeyThreadProc) {IsBackground = true, Name = "PSReadLine ReadKey Thread"};
+            _readKeyThread.Start();
         }
 
         private static void Chord(ConsoleKeyInfo? key = null, object arg = null)