EOperationKind::WriteTx (#976)

Author: azevaykin

ydb/core/tx/datashard/check_write_unit.cpp

@@ -41,7 +41,6 @@ EExecutionStatus TCheckWriteUnit::Execute(TOperation::TPtr op,
                                            TTransactionContext &,
                                            const TActorContext &ctx)
 {
-    Y_ABORT_UNLESS(op->IsDataTx() || op->IsReadTable());
     Y_ABORT_UNLESS(!op->IsAborted());
 
     if (CheckRejectDataTx(op, ctx)) {
@@ -50,8 +49,7 @@ EExecutionStatus TCheckWriteUnit::Execute(TOperation::TPtr op,
         return EExecutionStatus::Executed;
     }
 
-    TWriteOperation* writeOp = dynamic_cast<TWriteOperation*>(op.Get());
-    Y_VERIFY_S(writeOp, "cannot cast operation of kind " << op->GetKind());
+    TWriteOperation* writeOp = TWriteOperation::CastWriteOperation(op);
     auto writeTx = writeOp->GetWriteTx();
     Y_ABORT_UNLESS(writeTx);
     Y_ABORT_UNLESS(writeTx->Ready() || writeTx->RequirePrepare());

ydb/core/tx/datashard/datashard_pipeline.cpp

@@ -1563,7 +1563,7 @@ TOperation::TPtr TPipeline::BuildOperation(NEvents::TDataEvents::TEvWrite::TPtr&
                                            const TActorContext& ctx, NWilson::TSpan &&operationSpan)
 {
     const auto& rec = ev->Get()->Record;
-    TBasicOpInfo info(rec.GetTxId(), EOperationKind::DataTx, EvWrite::Convertor::GetProposeFlags(rec.GetTxMode()), 0, receivedAt, tieBreakerIndex);
+    TBasicOpInfo info(rec.GetTxId(), EOperationKind::WriteTx, EvWrite::Convertor::GetProposeFlags(rec.GetTxMode()), 0, receivedAt, tieBreakerIndex);
     auto writeOp = MakeIntrusive<TWriteOperation>(info, ev, Self, txc, ctx);
     writeOp->OperationSpan = std::move(operationSpan);
     auto writeTx = writeOp->GetWriteTx();

ydb/core/tx/datashard/datashard_write_operation.cpp

@@ -245,6 +245,7 @@ void TValidatedWriteTx::ComputeTxSize() {
 
 TWriteOperation* TWriteOperation::CastWriteOperation(TOperation::TPtr op)
 {
+    Y_ABORT_UNLESS(op->IsWriteTx());
     TWriteOperation* writeOp = dynamic_cast<TWriteOperation*>(op.Get());
     Y_ABORT_UNLESS(writeOp);
     return writeOp;
@@ -453,7 +454,7 @@ ERestoreDataStatus TWriteOperation::RestoreTxData(
 
 void TWriteOperation::FinalizeWriteTxPlan()
 {
-    Y_ABORT_UNLESS(IsDataTx());
+    Y_ABORT_UNLESS(IsWriteTx());
     Y_ABORT_UNLESS(!IsImmediate());
     Y_ABORT_UNLESS(!IsKqpScanTransaction());
 
@@ -494,8 +495,7 @@ class TFinalizeWriteTxPlanUnit: public TExecutionUnit {
         Y_UNUSED(txc);
         Y_UNUSED(ctx);
 
-        TWriteOperation* writeOp = dynamic_cast<TWriteOperation*>(op.Get());
-        Y_VERIFY_S(writeOp, "cannot cast operation of kind " << op->GetKind());
+        TWriteOperation* writeOp = TWriteOperation::CastWriteOperation(op);
 
         writeOp->FinalizeWriteTxPlan();
 

ydb/core/tx/datashard/operation.h

@@ -111,6 +111,7 @@ enum class EOperationKind : ui32 {
     // Values [100, inf) are used for internal kinds.
     DirectTx = 101,
     ReadTx = 102,
+    WriteTx = 103,
 };
 
 class TBasicOpInfo {
@@ -165,6 +166,7 @@ class TBasicOpInfo {
 
     bool IsDataTx() const { return Kind == EOperationKind::DataTx; }
     bool IsReadTx() const { return Kind == EOperationKind::ReadTx; }
+    bool IsWriteTx() const { return Kind == EOperationKind::WriteTx; }
     bool IsDirectTx() const { return Kind == EOperationKind::DirectTx; }
     bool IsSchemeTx() const { return Kind == EOperationKind::SchemeTx; }
     bool IsReadTable() const { return Kind == EOperationKind::ReadTable; }

ydb/core/tx/datashard/write_unit.cpp

@@ -98,8 +98,7 @@ class TWriteUnit : public TExecutionUnit {
     }
 
     EExecutionStatus Execute(TOperation::TPtr op, TTransactionContext& txc, const TActorContext& ctx) override {
-        TWriteOperation* writeOp = dynamic_cast<TWriteOperation*>(op.Get());
-        Y_ABORT_UNLESS(writeOp != nullptr);
+        TWriteOperation* writeOp = TWriteOperation::CastWriteOperation(op);
 
         LOG_DEBUG_S(ctx, NKikimrServices::TX_DATASHARD, "Executing write operation for " << *op << " at " << DataShard.TabletID());
 
@@ -164,8 +163,7 @@ class TWriteUnit : public TExecutionUnit {
         DataShard.EnqueueChangeRecords(std::move(op->ChangeRecords()));
         DataShard.EmitHeartbeats(ctx);
 
-        TWriteOperation* writeOp = dynamic_cast<TWriteOperation*>(op.Get());
-        Y_ABORT_UNLESS(writeOp != nullptr);
+        TWriteOperation* writeOp = TWriteOperation::CastWriteOperation(op);
 
         const auto& status = writeOp->GetWriteResult()->Record.status();
         LOG_DEBUG_S(ctx, NKikimrServices::TX_DATASHARD, "Completed write operation for " << *op << " at " << DataShard.TabletID() << ", status " << status);