[compiler-rt] ea624e6 - sanitizer_common: add Mutex::TryLock

[Dmitry Vyukov via llvm-commits]

Author: Dmitry Vyukov
Date: 2022-04-01T17:56:19+02:00
New Revision: ea624e697bbe1d2924b69208102c64c4c01e1400

URL: https://github.com/llvm/llvm-project/commit/ea624e697bbe1d2924b69208102c64c4c01e1400
DIFF: https://github.com/llvm/llvm-project/commit/ea624e697bbe1d2924b69208102c64c4c01e1400.diff

LOG: sanitizer_common: add Mutex::TryLock

Will be used in future changes.

Reviewed By: melver

Differential Revision: https://reviews.llvm.org/D122905

Added: 
    

Modified: 
    compiler-rt/lib/sanitizer_common/sanitizer_mutex.h
    compiler-rt/lib/sanitizer_common/tests/sanitizer_mutex_test.cpp

Removed: 
    


################################################################################
diff  --git a/compiler-rt/lib/sanitizer_common/sanitizer_mutex.h b/compiler-rt/lib/sanitizer_common/sanitizer_mutex.h
index c16f5cdc1d717..d2188a9e6d623 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_mutex.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_mutex.h
@@ -208,6 +208,20 @@ class SANITIZER_MUTEX Mutex : CheckedMutex {
     }
   }
 
+  bool TryLock() SANITIZER_TRY_ACQUIRE(true) {
+    u64 state = atomic_load_relaxed(&state_);
+    for (;;) {
+      if (UNLIKELY(state & (kWriterLock | kReaderLockMask)))
+        return false;
+      // The mutex is not read-/write-locked, try to lock.
+      if (LIKELY(atomic_compare_exchange_weak(
+              &state_, &state, state | kWriterLock, memory_order_acquire))) {
+        CheckedMutex::Lock();
+        return true;
+      }
+    }
+  }
+
   void Unlock() SANITIZER_RELEASE() {
     CheckedMutex::Unlock();
     bool wake_writer;

diff  --git a/compiler-rt/lib/sanitizer_common/tests/sanitizer_mutex_test.cpp b/compiler-rt/lib/sanitizer_common/tests/sanitizer_mutex_test.cpp
index 9955d7604c5eb..960427663480d 100644
--- a/compiler-rt/lib/sanitizer_common/tests/sanitizer_mutex_test.cpp
+++ b/compiler-rt/lib/sanitizer_common/tests/sanitizer_mutex_test.cpp
@@ -155,6 +155,15 @@ TEST(SanitizerCommon, Mutex) {
   for (int i = 0; i < kThreads; i++) PTHREAD_JOIN(threads[i], 0);
 }
 
+TEST(SanitizerCommon, MutexTry) {
+  Mutex mtx;
+  TestData<Mutex> data(&mtx);
+  pthread_t threads[kThreads];
+  for (int i = 0; i < kThreads; i++)
+    PTHREAD_CREATE(&threads[i], 0, try_thread<Mutex>, &data);
+  for (int i = 0; i < kThreads; i++) PTHREAD_JOIN(threads[i], 0);
+}
+
 struct SemaphoreData {
   Semaphore *sem;
   bool done;