diff --git a/ydb/public/lib/ydb_cli/commands/ydb_workload.cpp b/ydb/public/lib/ydb_cli/commands/ydb_workload.cpp index 98ec3cb4aaf8..02762fd5314a 100644 --- a/ydb/public/lib/ydb_cli/commands/ydb_workload.cpp +++ b/ydb/public/lib/ydb_cli/commands/ydb_workload.cpp @@ -58,6 +58,7 @@ TWorkloadCommand::TWorkloadCommand(const TString& name, const std::initializer_l : TYdbCommand(name, aliases, description) , TotalSec(0) , Threads(0) + , Rate(0) , ClientTimeoutMs(0) , OperationTimeoutMs(0) , CancelAfterTimeoutMs(0) @@ -67,6 +68,7 @@ TWorkloadCommand::TWorkloadCommand(const TString& name, const std::initializer_l , QueryExecuterType() , WindowHist(60000, 2) // highestTrackableValue 60000ms = 60s, precision 2 , TotalHist(60000, 2) + , TotalQueries(0) , TotalRetries(0) , WindowRetryCount(0) , TotalErrors(0) @@ -80,6 +82,17 @@ void TWorkloadCommand::Config(TConfig& config) { .DefaultValue(10).StoreResult(&TotalSec); config.Opts->AddLongOption('t', "threads", "Number of parallel threads in workload.") .DefaultValue(10).StoreResult(&Threads); + + const auto name = Parent->Parent->Name; + if (name == "kv") { + config.Opts->AddLongOption("rate", "Total rate for all threads (requests per second).") + .DefaultValue(0).StoreResult(&Rate); + } + else if (name == "stock") { + config.Opts->AddLongOption("rate", "Total rate for all threads (transactions per second).") + .DefaultValue(0).StoreResult(&Rate); + } + config.Opts->AddLongOption("quiet", "Quiet mode. Doesn't print statistics each second.") .StoreTrue(&Quiet); config.Opts->AddLongOption("print-timestamp", "Print timestamp each second with statistics.") @@ -206,10 +219,22 @@ void TWorkloadCommand::WorkerFn(int taskId, NYdbWorkload::IWorkloadQueryGenerato auto opStartTime = Now(); NYdbWorkload::TQueryInfoList::iterator it; - for (it = queryInfoList.begin(); it != queryInfoList.end(); ++it) { + for (it = queryInfoList.begin(); it != queryInfoList.end(); ) { + + if (Rate != 0) + { + const ui64 expectedQueries = (Now() - StartTime).SecondsFloat() * Rate; + if (TotalQueries > expectedQueries) { + Sleep(TDuration::MilliSeconds(1)); + continue; + } + } + queryInfo = *it; auto status = runQuery(); - if (!status.IsSuccess()) { + if (status.IsSuccess()) { + TotalQueries++; + } else { TotalErrors++; WindowErrors++; // if (status.GetStatus() != EStatus::ABORTED) { @@ -222,6 +247,8 @@ void TWorkloadCommand::WorkerFn(int taskId, NYdbWorkload::IWorkloadQueryGenerato WindowRetryCount += retryCount; } retryCount = -1; + + ++it; } if (it != queryInfoList.end()) { continue; diff --git a/ydb/public/lib/ydb_cli/commands/ydb_workload.h b/ydb/public/lib/ydb_cli/commands/ydb_workload.h index 15365b049e2f..105f310f9a73 100644 --- a/ydb/public/lib/ydb_cli/commands/ydb_workload.h +++ b/ydb/public/lib/ydb_cli/commands/ydb_workload.h @@ -43,6 +43,7 @@ class TWorkloadCommand : public TYdbCommand { size_t TotalSec; size_t Threads; + ui64 Rate; unsigned int ClientTimeoutMs; unsigned int OperationTimeoutMs; unsigned int CancelAfterTimeoutMs; @@ -60,6 +61,7 @@ class TWorkloadCommand : public TYdbCommand { NHdr::THistogram WindowHist; NHdr::THistogram TotalHist; + std::atomic_uint64_t TotalQueries; std::atomic_uint64_t TotalRetries; std::atomic_uint64_t WindowRetryCount; std::atomic_uint64_t TotalErrors;