Check for ChangeQueue overflow in EvWrite

ydb/core/tx/datashard/datashard_ut_write.cpp

@@ -13,12 +13,14 @@ using namespace NDataShardReadTableTest;
 Y_UNIT_TEST_SUITE(DataShardWrite) {
     const TString expectedTableState = "key = 0, value = 1\nkey = 2, value = 3\nkey = 4, value = 5\n";
 
-    std::tuple<TTestActorRuntime&, Tests::TServer::TPtr, TActorId> TestCreateServer() {
-        TPortManager pm;
-        TServerSettings serverSettings(pm.GetPort(2134));
-        serverSettings.SetDomainName("Root").SetUseRealThreads(false);
+    std::tuple<TTestActorRuntime&, Tests::TServer::TPtr, TActorId> TestCreateServer(std::optional<TServerSettings> serverSettings = {}) {
+        if (!serverSettings) {
+            TPortManager pm;
+            serverSettings = TServerSettings(pm.GetPort(2134));
+            serverSettings->SetDomainName("Root").SetUseRealThreads(false);
+        }
 
-        Tests::TServer::TPtr server = new TServer(serverSettings);
+        Tests::TServer::TPtr server = new TServer(serverSettings.value());
         auto& runtime = *server->GetRuntime();
         auto sender = runtime.AllocateEdgeActor();
 
@@ -318,6 +320,48 @@ Y_UNIT_TEST_SUITE(DataShardWrite) {
             UNIT_ASSERT_VALUES_EQUAL(tableState, expectedTableState);
         }
 
+    }
+
+    Y_UNIT_TEST(ShouldRejectOnChangeQueueOverflow) {
+        TPortManager pm;
+        TServerSettings serverSettings(pm.GetPort(2134));
+        serverSettings.SetDomainName("Root").SetUseRealThreads(false).SetChangesQueueItemsLimit(1);
+
+        auto [runtime, server, sender] = TestCreateServer(serverSettings);
+
+        auto opts = TShardedTableOptions()
+            .Columns({{"key", "Uint32", true, false},{"value", "Uint32", false, false}})
+            .Indexes({TShardedTableOptions::TIndex{"by_value", {"value"}, {}, NKikimrSchemeOp::EIndexTypeGlobalAsync}});
+
+        auto [shards, tableId] = CreateShardedTable(server, sender, "/Root", "table-1", opts);
+
+        TVector<THolder<IEventHandle>> blockedEnqueueRecords;
+        auto prevObserverFunc = runtime.SetObserverFunc([&](TAutoPtr<IEventHandle>& ev) {
+            if (ev->GetTypeRewrite() == NChangeExchange::TEvChangeExchange::EvEnqueueRecords) {
+                blockedEnqueueRecords.emplace_back(ev.Release());
+                return TTestActorRuntime::EEventAction::DROP;
+            }
+
+            return TTestActorRuntime::EEventAction::PROCESS;
+        });
+
+        const ui64 shard = shards[0];
+        const ui32 rowCount = 1;
+
+        ui64 txId = 100;
+
+        Cout << "========= Send immediate write, expecting success =========\n";
+        {
+            Write(runtime, sender, shard, tableId, opts.Columns_, rowCount, ++txId, NKikimrDataEvents::TEvWrite::MODE_IMMEDIATE);
+        }
+
+        UNIT_ASSERT_VALUES_EQUAL(blockedEnqueueRecords.size(), rowCount);
+
+        Cout << "========= Send immediate write, expecting overloaded =========\n";
+        {
+            Write(runtime, sender, shard, tableId, opts.Columns_, rowCount, ++txId, NKikimrDataEvents::TEvWrite::MODE_IMMEDIATE, NKikimrDataEvents::TEvWriteResult::STATUS_OVERLOADED);
+        }
+
     } // Y_UNIT_TEST
 
 } // Y_UNIT_TEST_SUITE

ydb/core/tx/datashard/datashard_write_operation.cpp

@@ -282,6 +282,16 @@ TWriteOperation* TWriteOperation::CastWriteOperation(TOperation::TPtr op)
     return writeOp;
 }
 
+TWriteOperation* TWriteOperation::TryCastWriteOperation(TOperation::TPtr op)
+{
+    if (!op->IsWriteTx())
+        return nullptr;
+
+    TWriteOperation* writeOp = dynamic_cast<TWriteOperation*>(op.Get());
+    Y_ABORT_UNLESS(writeOp);
+    return writeOp;
+}
+
 TWriteOperation::TWriteOperation(const TBasicOpInfo& op, ui64 tabletId)
     : TOperation(op)
     , TabletId(tabletId)

ydb/core/tx/datashard/datashard_write_operation.h

