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

[Kristof Beyls via llvm-commits]

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

>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 1/2] [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

>From faad7c7d61accc9cb93d25a456bc63600267d8fa Mon Sep 17 00:00:00 2001
From: Kristof Beyls <kristof.beyls at gmail.com>
Date: Thu, 21 Sep 2023 19:34:13 +0200
Subject: [PATCH 2/2] Update bolt/include/bolt/Core/MCPlusBuilder.h

Co-authored-by: Amir Ayupov <fads93 at gmail.com>
---
 bolt/include/bolt/Core/MCPlusBuilder.h | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index b67e172d43a38cc..4116424fffe2376 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -1791,8 +1791,7 @@ class MCPlusBuilder {
   /// \p Name wasn't registered previously.
   unsigned getOrCreateAnnotationIndex(StringRef Name) {
     {
-      std::optional<unsigned> Index = getAnnotationIndex(Name);
-      if (Index.has_value())
+      if (std::optional<unsigned> Index = getAnnotationIndex(Name))
         return *Index;
     }