Allow create_pq op to be a cancelable part of multipart operation (#4438)

Author: dcherednik

ydb/core/tx/schemeshard/schemeshard__operation_create_pq.cpp

@@ -243,7 +243,6 @@ void ApplySharding(TTxId txId,
     pqGroup->InitSplitMergeGraph();
 }
 
-
 class TCreatePQ: public TSubOperation {
     static TTxState::ETxState NextState() {
         return TTxState::CreateParts;
@@ -279,6 +278,11 @@ class TCreatePQ: public TSubOperation {
         }
     }
 
+    ui64 PqShardCountChange = 0;
+    ui64 StreamReservedThroughputChange = 0;
+    ui64 StreamReservedStorageChange = 0;
+    ui64 StreamShardsCountChange = 0;
+
 public:
     using TSubOperation::TSubOperation;
 
@@ -467,12 +471,10 @@ class TCreatePQ: public TSubOperation {
 
         ApplySharding(OperationId.GetTxId(), pathId, pqGroup, txState, tabletChannelsBinding, pqChannelsBinding, context.SS);
 
-        NIceDb::TNiceDb db(context.GetDB());
-
         for (auto& shard : pqGroup->Shards) {
             auto shardIdx = shard.first;
             for (const auto& pqInfo : shard.second->Partitions) {
-                context.SS->PersistPersQueue(db, pathId, shardIdx, *pqInfo);
+                context.DbChanges.PersistPersQueue(pathId, shardIdx, *pqInfo);
             }
         }
 
@@ -483,16 +485,17 @@ class TCreatePQ: public TSubOperation {
         context.SS->Topics[pathId]->AlterData = pqGroup;
         context.SS->IncrementPathDbRefCount(pathId);
 
-        context.SS->PersistPersQueueGroup(db, pathId, emptyGroup);
-        context.SS->PersistAddPersQueueGroupAlter(db, pathId, pqGroup);
+        context.DbChanges.PersistPersQueueGroup(pathId, emptyGroup);
+        context.DbChanges.PersistAddPersQueueGroupAlter(pathId, pqGroup);
 
         for (auto shard : txState.Shards) {
             Y_ABORT_UNLESS(shard.Operation == TTxState::CreateParts);
-            context.SS->PersistShardMapping(db, shard.Idx, InvalidTabletId, pathId, OperationId.GetTxId(), shard.TabletType);
-            context.SS->PersistChannelsBinding(db, shard.Idx, tabletChannelsBinding);
+            context.DbChanges.PersistShard(shard.Idx);
         }
+
         Y_ABORT_UNLESS(txState.Shards.size() == shardsToCreate);
-        context.SS->TabletCounters->Simple()[COUNTER_PQ_SHARD_COUNT].Add(shardsToCreate-1);
+        PqShardCountChange = shardsToCreate - 1;
+        context.SS->TabletCounters->Simple()[COUNTER_PQ_SHARD_COUNT].Add(PqShardCountChange);
         context.SS->TabletCounters->Simple()[COUNTER_PQ_RB_SHARD_COUNT].Add(1);
 
         dstPath.Base()->CreateTxId = OperationId.GetTxId();
@@ -517,22 +520,19 @@ class TCreatePQ: public TSubOperation {
             }
         }
 
-        context.SS->ChangeTxState(db, OperationId, TTxState::CreateParts);
         context.OnComplete.ActivateTx(OperationId);
 
-        context.SS->PersistTxState(db, OperationId);
-
+        context.DbChanges.PersistTxState(OperationId);
 
         if (!acl.empty()) {
             dstPath.Base()->ApplyACL(acl);
         }
-        context.SS->PersistPath(db, dstPath.Base()->PathId);
-
-        context.SS->PersistUpdateNextPathId(db);
-        context.SS->PersistUpdateNextShardIdx(db);
+        context.MemChanges.GrabPath(context.SS, dstPath.Base()->PathId);
+        context.DbChanges.PersistPath(dstPath.Base()->PathId);
 
         ++parentPath.Base()->DirAlterVersion;
-        context.SS->PersistPathDirAlterVersion(db, parentPath.Base());
+        context.MemChanges.GrabPath(context.SS, parentPath.Base()->PathId);
+        context.DbChanges.PersistPath(parentPath.Base()->PathId);
         context.SS->ClearDescribePathCaches(parentPath.Base());
         context.OnComplete.PublishToSchemeBoard(OperationId, parentPath.Base()->PathId);
 
@@ -544,10 +544,12 @@ class TCreatePQ: public TSubOperation {
         dstPath.DomainInfo()->IncPQPartitionsInside(partitionsToCreate);
         dstPath.DomainInfo()->IncPQReservedStorage(reserve.Storage);
 
-        context.SS->TabletCounters->Simple()[COUNTER_STREAM_RESERVED_THROUGHPUT].Add(reserve.Throughput);
-        context.SS->TabletCounters->Simple()[COUNTER_STREAM_RESERVED_STORAGE].Add(reserve.Storage);
-
-        context.SS->TabletCounters->Simple()[COUNTER_STREAM_SHARDS_COUNT].Add(partitionsToCreate);
+        StreamReservedThroughputChange = reserve.Throughput;
+        StreamReservedStorageChange = reserve.Storage;
+        StreamShardsCountChange = partitionsToCreate;
+        context.SS->TabletCounters->Simple()[COUNTER_STREAM_RESERVED_THROUGHPUT].Add(StreamReservedThroughputChange);
+        context.SS->TabletCounters->Simple()[COUNTER_STREAM_RESERVED_STORAGE].Add(StreamReservedStorageChange);
+        context.SS->TabletCounters->Simple()[COUNTER_STREAM_SHARDS_COUNT].Add(StreamShardsCountChange);
 
         dstPath.Base()->IncShardsInside(shardsToCreate);
         parentPath.Base()->IncAliveChildren();
@@ -556,8 +558,13 @@ class TCreatePQ: public TSubOperation {
         return result;
     }
 
-    void AbortPropose(TOperationContext&) override {
-        Y_ABORT("no AbortPropose for TCreatePQ");
+    void AbortPropose(TOperationContext& context) override {
+        context.SS->TabletCounters->Simple()[COUNTER_STREAM_RESERVED_THROUGHPUT].Sub(StreamReservedThroughputChange);
+        context.SS->TabletCounters->Simple()[COUNTER_STREAM_RESERVED_STORAGE].Sub(StreamReservedStorageChange);
+        context.SS->TabletCounters->Simple()[COUNTER_STREAM_SHARDS_COUNT].Sub(StreamShardsCountChange);
+
+        context.SS->TabletCounters->Simple()[COUNTER_PQ_SHARD_COUNT].Sub(PqShardCountChange);
+        context.SS->TabletCounters->Simple()[COUNTER_PQ_RB_SHARD_COUNT].Sub(1);
     }
 
     void AbortUnsafe(TTxId forceDropTxId, TOperationContext& context) override {

ydb/core/tx/schemeshard/schemeshard__operation_db_changes.cpp

@@ -85,6 +85,18 @@ void TStorageChanges::Apply(TSchemeShard* ss, NTabletFlatExecutor::TTransactionC
 
     ss->PersistUpdateNextPathId(db);
     ss->PersistUpdateNextShardIdx(db);
+
+    for (const auto& [pathId, shardIdx, pqInfo] : PersQueue) {
+        ss->PersistPersQueue(db, pathId, shardIdx, pqInfo);
+    }
+
+    for (const auto& [pathId, pqGroup] : PersQueueGroup) {
+        ss->PersistPersQueueGroup(db, pathId, pqGroup);
+    }
+
+    for (const auto& [pathId, pqGroup] : AddPersQueueGroupAlter) {
+        ss->PersistAddPersQueueGroupAlter(db, pathId, pqGroup);
+    }
 }
 
 }

ydb/core/tx/schemeshard/schemeshard__operation_db_changes.h

@@ -37,9 +37,26 @@ class TStorageChanges: public TSimpleRefCount<TStorageChanges> {
 
     TDeque<TPathId> Views;
 
+    //PQ part
+    TDeque<std::tuple<TPathId, TShardIdx, TTopicTabletInfo::TTopicPartitionInfo>> PersQueue;
+    TDeque<std::pair<TPathId, TTopicInfo::TPtr>> PersQueueGroup;
+    TDeque<std::pair<TPathId, TTopicInfo::TPtr>> AddPersQueueGroupAlter;
+
 public:
     ~TStorageChanges() = default;
 
+    void PersistPersQueue(const TPathId& pathId, const TShardIdx& shardIdx, const TTopicTabletInfo::TTopicPartitionInfo& pqInfo) {
+        PersQueue.emplace_back(pathId, shardIdx, pqInfo);
+    }
+
+    void PersistPersQueueGroup(const TPathId& pathId, const TTopicInfo::TPtr pqGroup) {
+        PersQueueGroup.emplace_back(pathId, pqGroup);
+    }
+
+    void PersistAddPersQueueGroupAlter(TPathId pathId, const TTopicInfo::TPtr alterData) {
+        AddPersQueueGroupAlter.emplace_back(pathId, alterData);
+    }
+
     void PersistPath(const TPathId& pathId) {
         Paths.push_back(pathId);
     }