[Basic] Avoid reentrant stat collection

Currently `UnifiedStatsReporter::flushTracesAndProfiles`
can kick off requests when computing the source
ranges for collected entities, which will try to
record additional stats about the requests.

This currently happens to work without issue,
but swiftlang#29289 bumped the counters down slightly such
that the vector storing the stats now performs a
re-allocation when we do a reentrant stat entry.
This then caused a use-after-free as we try to
continue iterating over the old buffer.

Fix this issue by refusing to record any new stats
while we're flushing out the ones we've already
recorded.

Author: hamishknight

include/swift/Basic/Statistic.h

@@ -162,6 +162,10 @@ class UnifiedStatsReporter {
   std::unique_ptr<StatsProfilers> EventProfilers;
   std::unique_ptr<StatsProfilers> EntityProfilers;
 
+  /// Whether we are currently flushing statistics and should not therefore
+  /// record any additional stats until we've finished.
+  bool IsFlushingTracesAndProfiles;
+
   void publishAlwaysOnStatsToLLVM();
   void printAlwaysOnStatsAndTimers(raw_ostream &OS);
 

lib/Basic/Statistic.cpp

@@ -25,6 +25,7 @@
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Process.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Support/SaveAndRestore.h"
 #include <chrono>
 #include <limits>
 
@@ -348,7 +349,8 @@ UnifiedStatsReporter::UnifiedStatsReporter(StringRef ProgramName,
                                         ProgramName, "Running Program")),
     SourceMgr(SM),
     ClangSourceMgr(CSM),
-    RecursiveTimers(llvm::make_unique<RecursionSafeTimers>())
+    RecursiveTimers(llvm::make_unique<RecursionSafeTimers>()),
+    IsFlushingTracesAndProfiles(false)
 {
   path::append(StatsFilename, makeStatsFileName(ProgramName, AuxName));
   path::append(TraceFilename, makeTraceFileName(ProgramName, AuxName));
@@ -557,6 +559,13 @@ UnifiedStatsReporter::saveAnyFrontendStatsEvents(
     bool IsEntry)
 {
   assert(MainThreadID == std::this_thread::get_id());
+
+  // Don't record any new stats if we're currently flushing the ones we've
+  // already recorded. This can happen when requests get kicked off when
+  // computing source ranges.
+  if (IsFlushingTracesAndProfiles)
+    return;
+
   // First make a note in the recursion-safe timers; these
   // are active anytime UnifiedStatsReporter is active.
   if (IsEntry) {
@@ -711,6 +720,10 @@ UnifiedStatsReporter::~UnifiedStatsReporter()
 
 void
 UnifiedStatsReporter::flushTracesAndProfiles() {
+  // Note that we're currently flushing statistics and shouldn't record any
+  // more until we've finished.
+  llvm::SaveAndRestore<bool> flushing(IsFlushingTracesAndProfiles, true);
+
   if (FrontendStatsEvents && SourceMgr) {
     std::error_code EC;
     raw_fd_ostream tstream(TraceFilename, EC, fs::F_Append | fs::F_Text);