main: add cooperative timeouts

in addition to the hard timeout that simply kills the process, we now also pass a timeout to the SMT solver directly so that we can have more graceful exits.

src/main.rs

@@ -28,7 +28,7 @@ use intrinsic::{distributions::init_distributions, list::init_lists};
 use opt::{boolify::Boolify, RemoveParens};
 use procs::add_default_specs;
 use proof_rules::init_encodings;
-use resource_limits::{await_with_resource_limits, LimitError};
+use resource_limits::{await_with_resource_limits, LimitError, LimitsRef};
 use thiserror::Error;
 use timing::DispatchBuilder;
 use tokio::task::JoinError;
@@ -262,25 +262,34 @@ pub async fn verify_files(
     files: &Arc<Mutex<Files>>,
     user_files: Vec<FileId>,
 ) -> Result<bool, VerifyError> {
-    let handle = {
+    let handle = |limits_ref: LimitsRef| {
         let options = options.clone();
         let files = files.clone();
-        tokio::task::spawn_blocking(move || verify_files_main(&options, &files, &user_files))
+        tokio::task::spawn_blocking(move || {
+            verify_files_main(&options, limits_ref, &files, &user_files)
+        })
     };
-    // Unpacking lots of results with `.await??` :-)
+    // Unpacking lots of Results with `.await??` :-)
     await_with_resource_limits(options.timeout, options.mem_limit, handle).await??
 }
 
 /// Synchronously verify the given source code. This is used for tests. The
 /// `--werr` option is enabled by default.
 #[cfg(test)]
 pub(crate) fn verify_test(source: &str) -> (Result<bool, VerifyError>, Files) {
+    use std::time::{Duration, Instant};
+
     let mut files = Files::new();
     let file_id = files.add(SourceFilePath::Builtin, source.to_owned()).id;
     let files_mutex = Mutex::new(files);
     let mut options = Options::default();
     options.werr = true;
-    let res = verify_files_main(&options, &files_mutex, &[file_id]);
+    let limits_ref = LimitsRef::new(
+        options
+            .timeout
+            .map(|seconds| Instant::now() + Duration::from_secs(seconds)),
+    );
+    let res = verify_files_main(&options, limits_ref, &files_mutex, &[file_id]);
     (res, files_mutex.into_inner().unwrap())
 }
 
@@ -355,6 +364,7 @@ pub(crate) fn single_desugar_test(source: &str) -> Result<String, VerifyError> {
 /// Synchronously verify the given files.
 fn verify_files_main(
     options: &Options,
+    limits_ref: LimitsRef,
     files_mutex: &Mutex<Files>,
     user_files: &[FileId],
 ) -> Result<bool, VerifyError> {
@@ -462,7 +472,7 @@ fn verify_files_main(
         let mut vc_expr = verify_unit.vcgen(&vcgen).unwrap();
 
         // 6. Unfolding
-        unfold_expr(options, &tcx, &mut vc_expr);
+        unfold_expr(options, &limits_ref, &tcx, &mut vc_expr)?;
 
         // 7. Quantifier elimination
         if !options.no_qelim {
@@ -515,7 +525,7 @@ fn verify_files_main(
         let sat_span = info_span!("SAT check");
         let sat_entered = sat_span.enter();
 
-        let mut prover = mk_valid_query_prover(&ctx, &smt_translate, &valid_query);
+        let mut prover = mk_valid_query_prover(&limits_ref, &ctx, &smt_translate, &valid_query);
         let smtlib = get_smtlib(options, &prover);
 
         let result = prover.check_proof();
@@ -539,9 +549,11 @@ fn verify_files_main(
             info!("Z3 tracing output will be written to `z3.log`.");
         }
 
-        // If the solver was interrupted from the keyboard, exit now.
-        if prover.get_reason_unknown() == Some(ReasonUnknown::Interrupted) {
-            return Err(VerifyError::Interrupted);
+        // Handle reasons to stop the verifier.
+        match prover.get_reason_unknown() {
+            Some(ReasonUnknown::Interrupted) => return Err(VerifyError::Interrupted),
+            Some(ReasonUnknown::Timeout) => return Err(LimitError::Timeout.into()),
+            _ => {}
         }
 
         all_proven = all_proven && result == ProveResult::Proof;
@@ -628,16 +640,22 @@ fn apply_boolify_opt(vc_expr_eq_infinity: &mut Expr) {
     (Boolify {}).visit_expr(vc_expr_eq_infinity).unwrap();
 }
 
-fn unfold_expr(options: &Options, tcx: &TyCtx, vc_expr: &mut VcUnit) {
+fn unfold_expr(
+    options: &Options,
+    limits_ref: &LimitsRef,
+    tcx: &TyCtx,
+    vc_expr: &mut VcUnit,
+) -> Result<(), LimitError> {
     let span = info_span!("unfolding");
     let _entered = span.enter();
     if !options.strict {
         let ctx = Context::new(&Config::default());
         let smt_ctx = SmtCtx::new(&ctx, tcx);
-        let mut unfolder = Unfolder::new(&smt_ctx);
-        unfolder.visit_expr(&mut vc_expr.expr).unwrap();
+        let mut unfolder = Unfolder::new(limits_ref.clone(), &smt_ctx);
+        unfolder.visit_expr(&mut vc_expr.expr)
     } else {
         apply_subst(tcx, &mut vc_expr.expr);
+        Ok(())
     }
 }
 
@@ -651,12 +669,17 @@ fn mk_z3_ctx(options: &Options) -> Context {
 }
 
 fn mk_valid_query_prover<'smt, 'ctx>(
+    limits_ref: &LimitsRef,
     ctx: &'ctx Context,
     smt_translate: &TranslateExprs<'smt, 'ctx>,
     valid_query: &Bool<'ctx>,
 ) -> Prover<'ctx> {
     // create the prover and set the params
     let mut prover = Prover::new(ctx);
+    if let Some(remaining) = limits_ref.time_left() {
+        prover.set_timeout(remaining);
+    }
+
     // add assumptions (from axioms and locals) to the prover
     smt_translate
         .ctx

src/opt/unfolder.rs

@@ -33,13 +33,16 @@ use crate::{
         BinOpKind, Expr, ExprBuilder, ExprData, ExprKind, Shared, Span, SpanVariant, Spanned,
         TyKind, UnOpKind,
     },
+    resource_limits::{LimitError, LimitsRef},
     smt::SmtCtx,
     vc::subst::Subst,
 };
 
 use crate::smt::translate_exprs::TranslateExprs;
 
 pub struct Unfolder<'smt, 'ctx> {
+    limits_ref: LimitsRef,
+
     /// The expressions may contain substitutions. We keep track of those.
     subst: Subst<'smt>,
 
@@ -54,8 +57,9 @@ pub struct Unfolder<'smt, 'ctx> {
 }
 
 impl<'smt, 'ctx> Unfolder<'smt, 'ctx> {
-    pub fn new(ctx: &'smt SmtCtx<'ctx>) -> Self {
+    pub fn new(limits_ref: LimitsRef, ctx: &'smt SmtCtx<'ctx>) -> Self {
         Unfolder {
+            limits_ref,
             subst: Subst::new(ctx.tcx()),
             translate: TranslateExprs::new(ctx),
             prover: Prover::new(ctx.ctx()),
@@ -147,9 +151,11 @@ impl<'smt, 'ctx> Unfolder<'smt, 'ctx> {
 }
 
 impl<'smt, 'ctx> VisitorMut for Unfolder<'smt, 'ctx> {
-    type Err = ();
+    type Err = LimitError;
 
     fn visit_expr(&mut self, e: &mut Expr) -> Result<(), Self::Err> {
+        self.limits_ref.check_limits()?;
+
         let span = e.span;
         let ty = e.ty.clone().unwrap();
         match &mut e.deref_mut().kind {
@@ -276,15 +282,19 @@ fn negate_expr(expr: Expr) -> Expr {
 #[cfg(test)]
 mod test {
     use super::Unfolder;
-    use crate::{ast::visit::VisitorMut, fuzz_expr_opt_test, opt::fuzz_test, smt::SmtCtx};
+    use crate::{
+        ast::visit::VisitorMut, fuzz_expr_opt_test, opt::fuzz_test, resource_limits::LimitsRef,
+        smt::SmtCtx,
+    };
 
     #[test]
     fn fuzz_unfolder() {
         fuzz_expr_opt_test!(|mut expr| {
             let tcx = fuzz_test::mk_tcx();
             let z3_ctx = z3::Context::new(&z3::Config::default());
             let smt_ctx = SmtCtx::new(&z3_ctx, &tcx);
-            let mut unfolder = Unfolder::new(&smt_ctx);
+            let limits_ref = LimitsRef::new(None);
+            let mut unfolder = Unfolder::new(limits_ref, &smt_ctx);
             unfolder.visit_expr(&mut expr).unwrap();
             expr
         })

src/resource_limits.rs

@@ -2,18 +2,23 @@
 
 use std::{
     future::{pending, Future},
-    time::Duration,
+    sync::{
+        atomic::{AtomicU8, Ordering},
+        Arc,
+    },
+    time::{Duration, Instant},
 };
 
 use simple_process_stats::ProcessStats;
 use thiserror::Error;
 use tokio::{
     select,
-    time::{interval, timeout, MissedTickBehavior},
+    time::{error::Elapsed, interval, timeout, MissedTickBehavior},
 };
 use tracing::error;
 
 const CHECK_MEM_USAGE_INTERVAL: Duration = Duration::from_millis(20);
+pub const HARD_TIMEOUT_SLACK: Duration = Duration::from_millis(500);
 
 /// A timeout or an out-of-memory condition.
 #[derive(Debug, Clone, Copy, Error)]
@@ -24,42 +29,63 @@ pub enum LimitError {
     Oom,
 }
 
-/// Wait for the given future `fut`. Returns [`LimitError::Timeout`] if the
-/// future took too long to complete. If the memory limit was exceeded, return
-/// [`LimitError::Oom`]. Otherwise, return the result wrapped in [`Result::Ok`].
+/// Wait for the future  returned by `fut`. Returns [`LimitError::Timeout`] if
+/// the future took too long to complete. If the memory limit was exceeded,
+/// return [`LimitError::Oom`]. Otherwise, return the result wrapped in
+/// [`Result::Ok`].
+///
+/// The `fut` is given a [`LimitsRef`] which it is supposed to check
+/// periodically. After the given duration plus [`HARD_TIMEOUT_SLACK`], the
+/// future will return a [`LimitError::Timeout`] in case `fut` did not catch the
+/// timeout itself.
 ///
 /// Note that the memory limit is checked for the whole process by a background
 /// thread. Therefore, the memory limit is not specific to the given future.
-pub async fn await_with_resource_limits<T: Unpin>(
+pub async fn await_with_resource_limits<T, F>(
     timeout_secs: Option<u64>,
     mem_limit_mb: Option<u64>,
-    fut: impl Future<Output = T>,
-) -> Result<T, LimitError> {
+    fut: impl FnOnce(LimitsRef) -> F,
+) -> Result<T, LimitError>
+where
+    T: Unpin,
+    F: Future<Output = T>,
+{
     if let Some(timeout_secs) = timeout_secs {
-        let fut = timeout(Duration::from_secs(timeout_secs), fut);
-        if let Some(mem_mbs) = mem_limit_mb {
+        let duration = Duration::from_secs(timeout_secs);
+        let limits_ref = LimitsRef::new(Some(Instant::now() + duration));
+
+        let hard_duration = duration + HARD_TIMEOUT_SLACK;
+        let fut = timeout(hard_duration, fut(limits_ref.clone()));
+        let res = if let Some(mem_mbs) = mem_limit_mb {
             select! {
                 _ = wait_for_oom(mem_mbs) => {
                     Err(LimitError::Oom)
                 }
                 res = fut => {
-                    res.map_err(|_| LimitError::Timeout)
+                    res.map_err(|_: Elapsed| LimitError::Timeout)
                 }
             }
         } else {
-            fut.await.map_err(|_| LimitError::Timeout)
+            fut.await.map_err(|_: Elapsed| LimitError::Timeout)
+        };
+        if let Err(err) = res {
+            limits_ref.set_error(err);
         }
+        res
     } else if let Some(mem_mbs) = mem_limit_mb {
+        let limits_ref = LimitsRef::new(None);
         select! {
             _ = wait_for_oom(mem_mbs) => {
+                limits_ref.set_error(LimitError::Oom);
                 Err(LimitError::Oom)
             }
-            res = fut => {
+            res = fut(limits_ref.clone()) => {
                 Ok(res)
             }
         }
     } else {
-        Ok(fut.await)
+        let limits_ref = LimitsRef::new(None);
+        Ok(fut(limits_ref).await)
     }
 }
 
@@ -92,3 +118,51 @@ async fn wait_for_oom(mem_limit_mb: u64) {
         }
     }
 }
+
+/// An object to pass around that allows to check whether the resource limits
+/// were exceeded and the task needs to be stopped as a consequence.
+#[derive(Debug, Clone)]
+pub struct LimitsRef(Arc<LimitsRefData>);
+
+#[derive(Debug)]
+struct LimitsRefData {
+    done: AtomicU8,
+    timeout: Option<Instant>,
+}
+
+impl LimitsRef {
+    pub fn new(timeout: Option<Instant>) -> Self {
+        LimitsRef(Arc::new(LimitsRefData {
+            done: AtomicU8::new(0),
+            timeout,
+        }))
+    }
+
+    pub fn check_limits(&self) -> Result<(), LimitError> {
+        match self.0.done.load(Ordering::Relaxed) {
+            0 => Ok(()),
+            1 => Err(LimitError::Timeout),
+            2 => Err(LimitError::Oom),
+            _ => unreachable!(),
+        }
+    }
+
+    /// Returns the time left or `None` if there was no timeout. Will return
+    /// zero if the timeout has elapsed.
+    pub fn time_left(&self) -> Option<Duration> {
+        Some(self.0.timeout?.duration_since(Instant::now()))
+    }
+
+    /// Sets an error. Will only store the first error, any subsequent errors
+    /// are discarded.
+    fn set_error(&self, err: LimitError) {
+        let new = match err {
+            LimitError::Timeout => 1,
+            LimitError::Oom => 2,
+        };
+        let _ = self
+            .0
+            .done
+            .compare_exchange(0, new, Ordering::Acquire, Ordering::Relaxed);
+    }
+}

z3rro/src/prover.rs

@@ -1,8 +1,10 @@
 //! Not a SAT solver, but a prover. There's a difference.
 
+use std::time::Duration;
+
 use z3::{ast::Bool, Context, Model, Params, SatResult, Solver};
 
-use crate::util::ReasonUnknown;
+use crate::util::{set_solver_timeout, ReasonUnknown};
 
 /// The result of a prove query.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -48,6 +50,11 @@ impl<'ctx> Prover<'ctx> {
         self.solver.set_params(params);
     }
 
+    /// Set a timeout using [`set_solver_timeout`].
+    pub fn set_timeout(&mut self, duration: Duration) {
+        set_solver_timeout(&self.solver, duration);
+    }
+
     /// Add an assumption to this prover.
     pub fn add_assumption(&mut self, value: &Bool<'ctx>) {
         self.solver.assert(value);