workload topic DirectWriteToPartition

6 files changed, 10 insertions, 1 deletions

diff --git a/ydb/public/lib/ydb_cli/commands/topic_operations_scenario.cpp b/ydb/public/lib/ydb_cli/commands/topic_operations_scenario.cpp
index 9a9e170937..4f424d70c1 100644
--- a/ydb/public/lib/ydb_cli/commands/topic_operations_scenario.cpp
+++ b/ydb/public/lib/ydb_cli/commands/topic_operations_scenario.cpp
@@ -212,6 +212,7 @@ void TTopicOperationsScenario::StartProducerThreads(std::vector<std::future<void
             .WriterIdx = writerIdx,
             .ProducerId = TGUID::CreateTimebased().AsGuidString(),
             .PartitionId = (partitionSeed + writerIdx) % partitionCount,
+            .Direct = Direct,
             .Codec = Codec
         };
 
diff --git a/ydb/public/lib/ydb_cli/commands/topic_operations_scenario.h b/ydb/public/lib/ydb_cli/commands/topic_operations_scenario.h
index f6e01bae30..b448701c61 100644
--- a/ydb/public/lib/ydb_cli/commands/topic_operations_scenario.h
+++ b/ydb/public/lib/ydb_cli/commands/topic_operations_scenario.h
@@ -52,6 +52,7 @@ public:
     ui32 ProducerThreadCount = 0;
     ui32 ConsumerThreadCount = 0;
     ui32 ConsumerCount = 0;
+    bool Direct = false;
     TString ConsumerPrefix;
     size_t MessageSize;
     size_t MessageRate;
diff --git a/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_run_full.cpp b/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_run_full.cpp
index e5d5693c07..04d0377f47 100644
--- a/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_run_full.cpp
+++ b/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_run_full.cpp
@@ -64,6 +64,9 @@ void TCommandWorkloadTopicRunFull::Config(TConfig& config)
         .Optional()
         .DefaultValue((TStringBuilder() << NTopic::ECodec::RAW))
         .StoreMappedResultT<TString>(&Scenario.Codec, &TCommandWorkloadTopicParams::StrToCodec);
+    config.Opts->AddLongOption("direct", "Direct write to a partition node.")
+        .Hidden()
+        .StoreTrue(&Scenario.Direct);
 
     config.Opts->MutuallyExclusive("message-rate", "byte-rate");
 
diff --git a/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_run_write.cpp b/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_run_write.cpp
index f6ec935273..4a8dba8167 100644
--- a/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_run_write.cpp
+++ b/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_run_write.cpp
@@ -55,7 +55,9 @@ void TCommandWorkloadTopicRunWrite::Config(TConfig& config)
         .Optional()
         .DefaultValue((TStringBuilder() << NTopic::ECodec::RAW))
         .StoreMappedResultT<TString>(&Scenario.Codec, &TCommandWorkloadTopicParams::StrToCodec);
-
+    config.Opts->AddLongOption("direct", "Direct write to a partition node.")
+        .Hidden()
+        .StoreTrue(&Scenario.Direct);
 
     config.Opts->MutuallyExclusive("message-rate", "byte-rate");
 
diff --git a/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_writer.cpp b/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_writer.cpp
index 341c716f26..1766a25add 100644
--- a/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_writer.cpp
+++ b/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_writer.cpp
@@ -225,6 +225,7 @@ void TTopicWorkloadWriterWorker::CreateWorker() {
     settings.Path(Params.TopicName);
     settings.ProducerId(Params.ProducerId);
     settings.PartitionId(Params.PartitionId);
+    settings.DirectWriteToPartition(Params.Direct);
     WriteSession = NYdb::NTopic::TTopicClient(Params.Driver).CreateWriteSession(settings);
 }
 
diff --git a/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_writer.h b/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_writer.h
index 9b53dcbf18..2d42f31e90 100644
--- a/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_writer.h
+++ b/ydb/public/lib/ydb_cli/commands/topic_workload/topic_workload_writer.h
@@ -27,6 +27,7 @@ namespace NYdb {
             ui32 WriterIdx;
             TString ProducerId;
             ui32 PartitionId;
+            bool Direct;
             ui32 Codec = 0;
         };