Add gflag to allow growing buffers that were created in different task

velox/exec/MemoryReclaimer.cpp

@@ -19,6 +19,11 @@
 #include "velox/exec/Driver.h"
 #include "velox/exec/Task.h"
 
+DEFINE_bool(
+    transferred_arbitration_allowed,
+    false,
+    "Whether to allow entering memory arbitration on memory pool that is from another task");
+
 namespace facebook::velox::exec {
 std::unique_ptr<memory::MemoryReclaimer> MemoryReclaimer::create() {
   return std::unique_ptr<memory::MemoryReclaimer>(new MemoryReclaimer());

velox/exec/MemoryReclaimer.h

@@ -18,6 +18,8 @@
 
 #include "velox/common/memory/MemoryArbitrator.h"
 
+DECLARE_bool(transferred_arbitration_allowed);
+
 namespace facebook::velox::exec {
 /// Provides the default leaf memory reclaimer implementation for velox task
 /// execution.

velox/exec/Operator.cpp

@@ -563,15 +563,19 @@ void Operator::MemoryReclaimer::enterArbitration() {
   }
 
   Driver* const runningDriver = driverThreadCtx->driverCtx.driver;
-  if (auto opDriver = ensureDriver()) {
-    // NOTE: the current running driver might not be the driver of the operator
-    // that requests memory arbitration. The reason is that an operator might
-    // extend the buffer allocated from the other operator either from the same
-    // or different drivers. But they must be from the same task.
-    VELOX_CHECK_EQ(
-        runningDriver->task()->taskId(),
-        opDriver->task()->taskId(),
-        "The current running driver and the request driver must be from the same task");
+  if (!FLAGS_transferred_arbitration_allowed) {
+    if (auto opDriver = ensureDriver()) {
+      // NOTE: the current running driver might not be the driver of the
+      // operator that requests memory arbitration. The reason is that an
+      // operator might extend the buffer allocated from the other operator
+      // either from the same or different drivers. But they must be from the
+      // same task as the following check. User could set
+      // FLAGS_transferred_arbitration_allowed=true to bypass this check.
+      VELOX_CHECK_EQ(
+          runningDriver->task()->taskId(),
+          opDriver->task()->taskId(),
+          "The current running driver and the request driver must be from the same task");
+    }
   }
   if (runningDriver->task()->enterSuspended(runningDriver->state()) !=
       StopReason::kNone) {
@@ -589,11 +593,13 @@ void Operator::MemoryReclaimer::leaveArbitration() noexcept {
     return;
   }
   Driver* const runningDriver = driverThreadCtx->driverCtx.driver;
-  if (auto opDriver = ensureDriver()) {
-    VELOX_CHECK_EQ(
-        runningDriver->task()->taskId(),
-        opDriver->task()->taskId(),
-        "The current running driver and the request driver must be from the same task");
+  if (!FLAGS_transferred_arbitration_allowed) {
+    if (auto opDriver = ensureDriver()) {
+      VELOX_CHECK_EQ(
+          runningDriver->task()->taskId(),
+          opDriver->task()->taskId(),
+          "The current running driver and the request driver must be from the same task");
+    }
   }
   runningDriver->task()->leaveSuspended(runningDriver->state());
 }