Refactor CancellationContext

This change makes the countdown period non-configurable (motivation: it needs to match the clock implementation, which is already non-configurable) and allows grouping multiple SoftCheck() calls into a single SoftCheck(n) call.

Benchmarks:
name                                  cpu/op
BM_CancellationContext_SoftCheck<1>   0.97ns ± 1%
BM_CancellationContext_SoftCheck<2>   1.45ns ± 1%
BM_CancellationContext_SoftCheck<4>   2.86ns ± 1%
BM_CancellationContext_SoftCheck<8>   5.61ns ± 1%
BM_CancellationContext_SoftCheck<-1>  11.4ns ± 1%
PiperOrigin-RevId: 721316028
Change-Id: I78bf9bcc7f9830dc6f32e4698ec597d86e7024bf

Author: apronchenkov

koladata/functor/BUILD

@@ -254,7 +254,6 @@ cc_test(
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/time",
         "@com_google_arolla//arolla/dense_array",
         "@com_google_arolla//arolla/expr",
         "@com_google_arolla//arolla/jagged_shape/testing",
@@ -297,7 +296,6 @@ cc_test(
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/time",
         "@com_google_arolla//arolla/expr",
         "@com_google_arolla//arolla/expr/operators/all",
         "@com_google_arolla//arolla/qexpr/operators/all",

koladata/functor/call_test.cc

@@ -26,7 +26,6 @@
 #include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
-#include "absl/time/time.h"
 #include "koladata/data_slice.h"
 #include "koladata/expr/expr_eval.h"
 #include "koladata/functor/functor.h"
@@ -244,37 +243,33 @@ TEST(CallTest, Cancellation) {
       }));
   ASSERT_OK_AND_ASSIGN(auto koda_signature,
                        CppSignatureToKodaSignature(signature));
-  // returns_expr = I.x + I.x + I.x
-  ASSERT_OK_AND_ASSIGN(
-      auto returns_expr,
-      WrapExpr(arolla::expr::CallOp(
-          "math.add", {arolla::expr::CallOp(
-                           "math.add", {CreateInput("a"), CreateInput("a")}),
-                       CreateInput("a")})));
-  ASSERT_OK_AND_ASSIGN(auto fn,
-                       CreateFunctor(returns_expr, koda_signature, {}));
-
+  const auto gen_returns_expr = [](int op_count) {
+    auto result = CreateInput("a");
+    for (int i = 0; i < op_count; ++i) {
+      result = arolla::expr::CallOp("math.add", {result, CreateInput("a")});
+    }
+    return WrapExpr(result);
+  };
   {
-    int op_count = 2;  // Stop after the second operator.
+    const int op_count = 512;  // Long enough computation.
+    ASSERT_OK_AND_ASSIGN(auto returns_expr, gen_returns_expr(op_count));
+    ASSERT_OK_AND_ASSIGN(auto fn,
+                         CreateFunctor(returns_expr, koda_signature, {}));
     auto cancel_ctx = arolla::CancellationContext::Make(
-        /*no cooldown*/ absl::Nanoseconds(-1),
-        /*no countdown*/ -1, [&op_count] {
-          return --op_count > 0 ? absl::OkStatus() : absl::CancelledError("");
-        });
+        /*no cooldown period*/ {}, [] { return absl::CancelledError(""); });
     expr::EvalOptions eval_options{.cancellation_context = cancel_ctx.get()};
     EXPECT_THAT(CallFunctorWithCompilationCache(
                     fn, /*args=*/{arolla::TypedRef::FromValue(1)},
                     /*kwnames=*/{}, eval_options),
                 StatusIs(absl::StatusCode::kCancelled));
   }
-  {
-    int op_count = 3;  // Should stop after the third operator;
-                       // however, there are only two operators.
+  {  // The computation is insufficiently long to detect cancellation.
+    const int op_count = 2;
+    ASSERT_OK_AND_ASSIGN(auto returns_expr, gen_returns_expr(op_count));
+    ASSERT_OK_AND_ASSIGN(auto fn,
+                         CreateFunctor(returns_expr, koda_signature, {}));
     auto cancel_ctx = arolla::CancellationContext::Make(
-        /*no cooldown*/ absl::Nanoseconds(-1),
-        /*no countdown*/ -1, [&op_count] {
-          return --op_count > 0 ? absl::OkStatus() : absl::CancelledError("");
-        });
+        /*no cooldown period*/ {}, [] { return absl::CancelledError(""); });
     expr::EvalOptions eval_options{.cancellation_context = cancel_ctx.get()};
     ASSERT_OK_AND_ASSIGN(auto result,
                          CallFunctorWithCompilationCache(

koladata/functor/map_test.cc

@@ -28,7 +28,6 @@
 #include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
-#include "absl/time/time.h"
 #include "koladata/data_bag.h"
 #include "koladata/data_slice.h"
 #include "koladata/expr/expr_eval.h"
@@ -171,11 +170,17 @@ TEST(MapTest, Cancellation) {
   ASSERT_OK_AND_ASSIGN(auto fn,
                        CreateFunctor(returns_expr, koda_signature, {}));
 
+  std::vector<int> items(512);  // Prepare a long enough test slice.
+  std::iota(items.begin(), items.end(), 0);
+  ASSERT_OK_AND_ASSIGN(
+      auto test_slice,
+      DataSlice::Create(
+          internal::DataSliceImpl::Create(arolla::CreateFullDenseArray(items)),
+          DataSlice::JaggedShape::FlatFromSize(items.size()),
+          internal::DataItem(schema::GetDType<int>())));
   auto cancel_ctx = arolla::CancellationContext::Make(
-      /*no cooldown*/ absl::Nanoseconds(-1),
-      /*no countdown*/ -1, [] { return absl::CancelledError(""); });
+      /*no cooldown*/ {}, [] { return absl::CancelledError(""); });
   expr::EvalOptions eval_options{.cancellation_context = cancel_ctx.get()};
-  auto test_slice = test::DataSlice<int>({1, 2, 3});
   EXPECT_THAT(MapFunctorWithCompilationCache(
                   fn, /*args=*/{test_slice},
                   /*kwnames=*/{}, /*include_missing=*/false, eval_options),