From 7b392b3510394a8c26757fd0bf4c00b5e7aec2a6 Mon Sep 17 00:00:00 2001
From: Saurabh Chalke <saurabhchalke@gmail.com>
Date: Wed, 8 May 2024 01:54:51 +0000
Subject: [PATCH] fix: interrupt solvers and cancel futures for early
 termination (#243)

The --early-exit flag is now properly handled to ensure that Halmos
terminates immediately after finding the first counterexample, avoiding
unnecessary solver computations.

The previous implementation allowed solver threads to continue running
even after a counterexample was found, leading to inefficient resource
usage and delayed termination.

The updated code introduces the following changes:

- A global list `solver_contexts` is introduced to store the solver
  contexts created in the `solve` function.

- The `solve` function is modified to append the solver context to the
  global list before returning.

- The list of futures is stored in a variable before submitting them to
  the thread pool.

- In the `future_callback` function, if `args.early_exit` is true and a
  counterexample is found:
  - All the solver contexts in the global list are interrupted using
    `ctx.interrupt()`.
  - The remaining futures are canceled using `future.cancel()` to stop
    unnecessary computations.
  - The global lists are cleared to prevent memory leaks.

- Exception handling is added to the interruption process to avoid
  segmentation faults.

With these changes, when `args.early_exit` is true and a counterexample
is found, the `future_callback` function interrupts all the running
solver contexts and cancels the remaining futures, effectively stopping
the unnecessary computations and allowing Halmos to terminate early.
---
 src/halmos/__main__.py | 45 ++++++++++++++++++++++++++++++++----------
 1 file changed, 35 insertions(+), 10 deletions(-)

diff --git a/src/halmos/__main__.py b/src/halmos/__main__.py
index 45c8331f6..885cd14f8 100644
--- a/src/halmos/__main__.py
+++ b/src/halmos/__main__.py
@@ -65,6 +65,8 @@
 VERBOSITY_TRACE_PATHS = 4
 VERBOSITY_TRACE_CONSTRUCTOR = 5
 
+solver_contexts = []
+
 
 @dataclass(frozen=True)
 class FunctionInfo:
@@ -657,7 +659,14 @@ def run(
     counterexamples = []
     traces = {}
 
+    early_exit = False
+
     def future_callback(future_model):
+        nonlocal early_exit
+
+        if early_exit:
+            return
+
         m = future_model.result()
         models.append(m)
 
@@ -665,17 +674,21 @@ def future_callback(future_model):
         if result == unsat:
             return
 
+        counterexample_found = False
+
         # model could be an empty dict here
         if model is not None:
             if is_valid:
                 print(red(f"Counterexample: {render_model(model)}"))
                 counterexamples.append(model)
+                counterexample_found = True
             else:
                 warn(
                     COUNTEREXAMPLE_INVALID,
                     f"Counterexample (potentially invalid): {render_model(model)}",
                 )
                 counterexamples.append(model)
+                counterexample_found = True
         else:
             warn(COUNTEREXAMPLE_UNKNOWN, f"Counterexample: {result}")
 
@@ -691,6 +704,26 @@ def future_callback(future_model):
             )
             print(traces[index], end="")
 
+        if args.early_exit and counterexample_found:
+            early_exit = True
+
+            # Interrupt all the solver contexts
+            for ctx in solver_contexts:
+                try:
+                    ctx.interrupt()
+                except Exception as e:
+                    if args.debug:
+                        traceback.print_exc()
+
+            # Cancel the remaining futures
+            for future in future_models:
+                if not future.done():
+                    future.cancel()
+
+            # Clear the global lists to prevent memory leaks
+            solver_contexts.clear()
+            future_models.clear()
+
     for idx, ex in enumerate(exs):
         result_exs.append(ex)
 
@@ -739,16 +772,7 @@ def future_callback(future_model):
             f"# of potential paths involving assertion violations: {len(future_models)} / {len(result_exs)}  (--solver-threads {args.solver_threads})"
         )
 
-    if args.early_exit:
-        while not (
-            len(counterexamples) > 0 or all([fm.done() for fm in future_models])
-        ):
-            time.sleep(1)
-
-        thread_pool.shutdown(wait=False, cancel_futures=True)
-
-    else:
-        thread_pool.shutdown(wait=True)
+    thread_pool.shutdown(wait=True)
 
     counter = Counter(str(m.result) for m in models)
     if counter["sat"] > 0:
@@ -1054,6 +1078,7 @@ def solve(
         solver.from_string(query)
         result = solver.check()
         model = copy_model(solver.model()) if result == sat else None
+        solver_contexts.append(solver.ctx)
         return result, model