From 44d3918e4f18e4e65233fef805aa03e7209d5815 Mon Sep 17 00:00:00 2001
From: "Justin C. Miller" <justin@devjustinian.com>
Date: Sat, 15 Jan 2022 17:51:55 -0800
Subject: [PATCH] [util] Add try_aquire to spinlock

Also added a scoped_trylock class that mirrors scoped_lock, but uses
try_aquire and has a scoped_lock::locked() accessor.
---
 src/libraries/util/include/util/spinlock.h | 32 +++++++++++++++++++++-
 src/libraries/util/spinlock.cpp            | 11 ++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/src/libraries/util/include/util/spinlock.h b/src/libraries/util/include/util/spinlock.h
index d9fc2d4..f22b19d 100644
--- a/src/libraries/util/include/util/spinlock.h
+++ b/src/libraries/util/include/util/spinlock.h
@@ -17,8 +17,10 @@ public:
     {
         bool blocked;
         waiter *next;
+        char const *where;
     };
 
+    bool try_acquire(waiter *w);
     void acquire(waiter *w);
     void release(waiter *w);
 
@@ -30,7 +32,10 @@ private:
 class scoped_lock
 {
 public:
-    inline scoped_lock(spinlock &lock) : m_lock(lock) {
+    inline scoped_lock(
+            spinlock &lock,
+            const char *where = __builtin_FUNCTION()
+            ) : m_lock(lock), m_waiter {false, nullptr, where} {
         m_lock.acquire(&m_waiter);
     }
 
@@ -43,4 +48,29 @@ private:
     spinlock::waiter m_waiter;
 };
 
+/// Scoped lock that owns a spinlock::waiter and calls
+/// spinlock::try_acquire
+class scoped_trylock
+{
+public:
+    inline scoped_trylock(
+            spinlock &lock,
+            const char *where = __builtin_FUNCTION()
+            ) : m_lock(lock), m_waiter {false, nullptr, where} {
+        m_locked = m_lock.try_acquire(&m_waiter);
+    }
+
+    inline ~scoped_trylock() {
+        if (m_locked)
+            m_lock.release(&m_waiter);
+    }
+
+    inline bool locked() const { return m_locked; }
+
+private:
+    bool m_locked;
+    spinlock &m_lock;
+    spinlock::waiter m_waiter;
+};
+
 } // namespace util
diff --git a/src/libraries/util/spinlock.cpp b/src/libraries/util/spinlock.cpp
index 4b62c25..5b54029 100644
--- a/src/libraries/util/spinlock.cpp
+++ b/src/libraries/util/spinlock.cpp
@@ -7,6 +7,17 @@ static constexpr int memorder = __ATOMIC_SEQ_CST;
 spinlock::spinlock() : m_lock {nullptr} {}
 spinlock::~spinlock() {}
 
+bool
+spinlock::try_acquire(waiter *w)
+{
+    w->next = nullptr;
+    w->blocked = false;
+
+    // Point this lock at the waiter only if it's empty
+    waiter *expected = nullptr;
+    return __atomic_compare_exchange_n(&m_lock, &expected, w, false, memorder, memorder);
+}
+
 void
 spinlock::acquire(waiter *w)
 {