Rate parameter (#1994)

Author: azevaykin

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;

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;