Wait writer (ydb-platform#2935)

GrigoriyPA · Mar 19, 2024 · 69bc320 · 69bc320
1 parent 7370f61
commit 69bc320
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/ydb/public/lib/ydb_cli/dump/restore_impl.cpp b/ydb/public/lib/ydb_cli/dump/restore_impl.cpp
@@ -315,6 +315,19 @@ 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) {
     if (desc.GetAttributes().contains(DOC_API_TABLE_VERSION_ATTR) && settings.SkipDocumentTables_) {
         return Result<TRestoreResult>();
@@ -354,6 +367,7 @@ TRestoreResult TRestoreClient::RestoreData(const TFsPath& fsPath, const TString&
         }
     }
 
+    TWriterWaiter waiter(*writer);
     ui32 dataFileId = 0;
     TFsPath dataFile = fsPath.Child(DataFileName(dataFileId));
 
@@ -389,8 +403,6 @@ TRestoreResult TRestoreClient::RestoreData(const TFsPath& fsPath, const TString&
         }
     }
 
-    writer->Wait();
-
     return Result<TRestoreResult>();
 }