From bf8a9074c904fd95d9b6e37e27e6dce7ecf17699 Mon Sep 17 00:00:00 2001
From: Andrii Grynenko <andrii@fb.com>
Date: Thu, 10 Dec 2015 18:16:51 -0800
Subject: [PATCH] Fix RefCountTest and RCURefCount race

Reviewed By: alikhtarov

Differential Revision: D2741459

fb-gh-sync-id: c4bd068cf735ae25364edba40960096fb35e8c43
---
 folly/experimental/RCURefCount.h         | 10 ++++++----
 folly/experimental/test/RefCountTest.cpp |  4 +++-
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/folly/experimental/RCURefCount.h b/folly/experimental/RCURefCount.h
index 567006ea..3814f819 100644
--- a/folly/experimental/RCURefCount.h
+++ b/folly/experimental/RCURefCount.h
@@ -40,12 +40,13 @@ class RCURefCount {
     auto& localCount = *localCount_;
 
     std::lock_guard<RCUReadLock> lg(RCUReadLock::instance());
+    auto state = state_.load();
 
-    if (LIKELY(state_ == State::LOCAL)) {
+    if (LIKELY(state == State::LOCAL)) {
       ++localCount;
 
       return 42;
-    } else if (state_ == State::GLOBAL_TRANSITION) {
+    } else if (state == State::GLOBAL_TRANSITION) {
       ++globalCount_;
 
       return 42;
@@ -67,15 +68,16 @@ class RCURefCount {
     auto& localCount = *localCount_;
 
     std::lock_guard<RCUReadLock> lg(RCUReadLock::instance());
+    auto state = state_.load();
 
-    if (LIKELY(state_ == State::LOCAL)) {
+    if (LIKELY(state == State::LOCAL)) {
       --localCount;
 
       return 42;
     } else {
       auto value = --globalCount_;
 
-      if (state_ == State::GLOBAL) {
+      if (state == State::GLOBAL) {
         assert(value >= 0);
         return value;
       } else {
diff --git a/folly/experimental/test/RefCountTest.cpp b/folly/experimental/test/RefCountTest.cpp
index 48ad80d9..6cd72474 100644
--- a/folly/experimental/test/RefCountTest.cpp
+++ b/folly/experimental/test/RefCountTest.cpp
@@ -69,7 +69,9 @@ void basicTest() {
   b.wait();
 
   count.useGlobal();
-  --count;
+  if (--count == 0) {
+    ++got0;
+  }
 
   for (auto& t: ts) {
     t.join();
-- 
2.34.1