[libc-commits] [libc] 8e488c3 - [libc] Add a multi-waiter mutex test.

[Siva Chandra Reddy via libc-commits]

Author: Siva Chandra Reddy
Date: 2021-08-23T05:47:30Z
New Revision: 8e488c3cc0da71afbe75963033e86bf32ee56c11

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

LOG: [libc] Add a multi-waiter mutex test.

A corresponding adjustment to mtx_lock has also been made.

Added: 
    

Modified: 
    libc/src/threads/linux/mtx_lock.cpp
    libc/test/src/threads/mtx_test.cpp

Removed: 
    


################################################################################
diff  --git a/libc/src/threads/linux/mtx_lock.cpp b/libc/src/threads/linux/mtx_lock.cpp
index 7ae30a78ba81f..39bd329e9e9a9 100644
--- a/libc/src/threads/linux/mtx_lock.cpp
+++ b/libc/src/threads/linux/mtx_lock.cpp
@@ -21,12 +21,16 @@ namespace __llvm_libc {
 // The implementation currently handles only plain mutexes.
 LLVM_LIBC_FUNCTION(int, mtx_lock, (mtx_t * mutex)) {
   FutexData *futex_data = reinterpret_cast<FutexData *>(mutex->__internal_data);
+  bool was_waiting = false;
   while (true) {
     uint32_t mutex_status = MS_Free;
     uint32_t locked_status = MS_Locked;
 
-    if (atomic_compare_exchange_strong(futex_data, &mutex_status, MS_Locked))
+    if (atomic_compare_exchange_strong(futex_data, &mutex_status, MS_Locked)) {
+      if (was_waiting)
+        atomic_store(futex_data, MS_Waiting);
       return thrd_success;
+    }
 
     switch (mutex_status) {
     case MS_Waiting:
@@ -35,6 +39,7 @@ LLVM_LIBC_FUNCTION(int, mtx_lock, (mtx_t * mutex)) {
       // 4th argument to the syscall function below.)
       __llvm_libc::syscall(SYS_futex, futex_data, FUTEX_WAIT_PRIVATE,
                            MS_Waiting, 0, 0, 0);
+      was_waiting = true;
       // Once woken up/unblocked, try everything all over.
       continue;
     case MS_Locked:
@@ -47,6 +52,7 @@ LLVM_LIBC_FUNCTION(int, mtx_lock, (mtx_t * mutex)) {
         // syscall will block only if the futex data is still `MS_Waiting`.
         __llvm_libc::syscall(SYS_futex, futex_data, FUTEX_WAIT_PRIVATE,
                              MS_Waiting, 0, 0, 0);
+        was_waiting = true;
       }
       continue;
     case MS_Free:

diff  --git a/libc/test/src/threads/mtx_test.cpp b/libc/test/src/threads/mtx_test.cpp
index 0f3df0ee03849..536ff520c6018 100644
--- a/libc/test/src/threads/mtx_test.cpp
+++ b/libc/test/src/threads/mtx_test.cpp
@@ -127,3 +127,56 @@ TEST(LlvmLibcMutexTest, WaitAndStep) {
   __llvm_libc::thrd_join(&thread, &retval);
   ASSERT_EQ(retval, 0);
 }
+
+static constexpr int THREAD_COUNT = 10;
+static mtx_t multiple_waiter_lock;
+static mtx_t counter_lock;
+static int wait_count = 0;
+
+int waiter_func(void *) {
+  __llvm_libc::mtx_lock(&counter_lock);
+  ++wait_count;
+  __llvm_libc::mtx_unlock(&counter_lock);
+
+  // Block on the waiter lock until the main
+  // thread unblocks.
+  __llvm_libc::mtx_lock(&multiple_waiter_lock);
+  __llvm_libc::mtx_unlock(&multiple_waiter_lock);
+
+  __llvm_libc::mtx_lock(&counter_lock);
+  --wait_count;
+  __llvm_libc::mtx_unlock(&counter_lock);
+
+  return 0;
+}
+
+TEST(LlvmLibcMutexTest, MultipleWaiters) {
+  __llvm_libc::mtx_init(&multiple_waiter_lock, mtx_plain);
+  __llvm_libc::mtx_init(&counter_lock, mtx_plain);
+
+  __llvm_libc::mtx_lock(&multiple_waiter_lock);
+  thrd_t waiters[THREAD_COUNT];
+  for (int i = 0; i < THREAD_COUNT; ++i) {
+    __llvm_libc::thrd_create(waiters + i, waiter_func, nullptr);
+  }
+
+  // Spin until the counter is incremented to the desired
+  // value.
+  while (true) {
+    __llvm_libc::mtx_lock(&counter_lock);
+    if (wait_count == THREAD_COUNT) {
+      __llvm_libc::mtx_unlock(&counter_lock);
+      break;
+    }
+    __llvm_libc::mtx_unlock(&counter_lock);
+  }
+
+  __llvm_libc::mtx_unlock(&multiple_waiter_lock);
+
+  int retval;
+  for (int i = 0; i < THREAD_COUNT; ++i) {
+    __llvm_libc::thrd_join(waiters + i, &retval);
+  }
+
+  ASSERT_EQ(wait_count, 0);
+}