let starlark evaluations during analysis be canceled

Summary:
Today, starlark evaluations must run to completion. This inserts an "infrequent instruction check" in the evaluator that is called every so often, in which we now check whether the `CancellationContext` has been notified. If so, we return with `EvaluatorError:Cancelled`.

#internal
This should allow eagerbuild's use of the `--premptible` flag to exit early from starlark evaluations and not hold up the next build in "synchronization and waiting".

This also provides a good spot for T219887296 to be addressed

Reviewed By: JakobDegen

Differential Revision: D73694716

fbshipit-source-id: 718253cf9e37fc8fee3de60f9f3075211f534b19

Author: Jeremy Braun

starlark/src/eval/bc/instr_impl.rs

@@ -1276,6 +1276,9 @@ impl BcInstr for InstrContinue {
             BcAddrOffset,
         ),
     ) -> InstrControl<'v, 'b> {
+        if let Err(e) = eval.run_infrequent_instr_checks() {
+            return InstrControl::Err(e);
+        }
         let iter = frame.get_bc_slot(*iter);
         let loop_depth = *loop_depth;
         let i = frame.get_iter_index(loop_depth);
@@ -1475,6 +1478,7 @@ impl BcFrozenCallable for FrozenValue {
         args: &Arguments<'v, '_>,
         eval: &mut Evaluator<'v, '_, '_>,
     ) -> crate::Result<Value<'v>> {
+        eval.run_infrequent_instr_checks()?;
         self.to_value().invoke_with_loc(Some(location), args, eval)
     }
 }
@@ -1487,6 +1491,7 @@ impl BcFrozenCallable for FrozenValueTyped<'static, FrozenDef> {
         args: &Arguments<'v, '_>,
         eval: &mut Evaluator<'v, '_, '_>,
     ) -> crate::Result<Value<'v>> {
+        eval.run_infrequent_instr_checks()?;
         eval.with_call_stack(self.to_value(), Some(location), |eval| {
             self.as_ref().invoke(self.to_value(), args, eval)
         })
@@ -1501,6 +1506,7 @@ impl BcFrozenCallable for BcNativeFunction {
         args: &Arguments<'v, '_>,
         eval: &mut Evaluator<'v, '_, '_>,
     ) -> crate::Result<Value<'v>> {
+        eval.run_infrequent_instr_checks()?;
         eval.with_call_stack(self.to_value(), Some(location), |eval| {
             self.invoke(args, eval)
         })
