From 02c117c2ea175663576a8af11b52f934fd9bb192 Mon Sep 17 00:00:00 2001
From: Daniil Demin <deminds@ydb.tech>
Date: Thu, 23 Jan 2025 18:09:14 +0000
Subject: [PATCH] export: correct cancellation state handling

---
 .../schemeshard_export__cancel.cpp            |  2 +-
 .../schemeshard_export__create.cpp            | 26 ++++++++++++++-----
 .../schemeshard_export__forget.cpp            |  8 ++++++
 3 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/ydb/core/tx/schemeshard/schemeshard_export__cancel.cpp b/ydb/core/tx/schemeshard/schemeshard_export__cancel.cpp
index f485618773e9..ec7635c7cd2a 100644
--- a/ydb/core/tx/schemeshard/schemeshard_export__cancel.cpp
+++ b/ydb/core/tx/schemeshard/schemeshard_export__cancel.cpp
@@ -60,7 +60,7 @@ struct TSchemeShard::TExport::TTxCancel: public TSchemeShard::TXxport::TTxBase {
         }
 
         LOG_D("TExport::TTxCancel, cancelling manually"
-            << ", info# " << exportInfo->ToString()
+            << ", info: " << exportInfo->ToString()
         );
 
         exportInfo->Issue = "Cancelled manually";
diff --git a/ydb/core/tx/schemeshard/schemeshard_export__create.cpp b/ydb/core/tx/schemeshard/schemeshard_export__create.cpp
index d883f5bad57e..4ba9bff6a045 100644
--- a/ydb/core/tx/schemeshard/schemeshard_export__create.cpp
+++ b/ydb/core/tx/schemeshard/schemeshard_export__create.cpp
@@ -659,7 +659,7 @@ struct TSchemeShard::TExport::TTxProgress: public TSchemeShard::TXxport::TTxBase
                 for (ui32 itemIdx : xrange(exportInfo->Items.size())) {
                     const auto& item = exportInfo->Items.at(itemIdx);
 
-                    if (item.State != EState::Dropping) {
+                    if (item.SourcePathType != NKikimrSchemeOp::EPathTypeTable || item.State != EState::Dropping) {
                         continue;
                     }
 
@@ -690,6 +690,15 @@ struct TSchemeShard::TExport::TTxProgress: public TSchemeShard::TXxport::TTxBase
         return itemIdx;
     }
 
+    void EndExport(TExportInfo::TPtr exportInfo, EState finalState, NIceDb::TNiceDb& db) {
+        exportInfo->State = finalState;
+        exportInfo->EndTime = TAppData::TimeProvider->Now();
+
+        Self->PersistExportState(db, exportInfo);
+        SendNotificationsIfFinished(exportInfo);
+        AuditLogExportEnd(*exportInfo.Get(), Self);
+    }
+
     void OnAllocateResult() {
         Y_ABORT_UNLESS(AllocateResult);
 
@@ -998,12 +1007,17 @@ struct TSchemeShard::TExport::TTxProgress: public TSchemeShard::TXxport::TTxBase
             Self->PersistExportItemState(db, exportInfo, itemIdx);
 
             if (AllOf(exportInfo->Items, &TExportInfo::TItem::IsDone)) {
-                exportInfo->State = EState::Done;
-                exportInfo->EndTime = TAppData::TimeProvider->Now();
+                EndExport(exportInfo, EState::Done, db);
+            }
+        } else if (exportInfo->State == EState::Cancellation) {
+            item.State = EState::Cancelled;
+            Self->PersistExportItemState(db, exportInfo, itemIdx);
 
-                Self->PersistExportState(db, exportInfo);
-                SendNotificationsIfFinished(exportInfo);
-                AuditLogExportEnd(*exportInfo.Get(), Self);
+            if (AllOf(exportInfo->Items, [](const TExportInfo::TItem& item) {
+                // on cancellation we wait only for transferring items
+                return item.State != EState::Transferring;
+            })) {
+                EndExport(exportInfo, EState::Cancelled, db);
             }
         }
     }
diff --git a/ydb/core/tx/schemeshard/schemeshard_export__forget.cpp b/ydb/core/tx/schemeshard/schemeshard_export__forget.cpp
index a829cba4fcd8..e5f8d9fc55ce 100644
--- a/ydb/core/tx/schemeshard/schemeshard_export__forget.cpp
+++ b/ydb/core/tx/schemeshard/schemeshard_export__forget.cpp
@@ -74,10 +74,14 @@ struct TSchemeShard::TExport::TTxForget: public TSchemeShard::TXxport::TTxBase {
             Self->Exports.erase(exportInfo->Id);
             Self->PersistRemoveExport(db, exportInfo);
         } else {
+            LOG_D("TExport::TTxForget, dropping export tables"
+                << ", info: " << exportInfo->ToString()
+            );
             exportInfo->WaitTxId = InvalidTxId;
             exportInfo->State = TExportInfo::EState::Dropping;
             Self->PersistExportState(db, exportInfo);
 
+            TVector<ui32> itemsToDrop;
             for (ui32 itemIdx : xrange(exportInfo->Items.size())) {
                 auto& item = exportInfo->Items.at(itemIdx);
 
@@ -87,10 +91,14 @@ struct TSchemeShard::TExport::TTxForget: public TSchemeShard::TXxport::TTxBase {
                 const TPath itemPath = TPath::Resolve(ExportItemPathName(Self, exportInfo, itemIdx), Self);
                 if (itemPath.IsResolved() && !itemPath.IsDeleted()) {
                     item.State = TExportInfo::EState::Dropping;
+                    itemsToDrop.emplace_back(itemIdx);
                 }
 
                 Self->PersistExportItemState(db, exportInfo, itemIdx);
             }
+            LOG_T("TExport::TTxForget, items to drop"
+                << ", items: " << JoinSeq(", ", itemsToDrop)
+            );
 
             Progress = true;
         }