Make keep_running thread-safe

PiperOrigin-RevId: 934411598
Change-Id: I25b77f5b6bbf21583f531d86e20cac0cc4448202

Author: andreas-abel

fleetbench/rpc/rpcperf_benchmark.cc

@@ -54,20 +54,29 @@ void BM_Rpc(benchmark::State &state, absl::string_view program) {
 
   absl::Mutex keep_running_mtx;
   int warmup = 1000;
+  bool finished = false;
   std::unique_ptr<fleetbench::rpc::GRPCClient> client =
       fleetbench::rpc::CreateAndStartClient(
           /*max_outstanding_rpcs=*/1, /*compress=*/false,
           /*checksum=*/false, /*skip_loopback=*/false,
           /*peers=*/peers, /*max_peers=*/-1,
           /*connections_per_peer=*/1, /*logstats_output_path=*/"",
           /*req_delay_us_dist=*/"", /*program=*/program,
-          /*keep_running=*/[&state, &keep_running_mtx, &warmup]() {
+          /*keep_running=*/
+          [&state, &keep_running_mtx, &warmup, &finished]() {
             absl::MutexLock l(keep_running_mtx);
+            if (finished) {
+              return false;
+            }
             if (warmup > 0) {
               --warmup;
               return true;
             }
-            return state.KeepRunning();
+            if (!state.KeepRunning()) {
+              finished = true;
+              return false;
+            }
+            return true;
           });
 
   fleetbench::rpc::Wait(std::move(server), std::move(client));