From 23533477cbf2282582867534bb9185fb80271327 Mon Sep 17 00:00:00 2001
From: Amirsaman Memaripour <amirsaman.memaripour@10gen.com>
Date: Thu, 16 Jan 2020 20:44:59 +0000
Subject: SERVER-45212 Add try_lock() to SpinLock

---
 src/mongo/db/client.h                         |  3 +
 src/mongo/util/concurrency/spin_lock.h        | 30 +++++++---
 src/mongo/util/concurrency/spin_lock_test.cpp | 79 +++++++++++++++++++--------
 3 files changed, 81 insertions(+), 31 deletions(-)

diff --git a/src/mongo/db/client.h b/src/mongo/db/client.h
index 6aea69b5c2a..64c0e7407c7 100644
--- a/src/mongo/db/client.h
+++ b/src/mongo/db/client.h
@@ -152,6 +152,9 @@ public:
     void unlock() {
         _lock.unlock();
     }
+    bool try_lock() {
+        return _lock.try_lock();
+    }
 
     /**
      * Makes a new operation context representing an operation on this client.  At most
diff --git a/src/mongo/util/concurrency/spin_lock.h b/src/mongo/util/concurrency/spin_lock.h
index 5c5a17b4b74..380ad9f7c4e 100644
--- a/src/mongo/util/concurrency/spin_lock.h
+++ b/src/mongo/util/concurrency/spin_lock.h
@@ -42,7 +42,7 @@
 namespace mongo {
 
 #if defined(_WIN32)
-class SpinLock {
+class SpinLock : public Latch {
     SpinLock(const SpinLock&) = delete;
     SpinLock& operator=(const SpinLock&) = delete;
 
@@ -55,14 +55,18 @@ public:
         DeleteCriticalSection(&_cs);
     }
 
-    void lock() {
+    void lock() override {
         EnterCriticalSection(&_cs);
     }
 
-    void unlock() {
+    void unlock() override {
         LeaveCriticalSection(&_cs);
     }
 
+    bool try_lock() override {
+        return TryEnterCriticalSection(&_cs);
+    }
+
 private:
     CRITICAL_SECTION _cs;
 };
@@ -70,44 +74,52 @@ private:
 #else
 
 #if MONGO_CONFIG_DEBUG_BUILD
-class SpinLock {
+class SpinLock : public Latch {
     SpinLock(const SpinLock&) = delete;
     SpinLock& operator=(const SpinLock&) = delete;
 
 public:
     SpinLock() = default;
 
-    void lock() {
+    void lock() override {
         _mutex.lock();
     }
 
-    void unlock() {
+    void unlock() override {
         _mutex.unlock();
     }
 
+    bool try_lock() override {
+        return _mutex.try_lock();
+    }
+
 private:
     stdx::mutex _mutex;  // NOLINT
 };
 
 #else
 
-class SpinLock {
+class SpinLock : public Latch {
     SpinLock(const SpinLock&) = delete;
     SpinLock& operator=(const SpinLock&) = delete;
 
 public:
     SpinLock() = default;
 
-    void unlock() {
+    void unlock() override {
         _locked.clear(std::memory_order_release);
     }
 
-    void lock() {
+    void lock() override {
         if (MONGO_likely(_tryLock()))
             return;
         _lockSlowPath();
     }
 
+    bool try_lock() override {
+        return _tryLock();
+    }
+
 private:
     bool _tryLock() {
         bool wasLocked = _locked.test_and_set(std::memory_order_acquire);
diff --git a/src/mongo/util/concurrency/spin_lock_test.cpp b/src/mongo/util/concurrency/spin_lock_test.cpp
index 62266221b2a..f9f760604ff 100644
--- a/src/mongo/util/concurrency/spin_lock_test.cpp
+++ b/src/mongo/util/concurrency/spin_lock_test.cpp
@@ -47,7 +47,7 @@ class LockTester {
 public:
     LockTester(SpinLock* spin, int* counter) : _spin(spin), _counter(counter), _requests(0) {}
 
-    ~LockTester() {
+    virtual ~LockTester() {
         delete _t;
     }
 
@@ -72,7 +72,7 @@ private:
 
     void test(int increments) {
         while (increments-- > 0) {
-            _spin->lock();
+            acquireLock();
             ++(*_counter);
             ++_requests;
             _spin->unlock();
@@ -81,35 +81,70 @@ private:
 
     LockTester(LockTester&);
     LockTester& operator=(LockTester&);
+
+protected:
+    SpinLock* spin() const {
+        return _spin;
+    }
+
+    virtual void acquireLock() {
+        spin()->lock();
+    }
 };
 
+class TryLockTester final : public LockTester {
+public:
+    TryLockTester(SpinLock* spin, int* counter) : LockTester(spin, counter) {}
 
-TEST(Concurrency, ConcurrentIncs) {
-    SpinLock spin;
-    int counter = 0;
+protected:
+    void acquireLock() override {
+        while (!spin()->try_lock()) {
+        }
+    }
+};
 
-    const int threads = 64;
-    const int incs = 50000;
-    LockTester* testers[threads];
+class SpinLockTester {
+public:
+    template <class T>
+    void run(std::string testName) {
+        SpinLock spin;
+        int counter = 0;
 
-    Timer timer;
+        std::vector<T*> testers(_threads);
+        Timer timer;
 
-    for (int i = 0; i < threads; i++) {
-        testers[i] = new LockTester(&spin, &counter);
-    }
-    for (int i = 0; i < threads; i++) {
-        testers[i]->start(incs);
-    }
-    for (int i = 0; i < threads; i++) {
-        testers[i]->join();
-        ASSERT_EQUALS(testers[i]->requests(), incs);
-        delete testers[i];
+        for (int i = 0; i < _threads; i++) {
+            testers[i] = new T(&spin, &counter);
+        }
+        for (int i = 0; i < _threads; i++) {
+            testers[i]->start(_incs);
+        }
+        for (int i = 0; i < _threads; i++) {
+            testers[i]->join();
+            ASSERT_EQUALS(testers[i]->requests(), _incs);
+            delete testers[i];
+        }
+
+        int ms = timer.millis();
+        mongo::unittest::log() << "spinlock " << testName << " time: " << ms << std::endl;
+
+        ASSERT_EQUALS(counter, _threads * _incs);
     }
 
-    int ms = timer.millis();
-    mongo::unittest::log() << "spinlock ConcurrentIncs time: " << ms << std::endl;
+private:
+    const int _threads = 64;
+    const int _incs = 50000;
+};
+
+
+TEST(Concurrency, ConcurrentIncs) {
+    SpinLockTester tester;
+    tester.run<LockTester>("ConcurrentIncs");
+}
 
-    ASSERT_EQUALS(counter, threads * incs);
+TEST(Concurrency, ConcurrentIncsWithTryLock) {
+    SpinLockTester tester;
+    tester.run<TryLockTester>("ConcurrentIncsWithTryLock");
 }
 
 }  // namespace
-- 
cgit v1.2.1