implement tpc-c mixed isolation mode (#46473)

5 files changed, 89 insertions, 4 deletions

diff --git a/ydb/library/workload/tpcc/runner.cpp b/ydb/library/workload/tpcc/runner.cpp
index 0c09831dee0..1623f7bdea3 100644
--- a/ydb/library/workload/tpcc/runner.cpp
+++ b/ydb/library/workload/tpcc/runner.cpp
@@ -265,6 +265,10 @@ TPCCRunner::TPCCRunner(const NConsoleClient::TClientCommand::TConfig& connection
             Config.SimulateTransactionMs,
             Config.SimulateTransactionSelect1Count,
             Config.TxMode,
+            Config.MixedTxMode,
+            Config.TxModeWeightSerializable,
+            Config.TxModeWeightSnapshot,
+            Config.TxModeWeightReadCommitted,
             TerminalsStopSource.get_token(),
             StopWarmup,
             PerThreadTerminalStats[i % threadCount],
diff --git a/ydb/library/workload/tpcc/runner.h b/ydb/library/workload/tpcc/runner.h
index 6580fc1f15b..5da9759a778 100644
--- a/ydb/library/workload/tpcc/runner.h
+++ b/ydb/library/workload/tpcc/runner.h
@@ -58,6 +58,12 @@ struct TRunConfig {
     EFormat Format = EFormat::Pretty;
     NQuery::TTxSettings TxMode = NQuery::TTxSettings::SerializableRW();
 
+    // mixed tx mode weights (used when MixedTxMode is true)
+    bool MixedTxMode = false;
+    double TxModeWeightSerializable = 0.0;
+    double TxModeWeightSnapshot = 0.0;
+    double TxModeWeightReadCommitted = 0.0;
+
     TString JsonResultPath;
 
     // advanced settings (normally, used by developer only)
diff --git a/ydb/library/workload/tpcc/terminal.cpp b/ydb/library/workload/tpcc/terminal.cpp
index 42524bacf10..fc4fc52b29e 100644
--- a/ydb/library/workload/tpcc/terminal.cpp
+++ b/ydb/library/workload/tpcc/terminal.cpp
@@ -80,6 +80,10 @@ TTerminal::TTerminal(size_t terminalID,
                      int simulateTransactionMs,
                      int simulateTransactionSelect1Count,
                      NQuery::TTxSettings txMode,
+                     bool mixedTxMode,
+                     double txModeWeightSerializable,
+                     double txModeWeightSnapshot,
+                     double txModeWeightReadCommitted,
                      std::stop_token stopToken,
                      std::atomic<bool>& stopWarmup,
                      std::shared_ptr<TTerminalStats>& stats,
@@ -87,12 +91,30 @@ TTerminal::TTerminal(size_t terminalID,
     : TaskQueue(taskQueue)
     , Context(terminalID, warehouseID, warehouseCount, TaskQueue, simulateTransactionMs, simulateTransactionSelect1Count, client, path, log, txMode)
     , NoDelays(noDelays)
+    , MixedTxMode(mixedTxMode)
+    , TxModeWeightSerializable(txModeWeightSerializable)
+    , TxModeWeightSnapshot(txModeWeightSnapshot)
+    , TxModeWeightReadCommitted(txModeWeightReadCommitted)
     , StopToken(stopToken)
     , StopWarmup(stopWarmup)
     , Stats(stats)
 {
 }
 
+NQuery::TTxSettings TTerminal::ChooseTxMode() const {
+    double totalWeight = TxModeWeightSerializable + TxModeWeightSnapshot + TxModeWeightReadCommitted;
+    double randomValue = ::RandomNumber<double>() * totalWeight; // randomValue in [0; totalWeight)
+    double cumulative = TxModeWeightSerializable;
+    if (randomValue <= cumulative) {
+        return NQuery::TTxSettings::SerializableRW();
+    }
+    cumulative += TxModeWeightSnapshot;
+    if (randomValue <= cumulative) {
+        return NQuery::TTxSettings::SnapshotRW();
+    }
+    return NQuery::TTxSettings::ReadCommittedRW();
+}
+
 void TTerminal::Start() {
     if (!Started) {
         TaskFuture = Run();
@@ -118,6 +140,10 @@ NThreading::TFuture<void> TTerminal::Run() {
         size_t txIndex = ChooseRandomTransactionIndex();
         auto& transaction = Transactions[txIndex];
 
+        if (MixedTxMode) {
+            Context.TxMode = ChooseTxMode();
+        }
+
         try {
             if (!NoDelays) {
                 LOG_T("Terminal " << Context.TerminalID << " keying time for " << transaction.Name << ": "
diff --git a/ydb/library/workload/tpcc/terminal.h b/ydb/library/workload/tpcc/terminal.h
index b95823ffe74..416900cbf16 100644
--- a/ydb/library/workload/tpcc/terminal.h
+++ b/ydb/library/workload/tpcc/terminal.h
@@ -163,6 +163,10 @@ public:
         int simulateTransactionMs,
         int simulateTransactionSelect1Count,
         NQuery::TTxSettings txMode,
+        bool mixedTxMode,
+        double txModeWeightSerializable,
+        double txModeWeightSnapshot,
+        double txModeWeightReadCommitted,
         std::stop_token stopToken,
         std::atomic<bool>& stopWarmup,
         std::shared_ptr<TTerminalStats>& stats,
@@ -184,11 +188,16 @@ public:
 
 private:
     NThreading::TFuture<void> Run();
+    NQuery::TTxSettings ChooseTxMode() const;
 
 private:
     ITaskQueue& TaskQueue;
     TTransactionContext Context;
     bool NoDelays;
+    bool MixedTxMode;
+    double TxModeWeightSerializable;
+    double TxModeWeightSnapshot;
+    double TxModeWeightReadCommitted;
     std::stop_token StopToken;
     std::atomic<bool>& StopWarmup;
     std::shared_ptr<TTerminalStats> Stats;
diff --git a/ydb/public/lib/ydb_cli/commands/ydb_workload_tpcc.cpp b/ydb/public/lib/ydb_cli/commands/ydb_workload_tpcc.cpp
index be1ea4300e7..f220f708419 100644
--- a/ydb/public/lib/ydb_cli/commands/ydb_workload_tpcc.cpp
+++ b/ydb/public/lib/ydb_cli/commands/ydb_workload_tpcc.cpp
@@ -256,7 +256,7 @@ void TCommandTPCCRun::Config(TConfig& config) {
 
     auto txModeOpt = config.Opts->AddLongOption(
         "tx-mode", TStringBuilder() << "Transaction mode: serializable-rw or snapshot-rw")
-            .OptionalArgument("STRING").StoreMappedResult(&RunConfig->TxMode, [](const TString& value) {
+            .OptionalArgument("STRING").StoreMappedResult(&RunConfig->TxMode, [runConfig = RunConfig](const TString& value) {
                 if (value == "serializable-rw") {
                     return NQuery::TTxSettings::SerializableRW();
                 } else if (value == "snapshot-rw") {
@@ -264,12 +264,30 @@ void TCommandTPCCRun::Config(TConfig& config) {
                 } else if (value == "read-committed-rw") {
                     // Experimental isolation level. Hidden from help at current time.
                     return NQuery::TTxSettings::ReadCommittedRW();
+                } else if (value == "mixed") {
+                    // This option is useful for stress tests. Hidden from help.
+                    runConfig->MixedTxMode = true;
+                    return NQuery::TTxSettings::SerializableRW();
                 }
                 throw yexception() << "Invalid transaction mode: " << value << ". Valid values are: serializable-rw, snapshot-rw";
             }).DefaultValue("serializable-rw")
-            .ChoicesWithCompletion({{"serializable-rw", "Serializable read-write"},
-                                    {"snapshot-rw", "Snapshot read-write"},
-                                    {"read-committed-rw", "Read Committed read-write"}});
+            .Completer(NLastGetopt::NComp::Choice({{"serializable-rw", "Serializable read-write"},
+                                                   {"snapshot-rw", "Snapshot read-write"}}));
+
+    auto txModeWeightSerializableOpt = config.Opts->AddLongOption(
+        "tx-mode-weight-serializable",
+        TStringBuilder() << "Weight for serializable-rw in mixed tx mode (default: 0)")
+            .RequiredArgument("FLOAT").StoreResult(&RunConfig->TxModeWeightSerializable).DefaultValue(0.0);
+
+    auto txModeWeightSnapshotOpt = config.Opts->AddLongOption(
+        "tx-mode-weight-snapshot",
+        TStringBuilder() << "Weight for snapshot-rw in mixed tx mode (default: 0)")
+            .RequiredArgument("FLOAT").StoreResult(&RunConfig->TxModeWeightSnapshot).DefaultValue(0.0);
+
+    auto txModeWeightReadCommittedOpt = config.Opts->AddLongOption(
+        "tx-mode-weight-read-committed",
+        TStringBuilder() << "Weight for read-committed-rw in mixed tx mode (default: 0)")
+            .RequiredArgument("FLOAT").StoreResult(&RunConfig->TxModeWeightReadCommitted).DefaultValue(0.0);
 
     auto simulateOpt = config.Opts->AddLongOption(
         "simulate", TStringBuilder() << "Simulate transaction execution (delay is simulated transaction latency ms)")
@@ -288,10 +306,32 @@ void TCommandTPCCRun::Config(TConfig& config) {
         noDelaysOpt.Hidden();
         simulateOpt.Hidden();
         simulateSelect1Opt.Hidden();
+        txModeWeightSerializableOpt.Hidden();
+        txModeWeightSnapshotOpt.Hidden();
+        txModeWeightReadCommittedOpt.Hidden();
     }
 }
 
 int TCommandTPCCRun::Run(TConfig& connectionConfig) {
+    if (RunConfig->TxModeWeightSerializable < 0.0
+        || RunConfig->TxModeWeightSnapshot < 0.0
+        || RunConfig->TxModeWeightReadCommitted < 0.0) {
+        Cout << "Run failed: --tx-mode-weight-* values must be non-negative" << Endl;
+        return 1;
+    }
+    double totalWeight = RunConfig->TxModeWeightSerializable
+        + RunConfig->TxModeWeightSnapshot
+        + RunConfig->TxModeWeightReadCommitted;
+    if (RunConfig->MixedTxMode) {
+        if (totalWeight <= 0.0) {
+            Cout << "Run failed: --tx-mode mixed requires at least one non-zero --tx-mode-weight-* option" << Endl;
+            return 1;
+        }
+    } else if (totalWeight > 0.0) {
+        Cout << "Run failed: --tx-mode-weight-* options are only valid with --tx-mode mixed" << Endl;
+        return 1;
+    }
+
     RunConfig->SetFullPath(connectionConfig);
     try {
         NTPCC::RunSync(connectionConfig, *RunConfig);