folly/test/SpinLockTest.cpp

   1 /*
   2  * Copyright 2017 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include <folly/SpinLock.h>
  18
  19 #include <folly/Random.h>
  20
  21 #include <thread>
  22
  23 #include <folly/portability/Asm.h>
  24 #include <folly/portability/GTest.h>
  25
  26 using folly::SpinLockGuardImpl;
  27
  28 namespace {
  29
  30 template <typename LOCK>
  31 struct LockedVal {
  32   int ar[1024];
  33   LOCK lock;
  34
  35   LockedVal() {
  36     memset(ar, 0, sizeof ar);
  37   }
  38 };
  39
  40 template <typename LOCK>
  41 void spinlockTestThread(LockedVal<LOCK>* v) {
  42   const int max = 1000;
  43   auto rng = folly::ThreadLocalPRNG();
  44   for (int i = 0; i < max; i++) {
  45     folly::asm_pause();
  46     SpinLockGuardImpl<LOCK> g(v->lock);
  47
  48     int first = v->ar[0];
  49     for (size_t j = 1; j < sizeof v->ar / sizeof j; ++j) {
  50       EXPECT_EQ(first, v->ar[j]);
  51     }
  52
  53     int byte = folly::Random::rand32(rng);
  54     memset(v->ar, char(byte), sizeof v->ar);
  55   }
  56 }
  57
  58 template <typename LOCK>
  59 struct TryLockState {
  60   LOCK lock1;
  61   LOCK lock2;
  62   bool locked{false};
  63   uint64_t obtained{0};
  64   uint64_t failed{0};
  65 };
  66
  67 template <typename LOCK>
  68 void trylockTestThread(TryLockState<LOCK>* state, size_t count) {
  69   while (true) {
  70     folly::asm_pause();
  71     SpinLockGuardImpl<LOCK> g(state->lock1);
  72     if (state->obtained >= count) {
  73       break;
  74     }
  75
  76     bool ret = state->lock2.try_lock();
  77     EXPECT_NE(state->locked, ret);
  78
  79     if (ret) {
  80       // We got lock2.
  81       ++state->obtained;
  82       state->locked = true;
  83
  84       // Release lock1 and wait until at least one other thread fails to
  85       // obtain the lock2 before continuing.
  86       auto oldFailed = state->failed;
  87       while (state->failed == oldFailed && state->obtained < count) {
  88         state->lock1.unlock();
  89         folly::asm_pause();
  90         state->lock1.lock();
  91       }
  92
  93       state->locked = false;
  94       state->lock2.unlock();
  95     } else {
  96       ++state->failed;
  97     }
  98   }
  99 }
 100
 101 template <typename LOCK>
 102 void correctnessTest() {
 103   int nthrs = sysconf(_SC_NPROCESSORS_ONLN) * 2;
 104   std::vector<std::thread> threads;
 105   LockedVal<LOCK> v;
 106   for (int i = 0; i < nthrs; ++i) {
 107     threads.push_back(std::thread(spinlockTestThread<LOCK>, &v));
 108   }
 109   for (auto& t : threads) {
 110     t.join();
 111   }
 112 }
 113
 114 template <typename LOCK>
 115 void trylockTest() {
 116   int nthrs = sysconf(_SC_NPROCESSORS_ONLN) + 4;
 117   std::vector<std::thread> threads;
 118   TryLockState<LOCK> state;
 119   size_t count = 100;
 120   for (int i = 0; i < nthrs; ++i) {
 121     threads.push_back(std::thread(trylockTestThread<LOCK>, &state, count));
 122   }
 123   for (auto& t : threads) {
 124     t.join();
 125   }
 126
 127   EXPECT_EQ(count, state.obtained);
 128   // Each time the code obtains lock2 it waits for another thread to fail
 129   // to acquire it.  The only time this might not happen is on the very last
 130   // loop when no other threads are left.
 131   EXPECT_GE(state.failed + 1, state.obtained);
 132 }
 133
 134 } // unnamed namespace
 135
 136 TEST(SpinLock, Correctness) {
 137   correctnessTest<folly::SpinLock>();
 138 }
 139 TEST(SpinLock, TryLock) {
 140   trylockTest<folly::SpinLock>();
 141 }