Fix tsan problem where the per-thread shared_ptr() can be locked right before the cache is destroyed causing a race where it tries to remove an entry from a destroyed cache.

[zhangqiaorjc]

This is originally from Parker Schuh. I'm just helping him creating this commit.

[rriddle]

Can you include more context in the description of the patch?

[zhangqiaorjc]

In D142394#4074500, @rriddle wrote:

Can you include more context in the description of the patch?

I updated the title to "Fix tsan problem where the per-thread shared_ptr() can be locked right before the cache is destroyed causing a race where it tries to remove an entry from a destroyed cache."

[zhangqiaorjc]

Could someone help review this patch? It's caught by tsan for jax tests and we verified the fix internally.

[zhangqiaorjc]

Thank you River for the review! I don't have commit access. @rriddle could you help us commit this patch? Thank you so much!

[zhangqiaorjc]

Thank you @bkramer !

[hctim]

Hey, looks like this broke the asan-aarch64 buildbot. https://lab.llvm.org/buildbot/#/builders/239/builds/771

Sorry it took so long to bisect this down. It seems like the relevant leaks happen because of:

Indirect leak of 104 byte(s) in 1 object(s) allocated from:
    #0 0xaaaad7bce55c in operator new(unsigned long) /home/mitchp/llvm-project/compiler-rt/lib/asan/asan_new_delete.cpp:95:3
    #1 0xaaaad7e81b44 in __libcpp_operator_new<unsigned long> /home/mitchp/build/libcxx_build_asan/include/c++/v1/new:266:10
    #2 0xaaaad7e81b44 in __libcpp_allocate /home/mitchp/build/libcxx_build_asan/include/c++/v1/new:292:10
    #3 0xaaaad7e81b44 in allocate /home/mitchp/build/libcxx_build_asan/include/c++/v1/__memory/allocator.h:115:38
    #4 0xaaaad7e81b44 in allocate /home/mitchp/build/libcxx_build_asan/include/c++/v1/__memory/allocator_traits.h:268:20
    #5 0xaaaad7e81b44 in __allocation_guard<std::__1::allocator<mlir::ThreadLocalCache<llvm::DenseSet<(anonymous namespace)::ParametricStorageUniquer::HashedStorage, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo> >::PerInstanceState> > /home/mitchp/build/libcxx_build_asan/include/c++/v1/__memory/allocation_guard.h:53:18
    #6 0xaaaad7e81b44 in allocate_shared<mlir::ThreadLocalCache<llvm::DenseSet<(anonymous namespace)::ParametricStorageUniquer::HashedStorage, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo> >::PerInstanceState, std::__1::allocator<mlir::ThreadLocalCache<llvm::DenseSet<(anonymous namespace)::ParametricStorageUniquer::HashedStorage, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo> >::PerInstanceState>, void> /home/mitchp/build/libcxx_build_asan/include/c++/v1/__memory/shared_ptr.h:947:48
    #7 0xaaaad7e81b44 in make_shared<mlir::ThreadLocalCache<llvm::DenseSet<(anonymous namespace)::ParametricStorageUniquer::HashedStorage, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo> >::PerInstanceState, void> /home/mitchp/build/libcxx_build_asan/include/c++/v1/__memory/shared_ptr.h:957:12
    #8 0xaaaad7e81b44 in ThreadLocalCache /home/mitchp/llvm-project/mlir/include/mlir/Support/ThreadLocalCache.h:123:7

And I'm guessing the rest of the leaks are cascading failures from that leak.

I've attached the rest of the logfile from ninja check-mlir inside of the sanitizer bot (https://github.com/google/sanitizers/wiki/SanitizerBotReproduceBuild, the buildscript that you'd want to use is buildbot_bootstrap_asan.sh instead of buildbot_fast.sh). It was generated with the following patch, which allows us to actually see the asan report:

diff --git a/mlir/test/CAPI/pdl.c b/mlir/test/CAPI/pdl.c
index 673ad8d6af7e..4dbada4060da 100644
--- a/mlir/test/CAPI/pdl.c
+++ b/mlir/test/CAPI/pdl.c
@@ -7,7 +7,9 @@
 //
 //===----------------------------------------------------------------------===//
 
-// RUN: mlir-capi-pdl-test 2>&1 | FileCheck %s
+// RUN: mlir-capi-pdl-test 2>&1
+//
+// | FileCheck %s
 
 #include "mlir-c/Dialect/PDL.h"
 #include "mlir-c/BuiltinTypes.h"
diff --git a/mlir/test/CAPI/transform.c b/mlir/test/CAPI/transform.c
index 1db5c2ecafe4..c69e40642c47 100644
--- a/mlir/test/CAPI/transform.c
+++ b/mlir/test/CAPI/transform.c
@@ -7,7 +7,9 @@
 //
 //===----------------------------------------------------------------------===//
 
-// RUN: mlir-capi-transform-test 2>&1 | FileCheck %s
+// RUN: mlir-capi-transform-test 
+//
+// 2>&1 | FileCheck %s
 
 #include "mlir-c/Dialect/Transform.h"
 #include "mlir-c/IR.h"

That patch is applied on top of bcc10817d5569172ee065015747e226280e9b698, which is the commit for this patch.

I'm not exactly sure why this doesn't reproduce on the x86_64 asan bot. I'm sorry about that.

aarch64-asan-log4 MBDownload

