[libc-commits] [libc] [libc] add internal rwlock implementation (PR #91142)

Schrodinger ZHU Yifan via libc-commits libc-commits at lists.llvm.org
Sun May 5 13:25:49 PDT 2024


https://github.com/SchrodingerZhu updated https://github.com/llvm/llvm-project/pull/91142

>From bed03df045b3f33d8b49a71e12fbc0e0845fa071 Mon Sep 17 00:00:00 2001
From: Schrodinger ZHU Yifan <yifanzhu at rochester.edu>
Date: Sun, 5 May 2024 16:22:22 -0400
Subject: [PATCH 1/3] [libc] add internal rwlock implementation

---
 libc/src/__support/CPP/atomic.h               |  20 +
 .../__support/threads/linux/CMakeLists.txt    |  28 +-
 libc/src/__support/threads/linux/futex_word.h |  50 ++-
 libc/src/__support/threads/linux/rwlock.h     | 412 ++++++++++++++++++
 libc/test/src/__support/CMakeLists.txt        |   1 +
 .../test/src/__support/threads/CMakeLists.txt |   3 +
 .../__support/threads/linux/CMakeLists.txt    |  13 +
 .../__support/threads/linux/rwlock_test.cpp   | 172 ++++++++
 8 files changed, 694 insertions(+), 5 deletions(-)
 create mode 100644 libc/src/__support/threads/linux/rwlock.h
 create mode 100644 libc/test/src/__support/threads/CMakeLists.txt
 create mode 100644 libc/test/src/__support/threads/linux/CMakeLists.txt
 create mode 100644 libc/test/src/__support/threads/linux/rwlock_test.cpp

diff --git a/libc/src/__support/CPP/atomic.h b/libc/src/__support/CPP/atomic.h
index 5e428940565b99..1392197321a80e 100644
--- a/libc/src/__support/CPP/atomic.h
+++ b/libc/src/__support/CPP/atomic.h
@@ -101,6 +101,26 @@ template <typename T> struct Atomic {
                                        int(mem_ord), int(mem_ord));
   }
 
