Add query abort injection test utility (facebookincubator#8948)

Summary:
Add a global scoped query abort injection utility to inject abort
during assert query results. Uses it in the SharedAbitrationTest
to make some query abort during memory arbitration.
The `TestScopedAbortInjection` is inspired by the PR facebookincubator#8921 .

Pull Request resolved: facebookincubator#8948

Reviewed By: tanjialiang

Differential Revision: D54569827

Pulled By: xiaoxmeng

fbshipit-source-id: 6251f84ba45e848636a8bb787f9c4fe05f488b5c

velox/common/memory/tests/SharedArbitratorTest.cpp

@@ -1073,6 +1073,7 @@ TEST_F(SharedArbitrationTest, concurrentArbitration) {
     const int maxNumZombieTasks = 8;
     std::vector<std::thread> queryThreads;
     queryThreads.reserve(numThreads);
+    TestScopedAbortInjection testScopedAbortInjection(10, numThreads);
     for (int i = 0; i < numThreads; ++i) {
       queryThreads.emplace_back([&, i]() {
         std::shared_ptr<Task> task;
@@ -1131,7 +1132,8 @@ TEST_F(SharedArbitrationTest, concurrentArbitration) {
         } catch (const VeloxException& e) {
           if (e.errorCode() != error_code::kMemCapExceeded.c_str() &&
               e.errorCode() != error_code::kMemAborted.c_str() &&
-              e.errorCode() != error_code::kMemAllocError.c_str()) {
+              e.errorCode() != error_code::kMemAllocError.c_str() &&
+              (e.message() != "Aborted for external error")) {
             std::rethrow_exception(std::current_exception());
           }
         }

velox/exec/tests/utils/QueryAssertions.cpp

@@ -19,6 +19,7 @@
 
 #include "duckdb/common/types.hpp" // @manual
 #include "velox/duckdb/conversion/DuckConversion.h"
+#include "velox/exec/tests/utils/ArbitratorTestUtil.h"
 #include "velox/exec/tests/utils/Cursor.h"
 #include "velox/exec/tests/utils/QueryAssertions.h"
 #include "velox/vector/VectorTypeUtils.h"
@@ -1316,6 +1317,43 @@ void assertResultsOrdered(
   }
 }
 
+tsan_atomic<int32_t>& testingAbortPct() {
+  static tsan_atomic<int32_t> abortPct = 0;
+  return abortPct;
+}
+
+tsan_atomic<int32_t>& testingAbortCounter() {
+  static tsan_atomic<int32_t> counter = 0;
+  return counter;
+}
+
+TestScopedAbortInjection::TestScopedAbortInjection(
+    int32_t abortPct,
+    int32_t maxInjections) {
+  testingAbortPct() = abortPct;
+  testingAbortCounter() = maxInjections;
+}
+
+TestScopedAbortInjection::~TestScopedAbortInjection() {
+  testingAbortPct() = 0;
+  testingAbortCounter() = 0;
+}
+
+bool testingMaybeTriggerAbort(exec::Task* task) {
+  if (testingAbortPct() <= 0 || testingAbortCounter() <= 0) {
+    return false;
+  }
+
+  if ((folly::Random::rand32() % 100) < testingAbortPct()) {
+    if (testingAbortCounter()-- > 0) {
+      task->requestAbort();
+      return true;
+    }
+  }
+
+  return false;
+}
+
 std::pair<std::unique_ptr<TaskCursor>, std::vector<RowVectorPtr>> readCursor(
     const CursorParameters& params,
     std::function<void(exec::Task*)> addSplits,
@@ -1329,6 +1367,7 @@ std::pair<std::unique_ptr<TaskCursor>, std::vector<RowVectorPtr>> readCursor(
   while (cursor->moveNext()) {
     result.push_back(cursor->current());
     addSplits(task);
+    testingMaybeTriggerAbort(task);
   }
 
   if (!waitForTaskCompletion(task, maxWaitMicros)) {

velox/exec/tests/utils/QueryAssertions.h

@@ -81,6 +81,27 @@ class DuckDbQueryRunner {
       std::function<void(std::vector<MaterializedRow>&)> resultCallback);
 };
 
+/// Scoped abort percentage utility that allows user to trigger abort during the
+///  query execution.
+/// 'abortPct' specifies the probability of of triggering abort. 100% means
+/// abort will always be triggered.
+/// 'maxInjections' indicates the max number of actual triggering, e.g. when
+/// 'abortPct' is 20 and 'maxInjections' is 10, continuous calls to
+/// testingMaybeTriggerAbort() will keep rolling the dice that has a chance of
+/// 20% triggering until 10 triggers have been invoked.
+class TestScopedAbortInjection {
+ public:
+  explicit TestScopedAbortInjection(
+      int32_t abortPct,
+      int32_t maxInjections = std::numeric_limits<int32_t>::max());
+
+  ~TestScopedAbortInjection();
+};
+
+/// Test utility that might trigger task abort. The function returns true if
+/// abortion triggers otherwise false.
+bool testingMaybeTriggerAbort(exec::Task* task);
+
 std::pair<std::unique_ptr<TaskCursor>, std::vector<RowVectorPtr>> readCursor(
     const CursorParameters& params,
     std::function<void(exec::Task*)> addSplits,