[llvm] [BOLT] Fix data race in MCPlusBuilder::getOrCreateAnnotationIndex (PR #67004)

[Kristof Beyls via llvm-commits]

https://github.com/kbeyls created https://github.com/llvm/llvm-project/pull/67004

MCPlusBuilder::getOrCreateAnnotationIndex(Name) can be called from different threads, for example when making use of
ParallelUtilities::runOnEachFunctionWithUniqueAllocId.

The race occurs when an Index for a particular annotation Name needs to be created for the first time.

For example, this can easily happen when multiple "copies" of an analysis pass run on different BinaryFunctions, and the analysis pass creates a new Annotation Index to be able to store analysis results as annotations.

This was found by using the ThreadSanitizer.

No regression test was added; I don't think there is good way to write regression tests that verify the absence of data races?

>From 3987fd2297ee9f13ef88c659294e08a514d62a80 Mon Sep 17 00:00:00 2001
From: Kristof Beyls <kristof.beyls at arm.com>
Date: Thu, 21 Sep 2023 12:05:06 +0200
Subject: [PATCH] [BOLT] Fix data race in
 MCPlusBuilder::getOrCreateAnnotationIndex

MCPlusBuilder::getAnnotationIndex(Name) can be called from different
threads, for example when making use of
ParallelUtilities::runOnEachFunctionWithUniqueAllocId.

The race occurs when an Index for a particular annotation Name needs to
be created for the first time.

For example, this can easily happen when multiple "copies" of an
analysis pass run on different BinaryFunctions, and the analysis pass
creates a new Annotation Index to be able to store analysis results as
annotations.

This was found by using the ThreadSanitizer.

No regression test was added; I don't think there is good way to write
regression tests that verify the absence of data races?
---
 bolt/include/bolt/Core/MCPlusBuilder.h | 27 ++++++++++++++++++--------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index f48cf21852dcbef..b67e172d43a38cc 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -29,6 +29,7 @@
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/ErrorOr.h"
+#include "llvm/Support/RWMutex.h"
 #include <cassert>
 #include <cstdint>
 #include <map>
@@ -166,6 +167,10 @@ class MCPlusBuilder {
   /// Names of non-standard annotations.
   SmallVector<std::string, 8> AnnotationNames;
 
+  /// A mutex that is used to control parallel accesses to
+  /// AnnotationNameIndexMap and AnnotationsNames.
+  mutable llvm::sys::RWMutex AnnotationNameMutex;
+
   /// Allocate the TailCall annotation value. Clients of the target-specific
   /// MCPlusBuilder classes must use convert/lower/create* interfaces instead.
   void setTailCall(MCInst &Inst);
@@ -1775,6 +1780,7 @@ class MCPlusBuilder {
 
   /// Return annotation index matching the \p Name.
   std::optional<unsigned> getAnnotationIndex(StringRef Name) const {
+    std::shared_lock<llvm::sys::RWMutex> Lock(AnnotationNameMutex);
     auto AI = AnnotationNameIndexMap.find(Name);
     if (AI != AnnotationNameIndexMap.end())
       return AI->second;
@@ -1784,15 +1790,20 @@ class MCPlusBuilder {
   /// Return annotation index matching the \p Name. Create a new index if the
   /// \p Name wasn't registered previously.
   unsigned getOrCreateAnnotationIndex(StringRef Name) {
-    auto AI = AnnotationNameIndexMap.find(Name);
-    if (AI != AnnotationNameIndexMap.end())
-      return AI->second;
+    {
+      std::optional<unsigned> Index = getAnnotationIndex(Name);
+      if (Index.has_value())
+        return *Index;
+    }
 
-    const unsigned Index =
-        AnnotationNameIndexMap.size() + MCPlus::MCAnnotation::kGeneric;
-    AnnotationNameIndexMap.insert(std::make_pair(Name, Index));
-    AnnotationNames.emplace_back(std::string(Name));
-    return Index;
+    {
+      std::unique_lock<llvm::sys::RWMutex> Lock(AnnotationNameMutex);
+      const unsigned Index =
+          AnnotationNameIndexMap.size() + MCPlus::MCAnnotation::kGeneric;
+      AnnotationNameIndexMap.insert(std::make_pair(Name, Index));
+      AnnotationNames.emplace_back(std::string(Name));
+      return Index;
+    }
   }
 
   /// Store an annotation value on an MCInst.  This assumes the annotation