From 2b662ef697797004703c0c2480c11293b1cff663 Mon Sep 17 00:00:00 2001
From: Remy Willems <rwillems@amazon.com>
Date: Fri, 9 Aug 2024 12:20:49 +0200
Subject: [PATCH] Add more cancellation

---
 Source/BoogieDriver/BoogieDriver.cs       | 11 +++++------
 Source/ExecutionEngine/ExecutionEngine.cs | 15 ++++++++++++---
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/Source/BoogieDriver/BoogieDriver.cs b/Source/BoogieDriver/BoogieDriver.cs
index 79b86bdd6..3e69ea093 100644
--- a/Source/BoogieDriver/BoogieDriver.cs
+++ b/Source/BoogieDriver/BoogieDriver.cs
@@ -21,6 +21,11 @@ public static int Main(string[] args)
       {
         return 1;
       }
+      var source = new CancellationTokenSource();
+      if (options.ProcessTimeLimit != 0)
+      {
+        source.CancelAfter(TimeSpan.FromSeconds(options.ProcessTimeLimit));
+      }
       using var executionEngine = ExecutionEngine.CreateWithoutSharedCache(options);
       
       if (options.ProcessInfoFlags())
@@ -64,12 +69,6 @@ public static int Main(string[] args)
 
       Helpers.ExtraTraceInformation(options, "Becoming sentient");
 
-      var source = new CancellationTokenSource();
-      if (options.ProcessTimeLimit != 0)
-      {
-        source.CancelAfter(TimeSpan.FromSeconds(options.ProcessTimeLimit));
-      }
-
       var success = executionEngine.ProcessFiles(Console.Out, fileList, cancellationToken: source.Token).Result;
       if (options.XmlSink != null)
       {
diff --git a/Source/ExecutionEngine/ExecutionEngine.cs b/Source/ExecutionEngine/ExecutionEngine.cs
index a1e026246..4f2529c1f 100644
--- a/Source/ExecutionEngine/ExecutionEngine.cs
+++ b/Source/ExecutionEngine/ExecutionEngine.cs
@@ -93,6 +93,7 @@ public async Task<bool> ProcessFiles(TextWriter output, IList<string> fileNames,
       if (Options.VerifySeparately && 1 < fileNames.Count) {
         var success = true;
         foreach (var fileName in fileNames) {
+          cancellationToken.ThrowIfCancellationRequested();
           success &= await ProcessFiles(output, new List<string> { fileName }, lookForSnapshots, fileName, cancellationToken: cancellationToken);
         }
         return success;
@@ -102,6 +103,7 @@ public async Task<bool> ProcessFiles(TextWriter output, IList<string> fileNames,
         var snapshotsByVersion = LookForSnapshots(fileNames);
         var success = true;
         foreach (var snapshots in snapshotsByVersion) {
+          cancellationToken.ThrowIfCancellationRequested();
           // BUG: Reusing checkers during snapshots doesn't work, even though it should. We create a new engine (and thus checker pool) to workaround this.
           using var engine = new ExecutionEngine(Options, Cache,
             CustomStackSizePoolTaskScheduler.Create(StackSize, Options.VcsCores), true);
@@ -110,14 +112,21 @@ public async Task<bool> ProcessFiles(TextWriter output, IList<string> fileNames,
         return success;
       }
 
+      return await ProcessProgramFiles(output, fileNames, programId, cancellationToken);
+    }
+
+    private Task<bool> ProcessProgramFiles(TextWriter output, IList<string> fileNames, string programId,
+      CancellationToken cancellationToken)
+    {
       using XmlFileScope xf = new XmlFileScope(Options.XmlSink, fileNames[^1]);
       Program program = ParseBoogieProgram(fileNames, false);
       var bplFileName = fileNames[^1];
-      if (program == null) {
-        return true;
+      if (program == null)
+      {
+        return Task.FromResult(true);
       }
 
-      return await ProcessProgram(output, program, bplFileName, programId, cancellationToken: cancellationToken);
+      return ProcessProgram(output, program, bplFileName, programId, cancellationToken: cancellationToken);
     }
 
     [Obsolete("Please inline this method call")]