From 811f2ec873e5d6d0e248cd011622be1037e1a56e Mon Sep 17 00:00:00 2001
From: Jialiang Tan <jtan6@meta.com>
Date: Thu, 16 May 2024 16:57:25 -0700
Subject: [PATCH] Fix flaky test in PeriodicStatsReport tests (#9846)

Summary:

The test reporter's map is not protected by lock, causing race conditions when counterMap is modified. This made the test flaky

Differential Revision: D57469593
---
 velox/common/base/tests/StatsReporterTest.cpp | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/velox/common/base/tests/StatsReporterTest.cpp b/velox/common/base/tests/StatsReporterTest.cpp
index debfde5390d29..b2d2899ab9cba 100644
--- a/velox/common/base/tests/StatsReporterTest.cpp
+++ b/velox/common/base/tests/StatsReporterTest.cpp
@@ -33,12 +33,14 @@ namespace facebook::velox {
 
 class TestReporter : public BaseStatsReporter {
  public:
+  mutable std::mutex m;
   mutable std::unordered_map<std::string, size_t> counterMap;
   mutable std::unordered_map<std::string, StatType> statTypeMap;
   mutable std::unordered_map<std::string, std::vector<int32_t>>
       histogramPercentilesMap;
 
   void clear() {
+    std::lock_guard<std::mutex> l(m);
     counterMap.clear();
     statTypeMap.clear();
     histogramPercentilesMap.clear();
@@ -74,28 +76,34 @@ class TestReporter : public BaseStatsReporter {
 
   void addMetricValue(const std::string& key, const size_t value)
       const override {
+    std::lock_guard<std::mutex> l(m);
     counterMap[key] += value;
   }
 
   void addMetricValue(const char* key, const size_t value) const override {
+    std::lock_guard<std::mutex> l(m);
     counterMap[key] += value;
   }
 
   void addMetricValue(folly::StringPiece key, size_t value) const override {
+    std::lock_guard<std::mutex> l(m);
     counterMap[key.str()] += value;
   }
 
   void addHistogramMetricValue(const std::string& key, size_t value)
       const override {
+    std::lock_guard<std::mutex> l(m);
     counterMap[key] = std::max(counterMap[key], value);
   }
 
   void addHistogramMetricValue(const char* key, size_t value) const override {
+    std::lock_guard<std::mutex> l(m);
     counterMap[key] = std::max(counterMap[key], value);
   }
 
   void addHistogramMetricValue(folly::StringPiece key, size_t value)
       const override {
+    std::lock_guard<std::mutex> l(m);
     counterMap[key.str()] = std::max(counterMap[key.str()], value);
   }
 };