[tsan] Fix a race during processUnitsAndWait()

After we added processUnitsAndWait() as a separte entry point that
ultimately called the onFilesChange() method, it created a race on the
DoneInitState, which was accessed directly in the unit event callback
before jumping to the unit processing queue. Simplify that logic and fix
the race by using processUnitsAndWait() during initialization if we
desire waiting, and get rid of DoneInitState.

The previous approach was more flexible about allowing you to wait at an
arbitrary point after creating the IndexDatastore, but in practice we
were not using that ability, and it made it harder to understand the
synchronization required.

Author: benlangmuir

include/IndexStoreDB/Index/IndexSystem.h

@@ -50,11 +50,10 @@ class INDEXSTOREDB_EXPORT IndexSystem {
                                              bool readonly,
                                              bool enableOutOfDateFileWatching,
                                              bool listenToUnitEvents,
+                                             bool waitUntilDoneInitializing,
                                              Optional<size_t> initialDBSize,
                                              std::string &Error);
 
-  void waitUntilDoneInitializing();
-
   bool isUnitOutOfDate(StringRef unitOutputPath, ArrayRef<StringRef> dirtyFiles);
   bool isUnitOutOfDate(StringRef unitOutputPath, llvm::sys::TimePoint<> outOfDateModTime);
 

lib/CIndexStoreDB/CIndexStoreDB.cpp

