libc/src/__support/threads/linux/raw_mutex.h

   1 //===--- Implementation of a Linux RawMutex class ---------------*- C++ -*-===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 #ifndef LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_RAW_MUTEX_H
   9 #define LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_RAW_MUTEX_H
  10
  11 #include "src/__support/CPP/optional.h"
  12 #include "src/__support/common.h"
  13 #include "src/__support/libc_assert.h"
  14 #include "src/__support/macros/attributes.h"
  15 #include "src/__support/macros/config.h"
  16 #include "src/__support/macros/optimization.h"
  17 #include "src/__support/threads/linux/futex_utils.h"
  18 #include "src/__support/threads/linux/futex_word.h"
  19 #include "src/__support/threads/sleep.h"
  20 #include "src/__support/time/linux/abs_timeout.h"
  21
  22 #ifndef LIBC_COPT_TIMEOUT_ENSURE_MONOTONICITY
  23 #define LIBC_COPT_TIMEOUT_ENSURE_MONOTONICITY 1
  24 #endif
  25
  26 #if LIBC_COPT_TIMEOUT_ENSURE_MONOTONICITY
  27 #include "src/__support/time/linux/monotonicity.h"
  28 #endif
  29
  30 #ifndef LIBC_COPT_RAW_MUTEX_DEFAULT_SPIN_COUNT
  31 #define LIBC_COPT_RAW_MUTEX_DEFAULT_SPIN_COUNT 100
  32 #endif
  33
  34 namespace LIBC_NAMESPACE_DECL {
  35 // Lock is a simple timable lock for internal usage.
  36 // This is separated from Mutex because this one does not need to consider
  37 // robustness and reentrancy. Also, this one has spin optimization for shorter
  38 // critical sections.
  39 class RawMutex {
  40 protected:
  41   Futex futex;
  42   LIBC_INLINE_VAR static constexpr FutexWordType UNLOCKED = 0b00;
  43   LIBC_INLINE_VAR static constexpr FutexWordType LOCKED = 0b01;
  44   LIBC_INLINE_VAR static constexpr FutexWordType IN_CONTENTION = 0b10;
  45
  46 private:
  47   LIBC_INLINE FutexWordType spin(unsigned spin_count) {
  48     FutexWordType result;
  49     for (;;) {
  50       result = futex.load(cpp::MemoryOrder::RELAXED);
  51       // spin until one of the following conditions is met:
  52       // - the mutex is unlocked
  53       // - the mutex is in contention
  54       // - the spin count reaches 0
  55       if (result != LOCKED || spin_count == 0u)
  56         return result;
  57       // Pause the pipeline to avoid extraneous memory operations due to
  58       // speculation.
  59       sleep_briefly();
  60       spin_count--;
  61     };
  62   }
  63
  64   // Return true if the lock is acquired. Return false if timeout happens before
  65   // the lock is acquired.
  66   LIBC_INLINE bool lock_slow(cpp::optional<Futex::Timeout> timeout,
  67                              bool is_pshared, unsigned spin_count) {
  68     FutexWordType state = spin(spin_count);
  69     // Before go into contention state, try to grab the lock.
  70     if (state == UNLOCKED &&
  71         futex.compare_exchange_strong(state, LOCKED, cpp::MemoryOrder::ACQUIRE,
  72                                       cpp::MemoryOrder::RELAXED))
  73       return true;
  74 #if LIBC_COPT_TIMEOUT_ENSURE_MONOTONICITY
  75     /* ADL should kick in */
  76     if (timeout)
  77       ensure_monotonicity(*timeout);
  78 #endif
  79     for (;;) {
  80       // Try to grab the lock if it is unlocked. Mark the contention flag if it
  81       // is locked.
  82       if (state != IN_CONTENTION &&
  83           futex.exchange(IN_CONTENTION, cpp::MemoryOrder::ACQUIRE) == UNLOCKED)
  84         return true;
  85       // Contention persists. Park the thread and wait for further notification.
  86       if (ETIMEDOUT == -futex.wait(IN_CONTENTION, timeout, is_pshared))
  87         return false;
  88       // Continue to spin after waking up.
  89       state = spin(spin_count);
  90     }
  91   }
  92
  93   LIBC_INLINE void wake(bool is_pshared) { futex.notify_one(is_pshared); }
  94
  95 public:
  96   LIBC_INLINE static void init(RawMutex *mutex) { mutex->futex = UNLOCKED; }
  97   LIBC_INLINE constexpr RawMutex() : futex(UNLOCKED) {}
  98   [[nodiscard]] LIBC_INLINE bool try_lock() {
  99     FutexWordType expected = UNLOCKED;
 100     // Use strong version since this is a one-time operation.
 101     return futex.compare_exchange_strong(
 102         expected, LOCKED, cpp::MemoryOrder::ACQUIRE, cpp::MemoryOrder::RELAXED);
 103   }
 104   LIBC_INLINE bool
 105   lock(cpp::optional<Futex::Timeout> timeout = cpp::nullopt,
 106        bool is_shared = false,
 107        unsigned spin_count = LIBC_COPT_RAW_MUTEX_DEFAULT_SPIN_COUNT) {
 108     // Timeout will not be checked if immediate lock is possible.
 109     if (LIBC_LIKELY(try_lock()))
 110       return true;
 111     return lock_slow(timeout, is_shared, spin_count);
 112   }
 113   LIBC_INLINE bool unlock(bool is_pshared = false) {
 114     FutexWordType prev = futex.exchange(UNLOCKED, cpp::MemoryOrder::RELEASE);
 115     // if there is someone waiting, wake them up
 116     if (LIBC_UNLIKELY(prev == IN_CONTENTION))
 117       wake(is_pshared);
 118     // Detect invalid unlock operation.
 119     return prev != UNLOCKED;
 120   }
 121   LIBC_INLINE void static destroy([[maybe_unused]] RawMutex *lock) {
 122     LIBC_ASSERT(lock->futex == UNLOCKED && "Mutex destroyed while used.");
 123   }
 124   LIBC_INLINE Futex &get_raw_futex() { return futex; }
 125   LIBC_INLINE void reset() { futex = UNLOCKED; }
 126 };
 127 } // namespace LIBC_NAMESPACE_DECL
 128
 129 #endif // LLVM_LIBC_SRC___SUPPORT_THREADS_LINUX_RAW_MUTEX_H