@@ -1545,6 +1551,7 @@ impl<A: BcCallArgs<Symbol>> InstrNoFlowImpl for InstrCallImpl<A> {
         _ip: BcPtrAddr,
         (this, args, span, target): &(BcSlotIn, A, FrozenRef<'static, FrameSpan>, BcSlotOut),
     ) -> crate::Result<()> {
+        eval.run_infrequent_instr_checks()?;
         let f = frame.get_bc_slot(*this);
         let arguments = Arguments(args.pop_from_stack(frame));
         let r = f.invoke_with_loc(Some(*span), &arguments, eval)?;
@@ -1565,6 +1572,7 @@ impl<F: BcFrozenCallable, A: BcCallArgs<Symbol>> InstrNoFlowImpl
         _ip: BcPtrAddr,
         (fun, args, span, target): &(F, A, FrozenRef<'static, FrameSpan>, BcSlotOut),
     ) -> crate::Result<()> {
+        eval.run_infrequent_instr_checks()?;
         let arguments = Arguments(args.pop_from_stack(frame));
         let r = fun.bc_invoke(*span, &arguments, eval)?;
         frame.set_bc_slot(*target, r);
@@ -1592,6 +1600,7 @@ impl<A: BcCallArgsForDef> InstrNoFlowImpl for InstrCallFrozenDefImpl<A> {
             BcSlotOut,
         ),
     ) -> crate::Result<()> {
+        eval.run_infrequent_instr_checks()?;
         let arguments = args.pop_from_stack(frame);
         let r = eval.with_call_stack(fun.to_value(), Some(*span), |eval| {
             fun.as_ref()
@@ -1613,6 +1622,7 @@ fn call_method_common<'v>(
     span: FrozenRef<'static, FrameSpan>,
     target: BcSlotOut,
 ) -> crate::Result<()> {
+    eval.run_infrequent_instr_checks()?;
     // TODO: wrong span: should be span of `object.method`, not of the whole expression
     let method = get_attr_hashed_raw(this, symbol, eval.heap())?;
     let r = method.invoke(this, span, arguments, eval)?;
@@ -1728,6 +1738,7 @@ impl InstrNoFlowImpl for InstrPossibleGcImpl {
         _ip: BcPtrAddr,
         (): &(),
     ) -> crate::Result<()> {
+        eval.run_infrequent_instr_checks()?;
         possible_gc(eval);
         Ok(())
     }

starlark/src/eval/runtime/evaluator.rs

@@ -104,11 +104,16 @@ enum EvaluatorError {
     CallstackSizeAlreadySet,
     #[error("Max callstack size cannot be zero")]
     ZeroCallstackSize,
+    #[error("Evaluation cancelled")]
+    Cancelled,
 }
 
 /// Number of bytes to allocate between GC's.
 pub(crate) const GC_THRESHOLD: usize = 100000;
 
+/// Number of instructions to execute before running "infrequent" checks
+const INFREQUENT_INSTRUCTION_CHECK_PERIOD: u32 = 1000;
+
 /// Default value for max starlark stack size
 pub(crate) const DEFAULT_STACK_SIZE: usize = 50;
 
@@ -164,6 +169,10 @@ pub struct Evaluator<'v, 'a, 'e> {
     // The Starlark-level call-stack of functions.
     // Must go last because it's quite a big structure
     pub(crate) call_stack: CheapCallStack<'v>,
+    /// Function to check if evaluation should be cancelled early
+    pub(crate) is_cancelled: Box<dyn Fn() -> bool + 'a>,
+    /// A counter to track when to perform "infrequent" checks like cancellation, timeouts, etc
+    pub(crate) infrequent_instr_check_counter: u32,
 }
 
 /// Just holds things that require using EvaluationCallbacksEnabled so that we can cache whether that needs to be enabled or not.
@@ -234,6 +243,8 @@ impl<'v, 'a, 'e: 'a> Evaluator<'v, 'a, 'e> {
             verbose_gc: false,
             static_typechecking: false,
             max_callstack_size: None,
+            is_cancelled: Box::new(|| false),
+            infrequent_instr_check_counter: 0,
         }
     }
 
@@ -449,6 +460,11 @@ impl<'v, 'a, 'e: 'a> Evaluator<'v, 'a, 'e> {
         self.soft_error_handler = handler;
     }
 
+    /// Set canceled-checking function. This function is called periodically to check if the evaluator should return early (with an error condition).
+    pub fn set_check_cancelled(&mut self, is_canceled: Box<dyn Fn() -> bool + 'a>) {
+        self.is_cancelled = is_canceled
+    }
+
     /// Called to add an entry to the call stack, by the function being invoked.
     /// Called for all types of function, including those written in Rust.
     #[inline(always)]
@@ -854,6 +870,20 @@ impl<'v, 'a, 'e: 'a> Evaluator<'v, 'a, 'e> {
         self.max_callstack_size = Some(stack_size);
         Ok(())
     }
+
+    #[inline(always)]
+    pub(crate) fn run_infrequent_instr_checks(&mut self) -> crate::Result<()> {
+        self.infrequent_instr_check_counter += 1;
+        if self.infrequent_instr_check_counter >= INFREQUENT_INSTRUCTION_CHECK_PERIOD {
+            if (self.is_cancelled)() {
+                return Err(crate::Error::new_other(EvaluatorError::Cancelled));
+            }
+            self.infrequent_instr_check_counter = 0
+        };
+
+        // TODO(T219887296): implement CPU-time-limiting checks here
+        Ok(())
+    }
 }
 
 pub(crate) trait EvaluationCallbacks {

starlark/src/tests/uncategorized.rs

@@ -582,6 +582,44 @@ fn test_module_visibility_preserved_by_evaluator() -> crate::Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_cancellation() -> crate::Result<()> {
+    // Make sure that when we use a module in the evaluator, the entering / exiting the
+    // module with ScopeData preserves the visibility of symbols.
+
+    let globals = Globals::standard();
+    let import = Module::new();
+
+    let mut eval = Evaluator::new(&import);
+    eval.set_check_cancelled(Box::new(|| true));
+
+    let ast = AstModule::parse(
+        "prelude.bzl",
+        // Note that the exact range here is unimportant, so long as it's small enough to not trigger the "infrequent" checks
+        "def loop():\n    for i in range(10):\n       pass\nloop()".to_owned(),
+        &Dialect::Standard,
+    )
+    .unwrap();
+    eval.eval_module(ast, &globals).unwrap();
+
+    let ast = AstModule::parse(
+        "prelude.bzl",
+        // Note that the exact range here is unimportant, so long as it's large enough to trigger the "infrequent" checks
+        "def loop():\n    for i in range(1000000):\n       pass\nloop()".to_owned(),
+        &Dialect::Standard,
+    )
+    .unwrap();
+    let err = eval.eval_module(ast, &globals);
+
+    let expected = "Evaluation cancelled";
+    let err_msg = format!("{:#?}", err);
+    if !err_msg.contains(expected) {
+        panic!("Error:\n{:#?}\nExpected:\n{:?}", err, expected)
+    }
+
+    Ok(())
+}
+
 #[test]
 fn test_load_did_you_mean() {
     let mut a = Assert::new();