From ebe396bad6f7c1c768c7f35216c91c0eaad80c3d Mon Sep 17 00:00:00 2001
From: Innokentii Mokin <innokentii@ydb.tech>
Date: Tue, 8 Oct 2024 19:32:42 +0300
Subject: [PATCH] Add op context to some TSubOperation calls (#10215)

---
 ydb/core/tx/schemeshard/schemeshard__init.cpp    |  3 ++-
 .../tx/schemeshard/schemeshard__operation.cpp    |  4 +++-
 ydb/core/tx/schemeshard/schemeshard__operation.h |  2 +-
 .../tx/schemeshard/schemeshard__operation_part.h | 16 ++++++++++++++++
 4 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/ydb/core/tx/schemeshard/schemeshard__init.cpp b/ydb/core/tx/schemeshard/schemeshard__init.cpp
index 263c6892cd35..30cc2265cc28 100644
--- a/ydb/core/tx/schemeshard/schemeshard__init.cpp
+++ b/ydb/core/tx/schemeshard/schemeshard__init.cpp
@@ -3560,7 +3560,8 @@ struct TSchemeShard::TTxInit : public TTransactionBase<TSchemeShard> {
 
                 TOperation::TPtr operation = Self->Operations.at(operationId.GetTxId());
                 Y_ABORT_UNLESS(operationId.GetSubTxId() == operation->Parts.size());
-                ISubOperation::TPtr part = operation->RestorePart(txState.TxType, txState.State);
+                TOperationContext context{Self, txc, ctx, OnComplete, MemChanges, DbChanges};
+                ISubOperation::TPtr part = operation->RestorePart(txState.TxType, txState.State, context);
                 operation->AddPart(part);
 
                 if (!txInFlightRowset.Next())
diff --git a/ydb/core/tx/schemeshard/schemeshard__operation.cpp b/ydb/core/tx/schemeshard/schemeshard__operation.cpp
index 337bd757e26f..2eedf9efb346 100644
--- a/ydb/core/tx/schemeshard/schemeshard__operation.cpp
+++ b/ydb/core/tx/schemeshard/schemeshard__operation.cpp
@@ -978,7 +978,7 @@ TOperation::TSplitTransactionsResult TOperation::SplitIntoTransactions(const TTx
     return result;
 }
 
-ISubOperation::TPtr TOperation::RestorePart(TTxState::ETxType txType, TTxState::ETxState txState) const {
+ISubOperation::TPtr TOperation::RestorePart(TTxState::ETxType txType, TTxState::ETxState txState, TOperationContext& context) const {
     switch (txType) {
     case TTxState::ETxType::TxMkDir:
         return CreateMkDir(NextPartId(), txState);
@@ -1197,6 +1197,8 @@ ISubOperation::TPtr TOperation::RestorePart(TTxState::ETxType txType, TTxState::
         Y_UNREACHABLE();
     }
 
+    Y_UNUSED(context); // TODO(Enjection): will be used by complex operations later
+
     Y_UNREACHABLE();
 }
 
diff --git a/ydb/core/tx/schemeshard/schemeshard__operation.h b/ydb/core/tx/schemeshard/schemeshard__operation.h
index d48f138b691c..45c132d56f56 100644
--- a/ydb/core/tx/schemeshard/schemeshard__operation.h
+++ b/ydb/core/tx/schemeshard/schemeshard__operation.h
@@ -79,7 +79,7 @@ struct TOperation: TSimpleRefCount<TOperation> {
     static TConsumeQuotaResult ConsumeQuota(const TTxTransaction& tx, TOperationContext& context);
     static TSplitTransactionsResult SplitIntoTransactions(const TTxTransaction& tx, const TOperationContext& context);
 
-    ISubOperation::TPtr RestorePart(TTxState::ETxType opType, TTxState::ETxState opState) const;
+    ISubOperation::TPtr RestorePart(TTxState::ETxType opType, TTxState::ETxState opState, TOperationContext& context) const;
     TVector<ISubOperation::TPtr> ConstructParts(const TTxTransaction& tx, TOperationContext& context) const;
     void AddPart(ISubOperation::TPtr part);
 
diff --git a/ydb/core/tx/schemeshard/schemeshard__operation_part.h b/ydb/core/tx/schemeshard/schemeshard__operation_part.h
index f3ecb60d5292..88e57aa3c38e 100644
--- a/ydb/core/tx/schemeshard/schemeshard__operation_part.h
+++ b/ydb/core/tx/schemeshard/schemeshard__operation_part.h
@@ -250,6 +250,10 @@ class TSubOperation: public TSubOperationBase {
     virtual TTxState::ETxState NextState(TTxState::ETxState state) const = 0;
     virtual TSubOperationState::TPtr SelectStateFunc(TTxState::ETxState state) = 0;
 
+    virtual TSubOperationState::TPtr SelectStateFunc(TTxState::ETxState state, TOperationContext&) {
+        return SelectStateFunc(state);
+    }
+
     virtual void StateDone(TOperationContext& context) {
         auto state = NextState(GetState());
         SetState(state);
@@ -272,6 +276,11 @@ class TSubOperation: public TSubOperationBase {
         return State;
     }
 
+    void SetState(TTxState::ETxState state, TOperationContext& context) {
+        State = state;
+        StateFunc = SelectStateFunc(state, context);
+    }
+
     void SetState(TTxState::ETxState state) {
         State = state;
         StateFunc = SelectStateFunc(state);
@@ -313,6 +322,13 @@ ISubOperation::TPtr MakeSubOperation(const TOperationId& id, const TTxTransactio
     return new T(id, tx, std::forward<Args>(args)...);
 }
 
+template <typename T, typename... Args>
+ISubOperation::TPtr MakeSubOperation(const TOperationId& id, TTxState::ETxState state, TOperationContext& context, Args&&... args) {
+    auto result = MakeHolder<T>(id, state, std::forward<Args>(args)...);
+    result->SetState(state, context);
+    return result.Release();
+}
+
 template <typename T, typename... Args>
 ISubOperation::TPtr MakeSubOperation(const TOperationId& id, TTxState::ETxState state, Args&&... args) {
     auto result = MakeHolder<T>(id, state, std::forward<Args>(args)...);