Remove adding new shards for data erasure from SetPartitioning (#19324)

Author: molotkov-and

ydb/core/tx/schemeshard/schemeshard__operation_common.cpp

@@ -1,4 +1,5 @@
 #include "schemeshard__operation_common.h"
+#include "schemeshard__data_erasure_manager.h"
 
 #include <ydb/core/blob_depot/events.h>
 #include <ydb/core/blockstore/core/blockstore.h>
@@ -928,8 +929,16 @@ void UpdatePartitioningForCopyTable(TOperationId operationId, TTxState &txState,
     TShardInfo datashardInfo = TShardInfo::DataShardInfo(operationId.GetTxId(), txState.TargetPathId);
     datashardInfo.BindedChannels = channelsBinding;
 
-    context.SS->SetPartitioning(txState.TargetPathId, dstTableInfo,
-        ApplyPartitioningCopyTable(datashardInfo, srcTableInfo, txState, context.SS));
+    auto newPartitioning = ApplyPartitioningCopyTable(datashardInfo, srcTableInfo, txState, context.SS);
+    TVector<TShardIdx> newShardsIdx;
+    newShardsIdx.reserve(newPartitioning.size());
+    for (const auto& part : newPartitioning) {
+        newShardsIdx.push_back(part.ShardIdx);
+    }
+    context.SS->SetPartitioning(txState.TargetPathId, dstTableInfo, std::move(newPartitioning));
+    if (context.SS->EnableDataErasure && context.SS->DataErasureManager->GetStatus() == EDataErasureStatus::IN_PROGRESS) {
+        context.OnComplete.Send(context.SS->SelfId(), new TEvPrivate::TEvAddNewShardToDataErasure(std::move(newShardsIdx)));
+    }
 
     ui32 newShardCout = dstTableInfo->GetPartitions().size();
 

ydb/core/tx/schemeshard/schemeshard__operation_copy_table.cpp

@@ -1,5 +1,6 @@
 #include "schemeshard__operation_part.h"
 #include "schemeshard__operation_common.h"
+#include "schemeshard__data_erasure_manager.h"
 #include "schemeshard_impl.h"
 #include "schemeshard_tx_infly.h"
 #include "schemeshard_cdc_stream_common.h"
@@ -670,11 +671,17 @@ class TCopyTable: public TSubOperation {
         TShardInfo datashardInfo = TShardInfo::DataShardInfo(OperationId.GetTxId(), newTable->PathId);
         datashardInfo.BindedChannels = channelsBinding;
         auto newPartition = NTableState::ApplyPartitioningCopyTable(datashardInfo, srcTableInfo, txState, context.SS);
+        TVector<TShardIdx> newShardsIdx;
+        newShardsIdx.reserve(newPartition.size());
         for (const auto& part: newPartition) {
             context.MemChanges.GrabNewShard(context.SS, part.ShardIdx);
             context.DbChanges.PersistShard(part.ShardIdx);
+            newShardsIdx.push_back(part.ShardIdx);
         }
         context.SS->SetPartitioning(newTable->PathId, tableInfo, std::move(newPartition));
+        if (context.SS->EnableDataErasure && context.SS->DataErasureManager->GetStatus() == EDataErasureStatus::IN_PROGRESS) {
+            context.OnComplete.Send(context.SS->SelfId(), new TEvPrivate::TEvAddNewShardToDataErasure(std::move(newShardsIdx)));
+        }
         for (const auto& shard : tableInfo->GetPartitions()) {
             Y_ABORT_UNLESS(context.SS->ShardInfos.contains(shard.ShardIdx), "shard info is set before");
             if (storePerShardConfig) {

ydb/core/tx/schemeshard/schemeshard__operation_split_merge.cpp

@@ -1,5 +1,6 @@
 #include "schemeshard__operation_part.h"
 #include "schemeshard__operation_common.h"
+#include "schemeshard__data_erasure_manager.h"
 #include "schemeshard_impl.h"
 
 #include <ydb/core/base/subdomain.h>
@@ -244,6 +245,7 @@ class TTransferData: public TSubOperationState {
 
         // Replace all Src datashard(s) with Dst datashard(s)
         TVector<TTableShardInfo> newPartitioning;
+        TVector<TShardIdx> newShardsIdx;
         THashSet<TShardIdx> allSrcShardIdxs;
         for (const auto& txShard : txState->Shards) {
             if (txShard.Operation == TTxState::TransferData)
@@ -280,6 +282,7 @@ class TTransferData: public TSubOperationState {
                     }
 
                     newPartitioning.push_back(dst);
+                    newShardsIdx.push_back(dst.ShardIdx);
                 }
 
                 dstAdded = true;
@@ -293,6 +296,9 @@ class TTransferData: public TSubOperationState {
         // Delete the whole old partitioning and persist the whole new partitioning as the indexes have changed
         context.SS->PersistTablePartitioningDeletion(db, tableId, tableInfo);
         context.SS->SetPartitioning(tableId, tableInfo, std::move(newPartitioning));
+        if (context.SS->EnableDataErasure && context.SS->DataErasureManager->GetStatus() == EDataErasureStatus::IN_PROGRESS) {
+            context.OnComplete.Send(context.SS->SelfId(), new TEvPrivate::TEvAddNewShardToDataErasure(std::move(newShardsIdx)));
+        }
         context.SS->PersistTablePartitioning(db, tableId, tableInfo);
         context.SS->PersistTablePartitionStats(db, tableId, tableInfo);
 

ydb/core/tx/schemeshard/schemeshard__tenant_data_erasure_manager.cpp

@@ -631,13 +631,13 @@ NTabletFlatExecutor::ITransaction* TSchemeShard::CreateTxCompleteDataErasureShar
 template NTabletFlatExecutor::ITransaction* TSchemeShard::CreateTxCompleteDataErasureShard<TEvDataShard::TEvForceDataCleanupResult::TPtr>(TEvDataShard::TEvForceDataCleanupResult::TPtr& ev);
 template NTabletFlatExecutor::ITransaction* TSchemeShard::CreateTxCompleteDataErasureShard<TEvKeyValue::TEvCleanUpDataResponse::TPtr>(TEvKeyValue::TEvCleanUpDataResponse::TPtr& ev);
 
-struct TSchemeShard::TTxAddEntryToDataErasure : public TSchemeShard::TRwTxBase {
-    const std::vector<TShardIdx> DataErasureShards;
+struct TSchemeShard::TTxAddNewShardToDataErasure : public TSchemeShard::TRwTxBase {
+    TEvPrivate::TEvAddNewShardToDataErasure::TPtr Ev;
     bool NeedResponseComplete = false;
 
-    TTxAddEntryToDataErasure(TSelf *self, const std::vector<TShardIdx>& dataErasureShards)
+    TTxAddNewShardToDataErasure(TSelf *self, TEvPrivate::TEvAddNewShardToDataErasure::TPtr& ev)
         : TRwTxBase(self)
-        , DataErasureShards(std::move(dataErasureShards))
+        , Ev(std::move(ev))
     {}
 
     TTxType GetTxType() const override { return TXTYPE_ADD_SHARDS_DATA_ERASURE; }
@@ -647,7 +647,7 @@ struct TSchemeShard::TTxAddEntryToDataErasure : public TSchemeShard::TRwTxBase {
             "TTxAddEntryToDataErasure Execute at schemestard: " << Self->TabletID());
 
         NIceDb::TNiceDb db(txc.DB);
-        Self->DataErasureManager->HandleNewPartitioning(DataErasureShards, db);
+        Self->DataErasureManager->HandleNewPartitioning(Ev->Get()->Shards, db);
         if (Self->DataErasureManager->GetStatus() == EDataErasureStatus::COMPLETED) {
             if (Self->DataErasureManager->IsRunning()) {
                 NeedResponseComplete = true;
@@ -664,8 +664,8 @@ struct TSchemeShard::TTxAddEntryToDataErasure : public TSchemeShard::TRwTxBase {
     }
 };
 
-NTabletFlatExecutor::ITransaction* TSchemeShard::CreateTxAddEntryToDataErasure(const std::vector<TShardIdx>& dataErasureShards) {
-    return new TTxAddEntryToDataErasure(this, dataErasureShards);
+NTabletFlatExecutor::ITransaction* TSchemeShard::CreateTxAddNewShardToDataErasure(TEvPrivate::TEvAddNewShardToDataErasure::TPtr& ev) {
+    return new TTxAddNewShardToDataErasure(this, ev);
 }
 
 struct TSchemeShard::TTxCancelDataErasureShards : public TSchemeShard::TRwTxBase {

ydb/core/tx/schemeshard/schemeshard_impl.cpp

@@ -4977,6 +4977,7 @@ void TSchemeShard::StateWork(STFUNC_SIG) {
         HFuncTraced(TEvSchemeShard::TEvTenantDataErasureRequest, Handle);
         HFuncTraced(TEvDataShard::TEvForceDataCleanupResult, Handle);
         HFuncTraced(TEvKeyValue::TEvCleanUpDataResponse, Handle);
+        HFuncTraced(TEvPrivate::TEvAddNewShardToDataErasure, Handle);
         HFuncTraced(TEvSchemeShard::TEvTenantDataErasureResponse, Handle);
         HFuncTraced(TEvSchemeShard::TEvDataErasureInfoRequest, Handle);
         HFuncTraced(TEvSchemeShard::TEvDataErasureManualStartupRequest, Handle);
@@ -7115,9 +7116,6 @@ void TSchemeShard::SetPartitioning(TPathId pathId, TTableInfo::TPtr tableInfo, T
                 newDataErasureShards.push_back(p.ShardIdx);
             }
         }
-        if (EnableDataErasure && DataErasureManager->GetStatus() == EDataErasureStatus::IN_PROGRESS) {
-            Execute(CreateTxAddEntryToDataErasure(newDataErasureShards), this->ActorContext());
-        }
 
         for (const auto& p: oldPartitioning) {
             if (!newPartitioningSet.contains(p.ShardIdx)) {
@@ -7678,6 +7676,10 @@ void TSchemeShard::Handle(TEvKeyValue::TEvCleanUpDataResponse::TPtr& ev, const T
     Execute(this->CreateTxCompleteDataErasureShard(ev), ctx);
 }
 
+void TSchemeShard::Handle(TEvPrivate::TEvAddNewShardToDataErasure::TPtr& ev, const TActorContext& ctx) {
+    Execute(CreateTxAddNewShardToDataErasure(ev), ctx);
+}
+
 void TSchemeShard::Handle(TEvSchemeShard::TEvTenantDataErasureResponse::TPtr& ev, const TActorContext& ctx) {
     Execute(CreateTxCompleteDataErasureTenant(ev), ctx);
 }

ydb/core/tx/schemeshard/schemeshard_impl.h

@@ -1048,8 +1048,8 @@ class TSchemeShard
     struct TTxRunDataErasure;
     NTabletFlatExecutor::ITransaction* CreateTxRunDataErasure(bool isNewDataErasure);
 
-    struct TTxAddEntryToDataErasure;
-    NTabletFlatExecutor::ITransaction* CreateTxAddEntryToDataErasure(const std::vector<TShardIdx>& dataErasureShards);
+    struct TTxAddNewShardToDataErasure;
+    NTabletFlatExecutor::ITransaction* CreateTxAddNewShardToDataErasure(TEvPrivate::TEvAddNewShardToDataErasure::TPtr& ev);
 
     struct TTxCancelDataErasureShards;
     NTabletFlatExecutor::ITransaction* CreateTxCancelDataErasureShards(const std::vector<TShardIdx>& oldShards);
@@ -1179,6 +1179,7 @@ class TSchemeShard
     void Handle(TEvDataShard::TEvForceDataCleanupResult::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvKeyValue::TEvCleanUpDataResponse__HandlePtr& ev, const TActorContext& ctx);
     void Handle(TEvSchemeShard::TEvTenantDataErasureResponse::TPtr& ev, const TActorContext& ctx);
+    void Handle(TEvPrivate::TEvAddNewShardToDataErasure::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvBlobStorage::TEvControllerShredResponse::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvSchemeShard::TEvDataErasureInfoRequest::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvSchemeShard::TEvDataErasureManualStartupRequest::TPtr& ev, const TActorContext& ctx);

ydb/core/tx/schemeshard/schemeshard_private.h

@@ -40,6 +40,7 @@ namespace TEvPrivate {
         EvRetryNodeSubscribe,
         EvRunDataErasure,
         EvRunTenantDataErasure,
+        EvAddNewShardToDataErasure,
         EvEnd
     };
 
@@ -238,6 +239,14 @@ namespace TEvPrivate {
             : NodeId(nodeId)
         { }
     };
+
+    struct TEvAddNewShardToDataErasure : public TEventLocal<TEvAddNewShardToDataErasure, EvAddNewShardToDataErasure> {
+        const std::vector<TShardIdx> Shards;
+
+        TEvAddNewShardToDataErasure(std::vector<TShardIdx>&& shards)
+            : Shards(std::move(shards))
+        {}
+    };
 }; // TEvPrivate
 
 } // NSchemeShard