[yql] Failure injector in mrrun/dqrun

YQL-16599

6 files changed, 26 insertions, 0 deletions

diff --git a/ydb/library/yql/tools/dqrun/CMakeLists.darwin-arm64.txt b/ydb/library/yql/tools/dqrun/CMakeLists.darwin-arm64.txt
index 605593ba8ab..d963c4b238a 100644
--- a/ydb/library/yql/tools/dqrun/CMakeLists.darwin-arm64.txt
+++ b/ydb/library/yql/tools/dqrun/CMakeLists.darwin-arm64.txt
@@ -60,6 +60,7 @@ target_link_libraries(dqrun PUBLIC
   yql-utils-backtrace
   yql-utils-bindings
   yql-utils-log
+  yql-utils-failure_injector
   yql-core-url_preprocessing
   yql-core-url_lister
   yt-comp_nodes-dq
diff --git a/ydb/library/yql/tools/dqrun/CMakeLists.darwin-x86_64.txt b/ydb/library/yql/tools/dqrun/CMakeLists.darwin-x86_64.txt
index c6fbeb546ee..895c759c512 100644
--- a/ydb/library/yql/tools/dqrun/CMakeLists.darwin-x86_64.txt
+++ b/ydb/library/yql/tools/dqrun/CMakeLists.darwin-x86_64.txt
@@ -61,6 +61,7 @@ target_link_libraries(dqrun PUBLIC
   yql-utils-backtrace
   yql-utils-bindings
   yql-utils-log
+  yql-utils-failure_injector
   yql-core-url_preprocessing
   yql-core-url_lister
   yt-comp_nodes-dq
diff --git a/ydb/library/yql/tools/dqrun/CMakeLists.linux-aarch64.txt b/ydb/library/yql/tools/dqrun/CMakeLists.linux-aarch64.txt
index 3cab1c92823..ecb0c3504bc 100644
--- a/ydb/library/yql/tools/dqrun/CMakeLists.linux-aarch64.txt
+++ b/ydb/library/yql/tools/dqrun/CMakeLists.linux-aarch64.txt
@@ -61,6 +61,7 @@ target_link_libraries(dqrun PUBLIC
   yql-utils-backtrace
   yql-utils-bindings
   yql-utils-log
+  yql-utils-failure_injector
   yql-core-url_preprocessing
   yql-core-url_lister
   yt-comp_nodes-dq
diff --git a/ydb/library/yql/tools/dqrun/CMakeLists.linux-x86_64.txt b/ydb/library/yql/tools/dqrun/CMakeLists.linux-x86_64.txt
index 3e069aad58b..d98045e536e 100644
--- a/ydb/library/yql/tools/dqrun/CMakeLists.linux-x86_64.txt
+++ b/ydb/library/yql/tools/dqrun/CMakeLists.linux-x86_64.txt
@@ -62,6 +62,7 @@ target_link_libraries(dqrun PUBLIC
   yql-utils-backtrace
   yql-utils-bindings
   yql-utils-log
+  yql-utils-failure_injector
   yql-core-url_preprocessing
   yql-core-url_lister
   yt-comp_nodes-dq
diff --git a/ydb/library/yql/tools/dqrun/dqrun.cpp b/ydb/library/yql/tools/dqrun/dqrun.cpp
index eaa8666a071..7d71897cc07 100644
--- a/ydb/library/yql/tools/dqrun/dqrun.cpp
+++ b/ydb/library/yql/tools/dqrun/dqrun.cpp
@@ -15,6 +15,7 @@
 #include <ydb/library/yql/utils/log/proto/logger_config.pb.h>
 #include <ydb/library/yql/core/url_preprocessing/url_preprocessing.h>
 #include <ydb/library/yql/utils/actor_system/manager.h>
+#include <ydb/library/yql/utils/failure_injector/failure_injector.h>
 
 #include <ydb/library/yql/parser/pg_wrapper/interface/comp_factory.h>
 #include <ydb/library/yql/providers/dq/provider/yql_dq_gateway.h>
@@ -423,6 +424,7 @@ int RunMain(int argc, const char* argv[])
         }
     }
     THashMap<TString, TString> customTokens;
+    THashMap<TString, std::pair<ui32, ui32>> failureInjections;
     TString folderId;
 
     TRunOptions runOptions;
@@ -545,6 +547,18 @@ int RunMain(int argc, const char* argv[])
     opts.AddLongOption("metrics-pusher-config", "Metrics Pusher Config")
         .StoreResult(&mestricsPusherConfig);
     opts.AddLongOption("enable-spilling", "Enable disk spilling").NoArgument();
+    opts.AddLongOption("failure-inject", "Activate failure injection")
+        .Optional()
+        .RequiredArgument("INJECTION_NAME=FAIL_COUNT or INJECTION_NAME=SKIP_COUNT/FAIL_COUNT")
+        .KVHandler([&failureInjections](TString key, TString value) {
+            TStringBuf fail = value;
+            TStringBuf skip;
+            if (TStringBuf(value).TrySplit('/', skip, fail)) {
+                failureInjections[key] = std::make_pair(FromString<ui32>(skip), FromString<ui32>(fail));
+            } else {
+                failureInjections[key] = std::make_pair(ui32(0), FromString<ui32>(fail));
+            }
+        });
     opts.AddHelpOption('h');
 
     opts.SetFreeArgsNum(0);
@@ -622,6 +636,13 @@ int RunMain(int argc, const char* argv[])
         return 1;
     }
 
+    if (!failureInjections.empty()) {
+        TFailureInjector::Activate();
+        for (auto& [name, count]: failureInjections) {
+            TFailureInjector::Set(name, count.first, count.second);
+        }
+    }
+
     runOptions.ResultsFormat =
             (format == TStringBuf("binary")) ? NYson::EYsonFormat::Binary
           : (format == TStringBuf("text")) ? NYson::EYsonFormat::Text
diff --git a/ydb/library/yql/tools/dqrun/ya.make b/ydb/library/yql/tools/dqrun/ya.make
index dbb1f3eb145..1b0ad5c9b0b 100644
--- a/ydb/library/yql/tools/dqrun/ya.make
+++ b/ydb/library/yql/tools/dqrun/ya.make
@@ -64,6 +64,7 @@ ENDIF()
         ydb/library/yql/utils/backtrace
         ydb/library/yql/utils/bindings
         ydb/library/yql/utils/log
+        ydb/library/yql/utils/failure_injector
         ydb/library/yql/core/url_preprocessing
         ydb/library/yql/core/url_lister
         ydb/library/yql/providers/yt/comp_nodes/dq