@@ -114,12 +114,9 @@ indexstoredb_index_create(const char *storePath, const char *databasePath,
   if (auto index =
           IndexSystem::create(storePath, databasePath, libProviderObj, delegate,
                               useExplicitOutputUnits, readonly,
-                              /*enableOutOfDateFileWatching=*/false, listenToUnitEvents,
+                              /*enableOutOfDateFileWatching=*/false, listenToUnitEvents, wait,
                               llvm::None, errMsg)) {
 
-    if (wait)
-      index->waitUntilDoneInitializing();
-
     return make_object(index);
 
   } else if (error) {

lib/Index/IndexDatastore.cpp

@@ -88,11 +88,6 @@ struct UnitEventInfo {
   : kind(kind), name(std::move(name)), isDependency(isDependency) {}
 };
 
-struct DoneInitState {
-  std::atomic<bool> DoneInit{false};
-  unsigned RemainingInitUnits = 0; // this is not accessed concurrently.
-};
-
 struct PollUnitsState {
   llvm::sys::Mutex pollMtx;
   llvm::StringMap<sys::TimePoint<>> knownUnits;
@@ -108,9 +103,6 @@ class StoreUnitRepo : public std::enable_shared_from_this<StoreUnitRepo> {
 
   std::shared_ptr<FilePathWatcher> PathWatcher;
 
-  DoneInitState InitializingState;
-  dispatch_semaphore_t InitSemaphore;
-
   PollUnitsState pollUnitsState;
 
   mutable llvm::sys::Mutex StateMtx;
@@ -129,22 +121,13 @@ class StoreUnitRepo : public std::enable_shared_from_this<StoreUnitRepo> {
     EnableOutOfDateFileWatching(enableOutOfDateFileWatching),
     Delegate(std::move(Delegate)),
     CanonPathCache(std::move(canonPathCache)) {
-    InitSemaphore = dispatch_semaphore_create(0);
-  }
-  ~StoreUnitRepo() {
-    dispatch_release(InitSemaphore);
   }
 
   void onFilesChange(std::vector<UnitEventInfo> evts,
                      std::shared_ptr<UnitProcessingSession> processSession,
                      function_ref<void(unsigned)> ReportCompleted,
                      function_ref<void()> DirectoryDeleted);
 
-  void setInitialUnitCount(unsigned count);
-  void processedInitialUnitCount(unsigned count);
-  void finishedUnitInitialization();
-  void waitUntilDoneInitializing();
-
   /// *For Testing* Poll for any changes to units and wait until they have been registered.
   void pollForUnitChangesAndWait();
 
@@ -185,9 +168,9 @@ class IndexDatastoreImpl {
             bool readonly,
             bool enableOutOfDateFileWatching,
             bool listenToUnitEvents,
+            bool waitUntilDoneInitializing,
             std::string &Error);
 
-  void waitUntilDoneInitializing();
   bool isUnitOutOfDate(StringRef unitOutputPath, ArrayRef<StringRef> dirtyFiles);
   bool isUnitOutOfDate(StringRef unitOutputPath, llvm::sys::TimePoint<> outOfDateModTime);
   void checkUnitContainingFileIsOutOfDate(StringRef file);
@@ -464,10 +447,6 @@ void StoreUnitRepo::onFilesChange(std::vector<UnitEventInfo> evts,
     };
     PathWatcher = std::make_shared<FilePathWatcher>(std::move(pathEventsReceiver));
   }
-
-  if (!InitializingState.DoneInit) {
-    processedInitialUnitCount(evts.size());
-  }
 }
 
 void StoreUnitRepo::registerUnit(StringRef unitName, std::shared_ptr<UnitProcessingSession> processSession) {
@@ -761,30 +740,6 @@ void StoreUnitRepo::purgeStaleData() {
   // IdxStore->purgeStaleRecords(ActiveRecNames);
 }
 
-void StoreUnitRepo::setInitialUnitCount(unsigned count) {
-  InitializingState.RemainingInitUnits = count;
-}
-
-void StoreUnitRepo::processedInitialUnitCount(unsigned count) {
-  assert(!InitializingState.DoneInit);
-  InitializingState.RemainingInitUnits -= std::min(count, InitializingState.RemainingInitUnits);
-  if (InitializingState.RemainingInitUnits == 0) {
-    finishedUnitInitialization();
-  }
-}
-
-void StoreUnitRepo::finishedUnitInitialization() {
-  assert(!InitializingState.DoneInit);
-  dispatch_semaphore_signal(InitSemaphore);
-  InitializingState.DoneInit = true;
-}
-
-void StoreUnitRepo::waitUntilDoneInitializing() {
-  if (InitializingState.DoneInit)
-    return;
-  dispatch_semaphore_wait(InitSemaphore, DISPATCH_TIME_FOREVER);
-}
-
 void StoreUnitRepo::pollForUnitChangesAndWait() {
   sys::ScopedLock L(pollUnitsState.pollMtx);
   std::vector<UnitEventInfo> events;
@@ -1081,6 +1036,7 @@ bool IndexDatastoreImpl::init(IndexStoreRef idxStore,
                               bool readonly,
                               bool enableOutOfDateFileWatching,
                               bool listenToUnitEvents,
+                              bool waitUntilDoneInitializing,
                               std::string &Error) {
   this->IdxStore = std::move(idxStore);
   if (!this->IdxStore)
@@ -1092,18 +1048,8 @@ bool IndexDatastoreImpl::init(IndexStoreRef idxStore,
   auto UnitRepo = std::make_shared<StoreUnitRepo>(this->IdxStore, SymIndex, useExplicitOutputUnits, enableOutOfDateFileWatching, Delegate, CanonPathCache);
   std::weak_ptr<StoreUnitRepo> WeakUnitRepo = UnitRepo;
   auto eventsDeque = std::make_shared<UnitEventInfoDeque>();
-  auto OnUnitsChange = [WeakUnitRepo, Delegate, eventsDeque](IndexStore::UnitEventNotification EventNote) {
-    if (EventNote.isInitial()) {
-      auto UnitRepo = WeakUnitRepo.lock();
-      if (!UnitRepo)
-        return;
-      size_t evtCount = EventNote.getEventsCount();
-      if (evtCount == 0) {
-        UnitRepo->finishedUnitInitialization();
-      } else {
-        UnitRepo->setInitialUnitCount(evtCount);
-      }
-    }
+  auto OnUnitsChange = [WeakUnitRepo, Delegate, eventsDeque, waitUntilDoneInitializing](IndexStore::UnitEventNotification EventNote) {
+    bool shouldWait = waitUntilDoneInitializing && EventNote.isInitial();
 
     std::vector<UnitEventInfo> evts;
     for (size_t i = 0, e = EventNote.getEventsCount(); i != e; ++i) {
@@ -1112,25 +1058,23 @@ bool IndexDatastoreImpl::init(IndexStoreRef idxStore,
     }
 
     auto session = std::make_shared<UnitProcessingSession>(eventsDeque, WeakUnitRepo, Delegate);
-    session->process(std::move(evts), /*waitForProcessing=*/false);
+    session->process(std::move(evts), shouldWait);
   };
 
+  this->UnitRepo = std::move(UnitRepo);
+
   if (listenToUnitEvents) {
     this->IdxStore->setUnitEventHandler(OnUnitsChange);
-    bool err = this->IdxStore->startEventListening(/*waitInitialSync=*/false, Error);
+    bool err = this->IdxStore->startEventListening(waitUntilDoneInitializing, Error);
     if (err)
       return true;
+  } else if (waitUntilDoneInitializing) {
+    pollForUnitChangesAndWait();
   }
 
-  this->UnitRepo = std::move(UnitRepo);
   return false;
 }
 
-void IndexDatastoreImpl::waitUntilDoneInitializing() {
-  if (UnitRepo)
-    UnitRepo->waitUntilDoneInitializing();
-}
-
 bool IndexDatastoreImpl::isUnitOutOfDate(StringRef unitOutputPath, ArrayRef<StringRef> dirtyFiles) {
   auto mostRecentFileAndTime = UnitMonitor::getMostRecentModTime(dirtyFiles);
   return isUnitOutOfDate(unitOutputPath, mostRecentFileAndTime.second);
@@ -1182,10 +1126,12 @@ IndexDatastore::create(IndexStoreRef idxStore,
                        bool readonly,
                        bool enableOutOfDateFileWatching,
                        bool listenToUnitEvents,
+                       bool waitUntilDoneInitializing,
                        std::string &Error) {
   std::unique_ptr<IndexDatastoreImpl> Impl(new IndexDatastoreImpl());
   bool Err = Impl->init(std::move(idxStore), std::move(SymIndex), std::move(Delegate), std::move(CanonPathCache),
-                        useExplicitOutputUnits, readonly, enableOutOfDateFileWatching, listenToUnitEvents, Error);
+                        useExplicitOutputUnits, readonly, enableOutOfDateFileWatching,
+                        listenToUnitEvents, waitUntilDoneInitializing, Error);
   if (Err)
     return nullptr;
 
@@ -1200,10 +1146,6 @@ IndexDatastore::~IndexDatastore() {
   delete IMPL;
 }
 
-void IndexDatastore::waitUntilDoneInitializing() {
-  return IMPL->waitUntilDoneInitializing();
-}
-
 bool IndexDatastore::isUnitOutOfDate(StringRef unitOutputPath, ArrayRef<StringRef> dirtyFiles) {
   return IMPL->isUnitOutOfDate(unitOutputPath, dirtyFiles);
 }

lib/Index/IndexDatastore.h

@@ -44,10 +44,9 @@ class IndexDatastore {
                                                 bool readonly,
                                                 bool enableOutOfDateFileWatching,
                                                 bool listenToUnitEvents,
+                                                bool waitUntilDoneInitializing,
                                                 std::string &Error);
 
-  void waitUntilDoneInitializing();
-
   bool isUnitOutOfDate(StringRef unitOutputPath, ArrayRef<StringRef> dirtyFiles);
   bool isUnitOutOfDate(StringRef unitOutputPath, llvm::sys::TimePoint<> outOfDateModTime);
 

lib/Index/IndexSystem.cpp

@@ -119,11 +119,10 @@ class IndexSystemImpl {
             std::shared_ptr<IndexSystemDelegate> Delegate,
             bool useExplicitOutputUnits, bool readonly,
             bool enableOutOfDateFileWatching, bool listenToUnitEvents,
+            bool waitUntilDoneInitializing,
             Optional<size_t> initialDBSize,
             std::string &Error);
 
-  void waitUntilDoneInitializing();
-
   bool isUnitOutOfDate(StringRef unitOutputPath, ArrayRef<StringRef> dirtyFiles);
   bool isUnitOutOfDate(StringRef unitOutputPath, llvm::sys::TimePoint<> outOfDateModTime);
   void checkUnitContainingFileIsOutOfDate(StringRef file);
@@ -210,6 +209,7 @@ bool IndexSystemImpl::init(StringRef StorePath,
                            std::shared_ptr<IndexSystemDelegate> Delegate,
                            bool useExplicitOutputUnits, bool readonly,
                            bool enableOutOfDateFileWatching, bool listenToUnitEvents,
+                           bool waitUntilDoneInitializing,
                            Optional<size_t> initialDBSize,
                            std::string &Error) {
   this->StorePath = StorePath;
@@ -252,17 +252,14 @@ bool IndexSystemImpl::init(StringRef StorePath,
                                             readonly,
                                             enableOutOfDateFileWatching,
                                             listenToUnitEvents,
+                                            waitUntilDoneInitializing,
                                             Error);
 
   if (!this->IndexStore)
     return true;
   return false;
 }
 
-void IndexSystemImpl::waitUntilDoneInitializing() {
-  IndexStore->waitUntilDoneInitializing();
-}
-
 bool IndexSystemImpl::isUnitOutOfDate(StringRef unitOutputPath, ArrayRef<StringRef> dirtyFiles) {
   return IndexStore->isUnitOutOfDate(unitOutputPath, dirtyFiles);
 }
@@ -611,12 +608,13 @@ IndexSystem::create(StringRef StorePath,
                     std::shared_ptr<IndexSystemDelegate> Delegate,
                     bool useExplicitOutputUnits, bool readonly,
                     bool enableOutOfDateFileWatching, bool listenToUnitEvents,
+                    bool waitUntilDoneInitializing,
                     Optional<size_t> initialDBSize,
                     std::string &Error) {
   std::unique_ptr<IndexSystemImpl> Impl(new IndexSystemImpl());
   bool Err = Impl->init(StorePath, dbasePath, std::move(storeLibProvider), std::move(Delegate),
                         useExplicitOutputUnits, readonly,
-                        enableOutOfDateFileWatching, listenToUnitEvents,
+                        enableOutOfDateFileWatching, listenToUnitEvents, waitUntilDoneInitializing,
                         initialDBSize, Error);
   if (Err)
     return nullptr;
@@ -632,10 +630,6 @@ IndexSystem::~IndexSystem() {
   delete IMPL;
 }
 
-void IndexSystem::waitUntilDoneInitializing() {
-  return IMPL->waitUntilDoneInitializing();
-}
-
 bool IndexSystem::isUnitOutOfDate(StringRef unitOutputPath, ArrayRef<StringRef> dirtyFiles) {
   return IMPL->isUnitOutOfDate(unitOutputPath, dirtyFiles);
 }