+  // Atomic compare exchange
+  LIBC_INLINE bool compare_exchange_strong(
+      T &expected, T desired, MemoryOrder success_order,
+      MemoryOrder failure_order,
+      [[maybe_unused]] MemoryScope mem_scope = cpp::MemoryScope::DEVICE) {
+    return __atomic_compare_exchange_n(&val, &expected, desired, false,
+                                       int(success_order), int(failure_order));
+  }
+
+  LIBC_INLINE bool compare_exchange_weak(
+      T &expected, T new_val, MemoryOrder success_order,
+      MemoryOrder failure_order,
+      [[maybe_unused]] MemoryScope mem_scope = cpp::MemoryScope::DEVICE) {
+    return __atomic_compare_exchange_n(
+        &val, &expected, new_val,
+        /* is_weak */ true,
+        /* success_memorder */ static_cast<int>(success_order),
+        /* failure_memorder */ static_cast<int>(failure_order));
+  }
+
   T exchange(T desired, MemoryOrder mem_ord = MemoryOrder::SEQ_CST,
              [[maybe_unused]] MemoryScope mem_scope = MemoryScope::DEVICE) {
 #if __has_builtin(__scoped_atomic_exchange_n)
diff --git a/libc/src/__support/threads/linux/CMakeLists.txt b/libc/src/__support/threads/linux/CMakeLists.txt
index 87a7a66ac6ea57..36edf99092792e 100644
--- a/libc/src/__support/threads/linux/CMakeLists.txt
+++ b/libc/src/__support/threads/linux/CMakeLists.txt
@@ -1,13 +1,21 @@
+if(NOT TARGET libc.src.__support.OSUtil.osutil)
+  return()
+endif()
+
 add_header_library(
   futex_word_type
   HDRS
     futex_word.h
+  DEPENDS
+    libc.src.__support.CPP.optional
+    libc.src.__support.CPP.atomic
+    libc.src.__support.CPP.limits
+    libc.src.__support.OSUtil.osutil
+    libc.include.sys_syscall
+    libc.src.time.clock_gettime
+    libc.src.errno.errno
 )
 
-if(NOT TARGET libc.src.__support.OSUtil.osutil)
-  return()
-endif()
-
 add_header_library(
   mutex
   HDRS
@@ -20,6 +28,18 @@ add_header_library(
     libc.src.__support.threads.mutex_common
 )
 
+add_header_library(
+  rwlock
+  HDRS
+    rwlock.h
+  DEPENDS
+    .futex_word_type
+    libc.include.sys_syscall
+    libc.src.__support.CPP.atomic
+    libc.src.__support.CPP.optional
+    libc.src.__support.libc_assert
+)
+
 add_object_library(
   thread
   SRCS
diff --git a/libc/src/__support/threads/linux/futex_word.h b/libc/src/__support/threads/linux/futex_word.h
index 67159b81b56132..aba650e2906bbe 100644
--- a/libc/src/__support/threads/linux/futex_word.h
+++ b/libc/src/__support/threads/linux/futex_word.h
@@ -9,9 +9,15 @@
 #ifndef LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_FUTEX_WORD_H
 #define LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_FUTEX_WORD_H
 
+#include "include/llvm-libc-types/struct_timespec.h"
+#include "src/__support/CPP/atomic.h"
+#include "src/__support/CPP/limits.h"
+#include "src/__support/CPP/optional.h"
+#include "src/__support/OSUtil/syscall.h"
+#include "src/errno/libc_errno.h"
+#include <linux/futex.h>
 #include <stdint.h>
 #include <sys/syscall.h>
-
 namespace LIBC_NAMESPACE {
 
 // Futexes are 32 bits in size on all platforms, including 64-bit platforms.
@@ -25,6 +31,48 @@ constexpr auto FUTEX_SYSCALL_ID = SYS_futex_time64;
 #error "futex and futex_time64 syscalls not available."
 #endif
 
+// Returns false on timeout, and true in all other cases.
+LIBC_INLINE bool futex_wait(cpp::Atomic<FutexWordType> &futex,
+                            FutexWordType expected,
+                            cpp::optional<::timespec> abs_timeout,
+                            bool is_shared = false) {
+  for (;;) {
+    if (futex.load(cpp::MemoryOrder::RELAXED) != expected)
+      return true;
+    // Use FUTEX_WAIT_BITSET rather than FUTEX_WAIT to be able to give an
+    // absolute time rather than a relative time.
+    long ret =
+        syscall_impl<long>(FUTEX_SYSCALL_ID, &futex,
+                           is_shared ? FUTEX_WAIT_BITSET
+                                     : (FUTEX_WAIT_BITSET | FUTEX_PRIVATE_FLAG),
+                           expected, abs_timeout ? &*abs_timeout : nullptr,
+                           nullptr, FUTEX_BITSET_MATCH_ANY);
+    switch (ret) {
+    case -EINTR:
+      continue;
+    case -ETIMEDOUT:
+      return false;
+    default:
+      return true;
+    }
+  }
+}
+
+LIBC_INLINE bool futex_wake_one(cpp::Atomic<FutexWordType> &futex,
+                                bool is_shared = false) {
+  long ret = syscall_impl<long>(FUTEX_SYSCALL_ID, &futex,
+                                is_shared ? FUTEX_WAKE : FUTEX_WAKE_PRIVATE, 1,
+                                0, 0, 0);
+  return ret > 0;
+}
+
+LIBC_INLINE void futex_wake_all(cpp::Atomic<FutexWordType> &futex,
+                                bool is_shared = false) {
+  syscall_impl<long>(FUTEX_SYSCALL_ID, &futex,
+                     is_shared ? FUTEX_WAKE : FUTEX_WAKE_PRIVATE,
+                     cpp::numeric_limits<int>::max(), 0, 0, 0);
+}
+
 } // namespace LIBC_NAMESPACE
 
 #endif // LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_FUTEX_WORD_H
diff --git a/libc/src/__support/threads/linux/rwlock.h b/libc/src/__support/threads/linux/rwlock.h
new file mode 100644
index 00000000000000..5a1c1f3cb518b2
--- /dev/null
+++ b/libc/src/__support/threads/linux/rwlock.h
@@ -0,0 +1,412 @@
+//===--- Implementation of a Linux rwlock class ------------------*- C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_MUTEX_H
+#define LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_MUTEX_H
+
+#include "src/__support/CPP/atomic.h"
+#include "src/__support/CPP/optional.h"
+#include "src/__support/libc_assert.h"
+#include "src/__support/macros/attributes.h"
+#include "src/__support/threads/linux/futex_word.h"
+
+#include "src/__support/threads/sleep.h"
+#include "src/errno/libc_errno.h"
+#include "src/time/clock_gettime.h"
+
+namespace LIBC_NAMESPACE {
+// A rwlock implementation using futexes.
+// Code is largely based on the Rust standard library implementation.
+// https://github.com/rust-lang/rust/blob/22a5267c83a3e17f2b763279eb24bb632c45dc6b/library/std/src/sys/sync/rwlock/futex.rs
+struct RwLock {
+  // The state consists of a 30-bit reader counter, a 'readers waiting' flag,
+  // and a 'writers waiting' flag. Bits 0..30:
+  //   0: Unlocked
+  //   1..=0x3FFF_FFFE: Locked by N readers
+  //   0x3FFF_FFFF: Write locked
+  // Bit 30: Readers are waiting on this futex.
+  // Bit 31: Writers are waiting on the writer_notify futex.
+  cpp::Atomic<FutexWordType> state;
+  // The 'condition variable' to notify writers through.
+  // Incremented on every signal.
+  cpp::Atomic<FutexWordType> writer_notify;
+  // If the rwlock is shared between processes.
+  bool is_shared;
+
+  LIBC_INLINE_VAR static constexpr FutexWordType READ_LOCKED = 1u;
+  LIBC_INLINE_VAR static constexpr FutexWordType MASK = (1u << 30) - 1u;
+  LIBC_INLINE_VAR static constexpr FutexWordType WRITE_LOCKED = MASK;
+  LIBC_INLINE_VAR static constexpr FutexWordType MAX_READERS = MASK - 1u;
+  LIBC_INLINE_VAR static constexpr FutexWordType READERS_WAITING = 1u << 30u;
+  LIBC_INLINE_VAR static constexpr FutexWordType WRITERS_WAITING = 1u << 31u;
+  LIBC_INLINE_VAR static constexpr int_fast32_t SPIN_LIMIT = 100u;
+  LIBC_INLINE static constexpr bool is_unlocked(FutexWordType state) {
+    return (state & MASK) == 0u;
+  }
+  LIBC_INLINE static constexpr bool is_write_locked(FutexWordType state) {
+    return (state & MASK) == WRITE_LOCKED;
+  }
+  LIBC_INLINE static constexpr bool has_readers_waiting(FutexWordType state) {
+    return (state & READERS_WAITING) != 0u;
+  }
+  LIBC_INLINE static constexpr bool has_writers_waiting(FutexWordType state) {
+    return (state & WRITERS_WAITING) != 0u;
+  }
+  // This also returns false if the counter could overflow if we tried to read
+  // lock it.
+  //
+  // We don't allow read-locking if there's readers waiting, even if the lock is
+  // unlocked and there's no writers waiting. The only situation when this
+  // happens is after unlocking, at which point the unlocking thread might be
+  // waking up writers, which have priority over readers. The unlocking thread
+  // will clear the readers waiting bit and wake up readers, if necessary.
+  LIBC_INLINE static constexpr bool is_read_lockable(FutexWordType state) {
+    return (state & MASK) < MAX_READERS && !has_readers_waiting(state) &&
+           !has_writers_waiting(state);
+  }
+  LIBC_INLINE static constexpr bool
+  has_reached_max_readers(FutexWordType state) {
+    return (state & MASK) == MAX_READERS;
+  }
+
+  // Convert a relative timeout to an absolute timespec.
+  LIBC_INLINE static void abs_timeout(cpp::optional<timespec> &timeout) {
+    if (!timeout)
+      return;
+
+    int errno_backup = libc_errno;
+    timespec now;
+    // if we failed to get time, we move timeout to infinity by setting it to
+    // nullopt.
+    if (LIBC_NAMESPACE::clock_gettime(CLOCK_MONOTONIC, &now) != 0) {
+      timeout = cpp::nullopt;
+      libc_errno = errno_backup;
+      return;
+    }
+    if (__builtin_add_overflow(now.tv_sec, timeout->tv_sec, &timeout->tv_sec)) {
+      timeout = cpp::nullopt;
+      return;
+    }
+    if (__builtin_add_overflow(now.tv_nsec, timeout->tv_nsec,
+                               &timeout->tv_nsec)) {
+      timeout = cpp::nullopt;
+      return;
+    }
+    if (now.tv_nsec >= 1'000'000'000) {
+      if (__builtin_add_overflow(now.tv_sec, 1, &timeout->tv_sec)) {
+        timeout = cpp::nullopt;
+        return;
+      }
+      timeout->tv_nsec -= 1'000'000'000;
+    }
+  }
+
+  template <typename F>
+  LIBC_INLINE bool fetch_update(cpp::Atomic<FutexWordType> &__restrict word,
+                                FutexWordType &__restrict prev,
+                                cpp::MemoryOrder set_order,
+                                cpp::MemoryOrder fetch_order, F &&f) {
+    prev = word.load(fetch_order);
+    // It is okay to have spurious failures here as we are in a loop.
+    while (cpp::optional<FutexWordType> new_val = f(prev))
+      if (word.compare_exchange_weak(prev, *new_val, set_order, fetch_order))
+        return true;
+
+    return false;
+  }
+
+  template <typename F>
+  LIBC_INLINE FutexWordType spin_until(int_fast32_t spin, F &&f) {
+    for (;;) {
+      FutexWordType state = this->state.load(cpp::MemoryOrder::RELAXED);
+      if (f(state) || spin == 0)
+        return state;
+      // pause the pipeline to avoid extra memory loads due to speculation
+      sleep_briefly();
+      --spin;
+    }
+  }
+
+  LIBC_INLINE static bool
+  is_timeout(const cpp::optional<timespec> &abs_timeout) {
+    if (!abs_timeout)
+      return false;
+    timespec now;
+    LIBC_NAMESPACE::clock_gettime(CLOCK_MONOTONIC, &now);
+    if (now.tv_sec > abs_timeout->tv_sec ||
+        (now.tv_sec == abs_timeout->tv_sec &&
+         now.tv_nsec >= abs_timeout->tv_nsec))
+      return true;
+    return false;
+  }
+
+  LIBC_INLINE FutexWordType spin_read(int_fast32_t spin) {
+    // Stop spinning when it's unlocked or read locked, or when there's waiting
+    // threads.
+    return spin_until(spin, [&](FutexWordType state) -> bool {
+      return !is_write_locked(state) || has_readers_waiting(state) ||
+             has_writers_waiting(state) || is_unlocked(state);
+    });
+  }
+
+  LIBC_INLINE FutexWordType spin_write(int_fast32_t spin) {
+    // Stop spinning when it's unlocked or when there's waiting writers, to keep
+    // things somewhat fair.
+    return spin_until(spin, [&](FutexWordType state) -> bool {
+      return is_unlocked(state) || has_writers_waiting(state);
+    });
+  }
+
+  [[gnu::cold]]
+  LIBC_INLINE bool
+  read_contented(cpp::optional<timespec> abs_timeout = cpp::nullopt) {
+    // Notice that the timeout is not checked during the fast spin loop.
+    FutexWordType prev = this->spin_read(SPIN_LIMIT);
+    int lockable_loop = 0;
+    for (;;) {
+      // We have a chance to lock it, go ahead and have a try.
+      if (is_read_lockable(prev)) {
+        if (this->state.compare_exchange_weak(prev, prev + READ_LOCKED,
+                                              cpp::MemoryOrder::ACQUIRE,
+                                              cpp::MemoryOrder::RELAXED))
+          return true;
+        // Timeout is checked the first time we enter the loop since there is
+        // always spin before this.
+        if (lockable_loop == 0 && is_timeout(abs_timeout)) {
+          libc_errno = ETIMEDOUT;
+          return false;
+        }
+        // In the rare case, we somehow stuck in the lockable loop for a long
+        // time. We should give a chance to allow timeout checking. We do this
+        // everytime we meet the SPIN_LIMIT.
+        lockable_loop = (lockable_loop + 1) % SPIN_LIMIT;
+        // we continue to spin if we failed to lock it.
+        // notice that prev is updated by the compare_exchange_weak.
+        continue;
+      }
+
+      if (has_reached_max_readers(prev)) {
+        // The read lock could not be acquired because the maximum
+        // number of read locks for rwlock has been exceeded.
+        libc_errno = EAGAIN;
+        return false;
+      }
+
+      // Make sure the readers waiting bit is set before we go to sleep.
+      // Strong CAS is required as this is a one-time operation.
+      if (!has_readers_waiting(prev) &&
+          !this->state.compare_exchange_strong(prev, prev | READERS_WAITING,
+                                               cpp::MemoryOrder::RELAXED))
+        continue;
+
+      if (!futex_wait(this->state, prev | READERS_WAITING, abs_timeout,
+                      is_shared)) {
+        // timeout happened.
+        libc_errno = ETIMEDOUT;
+        return false;
+      }
+
+      prev = this->spin_read(SPIN_LIMIT);
+      lockable_loop = 0;
+    }
+  }
+
+  /// Wake up a single writer.
+  LIBC_INLINE bool wake_writer() {
+    writer_notify.fetch_add(1, cpp::MemoryOrder::RELEASE);
+    return futex_wake_one(writer_notify);
+  }
+
+  /// Wake up waiting threads after unlocking.
+  ///
+  /// If both are waiting, this will wake up only one writer, but will fall
+  /// back to waking up readers if there was no writer to wake up.
+  [[gnu::cold]]
+  LIBC_INLINE void wake_writer_or_readers(FutexWordType prev) {
+    LIBC_ASSERT(is_unlocked(prev));
+    // The readers waiting bit might be turned on at any point now,
+    // since readers will block when there's anything waiting.
+    // Writers will just lock the lock though, regardless of the waiting bits,
+    // so we don't have to worry about the writer waiting bit.
+    //
+    // If the lock gets locked in the meantime, we don't have to do
+    // anything, because then the thread that locked the lock will take
+    // care of waking up waiters when it unlocks.
+
+    // If only writers are waiting, wake one of them up.
+    if (prev == WRITERS_WAITING)
+      if (this->state.compare_exchange_strong(
+              prev, 0, cpp::MemoryOrder::RELAXED, cpp::MemoryOrder::RELAXED))
+        return (void)wake_writer();
+    // notice that prev is updated.
+
+    // If both writers and readers are waiting, leave the readers waiting
+    // and only wake up one writer.
+    if (prev == (WRITERS_WAITING + READERS_WAITING)) {
+      // set READERS_WAITING and check if someone else locked the lock.
+      if (!this->state.compare_exchange_strong(prev, READERS_WAITING,
+                                               cpp::MemoryOrder::RELAXED,
+                                               cpp::MemoryOrder::RELAXED))
+        return; // someone else locked the lock.
+
+      if (wake_writer())
+        return;
+
+      // No writers were actually blocked on futex_wait, so we continue
+      // to wake up readers instead, since we can't be sure if we notified a
+      // writer.
+      prev = READERS_WAITING;
+    }
+
+    if (prev == READERS_WAITING &&
+        this->state.compare_exchange_strong(prev, 0, cpp::MemoryOrder::RELAXED,
+                                            cpp::MemoryOrder::RELAXED))
+      futex_wake_all(this->state);
+  }
+
+  [[gnu::cold]]
+  LIBC_INLINE bool write_contented(cpp::optional<timespec> abs_timeout) {
+    FutexWordType prev = this->spin_write(SPIN_LIMIT);
+    FutexWordType other_writers_waiting = 0;
+    int lockable_loop = 0;
+    for (;;) {
+      // We have a chance to lock it, go ahead and have a try.
+      if (is_unlocked(prev)) {
+        if (this->state.compare_exchange_weak(
+                prev, prev | WRITE_LOCKED | other_writers_waiting,
+                cpp::MemoryOrder::ACQUIRE, cpp::MemoryOrder::RELAXED))
+          return true;
+        // Timeout is checked the first time we enter the loop since there is
+        // always spin before this.
+        if (lockable_loop == 0 && is_timeout(abs_timeout)) {
+          libc_errno = ETIMEDOUT;
+          return false;
+        }
+        // In the rare case, we somehow stuck in the lockable loop for a long
+        // time. We should give a chance to allow timeout checking. We do this
+        // everytime we meet the SPIN_LIMIT.
+        lockable_loop = (lockable_loop + 1) % SPIN_LIMIT;
+        // we continue to spin if we failed to lock it.
+        // notice that prev is updated by the compare_exchange_weak.
+        continue;
+      }
+
+      // Set the waiting bit indicating that we're waiting on it.
+      if (!has_writers_waiting(prev) &&
+          !this->state.compare_exchange_strong(prev, prev | WRITERS_WAITING,
+                                               cpp::MemoryOrder::RELAXED))
+        continue;
+
+      // Other writers might be waiting now too, so we should make sure
+      // we keep that bit on once we manage lock it.
+      other_writers_waiting = WRITERS_WAITING;
+
+      // Examine the notification counter before we check if `state` has
+      // changed, to make sure we don't miss any notifications.
+      FutexWordType seq = this->writer_notify.load(cpp::MemoryOrder::ACQUIRE);
+
+      // Don't go to sleep if the lock has become available,
+      // or if the writers waiting bit is no longer set.
+      prev = this->state.load(cpp::MemoryOrder::RELAXED);
+      if (is_unlocked(prev) || !has_writers_waiting(prev))
+        continue;
+
+      // Wait for the state to change.
+      if (!futex_wait(this->writer_notify, seq, abs_timeout, is_shared)) {
+        // timeout happened.
+        libc_errno = ETIMEDOUT;
+        return false;
+      }
+
+      // Spin again after waking up.
+      prev = this->spin_write(SPIN_LIMIT);
+      lockable_loop = 0;
+    }
+  }
+
+public:
+  LIBC_INLINE explicit constexpr RwLock(bool shared = false)
+      : state(FutexWordType(0)), writer_notify(FutexWordType(0)),
+        is_shared(shared) {}
+
+  LIBC_INLINE bool try_read() {
+    FutexWordType prev;
+    // set order is the success order for the inner compare_exchange_weak.
+    // it effectively make the store ordering to be relaxed.
+    return fetch_update(this->state, prev, cpp::MemoryOrder::ACQUIRE,
+                        cpp::MemoryOrder::RELAXED,
+                        [](FutexWordType x) -> cpp::optional<FutexWordType> {
+                          if (is_read_lockable(x))
+                            return {x + READ_LOCKED};
+                          return cpp::nullopt;
+                        });
+  }
+
+  LIBC_INLINE bool read(cpp::optional<timespec> timeout = cpp::nullopt) {
+    // record absoulte time if timeout is provided.
+    abs_timeout(timeout);
+    FutexWordType prev = this->state.load(cpp::MemoryOrder::RELAXED);
+    // It is okay to have spurious failures here as we will fallback to
+    // contented routine.
+    if (!is_read_lockable(prev) ||
+        !this->state.compare_exchange_weak(prev, prev + READ_LOCKED,
+                                           cpp::MemoryOrder::ACQUIRE,
+                                           cpp::MemoryOrder::RELAXED))
+      return read_contented(timeout);
+    return true;
+  }
+
+  LIBC_INLINE void read_unlock() {
+    FutexWordType prev =
+        this->state.fetch_sub(READ_LOCKED, cpp::MemoryOrder::RELEASE) -
+        READ_LOCKED;
+    // It's impossible for a reader to be waiting on a read-locked RwLock,
+    // except if there is also a writer waiting.
+    LIBC_ASSERT(!has_readers_waiting(state) || has_writers_waiting(state));
+
+    // Wake up a writer if we were the last reader and there's a writer waiting.
+    if (is_unlocked(prev) && has_writers_waiting(prev))
+      wake_writer_or_readers(prev);
+  }
+  LIBC_INLINE bool try_write() {
+    FutexWordType prev;
+    return fetch_update(this->state, prev, cpp::MemoryOrder::ACQUIRE,
+                        cpp::MemoryOrder::RELAXED,
+                        [](FutexWordType s) -> cpp::optional<FutexWordType> {
+                          if (is_unlocked(s))
+                            return s + WRITE_LOCKED;
+                          return cpp::nullopt;
+                        });
+  }
+  LIBC_INLINE bool write(cpp::optional<timespec> timeout = cpp::nullopt) {
+    // record absoulte time if timeout is provided.
+    abs_timeout(timeout);
+    FutexWordType prev = 0;
+    // It is okay to have spurious failures here as we will fallback to
+    // contented routine.
+    if (!this->state.compare_exchange_weak(prev, WRITE_LOCKED,
+                                           cpp::MemoryOrder::ACQUIRE,
+                                           cpp::MemoryOrder::RELAXED))
+      return write_contented(timeout);
+    return true;
+  }
+  LIBC_INLINE void write_unlock() {
+    FutexWordType prev =
+        this->state.fetch_sub(WRITE_LOCKED, cpp::MemoryOrder::RELEASE) -
+        WRITE_LOCKED;
+
+    LIBC_ASSERT(is_unlocked(prev));
+
+    if (has_writers_waiting(prev) || has_readers_waiting(prev))
+      wake_writer_or_readers(prev);
+  }
+};
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_MUTEX_H
diff --git a/libc/test/src/__support/CMakeLists.txt b/libc/test/src/__support/CMakeLists.txt
index 5d1230f5f3a70f..7c084b5d284643 100644
--- a/libc/test/src/__support/CMakeLists.txt
+++ b/libc/test/src/__support/CMakeLists.txt
@@ -206,3 +206,4 @@ add_subdirectory(OSUtil)
 add_subdirectory(FPUtil)
 add_subdirectory(fixed_point)
 add_subdirectory(HashTable)
+add_subdirectory(threads)
diff --git a/libc/test/src/__support/threads/CMakeLists.txt b/libc/test/src/__support/threads/CMakeLists.txt
new file mode 100644
index 00000000000000..b4bbe81c92ff2e
--- /dev/null
+++ b/libc/test/src/__support/threads/CMakeLists.txt
@@ -0,0 +1,3 @@
+if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${LIBC_TARGET_OS})
+  add_subdirectory(${LIBC_TARGET_OS})
+endif()
diff --git a/libc/test/src/__support/threads/linux/CMakeLists.txt b/libc/test/src/__support/threads/linux/CMakeLists.txt
new file mode 100644
index 00000000000000..ba332dcb0015fe
--- /dev/null
+++ b/libc/test/src/__support/threads/linux/CMakeLists.txt
@@ -0,0 +1,13 @@
+add_custom_target(libc___support_threads_linux_rwlock_unittests)
+add_libc_unittest(
+  rwlock_test
+  SUITE
+    libc___support_threads_linux_rwlock_unittests
+  SRCS
+    rwlock_test.cpp
+  DEPENDS
+    libc.src.__support.threads.linux.rwlock
+    libc.test.UnitTest.ErrnoSetterMatcher
+    libc.src.sys.mman.mmap
+    libc.src.sys.mman.munmap
+)
diff --git a/libc/test/src/__support/threads/linux/rwlock_test.cpp b/libc/test/src/__support/threads/linux/rwlock_test.cpp
new file mode 100644
index 00000000000000..c259cc7f9511f2
--- /dev/null
+++ b/libc/test/src/__support/threads/linux/rwlock_test.cpp
@@ -0,0 +1,172 @@
+#include "src/__support/CPP/atomic.h"
+#include "src/__support/libc_assert.h"
+#include "src/__support/threads/linux/rwlock.h"
+#include "src/errno/libc_errno.h"
+#include "src/sys/mman/mmap.h"
+#include "src/sys/mman/munmap.h"
+#include "test/UnitTest/Test.h"
+#include <linux/mman.h>
+#include <llvm-libc-types/pid_t.h>
+#include <pthread.h>
+extern "C" pid_t fork(void);
+extern "C" [[noreturn]] void exit(int);
+namespace LIBC_NAMESPACE {
+TEST(LlvmLibcRwLock, Smoke) {
+  RwLock l{};
+  l.read();
+  l.read_unlock();
+  l.write();
+  l.write_unlock();
+  l.read();
+  l.read();
+  l.read_unlock();
+  l.read_unlock();
+  l.write();
+  l.write_unlock();
+}
+
+TEST(LlvmLibcRwLock, TrivialReadLock) {
+  RwLock l{};
+  pthread_t t[100];
+  for (int i = 0; i < 100; i++) {
+    pthread_create(
+        &t[i], nullptr,
+        [](void *arg) -> void * {
+          RwLock *l = static_cast<RwLock *>(arg);
+          if (!l->try_read())
+            l->read();
+          l->read_unlock();
+          return nullptr;
+        },
+        &l);
+  }
+  for (int i = 0; i < 100; i++) {
+    pthread_join(t[i], nullptr);
+  }
+}
+
+TEST(LlvmLibcRwLock, TrivialWriteLock) {
+  struct Data {
+    RwLock lock{};
+    cpp::Atomic<int> i = 0;
+  } data;
+
+  pthread_t t[100];
+  for (int i = 0; i < 100; i++) {
+    pthread_create(
+        &t[i], nullptr,
+        [](void *arg) -> void * {
+          Data *data = static_cast<Data *>(arg);
+          data->lock.write();
+          int x = data->i.load(cpp::MemoryOrder::RELAXED);
+          x++;
+          data->i.store(x, cpp::MemoryOrder::RELAXED);
+          data->lock.write_unlock();
+          return nullptr;
+        },
+        &data);
+  }
+  for (int i = 0; i < 100; i++) {
+    pthread_join(t[i], nullptr);
+  }
+  ASSERT_EQ(data.i.load(cpp::MemoryOrder::RELAXED), 100);
+}
+
+TEST(LlvmLibcRwLock, DeadWriteTimeoutRead) {
+  RwLock l{};
+  l.write();
+  pthread_t t[100];
+  for (int i = 0; i < 100; i++) {
+    pthread_create(
+        &t[i], nullptr,
+        [](void *arg) -> void * {
+          RwLock *l = static_cast<RwLock *>(arg);
+          if (l->try_read())
+            __builtin_trap();
+          if (l->try_write())
+            __builtin_trap();
+          ::timespec ts{};
+          ts.tv_sec = 1;
+          if (l->read(ts))
+            __builtin_trap();
+          if (libc_errno != ETIMEDOUT)
+            __builtin_trap();
+          return nullptr;
+        },
+        &l);
+  }
+  for (int i = 0; i < 100; i++) {
+    pthread_join(t[i], nullptr);
+  }
+}
+
+TEST(LlvmLibcRwLock, DeadReadTimeoutWrite) {
+  RwLock l{};
+  l.read();
+  pthread_t t[100];
+  for (int i = 0; i < 100; i++) {
+    pthread_create(
+        &t[i], nullptr,
+        [](void *arg) -> void * {
+          RwLock *l = static_cast<RwLock *>(arg);
+          // l->read_unlock();
+          ::timespec ts{};
+          ts.tv_sec = 1;
+          if (l->try_write())
+            __builtin_trap();
+          if (l->write(ts))
+            __builtin_trap();
+          if (libc_errno != ETIMEDOUT)
+            __builtin_trap();
+          return nullptr;
+        },
+        &l);
+  }
+  for (int i = 0; i < 100; i++) {
+    pthread_join(t[i], nullptr);
+  }
+}
+
+TEST(LlvmLibcRwLock, Fork) {
+  static constexpr struct Data {
+    RwLock lock{true};
+    cpp::Atomic<int> i = 0;
+    cpp::Atomic<int> pending = 2;
+  } DEFAULT;
+
+  void *mmap_addr = mmap(nullptr, sizeof(Data), PROT_READ | PROT_WRITE,
+                         MAP_ANONYMOUS | MAP_SHARED, -1, 0);
+  ASSERT_NE(mmap_addr, MAP_FAILED);
+  Data *data = static_cast<Data *>(mmap_addr);
+  // placement new is not available as we cannot include <new>
+  __builtin_memcpy(data, &DEFAULT, sizeof(Data));
+  int pid = ::fork();
+
+  pthread_t t[100];
+  for (int i = 0; i < 100; i++) {
+    pthread_create(
+        &t[i], nullptr,
+        [](void *arg) -> void * {
+          Data *data = static_cast<Data *>(arg);
+          data->lock.write();
+          int x = data->i.load(cpp::MemoryOrder::RELAXED);
+          x++;
+          data->i.store(x, cpp::MemoryOrder::RELAXED);
+          data->lock.write_unlock();
+          return nullptr;
+        },
+        data);
+  }
+  for (int i = 0; i < 100; i++) {
+    pthread_join(t[i], nullptr);
+  }
+  data->pending.fetch_sub(1);
+  while (data->pending)
+    /* not UB as operations are atomic */;
+  // early exit to avoid pollute the test output
+  if (pid == 0)
+    ::exit(0);
+  ASSERT_EQ(data->i.load(cpp::MemoryOrder::RELAXED), 200);
+  ASSERT_EQ(munmap(mmap_addr, sizeof(Data)), 0);
+}
+} // namespace LIBC_NAMESPACE
\ No newline at end of file

>From 06b04cee62d1338b6613d4477a9322fff8fbbb3e Mon Sep 17 00:00:00 2001
From: Schrodinger ZHU Yifan <yifanzhu at rochester.edu>
Date: Sun, 5 May 2024 16:24:25 -0400
Subject: [PATCH 2/3] [libc] correct format

---
 libc/src/__support/threads/linux/rwlock.h             | 6 +++---
 libc/test/src/__support/threads/linux/rwlock_test.cpp | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/libc/src/__support/threads/linux/rwlock.h b/libc/src/__support/threads/linux/rwlock.h
index 5a1c1f3cb518b2..4098be7d133014 100644
--- a/libc/src/__support/threads/linux/rwlock.h
+++ b/libc/src/__support/threads/linux/rwlock.h
@@ -6,8 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_MUTEX_H
-#define LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_MUTEX_H
+#ifndef LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_RWLOCK_H
+#define LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_RWLOCK_H
 
 #include "src/__support/CPP/atomic.h"
 #include "src/__support/CPP/optional.h"
@@ -409,4 +409,4 @@ struct RwLock {
 
 } // namespace LIBC_NAMESPACE
 
-#endif // LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_MUTEX_H
+#endif // LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_RWLOCK_H
diff --git a/libc/test/src/__support/threads/linux/rwlock_test.cpp b/libc/test/src/__support/threads/linux/rwlock_test.cpp
index c259cc7f9511f2..464bb9218db91b 100644
--- a/libc/test/src/__support/threads/linux/rwlock_test.cpp
+++ b/libc/test/src/__support/threads/linux/rwlock_test.cpp
@@ -169,4 +169,4 @@ TEST(LlvmLibcRwLock, Fork) {
   ASSERT_EQ(data->i.load(cpp::MemoryOrder::RELAXED), 200);
   ASSERT_EQ(munmap(mmap_addr, sizeof(Data)), 0);
 }
-} // namespace LIBC_NAMESPACE
\ No newline at end of file
+} // namespace LIBC_NAMESPACE

>From 462f03d507ebe5a0361fcad88be138d1ea8002c8 Mon Sep 17 00:00:00 2001
From: Schrodinger ZHU Yifan <yifanzhu at rochester.edu>
Date: Sun, 5 May 2024 16:25:38 -0400
Subject: [PATCH 3/3] [libc] adjust check ordering

---
 libc/test/src/__support/threads/linux/rwlock_test.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/libc/test/src/__support/threads/linux/rwlock_test.cpp b/libc/test/src/__support/threads/linux/rwlock_test.cpp
index 464bb9218db91b..7984d92daa68c9 100644
--- a/libc/test/src/__support/threads/linux/rwlock_test.cpp
+++ b/libc/test/src/__support/threads/linux/rwlock_test.cpp
@@ -163,10 +163,10 @@ TEST(LlvmLibcRwLock, Fork) {
   data->pending.fetch_sub(1);
   while (data->pending)
     /* not UB as operations are atomic */;
+  ASSERT_EQ(data->i.load(cpp::MemoryOrder::RELAXED), 200);
   // early exit to avoid pollute the test output
   if (pid == 0)
     ::exit(0);
-  ASSERT_EQ(data->i.load(cpp::MemoryOrder::RELAXED), 200);
   ASSERT_EQ(munmap(mmap_addr, sizeof(Data)), 0);
 }
 } // namespace LIBC_NAMESPACE



More information about the libc-commits mailing list