// Copyright 2020 The IREE Authors
//
// Licensed 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

// NOTE: the best kind of synchronization is no synchronization; always try to
// design your algorithm so that you don't need anything from this file :)
// See https://travisdowns.github.io/blog/2020/07/06/concurrency-costs.html

#ifndef IREE_BASE_THREADING_MUTEX_H_
#define IREE_BASE_THREADING_MUTEX_H_

#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>

#include "iree/base/api.h"
#include "iree/base/internal/atomics.h"
#include "iree/base/threading/futex.h"

// Clang thread-safety analysis is a C++ language extension. The capability
// model requires capabilities to be C++ records with annotated acquire/release
// member functions, and GUARDED_BY() uses C++ name lookup to resolve the
// capability object. These assumptions do not hold in C: the attribute may be
// ignored or mis-parsed by the C frontend, and free-function lock/unlock
// primitives lack the this-based capability tracking that the analysis relies
// on. The analysis therefore only fires meaningfully in C++.
// See https://clang.llvm.org/docs/ThreadSafetyAnalysis.html ("Introduction").
// Enable the annotations when compiling C++ with Clang so that callers can use
// IREE_GUARDED_BY et al. and iree_slim_mutex_t carries the required capability
// annotation.
#if defined(__cplusplus) && defined(__clang__)
#define IREE_THREAD_ANNOTATION_ATTRIBUTE(x) __attribute__((x))
#else
#define IREE_THREAD_ANNOTATION_ATTRIBUTE(x)
#endif  // __cplusplus && __clang__

#ifdef __cplusplus
// Documents if a shared field or global variable needs to be protected by a
// mutex. IREE_GUARDED_BY() allows the user to specify a particular mutex that
// should be held when accessing the annotated variable.
#define IREE_GUARDED_BY(x) IREE_THREAD_ANNOTATION_ATTRIBUTE(guarded_by(x))
#else
#define IREE_GUARDED_BY(x)
#endif  // __cplusplus

#ifdef __cplusplus
// Like IREE_GUARDED_BY but specifies that the contents of a pointer are guarded
// by a mutex instead of the pointer itself.
#define IREE_PTR_GUARDED_BY(x) \
  IREE_THREAD_ANNOTATION_ATTRIBUTE(pt_guarded_by(x))
#else
#define IREE_PTR_GUARDED_BY(x)
#endif  // __cplusplus

// Disables thread-safety analysis for a function. Use on init/destroy functions
// that access mutex-guarded data before/after the mutex is live, where single-
// threaded access is guaranteed by the caller's lifecycle contract.
#define IREE_NO_THREAD_SAFETY_ANALYSIS \
  IREE_THREAD_ANNOTATION_ATTRIBUTE(no_thread_safety_analysis)

// Allow users to fully disable all synchronization for systems that are known
// to never need it. This removes our dependency on pthreads.
#if !IREE_SYNCHRONIZATION_DISABLE_UNSAFE

#if defined(IREE_PLATFORM_APPLE)
#include <os/lock.h>
#endif  // IREE_PLATFORM_APPLE

#if !defined(IREE_PLATFORM_WINDOWS)
#include <pthread.h>
#endif  // !IREE_PLATFORM_WINDOWS

// We have the CRITICAL_SECTION path for now but Slim Reader/Writer lock (SRW)
// is much better (and what std::mutex uses). SRW doesn't spin, though, and has
// some other implications that don't quite line up with pthread_mutex_t on most
// platforms. Once we have larger end-to-end benchmarks we should choose based
// on workloads.
#define IREE_MUTEX_USE_WIN32_SRW 1

#endif  // !IREE_SYNCHRONIZATION_DISABLE_UNSAFE

