Fix TSAN data race flaky test

Summary: There were some race conditions that I'm fixing.

Differential Revision: D55817149

velox/dwio/common/tests/utils/UnitLoaderTestTools.cpp

@@ -12,7 +12,7 @@ ReaderMock::ReaderMock(
     UnitLoaderFactory& factory)
     : rowsPerUnit_{std::move(rowsPerUnit)},
       ioSizes_{std::move(ioSizes)},
-      unitsLoaded_(rowsPerUnit_.size(), false),
+      unitsLoaded_(std::vector<bool>(rowsPerUnit_.size())),
       loader_{factory.create(getUnits())} {
   VELOX_CHECK(rowsPerUnit_.size() == ioSizes_.size());
 }

velox/dwio/common/tests/utils/UnitLoaderTestTools.h

@@ -6,6 +6,7 @@
 #include <memory>
 #include <vector>
 
+#include "folly/Synchronized.h"
 #include "velox/common/base/Exceptions.h"
 #include "velox/dwio/common/UnitLoader.h"
 
@@ -16,7 +17,7 @@ class LoadUnitMock : public LoadUnit {
   LoadUnitMock(
       uint64_t rowCount,
       uint64_t ioSize,
-      std::vector<bool>& unitsLoaded,
+      folly::Synchronized<std::vector<bool>>& unitsLoaded,
       size_t unitId)
       : rowCount_{rowCount},
         ioSize_{ioSize},
@@ -27,12 +28,12 @@ class LoadUnitMock : public LoadUnit {
 
   void load() override {
     VELOX_CHECK(!isLoaded());
-    unitsLoaded_[unitId_] = true;
+    (*unitsLoaded_.wlock())[unitId_] = true;
   }
 
   void unload() override {
     VELOX_CHECK(isLoaded());
-    unitsLoaded_[unitId_] = false;
+    (*unitsLoaded_.wlock())[unitId_] = false;
   }
 
   uint64_t getNumRows() override {
@@ -44,13 +45,13 @@ class LoadUnitMock : public LoadUnit {
   }
 
   bool isLoaded() const {
-    return unitsLoaded_[unitId_];
+    return (*unitsLoaded_.rlock())[unitId_];
   }
 
  private:
   uint64_t rowCount_;
   uint64_t ioSize_;
-  std::vector<bool>& unitsLoaded_;
+  folly::Synchronized<std::vector<bool>>& unitsLoaded_;
   size_t unitId_;
 };
 
@@ -63,8 +64,8 @@ class ReaderMock {
 
   bool read(uint64_t maxRows);
 
-  const std::vector<bool>& unitsLoaded() const {
-    return unitsLoaded_;
+  std::vector<bool> unitsLoaded() const {
+    return unitsLoaded_.withRLock([](auto unitsLoaded) { return unitsLoaded; });
   }
 
  private:
@@ -74,7 +75,7 @@ class ReaderMock {
 
   std::vector<uint64_t> rowsPerUnit_;
   std::vector<uint64_t> ioSizes_;
-  std::vector<bool> unitsLoaded_;
+  folly::Synchronized<std::vector<bool>> unitsLoaded_;
   std::unique_ptr<UnitLoader> loader_;
   size_t currentUnit_{0};
   size_t currentRowInUnit_{0};