Enhanced parallelism of data restoring in ydb tools restore (#12203)

Author: pixcc

ydb/apps/ydb/CHANGELOG.md

@@ -1,4 +1,6 @@
 
+* Enhanced parallelism of data restoring in `ydb tools restore`
+
 ## 2.16.0 ##
 
 * Improved throughput of `ydb import file csv` command. It is now approximately x3 times faster

ydb/public/lib/ydb_cli/dump/restore_compat.cpp

@@ -95,7 +95,10 @@ class TDataAccumulator: public NPrivate::IDataAccumulator, protected TQueryBuild
         Bytes = 0;
         RequestUnitsX2 = 0;
 
-        return {}; // Writer gets data directly from accumulator
+        // Writer gets data directly from accumulator
+        NPrivate::TBatch batch;
+        batch.SetOriginAccumulator(this);
+        return batch; 
     }
 
 private:
@@ -115,35 +118,37 @@ class TDataWriter: public NPrivate::IDataWriter {
     explicit TDataWriter(
             const TString& path,
             TTableClient& tableClient,
-            NPrivate::IDataAccumulator* accumulator,
+            const NPrivate::IDataAccumulator* accumulator,
             const TRestoreSettings& settings)
         : Path(path)
         , TableClient(tableClient)
-        , Accumulator(dynamic_cast<TDataAccumulator*>(accumulator))
         , UseBulkUpsert(settings.Mode_ == TRestoreSettings::EMode::BulkUpsert)
-    {
-        Y_ENSURE(Accumulator);
+    {   
+        const auto* dataAccumulator = dynamic_cast<const TDataAccumulator*>(accumulator);
+        Y_ENSURE(dataAccumulator);
 
         TUploader::TOptions opts;
         opts.InFly = settings.InFly_;
         opts.Rate = settings.RateLimiterSettings_.Rate_;
         opts.Interval = settings.RateLimiterSettings_.Interval_;
         opts.ReactionTime = settings.RateLimiterSettings_.ReactionTime_;
 
-        Uploader = MakeHolder<TUploader>(opts, TableClient, Accumulator->GetQueryString());
+        Uploader = MakeHolder<TUploader>(opts, TableClient, dataAccumulator->GetQueryString());
     }
 
-    bool Push(NPrivate::TBatch&&) override {
-        bool ok;
+    bool Push(NPrivate::TBatch&& batch) override {
+        auto* accumulator = dynamic_cast<TDataAccumulator*>(batch.GetOriginAccumulator());
+        Y_ENSURE(accumulator);
 
+        bool ok;
         if (UseBulkUpsert) {
-            ok = Uploader->Push(Path, Accumulator->EndAndGetResultingValue());
+            ok = Uploader->Push(Path, accumulator->EndAndGetResultingValue());
         } else {
-            ok = Uploader->Push(Accumulator->EndAndGetResultingParams());
+            ok = Uploader->Push(accumulator->EndAndGetResultingParams());
         }
 
         if (ok) {
-            Accumulator->Begin();
+            accumulator->Begin();
         }
 
         return ok;
@@ -156,7 +161,6 @@ class TDataWriter: public NPrivate::IDataWriter {
 private:
     const TString Path;
     TTableClient& TableClient;
-    TDataAccumulator* Accumulator;
     const bool UseBulkUpsert;
     THolder<TUploader> Uploader;
 
@@ -174,7 +178,7 @@ NPrivate::IDataAccumulator* CreateCompatAccumulator(
 NPrivate::IDataWriter* CreateCompatWriter(
         const TString& path,
         TTableClient& tableClient,
-        NPrivate::IDataAccumulator* accumulator,
+        const NPrivate::IDataAccumulator* accumulator,
         const TRestoreSettings& settings) {
     return new TDataWriter(path, tableClient, accumulator, settings);
 }

ydb/public/lib/ydb_cli/dump/restore_compat.h

@@ -13,7 +13,7 @@ NPrivate::IDataAccumulator* CreateCompatAccumulator(
 NPrivate::IDataWriter* CreateCompatWriter(
     const TString& path,
     NTable::TTableClient& tableClient,
-    NPrivate::IDataAccumulator* accumulator,
+    const NPrivate::IDataAccumulator* accumulator,
     const TRestoreSettings& settings);
 
 } // NDump

ydb/public/lib/ydb_cli/dump/restore_impl.cpp

@@ -80,6 +80,24 @@ bool IsOperationStarted(TStatus operationStatus) {
     return operationStatus.IsSuccess() || operationStatus.GetStatus() == EStatus::STATUS_UNDEFINED;
 }
 
+ui32 CountDataFiles(const TFsPath& fsPath) {
+    ui32 dataFileId = 0;
+    TFsPath dataFile = fsPath.Child(DataFileName(dataFileId));
+    while (dataFile.Exists()) {
+        dataFile = fsPath.Child(DataFileName(++dataFileId));
+    }
+    return dataFileId;
+}
+
+TRestoreResult CombineResults(const TVector<TRestoreResult>& results) {
+    for (auto result : results) {
+        if (!result.IsSuccess()) {
+            return result;
+        }
+    }
+    return Result<TRestoreResult>();
+}
+
 } // anonymous
 
 namespace NPrivate {
@@ -392,95 +410,164 @@ TRestoreResult TRestoreClient::CheckSchema(const TString& dbPath, const TTableDe
     return Result<TRestoreResult>();
 }
 
-struct TWriterWaiter {
-    NPrivate::IDataWriter& Writer;
-
-    TWriterWaiter(NPrivate::IDataWriter& writer)
-        : Writer(writer)
-    {
-    }
-
-    ~TWriterWaiter() {
-        Writer.Wait();
-    }
-};
-
-TRestoreResult TRestoreClient::RestoreData(const TFsPath& fsPath, const TString& dbPath, const TRestoreSettings& settings, const TTableDescription& desc) {
-    THolder<NPrivate::IDataAccumulator> accumulator;
+THolder<NPrivate::IDataWriter> TRestoreClient::CreateDataWriter(const TString& dbPath, const TRestoreSettings& settings,
+    const TTableDescription& desc, const TVector<THolder<NPrivate::IDataAccumulator>>& accumulators)
+{   
     THolder<NPrivate::IDataWriter> writer;
-
     switch (settings.Mode_) {
         case TRestoreSettings::EMode::Yql:
         case TRestoreSettings::EMode::BulkUpsert: {
-            accumulator.Reset(CreateCompatAccumulator(dbPath, desc, settings));
-            writer.Reset(CreateCompatWriter(dbPath, TableClient, accumulator.Get(), settings));
+            // Need only one accumulator to initialize query string
+            writer.Reset(CreateCompatWriter(dbPath, TableClient, accumulators[0].Get(), settings));
+            break;
+        }
 
+        case TRestoreSettings::EMode::ImportData: {
+            writer.Reset(CreateImportDataWriter(dbPath, desc, ImportClient, TableClient, accumulators, settings, Log));
             break;
         }
+    }
+    return writer;
+}
+
+TRestoreResult TRestoreClient::CreateDataAccumulators(TVector<THolder<NPrivate::IDataAccumulator>>& outAccumulators,
+    const TString& dbPath, const TRestoreSettings& settings, const NTable::TTableDescription& desc, ui32 dataFilesCount)
+{
+    const ui32 accumulatorsCount = std::min(settings.InFly_, dataFilesCount);
+    outAccumulators.resize(accumulatorsCount);
+   
+    switch (settings.Mode_) {
+        case TRestoreSettings::EMode::Yql:
+        case TRestoreSettings::EMode::BulkUpsert:
+            for (size_t i = 0; i < accumulatorsCount; ++i) {
+                outAccumulators[i].Reset(CreateCompatAccumulator(dbPath, desc, settings));
+            }
+            break;
 
         case TRestoreSettings::EMode::ImportData: {
             TMaybe<TTableDescription> actualDesc;
             auto descResult = DescribeTable(TableClient, dbPath, actualDesc);
             if (!descResult.IsSuccess()) {
                 return Result<TRestoreResult>(dbPath, std::move(descResult));
             }
-
-            accumulator.Reset(CreateImportDataAccumulator(desc, *actualDesc, settings, Log));
-            writer.Reset(CreateImportDataWriter(dbPath, desc, ImportClient, TableClient, accumulator.Get(), settings, Log));
-
+            for (size_t i = 0; i < accumulatorsCount; ++i) {
+                outAccumulators[i].Reset(CreateImportDataAccumulator(desc, *actualDesc, settings, Log));
+            }
             break;
         }
     }
+    return Result<TRestoreResult>();
+}
 
-    TWriterWaiter waiter(*writer);
-
-    ui32 dataFileId = 0;
-    TFsPath dataFile = fsPath.Child(DataFileName(dataFileId));
-    TVector<TString> dataFileNames;
-
-    while (dataFile.Exists()) {
-        LOG_D("Read data from " << dataFile.GetPath().Quote());
-
-        dataFileNames.push_back(dataFile);
-        TFileInput input(dataFile, settings.FileBufferSize_);
-        TString line;
-        ui64 lineNo = 0;
-
-        while (input.ReadLine(line)) {
-            auto l = NPrivate::TLine(std::move(line), dataFileNames.back(), ++lineNo);
-            for (auto status = accumulator->Check(l); status != NPrivate::IDataAccumulator::OK; status = accumulator->Check(l)) {
-                if (status == NPrivate::IDataAccumulator::ERROR) {
-                    return Result<TRestoreResult>(dbPath, EStatus::GENERIC_ERROR,
-                        TStringBuilder() << "Invalid data: " << l.GetLocation());
-                }
-
-                if (!accumulator->Ready(true)) {
-                    LOG_E("Error reading data from " << dataFile.GetPath().Quote());
-                    return Result<TRestoreResult>(dbPath, EStatus::INTERNAL_ERROR, "Data is not ready");
-                }
+TRestoreResult TRestoreClient::RestoreData(const TFsPath& fsPath, const TString& dbPath, const TRestoreSettings& settings, const TTableDescription& desc) {
+    const ui32 dataFilesCount = CountDataFiles(fsPath);
+    if (dataFilesCount == 0) {
+        return Result<TRestoreResult>();
+    }
 
-                if (!writer->Push(accumulator->GetData(true))) {
-                    LOG_E("Error writing data to " << dbPath.Quote() << ", file: " << dataFile.GetPath().Quote());
-                    return Result<TRestoreResult>(dbPath, EStatus::GENERIC_ERROR, "Cannot write data #1");
+    TVector<THolder<NPrivate::IDataAccumulator>> accumulators;
+    if (auto res = CreateDataAccumulators(accumulators, dbPath, settings, desc, dataFilesCount); !res.IsSuccess()) {
+        return res;
+    }
+    
+    THolder<NPrivate::IDataWriter> writer = CreateDataWriter(dbPath, settings, desc, accumulators);
+
+    TVector<TRestoreResult> accumulatorWorkersResults(accumulators.size(), Result<TRestoreResult>());
+    TThreadPool accumulatorWorkers(TThreadPool::TParams().SetBlocking(true));
+    accumulatorWorkers.Start(accumulators.size(), accumulators.size());  
+
+    const ui32 dataFilesPerAccumulator = dataFilesCount / accumulators.size();
+    const ui32 dataFilesPerAccumulatorRemainder = dataFilesCount % accumulators.size();
+    for (ui32 i = 0; i < accumulators.size(); ++i) {
+        auto* accumulator = accumulators[i].Get();
+
+        ui32 dataFileIdStart = dataFilesPerAccumulator * i + std::min(i, dataFilesPerAccumulatorRemainder);
+        ui32 dataFileIdEnd = dataFilesPerAccumulator * (i + 1) + std::min(i + 1, dataFilesPerAccumulatorRemainder);
+        auto func = [&, i, dataFileIdStart, dataFileIdEnd, accumulator]() {
+            for (size_t id = dataFileIdStart; id < dataFileIdEnd; ++id) {
+                TFsPath dataFile = fsPath.Child(DataFileName(id));
+
+                LOG_D("Read data from " << dataFile.GetPath().Quote());
+
+                TFileInput input(dataFile, settings.FileBufferSize_);
+                TString line;
+                ui64 lineNo = 0;
+
+                while (input.ReadLine(line)) {
+                    auto l = NPrivate::TLine(std::move(line), dataFile.GetPath(), ++lineNo);
+
+                    for (auto status = accumulator->Check(l); status != NPrivate::IDataAccumulator::OK; status = accumulator->Check(l)) {
+                        if (status == NPrivate::IDataAccumulator::ERROR) {
+                            accumulatorWorkersResults[i] = Result<TRestoreResult>(dbPath, EStatus::GENERIC_ERROR,
+                                TStringBuilder() << "Invalid data: " << l.GetLocation());
+                            return;
+                        }
+
+                        if (!accumulator->Ready(true)) {
+                            LOG_E("Error reading data from " << dataFile.GetPath().Quote());
+                            accumulatorWorkersResults[i] = Result<TRestoreResult>(dbPath, EStatus::INTERNAL_ERROR, "Data is not ready");
+                            return;
+                        }
+
+                        if (!writer->Push(accumulator->GetData(true))) {
+                            LOG_E("Error writing data to " << dbPath.Quote() << ", file: " << dataFile.GetPath().Quote());
+                            accumulatorWorkersResults[i] = Result<TRestoreResult>(dbPath, EStatus::GENERIC_ERROR, "Cannot write data #1");
+                            return;
+                        }
+                    }
+
+                    accumulator->Feed(std::move(l));
+                    if (accumulator->Ready()) {
+                        if (!writer->Push(accumulator->GetData())) {
+                            LOG_E("Error writing data to " << dbPath.Quote() << ", file: " << dataFile.GetPath().Quote());
+                            accumulatorWorkersResults[i] = Result<TRestoreResult>(dbPath, EStatus::GENERIC_ERROR, "Cannot write data #2");
+                            return;
+                        }
+                    }
                 }
             }
 
-            accumulator->Feed(std::move(l));
-            if (accumulator->Ready()) {
-                if (!writer->Push(accumulator->GetData())) {
-                    LOG_E("Error writing data to " << dbPath.Quote() << ", file: " << dataFile.GetPath().Quote());
-                    return Result<TRestoreResult>(dbPath, EStatus::GENERIC_ERROR, "Cannot write data #2");
+            while (accumulator->Ready(true)) {
+                if (!writer->Push(accumulator->GetData(true))) {
+                    accumulatorWorkersResults[i] = Result<TRestoreResult>(dbPath, EStatus::GENERIC_ERROR, "Cannot write data #3");
+                    return;
                 }
             }
+        };
+    
+        if (!accumulatorWorkers.AddFunc(std::move(func))) {
+            return Result<TRestoreResult>(dbPath, EStatus::GENERIC_ERROR, "Can't start restoring data: queue is full or shutting down");
         }
+    }
 
-        dataFile = fsPath.Child(DataFileName(++dataFileId));
+    accumulatorWorkers.Stop();
+    if (auto res = CombineResults(accumulatorWorkersResults); !res.IsSuccess()) {
+        return res;
     }
 
-    while (accumulator->Ready(true)) {
-        if (!writer->Push(accumulator->GetData(true))) {
-            LOG_E("Error writing data to " << dbPath.Quote() << ", file: " << dataFile.GetPath().Quote());
-            return Result<TRestoreResult>(dbPath, EStatus::GENERIC_ERROR, "Cannot write data #3");
+    // ensure that all data is restored
+    while (true) {
+        writer->Wait();
+
+        bool dataFound = false;
+        for (auto& acc : accumulators) {
+            if (acc->Ready(true)) {
+                dataFound = true;
+                break;
+            }
+        }
+
+        if (dataFound) {
+            writer = CreateDataWriter(dbPath, settings, desc, accumulators);
+            for (auto& acc : accumulators) {
+                while (acc->Ready(true)) {
+                    if (!writer->Push(acc->GetData(true))) {
+                        return Result<TRestoreResult>(dbPath, EStatus::GENERIC_ERROR, "Cannot write data #4");
+                    }
+                }
+            }
+        } else {
+            break;
         }
     }
 

ydb/public/lib/ydb_cli/dump/restore_impl.h

@@ -61,9 +61,12 @@ class TLine {
     }
 };
 
+class IDataAccumulator;
+
 class TBatch {
     TStringBuilder Data;
     TVector<TLocation> Locations;
+    IDataAccumulator* OriginAccumulator;
 
 public:
     void Add(const TLine& line);
@@ -84,6 +87,14 @@ class TBatch {
     inline auto size() const {
         return Data.size();
     }
+
+    inline void SetOriginAccumulator(IDataAccumulator* originAccumulator) {
+        OriginAccumulator = originAccumulator;
+    }
+
+    inline IDataAccumulator* GetOriginAccumulator() const {
+        return OriginAccumulator;
+    }
 };
 
 class IDataAccumulator {
@@ -120,6 +131,12 @@ class TRestoreClient {
     TRestoreResult RestoreIndexes(const TString& dbPath, const NTable::TTableDescription& desc);
     TRestoreResult RestorePermissions(const TFsPath& fsPath, const TString& dbPath, const TRestoreSettings& settings, const THashSet<TString>& oldEntries);
 
+    THolder<NPrivate::IDataWriter> CreateDataWriter(const TString& dbPath, const TRestoreSettings& settings,
+        const NTable::TTableDescription& desc, const TVector<THolder<NPrivate::IDataAccumulator>>& accumulators);
+    TRestoreResult CreateDataAccumulators(TVector<THolder<NPrivate::IDataAccumulator>>& outAccumulators,
+        const TString& dbPath, const TRestoreSettings& settings, const NTable::TTableDescription& desc,
+        ui32 dataFilesCount);
+
 public:
     explicit TRestoreClient(const TDriver& driver, const std::shared_ptr<TLog>& log);