Fix console cancellation

#423 #393
cyanfish · Sep 2, 2024 · 58a392f · 58a392f
1 parent 802656c
commit 58a392f
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 6 deletions.
diff --git a/NAPS2.Lib/Automation/AutomatedScanning.cs b/NAPS2.Lib/Automation/AutomatedScanning.cs
@@ -33,6 +33,7 @@ internal class AutomatedScanning
 
     private readonly ConsoleOutput _output;
     private readonly AutomatedScanningOptions _options;
+    private readonly CancellationTokenSource _cts = new();
     private PdfEncryption _parsedEncryptConfigOption = null!;
     private List<List<ProcessedImage>> _scanList = null!;
     private int _pagesScanned;
@@ -78,6 +79,15 @@ private void OutputVerbose(string value, params object[] args)
 
     public async Task Execute()
     {
+        Console.CancelKeyPress += (_, e) =>
+        {
+            if (!_cts.IsCancellationRequested)
+            {
+                Console.WriteLine(ConsoleResources.Cancelling);
+                _cts.Cancel();
+                e.Cancel = true;
+            }
+        };
         bool hasUnexpectedException = false;
         try
         {
@@ -115,6 +125,8 @@ public async Task Execute()
                 await ImportImages();
             }
 
+            if (_cts.IsCancellationRequested) return;
+
             ConfigureOcr();
 
             if (_options.Number > 0)
@@ -130,6 +142,8 @@ public async Task Execute()
                 await PerformScan(profile);
             }
 
+            if (_cts.IsCancellationRequested) return;
+
             ReorderScannedImages();
 
             if (_options.OutputPath != null)
@@ -245,9 +259,7 @@ private async Task ListDevices()
         };
         await SetProfileOverrides(profile);
 
-        var cts = new CancellationTokenSource();
-        Console.CancelKeyPress += (_, _) => cts.Cancel();
-        await foreach (var device in _scanPerformer.GetDevices(profile, cts.Token))
+        await foreach (var device in _scanPerformer.GetDevices(profile, _cts.Token))
         {
             _output.Writer.WriteLine(device.Name);
         }
@@ -332,7 +344,7 @@ private async Task ImportImages()
                         PatchTOnly = true
                     }
                 };
-                await foreach (var image in _fileImporter.Import(actualPath, importParams))
+                await foreach (var image in _fileImporter.Import(actualPath, importParams, _cts.Token))
                 {
                     scan.Add(PostProcessImportedImage(image));
                 }
@@ -443,7 +455,7 @@ private async Task EmailScannedImages()
             }
 
             OutputVerbose(ConsoleResources.SendingEmail);
-            if (await _emailProviderFactory.Default.SendEmail(message))
+            if (await _emailProviderFactory.Default.SendEmail(message, _cts.Token))
             {
                 OutputVerbose(ConsoleResources.EmailSent);
             }
@@ -590,6 +602,7 @@ private async Task DoExportToImageFiles(string outputPath)
         foreach (var scan in _scanList)
         {
             var op = _operationFactory.Create<SaveImagesOperation>();
+            _cts.Token.Register(op.Cancel);
             int i = -1;
             op.StatusChanged += (sender, args) =>
             {
@@ -673,6 +686,7 @@ private async Task<bool> DoExportToPdf(string path, bool email)
         foreach (var fileContents in _scanList)
         {
             var op = _operationFactory.Create<SavePdfOperation>();
+            _cts.Token.Register(op.Cancel);
             int i = -1;
             op.StatusChanged += (sender, args) =>
             {
@@ -731,7 +745,7 @@ private async Task PerformScan(ScanProfile profile)
                 DetectPatchT = _options.SplitPatchT,
                 OcrParams = _ocrParams
             };
-            await foreach (var image in _scanPerformer.PerformScan(profile, scanParams))
+            await foreach (var image in _scanPerformer.PerformScan(profile, scanParams, cancelToken: _cts.Token))
             {
                 ReceiveScannedImage(image);
             }

diff --git a/NAPS2.Lib/Lang/ConsoleResources/ConsoleResources.Designer.cs b/NAPS2.Lib/Lang/ConsoleResources/ConsoleResources.Designer.cs
diff --git a/NAPS2.Lib/Lang/ConsoleResources/ConsoleResources.resx b/NAPS2.Lib/Lang/ConsoleResources/ConsoleResources.resx
@@ -263,4 +263,7 @@ Use the "--importpassword" option.</value>
   <data name="DriverRequired" xml:space="preserve">
     <value>The --driver option must be set to use the --device or --listdevices option. Possible values: {0}</value>
   </data>
+  <data name="Cancelling" xml:space="preserve">
+    <value>Cancelling... Press Ctrl+C again to terminate.</value>
+  </data>
 </root>
diff --git a/NAPS2.Sdk/Scan/Internal/Twain/RemoteTwainController.cs b/NAPS2.Sdk/Scan/Internal/Twain/RemoteTwainController.cs
@@ -37,6 +37,9 @@ public async Task StartScan(ScanOptions options, ITwainEvents twainEvents, Cance
             // We need to report cancellation so that TwainImageProcessor doesn't return a partial image
             _scanningContext.Logger.LogDebug("NAPS2.TW - Sending cancel event");
             twainEvents.TransferCanceled(new TwainTransferCanceled());
+            // We also want to wait until the worker closes so we guarantee the parent process doesn't die before the
+            // TWAIN process has a chance to clean up
+            await workerContext.Stop();
         }
     }