From 69bc3204f034b7b3965c156c9ee9ae724bb25136 Mon Sep 17 00:00:00 2001
From: Ilnaz Nizametdinov <ilnaz@ydb.tech>
Date: Tue, 19 Mar 2024 20:02:38 +0300
Subject: [PATCH] Wait writer (#2935)

---
 ydb/public/lib/ydb_cli/dump/restore_impl.cpp | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/ydb/public/lib/ydb_cli/dump/restore_impl.cpp b/ydb/public/lib/ydb_cli/dump/restore_impl.cpp
index 84ba446aea5e..88a2411202bf 100644
--- 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>();
 }