Make RPC buffer selection thread-safe

PiperOrigin-RevId: 934366192
Change-Id: Iedee832eeaf631157c0010ead46ec17587df5c61

Author: andreas-abel

fleetbench/common/BUILD

@@ -12,11 +12,13 @@ cc_library(
     linkstatic = 1,
     deps = [
         "@bazel_tools//tools/cpp/runfiles",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
         "@com_google_benchmark//:benchmark",
     ],
     alwayslink = 1,

fleetbench/common/common.cc

@@ -26,6 +26,7 @@
 #include <vector>
 
 #include "tools/cpp/runfiles/runfiles.h"
+#include "absl/base/const_init.h"
 #include "absl/container/btree_map.h"
 #include "absl/flags/flag.h"
 #include "absl/log/check.h"
@@ -35,6 +36,7 @@
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
 #include "benchmark/benchmark.h"
 
 #include "absl/log/log.h"
@@ -89,6 +91,15 @@ void Random::Reset() {
   }
 }
 
+std::default_random_engine& GetThreadLocalRNG() {
+  static absl::Mutex mutex(absl::kConstInit);
+  thread_local std::default_random_engine local_rng = [&]() {
+    absl::MutexLock l(mutex);
+    return std::default_random_engine(GetRNG()());
+  }();
+  return local_rng;
+}
+
 std::vector<std::filesystem::path> GetMatchingFiles(absl::string_view dir,
                                                     absl::string_view prefix) {
   std::vector<std::filesystem::path> files;

fleetbench/common/common.h

@@ -62,6 +62,10 @@ class Random {
 
 inline std::default_random_engine& GetRNG() { return Random::instance().rng(); }
 
+// Returns a thread-local random engine, seeded from GetRNG().
+// This is thread-safe.
+std::default_random_engine& GetThreadLocalRNG();
+
 // Returns a sorted list of the files in directory 'dir' whose filenames start
 // with 'prefix'.
 std::vector<std::filesystem::path> GetMatchingFiles(absl::string_view dir,

fleetbench/common/common_test.cc

@@ -116,6 +116,12 @@ TEST(RandomTest, Seed) {
   EXPECT_EQ(random1, random2);
 }
 
+TEST(RandomTest, GetThreadLocalRNG) {
+  auto& rng1 = GetThreadLocalRNG();
+  auto& rng2 = GetThreadLocalRNG();
+  EXPECT_EQ(&rng1, &rng2);
+}
+
 TEST(DeathTest, SeedFlags) {
   ASSERT_DEATH(
       {

fleetbench/rpc/grpc_server.cc

@@ -17,6 +17,7 @@
 #include <chrono>  // NOLINT
 #include <cstddef>
 #include <memory>
+#include <random>
 #include <utility>
 
 #include "absl/log/check.h"
@@ -103,7 +104,7 @@ void GRPCServer::Start(absl::string_view filepath,
 }
 
 const fleetbench::rpc::ResponseMessage& GRPCServer::Buffer() const {
-  return message_buffers_[GetRNG()() % message_buffers_.size()];
+  return message_buffers_[GetThreadLocalRNG()() % message_buffers_.size()];
 }
 
 void GRPCServer::Delay() {