fix use-after-free in addFunctionOnce

[folly.git] / folly / experimental / test / RefCountTest.cpp

diff --git a/folly/experimental/test/RefCountTest.cpp b/folly/experimental/test/RefCountTest.cpp
index b9ef475bbdda32c8c6a1f6c86dc9e6a453b07382..7a6fc8103d2fa68ec1d0f4cd856f71fe8899c3a6 100644 (file)
--- a/folly/experimental/test/RefCountTest.cpp
+++ b/folly/experimental/test/RefCountTest.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright 2016 Facebook, Inc.
+ * Copyright 2017 Facebook, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,8 +18,7 @@
 #include <folly/Baton.h>
 #include <folly/experimental/RCURefCount.h>
 #include <folly/experimental/TLRefCount.h>
-
-#include <gtest/gtest.h>
+#include <folly/portability/GTest.h>
 
 namespace folly {
 
@@ -83,6 +82,40 @@ void basicTest() {
   EXPECT_EQ(0, ++count);
 }
 
+template <typename RefCount>
+void stressTest() {
+  constexpr size_t kItersCount = 10000;
+
+  for (size_t i = 0; i < kItersCount; ++i) {
+    RefCount count;
+    std::mutex mutex;
+    int a{1};
+
+    std::thread t1([&]() {
+      if (++count) {
+        {
+          std::lock_guard<std::mutex> lg(mutex);
+          EXPECT_EQ(1, a);
+        }
+        --count;
+      }
+    });
+
+    std::thread t2([&]() {
+      count.useGlobal();
+      if (--count == 0) {
+        std::lock_guard<std::mutex> lg(mutex);
+        a = 0;
+      }
+    });
+
+    t1.join();
+    t2.join();
+
+    EXPECT_EQ(0, ++count);
+  }
+}
+
 TEST(RCURefCount, Basic) {
   basicTest<RCURefCount>();
 }
@@ -91,4 +124,11 @@ TEST(TLRefCount, Basic) {
   basicTest<TLRefCount>();
 }
 
+TEST(RCURefCount, Stress) {
+  stressTest<TLRefCount>();
+}
+
+TEST(TLRefCount, Stress) {
+  stressTest<TLRefCount>();
+}
 }