Ускоряем рекурсивную операцию ydb schem ls (#2023)

Author: iddqdex

ydb/public/lib/ydb_cli/commands/ydb_service_scheme.cpp

@@ -858,6 +858,8 @@ void TCommandList::Config(TConfig& config) {
         .StoreTrue(&Recursive);
     config.Opts->AddCharOption('1', "List one object per line")
         .StoreTrue(&FromNewLine);
+    config.Opts->AddCharOption('m', "Multithread recursive request")
+        .StoreTrue(&Multithread);
     AddFormats(config, { EOutputFormat::Pretty, EOutputFormat::Json });
     config.SetFreeArgsMax(1);
     SetFreeArgTitle(0, "<path>", "Path to list");
@@ -877,6 +879,7 @@ int TCommandList::Run(TConfig& config) {
     ISchemePrinter::TSettings settings = {
         Path,
         Recursive,
+        Multithread,
         FromNewLine,
         FillSettings(NScheme::TListDirectorySettings()),
         FillSettings(NTable::TDescribeTableSettings().WithTableStatistics(true))

ydb/public/lib/ydb_cli/commands/ydb_service_scheme.h

@@ -102,6 +102,7 @@ class TCommandList : public TYdbOperationCommand, public TCommandWithPath, publi
     bool AdvancedMode = false;
     bool Recursive = false;
     bool FromNewLine = false;
+    bool Multithread = false;
 };
 
 class TCommandPermissions : public TClientCommandTree {

ydb/public/lib/ydb_cli/common/scheme_printers.cpp

@@ -15,7 +15,7 @@ TSchemePrinterBase::TSchemePrinterBase(const TDriver& driver, TSettings&& settin
 {}
 
 void TSchemePrinterBase::Print() {
-    PrintDirectoryRecursive(Settings.Path, "");
+    PrintDirectoryRecursive(Settings.Path, "").GetValueSync();
 }
 
 bool TSchemePrinterBase::IsDirectoryLike(const NScheme::TSchemeEntry& entry) {
@@ -24,30 +24,41 @@ bool TSchemePrinterBase::IsDirectoryLike(const NScheme::TSchemeEntry& entry) {
         || entry.Type == NScheme::ESchemeEntryType::ColumnStore;
 }
 
-void TSchemePrinterBase::PrintDirectoryRecursive(const TString& fullPath, const TString& relativePath) {
-    NScheme::TListDirectoryResult result = SchemeClient.ListDirectory(
+NThreading::TFuture<void> TSchemePrinterBase::PrintDirectoryRecursive(const TString& fullPath, const TString& relativePath) {
+    return SchemeClient.ListDirectory(
         fullPath,
         Settings.ListDirectorySettings
-    ).GetValueSync();
-    ThrowOnError(result);
+    ).Apply([this, fullPath, relativePath](const NScheme::TAsyncListDirectoryResult& resultFuture) {
+        const auto& result = resultFuture.GetValueSync();
+        ThrowOnError(result);
 
-    if (relativePath || IsDirectoryLike(result.GetEntry())) {
-        PrintDirectory(relativePath, result);
-    } else {
-        PrintEntry(relativePath, result.GetEntry());
-    }
+        if (relativePath || IsDirectoryLike(result.GetEntry())) {
+            std::lock_guard g(Lock);
+            PrintDirectory(relativePath, result);
+        } else {
+            std::lock_guard g(Lock);
+            PrintEntry(relativePath, result.GetEntry());
+        }
 
-    if (Settings.Recursive) {
-        for (const auto& child : result.GetChildren()) {
-            TString childRelativePath = relativePath + (relativePath ? "/" : "") + child.Name;
-            TString childFullPath = fullPath + "/" + child.Name;
-            if (IsDirectoryLike(child)) {
-                PrintDirectoryRecursive(childFullPath, childRelativePath);
-            } else {
-                PrintEntry(childRelativePath, child);
+        TVector<NThreading::TFuture<void>> childFutures;
+        if (Settings.Recursive) {
+            for (const auto& child : result.GetChildren()) {
+                TString childRelativePath = relativePath + (relativePath ? "/" : "") + child.Name;
+                TString childFullPath = fullPath + "/" + child.Name;
+                if (IsDirectoryLike(child)) {
+                    childFutures.push_back(PrintDirectoryRecursive(childFullPath, childRelativePath));
+                    if (!Settings.Multithread) {
+                        childFutures.back().Wait();
+                        childFutures.back().TryRethrow();
+                    }
+                } else {
+                    std::lock_guard g(Lock);
+                    PrintEntry(childRelativePath, child);
+                }
             }
         }
-    }
+        return NThreading::WaitExceptionOrAll(childFutures);
+    });
 }
 
 NTable::TDescribeTableResult TSchemePrinterBase::DescribeTable(const TString& relativePath) {

ydb/public/lib/ydb_cli/common/scheme_printers.h

@@ -13,6 +13,7 @@ class ISchemePrinter {
     struct TSettings {
         TString Path;
         bool Recursive;
+        bool Multithread;
         bool FromNewLine;
         NScheme::TListDirectorySettings ListDirectorySettings;
         NTable::TDescribeTableSettings DescribeTableSettings;
@@ -37,13 +38,14 @@ class TSchemePrinterBase : public ISchemePrinter {
     NTable::TDescribeTableResult DescribeTable(const TString& relativePath);
 
 private:
-    void PrintDirectoryRecursive(const TString& fullPath, const TString& relativePath);
+    NThreading::TFuture<void> PrintDirectoryRecursive(const TString& fullPath, const TString& relativePath);
     static bool IsDirectoryLike(const NScheme::TSchemeEntry& entry);
 
 protected:
     NTable::TTableClient TableClient;
     NScheme::TSchemeClient SchemeClient;
     const TSettings Settings;
+    std::mutex Lock;
 };
 
 class TDefaultSchemePrinter : public TSchemePrinterBase {