mlir/include/mlir/Support/ThreadLocalCache.h

 #include "llvm/Support/Mutex.h"
 
 namespace mlir {
 /// This class provides support for defining a thread local object with non
 /// static storage duration. This is very useful for situations in which a data
 /// cache has very large lock contention.
 template <typename ValueT>
 class ThreadLocalCache {
+  // Keep a separate shared_ptr protected state that can be acquired atomically
+  // instead of using shared_ptr's for each value. This avoids a problem
+  // where the instance shared_ptr is locked() successfully, and then the
+  // ThreadLocalCache gets destroyed before remove() can be called successfully.
+  struct PerInstanceState {
+    /// Remove the given value entry. This is generally called when a thread
+    /// local cache is destructing.
+    void remove(ValueT *value) {
+      // Erase the found value directly, because it is guaranteed to be in the
+      // list.
+      llvm::sys::SmartScopedLock<true> threadInstanceLock(instanceMutex);
+      auto it =
+          llvm::find_if(instances, [&](std::unique_ptr<ValueT> &instance) {
+            return instance.get() == value;
+          });
+      assert(it != instances.end() && "expected value to exist in cache");
+      instances.erase(it);
+    }
+
+    /// Owning pointers to all of the values that have been constructed for this
+    /// object in the static cache.
+    SmallVector<std::unique_ptr<ValueT>, 1> instances;
+
+    /// A mutex used when a new thread instance has been added to the cache for
+    /// this object.
+    llvm::sys::SmartMutex<true> instanceMutex;
+  };
+
   /// The type used for the static thread_local cache. This is a map between an
   /// instance of the non-static cache and a weak reference to an instance of
   /// ValueT. We use a weak reference here so that the object can be destroyed
   /// without needing to lock access to the cache itself.
-  struct CacheType : public llvm::SmallDenseMap<ThreadLocalCache<ValueT> *,
-                                                std::weak_ptr<ValueT>> {
+  struct CacheType
+      : public llvm::SmallDenseMap<PerInstanceState *, std::weak_ptr<ValueT>> {
     ~CacheType() {
       // Remove the values of this cache that haven't already expired.
       for (auto &it : *this)
         if (std::shared_ptr<ValueT> value = it.second.lock())
           it.first->remove(value.get());
     }
 
     /// Clear out any unused entries within the map. This method is not
~ThreadLocalCache() {
     // No cleanup is necessary here as the shared_pointer memory will go out of
     // scope and invalidate the weak pointers held by the thread_local caches.
   }
 
   /// Return an instance of the value type for the current thread.
   ValueT &get() {
     // Check for an already existing instance for this thread.
     CacheType &staticCache = getStaticCache();
-    std::weak_ptr<ValueT> &threadInstance = staticCache[this];
+    std::weak_ptr<ValueT> &threadInstance = staticCache[perInstanceState.get()];
     if (std::shared_ptr<ValueT> value = threadInstance.lock())
       return *value;
 
     // Otherwise, create a new instance for this thread.
-    llvm::sys::SmartScopedLock<true> threadInstanceLock(instanceMutex);
-    instances.push_back(std::make_shared<ValueT>());
-    std::shared_ptr<ValueT> &instance = instances.back();
-    threadInstance = instance;
+    llvm::sys::SmartScopedLock<true> threadInstanceLock(
+        perInstanceState->instanceMutex);
+    perInstanceState->instances.push_back(std::make_unique<ValueT>());
+    ValueT *instance = perInstanceState->instances.back().get();
+    threadInstance = std::shared_ptr<ValueT>(perInstanceState, instance);
 
     // Before returning the new instance, take the chance to clear out any used
     // entries in the static map. The cache is only cleared within the same
     // thread to remove the need to lock the cache itself.
     staticCache.clearExpiredEntries();
     return *instance;
   }
   ValueT &operator*() { return get(); }
   ValueT *operator->() { return &get(); }
 
 private:
   ThreadLocalCache(ThreadLocalCache &&) = delete;
   ThreadLocalCache(const ThreadLocalCache &) = delete;
   ThreadLocalCache &operator=(const ThreadLocalCache &) = delete;
 
   /// Return the static thread local instance of the cache type.
   static CacheType &getStaticCache() {
     static LLVM_THREAD_LOCAL CacheType cache;
     return cache;
   }
 
-  /// Remove the given value entry. This is generally called when a thread local
-  /// cache is destructing.
-  void remove(ValueT *value) {
-    // Erase the found value directly, because it is guaranteed to be in the
-    // list.
-    llvm::sys::SmartScopedLock<true> threadInstanceLock(instanceMutex);
-    auto it = llvm::find_if(instances, [&](std::shared_ptr<ValueT> &instance) {
-      return instance.get() == value;
-    });
-    assert(it != instances.end() && "expected value to exist in cache");
-    instances.erase(it);
-  }
-
-  /// Owning pointers to all of the values that have been constructed for this
-  /// object in the static cache.
-  SmallVector<std::shared_ptr<ValueT>, 1> instances;
-
-  /// A mutex used when a new thread instance has been added to the cache for
-  /// this object.
-  llvm::sys::SmartMutex<true> instanceMutex;
+  std::shared_ptr<PerInstanceState> perInstanceState =
+      std::make_shared<PerInstanceState>();
 };
 } // namespace mlir
 
 #endif // MLIR_SUPPORT_THREADLOCALCACHE_H