Fix generator state work (#6219)

Author: iddqdex

ydb/library/workload/abstract/workload_query_generator.h

@@ -52,7 +52,8 @@ class IBulkDataGenerator {
         , Size(size)
     {}
 
-    struct TDataPortion: public TAtomicRefCount<TDataPortion>, TMoveOnly {
+    class TDataPortion: public TAtomicRefCount<TDataPortion>, TMoveOnly {
+    public:
         struct TCsv {
             TCsv(TString&& data, const TString& formatString = TString())
                 : Data(std::move(data))
@@ -71,16 +72,26 @@ class IBulkDataGenerator {
             TString Schema;
         };
 
+        using TDataType = std::variant<NYdb::TValue, TCsv, TArrow>;
+
         template<class T>
         TDataPortion(const TString& table, T&& data, ui64 size)
             : Table(table)
-            , Data(std::move(data))
             , Size(size)
+            , Data(std::move(data))
         {}
 
-        TString Table;
-        std::variant<NYdb::TValue, TCsv, TArrow> Data;
-        ui64 Size;
+        virtual ~TDataPortion() = default;
+        virtual void SetSendResult(const NYdb::TStatus& status) {
+            Y_UNUSED(status);
+        }
+        TDataType& MutableData() {
+            return Data;
+        }
+        YDB_READONLY_DEF(TString, Table);
+        YDB_READONLY(ui64, Size, 0);
+    private:
+        TDataType Data;
     };
 
     using TDataPortionPtr = TIntrusivePtr<TDataPortion>;

ydb/library/workload/benchmark_base/state.cpp

@@ -3,7 +3,6 @@
 #include <library/cpp/json/json_writer.h>
 #include <library/cpp/json/json_reader.h>
 #include <util/stream/file.h>
-#include <thread>
 
 namespace NYdbWorkload {
 
@@ -25,48 +24,28 @@ TGeneratorStateProcessor::TGeneratorStateProcessor(const TFsPath& path, bool cle
     }
 }
 
-void TGeneratorStateProcessor::AddPortion(const TString& source, ui64 from, ui64 size) {
-    InProcess.Get().emplace_back(TInProcessPortion{source, from, size});
-}
-
-void TGeneratorStateProcessor::FinishPortions() {
-    bool needSave = false;
+void TGeneratorStateProcessor::FinishPortion(const TString& source, ui64 from, ui64 size) {
     auto g = Guard(Lock);
-    for (const auto& p: InProcess.Get()) {
-        needSave |= StateImpl[p.Source].FinishPortion(p.From, p.Size, State[p.Source].Position);
-    }
-    InProcess.Get().clear();
+    bool needSave = StateImpl[source].FinishPortion(from, size, State[source].Position);
     if (needSave) {
         Save();
     }
 }
 
 bool TGeneratorStateProcessor::TSourceStateImpl::FinishPortion(ui64 from, ui64 size, ui64& position) {
-    const ui64 threadId = std::hash<std::thread::id>()(std::this_thread::get_id());
-    ThreadsState[threadId].FinishedPortions.emplace_back(TSourcePortion{from, size});
-    TVector<TThreadSourceState::TFinishedPortions*> portions;
-    for (auto& [t, ss]: ThreadsState) {
-        while (!ss.FinishedPortions.empty() && ss.FinishedPortions.front().first < position) {
-            ss.FinishedPortions.pop_front();
-        }
-        if (!ss.FinishedPortions.empty()) {
-            portions.push_back(&ss.FinishedPortions);
-        }
+    auto i = FinishedPortions.begin();
+    while (i != FinishedPortions.end() && i->first < from) {
+         ++i;
     }
-    Y_VERIFY(!portions.empty());
-    auto portionsCmp = [](auto l, auto r) {return l->front().first > r->front().first;};
-    std::make_heap(portions.begin(), portions.end(), portionsCmp);
+    FinishedPortions.emplace(i, TSourcePortion{from, size});
     bool result = false;
-    while (!portions.empty() && portions.front()->front().first == position) {
+    while (!FinishedPortions.empty() && FinishedPortions.front().first < position) {
+        FinishedPortions.pop_front();
+    }
+    while (!FinishedPortions.empty() && FinishedPortions.front().first == position) {
         result = true;
-        position += portions.front()->front().second;
-        std::pop_heap(portions.begin(), portions.end(), portionsCmp);
-        portions.back()->pop_front();
-        if (portions.back()->empty()) {
-            portions.pop_back();
-        } else {
-            std::push_heap(portions.begin(), portions.end(), portionsCmp);
-        }
+        position += FinishedPortions.front().second;
+        FinishedPortions.pop_front();
     }
     return result;
 }

ydb/library/workload/benchmark_base/state.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <ydb/library/workload/abstract/workload_query_generator.h>
 #include <ydb/library/accessor/accessor.h>
 #include <util/generic/deque.h>
 #include <util/generic/map.h>
@@ -18,29 +19,17 @@ class TGeneratorStateProcessor {
 
 public:
     TGeneratorStateProcessor(const TFsPath& path, bool clear);
-    void AddPortion(const TString& source, ui64 from, ui64 size);
-    void FinishPortions();
+    void FinishPortion(const TString& source, ui64 from, ui64 size);
     YDB_READONLY_DEF(TState, State);
 
 private:
     using TSourcePortion = std::pair<ui64, ui64>;
-    struct TThreadSourceState {
-        using TFinishedPortions = TDeque<TSourcePortion>;
-        TFinishedPortions FinishedPortions;
-    };
-
     class TSourceStateImpl {
     public:
         bool FinishPortion(ui64 from, ui64 size, ui64& position);
 
     private:
-        TMap<ui64, TThreadSourceState> ThreadsState;
-    };
-
-    struct TInProcessPortion {
-        TString Source;
-        ui64 From;
-        ui64 Size;
+        TList<TSourcePortion> FinishedPortions;
     };
 
 private:
@@ -52,7 +41,27 @@ class TGeneratorStateProcessor {
     TFsPath TmpPath;
     TMap<TString, TSourceStateImpl> StateImpl;
     TAdaptiveLock Lock;
-    Y_THREAD(TVector<TInProcessPortion>) InProcess;
+};
+
+class TDataPortionWithState: public IBulkDataGenerator::TDataPortion {
+public:
+    template<class T>
+    TDataPortionWithState(TGeneratorStateProcessor* stateProcessor, const TString& table, const TString& stateSource, T&& data, ui64 position, ui64 size)
+        : TDataPortion(table, std::move(data), size)
+        , StateSource(stateSource)
+        , Position(position)
+        , StateProcessor(stateProcessor)
+    {}
+
+    virtual void SetSendResult(const NYdb::TStatus& status) {
+        if (StateProcessor && status.IsSuccess()) {
+            StateProcessor->FinishPortion(StateSource, Position, GetSize());
+        }
+    }
+private:
+    TString StateSource;
+    ui64 Position;
+    TGeneratorStateProcessor* StateProcessor;
 };
 
 }

ydb/library/workload/benchmark_base/state_ut.cpp

@@ -6,38 +6,34 @@
 Y_UNIT_TEST_SUITE(DataGeneratorState) {
     Y_UNIT_TEST(PortionProcessing) {
         TFsPath path;
-        NYdbWorkload::TGeneratorStateProcessor proccessor(path);
-        proccessor.AddPortion("test1", 0, 10);
+        NYdbWorkload::TGeneratorStateProcessor proccessor(path, true);
         UNIT_ASSERT(!proccessor.GetState().contains("test1"));
-        proccessor.FinishPortions();
+        proccessor.FinishPortion("test1", 0, 10);
         UNIT_ASSERT_EQUAL(proccessor.GetState().at("test1").Position, 10);
         TThread t([&proccessor](){
-            proccessor.AddPortion("test1", 10, 10);
-            proccessor.AddPortion("test1", 50, 10);
-            proccessor.FinishPortions();
+            proccessor.FinishPortion("test1", 10, 10);
+            proccessor.FinishPortion("test1", 50, 10);
         });
         t.Start();
         t.Join();
         UNIT_ASSERT_EQUAL(proccessor.GetState().at("test1").Position, 20);
-        proccessor.AddPortion("test1", 20, 30);
-        proccessor.FinishPortions();
+        proccessor.FinishPortion("test1", 20, 30);
         UNIT_ASSERT_EQUAL(proccessor.GetState().at("test1").Position, 60);
     };
 
     Y_UNIT_TEST(SaveLoad) {
         TFsPath path("test_state.json");
         path.DeleteIfExists();
 
-        NYdbWorkload::TGeneratorStateProcessor proccessor1(path);
-        proccessor1.AddPortion("test1", 0, 10);
-        proccessor1.AddPortion("test2", 0, 15);
-        proccessor1.AddPortion("test3", 0, 17);
-        proccessor1.FinishPortions();
+        NYdbWorkload::TGeneratorStateProcessor proccessor1(path, true);
+        proccessor1.FinishPortion("test1", 0, 10);
+        proccessor1.FinishPortion("test2", 0, 15);
+        proccessor1.FinishPortion("test3", 0, 17);
         UNIT_ASSERT_EQUAL(proccessor1.GetState().at("test1").Position, 10);
         UNIT_ASSERT_EQUAL(proccessor1.GetState().at("test2").Position, 15);
         UNIT_ASSERT_EQUAL(proccessor1.GetState().at("test3").Position, 17);
 
-        NYdbWorkload::TGeneratorStateProcessor proccessor2(path);
+        NYdbWorkload::TGeneratorStateProcessor proccessor2(path, false);
         UNIT_ASSERT_EQUAL(proccessor2.GetState().at("test1").Position, 10);
         UNIT_ASSERT_EQUAL(proccessor2.GetState().at("test2").Position, 15);
         UNIT_ASSERT_EQUAL(proccessor2.GetState().at("test3").Position, 17);

ydb/library/workload/benchmark_base/ut/ya.make

@@ -1,4 +1,4 @@
-UNITTEST_FOR(kikimr/tests/acceptance/olap/cli/workload_base)
+UNITTEST_FOR(ydb/library/workload/benchmark_base)
     SRCS(
         state_ut.cpp
     )

ydb/library/workload/benchmark_base/ya.make

@@ -16,3 +16,7 @@ PEERDIR(
 GENERATE_ENUM_SERIALIZATION(workload.h)
 
 END()
+
+RECURSE_FOR_TESTS(
+    ut
+)

ydb/library/workload/clickbench/data_generator.cpp

@@ -45,9 +45,6 @@ TClickbenchWorkloadDataInitializerGenerator::TDataGenerartor::TDataGenerartor(co
 }
 
 IBulkDataGenerator::TDataPortions TClickbenchWorkloadDataInitializerGenerator::TDataGenerartor::GenerateDataPortion() {
-    if (Owner.StateProcessor) {
-        Owner.StateProcessor->FinishPortions();
-    }
     while (true) {
         size_t index;
         TFile::TPtr file;
@@ -139,13 +136,14 @@ class TClickbenchWorkloadDataInitializerGenerator::TDataGenerartor::TCsvFileBase
         data.reserve(lines.size() * 10000);
         data << Header << Endl;
         data << JoinSeq("\n", lines) << Endl;
-        if (Owner.Owner.StateProcessor) {
-            Owner.Owner.StateProcessor->AddPortion(Path, Readed, lines.size());
-        }
+        const ui64 position = Readed;
         Readed += lines.size();
-        return MakeIntrusive<TDataPortion>(
+        return MakeIntrusive<TDataPortionWithState>(
+            Owner.Owner.StateProcessor.Get(),
             Owner.Owner.Params.GetFullTableName(nullptr),
+            Path,
             TDataPortion::TCsv(std::move(data), Foramt),
+            position,
             lines.size()
         );
     }

ydb/library/workload/tpch/data_generator.cpp

@@ -64,9 +64,10 @@ ui64 TTpchWorkloadDataInitializerGenerator::TBulkDataGenerator::CalcCountToGener
     return rowCount / owner.GetProcessCount() + extraRows - position;
 }
 
-TTpchWorkloadDataInitializerGenerator::TBulkDataGenerator::TContext::TContext(const TBulkDataGenerator& owner, int tableNum)
+TTpchWorkloadDataInitializerGenerator::TBulkDataGenerator::TContext::TContext(const TBulkDataGenerator& owner, int tableNum, TGeneratorStateProcessor* state)
     : Owner(owner)
     , TableNum(tableNum)
+    , State(state)
 {}
 
 NYdb::TValueBuilder& TTpchWorkloadDataInitializerGenerator::TBulkDataGenerator::TContext::GetBuilder() {
@@ -85,9 +86,23 @@ void TTpchWorkloadDataInitializerGenerator::TBulkDataGenerator::TContext::Append
     const auto path = Owner.GetFullTableName(tdefs[TableNum].name);
     if (Builder) {
         Builder->EndList();
-        result.push_back(MakeIntrusive<TDataPortion>(path, Builder->Build(), Count));
+        result.push_back(MakeIntrusive<TDataPortionWithState>(
+            State,
+            path,
+            tdefs[TableNum].name,
+            Builder->Build(),
+            Start - 1,
+            Count
+        ));
     } else if (Csv) {
-        result.push_back(MakeIntrusive<TDataPortion>(path, TDataPortion::TCsv(std::move(Csv), TWorkloadGeneratorBase::TsvFormatString), Count));
+        result.push_back(MakeIntrusive<TDataPortionWithState>(
+            State,
+            path,
+            tdefs[TableNum].name,
+            TDataPortion::TCsv(std::move(Csv), TWorkloadGeneratorBase::TsvFormatString),
+            Start - 1,
+            Count
+        ));
     }
 }
 
@@ -103,27 +118,26 @@ TTpchWorkloadDataInitializerGenerator::TBulkDataGenerator::TBulkDataGenerator(co
 {}
 
 TTpchWorkloadDataInitializerGenerator::TBulkDataGenerator::TDataPortions TTpchWorkloadDataInitializerGenerator::TBulkDataGenerator::GenerateDataPortion() {
-    if (Owner.StateProcessor) {
-        Owner.StateProcessor->FinishPortions();
-    }
     TDataPortions result;
     if (TableSize == 0) {
         return result;
     }
     TContexts ctxs;
-    ctxs.emplace_back(*this, TableNum);
+    ctxs.emplace_back(*this, TableNum, Owner.StateProcessor.Get());
     if (tdefs[TableNum].child != NONE) {
-        ctxs.emplace_back(*this, tdefs[TableNum].child);
+        ctxs.emplace_back(*this, tdefs[TableNum].child, nullptr);
     }
     with_lock(NumbersLock) {
         if (!Generated) {
             if (Owner.GetProcessCount() > 1) {
                 DSS_HUGE e;
                 set_state(TableNum, Owner.GetScale(), Owner.GetProcessCount(), Owner.GetProcessIndex() + 1, &e);
             }
-            if (!!Owner.StateProcessor && Owner.StateProcessor->GetState().contains(GetName())) {
-                Generated = Owner.StateProcessor->GetState().at(TString(GetName())).Position;
-                GenSeed(TableNum, Generated);
+            if (!!Owner.StateProcessor) {
+                if (const auto* state = MapFindPtr(Owner.StateProcessor->GetState(), GetName())) {
+                    Generated = state->Position;
+                    GenSeed(TableNum, Generated);
+                }
             }
         }
         const auto count = TableSize > Generated ? std::min(ui64(TableSize - Generated), Owner.Params.BulkSize) : 0;
@@ -132,9 +146,6 @@ TTpchWorkloadDataInitializerGenerator::TBulkDataGenerator::TDataPortions TTpchWo
         }
         ctxs.front().SetCount(count);
         ctxs.front().SetStart((tdefs[TableNum].base * Owner.GetScale() / Owner.GetProcessCount()) * Owner.GetProcessIndex() + Generated + 1);
-        if (Owner.StateProcessor) {
-            Owner.StateProcessor->AddPortion(TString(GetName()), Generated, count);
-        }
         Generated += count;
     }
     GenerateRows(ctxs);

ydb/library/workload/tpch/data_generator.h

@@ -31,7 +31,7 @@ class TTpchWorkloadDataInitializerGenerator: public TWorkloadDataInitializerBase
     protected:
         class TContext {
         public:
-            TContext(const TBulkDataGenerator& owner, int tableNum);
+            TContext(const TBulkDataGenerator& owner, int tableNum, TGeneratorStateProcessor* state);
             NYdb::TValueBuilder& GetBuilder();
             TStringBuilder& GetCsv();
             void AppendPortions(TDataPortions& result);
@@ -43,6 +43,7 @@ class TTpchWorkloadDataInitializerGenerator: public TWorkloadDataInitializerBase
             TStringBuilder Csv;
             const TBulkDataGenerator& Owner;
             int TableNum;
+            TGeneratorStateProcessor* State;
         };
 
         using TContexts = TVector<TContext>;