#ifdef __cplusplus
extern "C" {
#endif

//==============================================================================
// iree_mutex_t
//==============================================================================

// A normal fat mutex (ala std::mutex).
// This may be implemented as a slim mutex on certain platforms but in the worst
// case will be the native platform primitive (like pthread_mutex_t) and as such
// should not be embedded in structures meant to be kept small.
//
// Windows: Slim Reader/Writer (SRW) Locks
// All others: pthread_mutex_t
typedef struct iree_mutex_t {
#if IREE_SYNCHRONIZATION_DISABLE_UNSAFE
  int reserved;
#elif defined(IREE_PLATFORM_WINDOWS) && defined(IREE_MUTEX_USE_WIN32_SRW)
  SRWLOCK value;
#elif defined(IREE_PLATFORM_WINDOWS)
  CRITICAL_SECTION value;
#else
  pthread_mutex_t value;
#endif  // IREE_PLATFORM_*
#if (IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_SLOW_LOCKS)
  uint32_t lock_id;
#endif  // IREE_TRACING_FEATURE_SLOW_LOCKS
} IREE_THREAD_ANNOTATION_ATTRIBUTE(capability("mutex")) iree_mutex_t;

#if (IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_SLOW_LOCKS)
// Initializes |out_mutex| to the well-defined unlocked contents.
// Must be called prior to using any other iree_mutex_* method.
#define iree_mutex_initialize(out_mutex)                                      \
  static const iree_tracing_location_t TracyConcat(                           \
      __tracy_source_location, __LINE__) = {NULL, __FUNCTION__, __FILE__,     \
                                            (uint32_t)__LINE__, 0};           \
  iree_mutex_initialize_impl(&TracyConcat(__tracy_source_location, __LINE__), \
                             out_mutex);
IREE_API_EXPORT void iree_mutex_initialize_impl(
    const iree_tracing_location_t* src_loc, iree_mutex_t* out_mutex);
#else
// Initializes |out_mutex| to the well-defined unlocked contents.
// Must be called prior to using any other iree_mutex_* method.
IREE_API_EXPORT void iree_mutex_initialize(iree_mutex_t* out_mutex);
#endif  // IREE_TRACING_FEATURE_SLOW_LOCKS

// Deinitializes |mutex| (after a prior call to iree_mutex_initialize).
// The mutex must not be held by any thread.
IREE_API_EXPORT void iree_mutex_deinitialize(iree_mutex_t* mutex)
    IREE_THREAD_ANNOTATION_ATTRIBUTE(locks_excluded(mutex));

// Locks the |mutex| and returns when held by the caller.
IREE_API_EXPORT void iree_mutex_lock(iree_mutex_t* mutex)
    IREE_THREAD_ANNOTATION_ATTRIBUTE(acquire_capability(mutex));

// Tries to lock the |mutex| and returns true if the caller holds the lock.
IREE_API_EXPORT bool iree_mutex_try_lock(iree_mutex_t* mutex)
    IREE_THREAD_ANNOTATION_ATTRIBUTE(try_acquire_capability(true, mutex));

// Unlocks the |mutex|, which must be held by the caller.
IREE_API_EXPORT void iree_mutex_unlock(iree_mutex_t* mutex)
    IREE_THREAD_ANNOTATION_ATTRIBUTE(release_capability(mutex));

//==============================================================================
// iree_slim_mutex_t
//==============================================================================

// A lightweight unfair lock.
// Depending on platform this is significantly smaller than a mutex (4-8 bytes
// vs 64+ bytes), can always be statically initialized/requires no allocations,
// and performs the minimal amount of work possible while still playing nicely
// with the OS thread scheduler.
//
// Unlike a full mutex these don't have the ability to be shared across
// processes (not something we care about), don't have a way to define timeouts,
// and have only a binary held/unheld state. They are often an order of
// magnitude faster in uncontended/lightly-contended code and the same
// performance in highly-contended code, though, so it's worth it for locks that
// be guarding small data structures (queue pointers, etc) and touched from many
// threads. Since they are so lightweight it's possible to embed them per-object
// instead of per-manager and change from a single highly-contended lock to
// thousands of almost completely uncontended slim locks.
//
// Though these locks have lightweight return paths (using atomics ops, see
// the below paragraph, with or without spinning), they always have a
// heavyweight fallback path that ends up calling into the kernel to properly
// let the thread wait. This is critical to avoid pathological cases under
// contention and allowing for thread priority inheritance when there are
// multiple threads competing that may otherwise be scheduled in a potentially
// livelocking order.
//
// The "unfair" here comes from the fact that it's possible on certain platforms
// for certain threads to never be able to acquire the lock in cases of
// extremely high contention or widely disparate thread priority levels. This is
// mitigated by ensuring only very small regions of code are guarded and that
// there's enough work happening outside of the lock on any particular thread to
// ensure that there's some chance of other threads being able to acquire it.
//
// Notes on weakly-ordered atomics (for the lightweight return paths)
// ------------------------------------------------------------------
//
// The lightweight return paths (avoiding waiting in the kernel) are typically
// implemented using acquire/release weakly-ordered atomics. When these code
// paths are taken:
//
//   * iree_slim_mutex_lock is a read-modify-write operation on the
//     mutex with memory_order_acquire.
//   * iree_slim_mutex_unlock is a read-modify-write operation on the
//     mutex with memory_order_release.
//
// This means the following guarantee for the caller of iree_slim_mutex_lock:
//
//   When iree_slim_mutex_lock returns on this thread T1 from having waited on
//   another thread T2 calling iree_slim_mutex_unlock, all memory read and write
//   operations performed on thread T2 prior to calling iree_slim_mutex_unlock
//   are guaranteed to "happen-before" any subsequent memory read or write on
//   thread T1.
//
// This is meant to be an implementation detail within the standard contract
// for anything called a "mutex". The C++ standard is a good reference for what
// it means to be a mutex in the context of the C11/C++11 memory model:
// https://eel.is/c++draft/thread.mutex#requirements.mutex
//
// OS-specific implementation aspects (for the heavyweight return paths)
// ---------------------------------------------------------------------
//
// MacOS/iOS: os_unfair_lock
//   Spins and after a short backoff drops to a futex-like behavior of waiting
//   in the kernel. Unfortunately real futexes aren't supported.
// See:
//   https://developer.apple.com/documentation/os/synchronization
//   https://opensource.apple.com/source/libplatform/libplatform-125/src/os/lock.c.auto.html
//
// Wasm: memory.atomic.wait32/memory.atomic.notify
//   Uses Wasm thread builtins (__builtin_wasm_memory_atomic_wait32 and
//   __builtin_wasm_memory_atomic_notify) which map to Atomics.wait/notify
//   in the JS host. Spins and drops to a futex-like wait.
// See:
//   https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Atomics/wait
//   https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Atomics/notify
//
// Windows: WaitOnAddress/WakeByAddress*
//   Spins and after a short backoff drops to a futex and waits in the kernel.
// See:
//   https://docs.microsoft.com/en-us/windows/win32/api/synchapi/nf-synchapi-waitonaddress
//   https://devblogs.microsoft.com/oldnewthing/20170601-00/?p=96265
//
// Linux/Android/others: futex
//   Spins and after a short backoff drops to a futex and waits in the kernel.
// See:
//   http://locklessinc.com/articles/futex_cheat_sheet/
//   https://man7.org/linux/man-pages/man2/futex.2.html
//   https://eli.thegreenplace.net/2018/basics-of-futexes/
//   https://bartoszmilewski.com/2008/09/01/thin-lock-vs-futex/
typedef struct iree_slim_mutex_t {
#if IREE_SYNCHRONIZATION_DISABLE_UNSAFE
  int reserved;
#elif (IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_FAST_LOCKS)
  iree_mutex_t impl;  // re-route to slow mutex
#elif defined(IREE_PLATFORM_APPLE)
  os_unfair_lock value;
#elif defined(IREE_PLATFORM_WINDOWS) && defined(IREE_MUTEX_USE_WIN32_SRW)
  SRWLOCK value;
#elif defined(IREE_RUNTIME_USE_FUTEX)
  iree_atomic_int32_t value;
#else
  iree_mutex_t impl;  // fallback
#endif  // IREE_PLATFORM_*
} IREE_THREAD_ANNOTATION_ATTRIBUTE(capability("mutex")) iree_slim_mutex_t;

#if (IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_FAST_LOCKS)
// Initializes |out_mutex| to the well-defined unlocked contents.
// Must be called prior to using any other iree_slim_mutex_* method.
#define iree_slim_mutex_initialize(out_mutex)                             \
  static const iree_tracing_location_t TracyConcat(                       \
      __tracy_source_location, __LINE__) = {NULL, __FUNCTION__, __FILE__, \
                                            (uint32_t)__LINE__, 0};       \
  iree_slim_mutex_initialize_impl(                                        \
      &TracyConcat(__tracy_source_location, __LINE__), out_mutex);
IREE_API_EXPORT void iree_slim_mutex_initialize_impl(
    const iree_tracing_location_t* src_loc, iree_slim_mutex_t* out_mutex);
#else
// Initializes |out_mutex| to the well-defined unlocked contents.
// Must be called prior to using any other iree_slim_mutex_* method.
//
// Though optional (static initialization is fine) this is required to support
// lock tracing. Assume it's (mostly) free and always call it if possible. This
// also allows us to swap in a non-slim lock for enhanced debugging if we run
// into threading issues.
IREE_API_EXPORT void iree_slim_mutex_initialize(iree_slim_mutex_t* out_mutex);
#endif  // IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_FAST_LOCKS

// Deinitializes |mutex| (after a prior call to iree_slim_mutex_initialize).
// The mutex must not be held by any thread.
IREE_API_EXPORT void iree_slim_mutex_deinitialize(iree_slim_mutex_t* mutex)
    IREE_THREAD_ANNOTATION_ATTRIBUTE(locks_excluded(mutex));

// Locks the |mutex| and returns when held by the caller.
IREE_API_EXPORT void iree_slim_mutex_lock(iree_slim_mutex_t* mutex)
    IREE_THREAD_ANNOTATION_ATTRIBUTE(acquire_capability(mutex));

// Tries to lock the |mutex| and returns true if the caller holds the lock.
//
// This function is allowed to fail spuriously, i.e. even if the lock isn't
// held by another thread.
IREE_API_EXPORT bool iree_slim_mutex_try_lock(iree_slim_mutex_t* mutex)
    IREE_THREAD_ANNOTATION_ATTRIBUTE(try_acquire_capability(true, mutex));

// Unlocks the |mutex|, which must be held by the caller.
IREE_API_EXPORT void iree_slim_mutex_unlock(iree_slim_mutex_t* mutex)
    IREE_THREAD_ANNOTATION_ATTRIBUTE(release_capability(mutex));

#ifdef __cplusplus
}  // extern "C"
#endif

#endif  // IREE_BASE_THREADING_MUTEX_H_
