fix: interrupt solvers for early termination with --early-exit (a16z#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.

- 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()`.

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

- After interrupting the solvers, the global list is cleared to prevent
  memory leaks.

With these changes, when `args.early_exit` is true and a counterexample
is found, the `future_callback` function interrupts all the running
solver contexts, effectively stopping the unnecessary computations and
allowing Halmos to terminate early.

src/halmos/__main__.py

@@ -65,6 +65,8 @@
 VERBOSITY_TRACE_PATHS = 4
 VERBOSITY_TRACE_CONSTRUCTOR = 5
 
+solver_contexts = []
+
 
 @dataclass(frozen=True)
 class FunctionInfo:
@@ -657,25 +659,36 @@ 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)
 
         model, is_valid, index, result = m.model, m.is_valid, m.index, m.result
         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,20 @@ 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()
+
+            # Clear the global list to prevent memory leaks
+            solver_contexts.clear()
+
     for idx, ex in enumerate(exs):
         result_exs.append(ex)
 
@@ -739,16 +766,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 +1072,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