Use thread safe random

3 files changed, 7 insertions, 11 deletions

diff --git a/ydb/library/workload/kv_workload.cpp b/ydb/library/workload/kv_workload.cpp
index 58ba99e1bbd..70caf9d2b12 100644
--- a/ydb/library/workload/kv_workload.cpp
+++ b/ydb/library/workload/kv_workload.cpp
@@ -1,4 +1,5 @@
 #include "kv_workload.h"
+#include "util/random/random.h"
 
 #include <util/datetime/base.h>
 
@@ -23,11 +24,7 @@ TKvWorkloadGenerator::TKvWorkloadGenerator(const TKvWorkloadParams* params)
     : DbPath(params->DbPath)
     , Params(*params)
     , BigString(NKikimr::GenDataForLZ4(Params.StringLen))
-    , Rd()
-    , Gen(Rd())
-    , KeyUniformDistGen(0, Params.MaxFirstKey)
 {
-    Gen.seed(Now().MicroSeconds());
 }
 
 TKvWorkloadParams* TKvWorkloadGenerator::GetParams() {
@@ -84,7 +81,7 @@ TQueryInfoList TKvWorkloadGenerator::AddOperation(TString operation) {
     for (size_t row = 0; row < Params.RowsCnt; ++row) {
         TString pkname = "$r" + std::to_string(row);
         ss << "DECLARE " << pkname << " AS Uint64;\n";
-        paramsBuilder.AddParam(pkname).Uint64(KeyUniformDistGen(Gen)).Build();
+        paramsBuilder.AddParam(pkname).Uint64(RandomNumber<ui64>(Params.MaxFirstKey)).Build();
 
         for (size_t col = 1; col < Params.ColumnsCnt; ++col) {
             TString cname = "$c" + std::to_string(row) + std::to_string(col);
@@ -143,7 +140,7 @@ TQueryInfoList TKvWorkloadGenerator::SelectRandom() {
 
     for (size_t row = 0; row < Params.RowsCnt; ++row) {
         ss << "DECLARE $r" << row << " AS Uint64;\n";
-        paramsBuilder.AddParam("$r" + std::to_string(row)).Uint64(KeyUniformDistGen(Gen)).Build();
+        paramsBuilder.AddParam("$r" + std::to_string(row)).Uint64(RandomNumber<ui64>(Params.MaxFirstKey)).Build();
     }
 
     ss << "SELECT ";
@@ -174,7 +171,7 @@ TQueryInfoList TKvWorkloadGenerator::ReadRowsRandom() {
     for (size_t i = 0; i < Params.RowsCnt; ++i) {
         keys.AddListItem()
             .BeginStruct()
-                .AddMember("c0").Uint64(KeyUniformDistGen(Gen))
+                .AddMember("c0").Uint64(RandomNumber<ui64>(Params.MaxFirstKey))
             .EndStruct();
     }
     keys.EndList();
diff --git a/ydb/library/workload/kv_workload.h b/ydb/library/workload/kv_workload.h
index 2abfd0d6d21..011082f9e22 100644
--- a/ydb/library/workload/kv_workload.h
+++ b/ydb/library/workload/kv_workload.h
@@ -68,10 +68,6 @@ private:
     std::string DbPath;
     TKvWorkloadParams Params;
     TString BigString;
-
-    std::random_device Rd;
-    std::mt19937_64 Gen;
-    std::uniform_int_distribution<ui64> KeyUniformDistGen;
 };
 
 } // namespace NYdbWorkload
diff --git a/ydb/public/lib/ydb_cli/commands/ydb_workload.cpp b/ydb/public/lib/ydb_cli/commands/ydb_workload.cpp
index 3249bc67336..f150b4afd88 100644
--- a/ydb/public/lib/ydb_cli/commands/ydb_workload.cpp
+++ b/ydb/public/lib/ydb_cli/commands/ydb_workload.cpp
@@ -6,6 +6,7 @@
 #include "tpch.h"
 #include "topic_workload/topic_workload.h"
 #include "transfer_workload/transfer_workload.h"
+#include "util/random/random.h"
 
 #include <ydb/library/workload/workload_factory.h>
 #include <ydb/public/lib/ydb_cli/commands/ydb_common.h>
@@ -86,6 +87,8 @@ void TWorkloadCommand::Config(TConfig& config) {
 }
 
 void TWorkloadCommand::PrepareForRun(TConfig& config) {
+    SetRandomSeed(Now().MicroSeconds());
+
     auto driverConfig = TDriverConfig()
         .SetEndpoint(config.Address)
         .SetDatabase(config.Database)