@@ -134,6 +134,7 @@ class TWriteOperation : public TOperation {
     using TPtr = TIntrusivePtr<TWriteOperation>;
 
     static TWriteOperation* CastWriteOperation(TOperation::TPtr op);
+    static TWriteOperation* TryCastWriteOperation(TOperation::TPtr op);
 
     explicit TWriteOperation(const TBasicOpInfo& op, ui64 tabletId);
     explicit TWriteOperation(const TBasicOpInfo& op, NEvents::TDataEvents::TEvWrite::TPtr&& ev, TDataShard* self);

ydb/core/tx/datashard/execution_unit.cpp

@@ -158,6 +158,8 @@ THolder<TExecutionUnit> CreateExecutionUnit(EExecutionUnitKind kind,
 }
 
 bool TExecutionUnit::CheckRejectDataTx(TOperation::TPtr op, const TActorContext& ctx) {
+    TWriteOperation* writeOp = TWriteOperation::TryCastWriteOperation(op);
+
     // Reject operations after receiving EvSplit
     // This is to avoid races when split is in progress
     if (DataShard.GetState() == TShardState::SplitSrcWaitForNoTxInFlight ||
@@ -168,8 +170,13 @@ bool TExecutionUnit::CheckRejectDataTx(TOperation::TPtr op, const TActorContext&
         TString err = TStringBuilder()
             << "Wrong shard state: " << (TShardState)DataShard.GetState()
             << " tablet id: " << DataShard.TabletID();
-        BuildResult(op, NKikimrTxDataShard::TEvProposeTransactionResult::OVERLOADED)
-            ->AddError(NKikimrTxDataShard::TError::WRONG_SHARD_STATE, err);
+
+        if (writeOp) {
+            writeOp->SetError(NKikimrDataEvents::TEvWriteResult::STATUS_OVERLOADED, err);
+        } else {
+            BuildResult(op, NKikimrTxDataShard::TEvProposeTransactionResult::OVERLOADED)
+                ->AddError(NKikimrTxDataShard::TError::WRONG_SHARD_STATE, err);
+        }
 
         LOG_NOTICE_S(ctx, NKikimrServices::TX_DATASHARD,
             "Tablet " << DataShard.TabletID() << " rejecting tx due to split");
@@ -187,7 +194,11 @@ bool TExecutionUnit::CheckRejectDataTx(TOperation::TPtr op, const TActorContext&
             << "Wrong shard state: " << (TShardState)DataShard.GetState()
             << " tablet id: " << DataShard.TabletID();
         // TODO: Return SCHEME_CHANGED if the shard has been split
-        BuildResult(op)->AddError(NKikimrTxDataShard::TError::WRONG_SHARD_STATE, err);
+        if (writeOp) {
+            writeOp->SetError(NKikimrDataEvents::TEvWriteResult::STATUS_INTERNAL_ERROR, err);
+        } else {        
+            BuildResult(op)->AddError(NKikimrTxDataShard::TError::WRONG_SHARD_STATE, err);
+        }
 
         LOG_NOTICE_S(ctx, NKikimrServices::TX_DATASHARD, err);
 
@@ -198,8 +209,13 @@ bool TExecutionUnit::CheckRejectDataTx(TOperation::TPtr op, const TActorContext&
     if (DataShard.IsStopping()) {
         TString err = TStringBuilder()
             << "Tablet " << DataShard.TabletID() << " is restarting";
-        BuildResult(op, NKikimrTxDataShard::TEvProposeTransactionResult::OVERLOADED)
-            ->AddError(NKikimrTxDataShard::TError::WRONG_SHARD_STATE, err);
+
+        if (writeOp) {
+            writeOp->SetError(NKikimrDataEvents::TEvWriteResult::STATUS_OVERLOADED, err);
+        } else {
+            BuildResult(op, NKikimrTxDataShard::TEvProposeTransactionResult::OVERLOADED)
+                ->AddError(NKikimrTxDataShard::TError::WRONG_SHARD_STATE, err);
+        }
 
         LOG_NOTICE_S(ctx, NKikimrServices::TX_DATASHARD, err);
 
@@ -211,8 +227,13 @@ bool TExecutionUnit::CheckRejectDataTx(TOperation::TPtr op, const TActorContext&
         TString err = TStringBuilder()
                 << "Wrong shard state: " << (TShardState)DataShard.GetState()
                 << " tablet id: " << DataShard.TabletID();
-        BuildResult(op, NKikimrTxDataShard::TEvProposeTransactionResult::OVERLOADED)
+
+        if (writeOp) {
+            writeOp->SetError(NKikimrDataEvents::TEvWriteResult::STATUS_OVERLOADED, err);
+        } else {
+            BuildResult(op, NKikimrTxDataShard::TEvProposeTransactionResult::OVERLOADED)
                 ->AddError(NKikimrTxDataShard::TError::WRONG_SHARD_STATE, err);
+        }
 
         LOG_NOTICE_S(ctx, NKikimrServices::TX_DATASHARD,
                      "Tablet " << DataShard.TabletID() << " rejecting tx due to mvcc state change");
@@ -224,8 +245,13 @@ bool TExecutionUnit::CheckRejectDataTx(TOperation::TPtr op, const TActorContext&
     if (!op->IsReadOnly() && DataShard.CheckChangesQueueOverflow()) {
         TString err = TStringBuilder()
                 << "Can't execute at blocked shard: " << " tablet id: " << DataShard.TabletID();
-        BuildResult(op, NKikimrTxDataShard::TEvProposeTransactionResult::OVERLOADED)
-                ->AddError(NKikimrTxDataShard::TError::SHARD_IS_BLOCKED, err);
+
+        if (writeOp) {
+            writeOp->SetError(NKikimrDataEvents::TEvWriteResult::STATUS_OVERLOADED, err);
+        } else {                
+            BuildResult(op, NKikimrTxDataShard::TEvProposeTransactionResult::OVERLOADED)
+                    ->AddError(NKikimrTxDataShard::TError::SHARD_IS_BLOCKED, err);
+        }
 
         LOG_NOTICE_S(ctx, NKikimrServices::TX_DATASHARD,
                      "Tablet " << DataShard.TabletID() << " rejecting tx due to changes queue overflow");
@@ -238,8 +264,13 @@ bool TExecutionUnit::CheckRejectDataTx(TOperation::TPtr op, const TActorContext&
         TString err = TStringBuilder()
             << "Can't execute write tx at replicated table:"
             << " tablet id: " << DataShard.TabletID();
-        BuildResult(op, NKikimrTxDataShard::TEvProposeTransactionResult::EXEC_ERROR)
+
+        if (writeOp) {
+            writeOp->SetError(NKikimrDataEvents::TEvWriteResult::STATUS_BAD_REQUEST, err);
+        } else {            
+            BuildResult(op, NKikimrTxDataShard::TEvProposeTransactionResult::EXEC_ERROR)
                 ->AddError(NKikimrTxDataShard::TError::WRONG_SHARD_STATE, err);
+        }
 
         LOG_NOTICE_S(ctx, NKikimrServices::TX_DATASHARD, err);
 

ydb/core/tx/datashard/load_write_details_unit.cpp

@@ -40,8 +40,7 @@ EExecutionStatus TLoadWriteDetailsUnit::Execute(TOperation::TPtr op,
                                              TTransactionContext &txc,
                                              const TActorContext &ctx)
 {
-    TWriteOperation::TPtr writeOp = dynamic_cast<TWriteOperation*>(op.Get());
-    Y_VERIFY_S(writeOp, "cannot cast operation of kind " << op->GetKind());
+    TWriteOperation* writeOp = TWriteOperation::CastWriteOperation(op);
 
     if (!Pipeline.LoadWriteDetails(txc, ctx, writeOp))
         return EExecutionStatus::Restart;

ydb/core/tx/datashard/prepare_write_tx_in_rs_unit.cpp

@@ -32,8 +32,7 @@ bool TPrepareWriteTxInRSUnit::IsReadyToExecute(TOperation::TPtr) const {
 EExecutionStatus TPrepareWriteTxInRSUnit::Execute(TOperation::TPtr op, TTransactionContext &txc,
     const TActorContext &ctx)
 {
-    TWriteOperation* writeOp = dynamic_cast<TWriteOperation*>(op.Get());
-    Y_VERIFY_S(writeOp, "cannot cast operation of kind " << op->GetKind());
+    TWriteOperation* writeOp = TWriteOperation::CastWriteOperation(op);
 
     const TValidatedWriteTx::TPtr& writeTx = writeOp->GetWriteTx();
 

ydb/core/tx/datashard/store_and_send_write_out_rs_unit.cpp

@@ -43,8 +43,7 @@ EExecutionStatus TStoreAndSendWriteOutRSUnit::Execute(TOperation::TPtr op,
                                                  TTransactionContext &txc,
                                                  const TActorContext &ctx)
 {
-    TWriteOperation* writeOp = dynamic_cast<TWriteOperation*>(op.Get());
-    Y_VERIFY_S(writeOp, "cannot cast operation of kind " << op->GetKind());
+    TWriteOperation* writeOp = TWriteOperation::CastWriteOperation(op);
 
     bool newArtifact = false;
     // TODO: move artifact flags into operation flags.