Add op context to some TSubOperation calls (#10215)

Author: Enjection

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())

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();
 }
 

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);
 

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)...);