From 39e93deac6e57510954303c222cdd5fb07e6cc9f Mon Sep 17 00:00:00 2001
From: ivanmorozov333 <ivanmorozov@ydb.tech>
Date: Wed, 8 Jan 2025 10:25:08 +0300
Subject: [PATCH] simultaneous writing session for tests (#13187)

---
 ydb/core/kqp/ut/olap/helpers/typed_local.cpp |  6 ++++--
 ydb/core/kqp/ut/olap/helpers/typed_local.h   | 16 +++++++---------
 ydb/core/kqp/ut/olap/write_ut.cpp            | 10 +++++-----
 3 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/ydb/core/kqp/ut/olap/helpers/typed_local.cpp b/ydb/core/kqp/ut/olap/helpers/typed_local.cpp
index 6ead7bdd6624..492bc7f4d2c0 100644
--- a/ydb/core/kqp/ut/olap/helpers/typed_local.cpp
+++ b/ydb/core/kqp/ut/olap/helpers/typed_local.cpp
@@ -191,7 +191,7 @@ void TTypedLocalHelper::GetStats(std::vector<NJson::TJsonValue>& stats, const bo
     }
 }
 
-void TTypedLocalHelper::TWritingGuard::SendDataViaActorSystem(TString testTable, std::shared_ptr<arrow::RecordBatch> batch,
+void TTypedLocalHelper::TSimultaneousWritingSession::SendDataViaActorSystem(TString testTable, std::shared_ptr<arrow::RecordBatch> batch,
     const Ydb::StatusIds_StatusCode expectedStatus /*= = Ydb::StatusIds::SUCCESS*/) const {
     auto* runtime = KikimrRunner.GetTestServer().GetRuntime();
 
@@ -224,7 +224,9 @@ void TTypedLocalHelper::TWritingGuard::SendDataViaActorSystem(TString testTable,
     });
 }
 
-void TTypedLocalHelper::TWritingGuard::WaitWritings() {
+void TTypedLocalHelper::TSimultaneousWritingSession::Finalize() {
+    AFL_VERIFY(!Finished);
+    Finished = true;
     auto* runtime = KikimrRunner.GetTestServer().GetRuntime();
     TDispatchOptions options;
     options.CustomFinalCondition = [&]() {
diff --git a/ydb/core/kqp/ut/olap/helpers/typed_local.h b/ydb/core/kqp/ut/olap/helpers/typed_local.h
index 0753cd72b02b..a0ead56e8769 100644
--- a/ydb/core/kqp/ut/olap/helpers/typed_local.h
+++ b/ydb/core/kqp/ut/olap/helpers/typed_local.h
@@ -37,18 +37,17 @@ class TTypedLocalHelper: public Tests::NCS::THelper {
         SetShardingMethod("HASH_FUNCTION_CONSISTENCY_64");
     }
 
-    class TWritingGuard {
+    class TSimultaneousWritingSession {
     private:
+        bool Finished = false;
         TKikimrRunner& KikimrRunner;
         const TString TablePath;
         mutable std::atomic<size_t> Responses = 0;
         void SendDataViaActorSystem(TString testTable, std::shared_ptr<arrow::RecordBatch> batch,
             const Ydb::StatusIds_StatusCode expectedStatus = Ydb::StatusIds::SUCCESS) const;
 
-        void WaitWritings();
-
     public:
-        TWritingGuard(TKikimrRunner& kikimrRunner, const TString& tablePath)
+        TSimultaneousWritingSession(TKikimrRunner& kikimrRunner, const TString& tablePath)
             : KikimrRunner(kikimrRunner)
             , TablePath(tablePath)
         {
@@ -56,6 +55,7 @@ class TTypedLocalHelper: public Tests::NCS::THelper {
 
         template <class TFiller>
         void FillTable(const TFiller& fillPolicy, const double pkKff = 0, const ui32 numRows = 800000) const {
+            AFL_VERIFY(!Finished);
             std::vector<NArrow::NConstruction::IArrayBuilder::TPtr> builders;
             builders.emplace_back(
                 NArrow::NConstruction::TSimpleArrayConstructor<NArrow::NConstruction::TIntSeqFiller<arrow::Int64Type>>::BuildNotNullable(
@@ -66,13 +66,11 @@ class TTypedLocalHelper: public Tests::NCS::THelper {
             SendDataViaActorSystem(TablePath, batch, Ydb::StatusIds::SUCCESS);
         }
 
-        void Finalize() {
-            WaitWritings();
-        }
+        void Finalize();
     };
 
-    TWritingGuard StartWriting(const TString& tablePath) {
-        return TWritingGuard(KikimrRunner, tablePath);
+    TSimultaneousWritingSession StartWriting(const TString& tablePath) {
+        return TSimultaneousWritingSession(KikimrRunner, tablePath);
     }
 
     void ExecuteSchemeQuery(const TString& alterQuery, const NYdb::EStatus expectedStatus = NYdb::EStatus::SUCCESS) const;
diff --git a/ydb/core/kqp/ut/olap/write_ut.cpp b/ydb/core/kqp/ut/olap/write_ut.cpp
index 83c594db5db6..470bd0e9a0bc 100644
--- a/ydb/core/kqp/ut/olap/write_ut.cpp
+++ b/ydb/core/kqp/ut/olap/write_ut.cpp
@@ -201,14 +201,14 @@ Y_UNIT_TEST_SUITE(KqpOlapWrite) {
         Tests::NCommon::TLoggerInit(kikimr).Initialize();
         TTypedLocalHelper helper("Utf8", kikimr);
         helper.CreateTestOlapTable();
-        auto writeGuard = helper.StartWriting("/Root/olapStore/olapTable");
-        writeGuard.FillTable(NArrow::NConstruction::TStringPoolFiller(1, 1, "aaa", 1), 0, 800000);
+        auto writeSession = helper.StartWriting("/Root/olapStore/olapTable");
+        writeSession.FillTable(NArrow::NConstruction::TStringPoolFiller(1, 1, "aaa", 1), 0, 800000);
         Sleep(TDuration::Seconds(1));
-        writeGuard.FillTable(NArrow::NConstruction::TStringPoolFiller(1, 1, "bbb", 1), 0.5, 800000);
+        writeSession.FillTable(NArrow::NConstruction::TStringPoolFiller(1, 1, "bbb", 1), 0.5, 800000);
         Sleep(TDuration::Seconds(1));
-        writeGuard.FillTable(NArrow::NConstruction::TStringPoolFiller(1, 1, "ccc", 1), 0.75, 800000);
+        writeSession.FillTable(NArrow::NConstruction::TStringPoolFiller(1, 1, "ccc", 1), 0.75, 800000);
         Sleep(TDuration::Seconds(1));
-        writeGuard.Finalize();
+        writeSession.Finalize();
 
         auto selectQuery = TString(R"(
                 SELECT