/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <thread>
#include <folly/concurrency/memory/TLRefCount.h>
#include <folly/portability/GTest.h>
#include <folly/synchronization/Baton.h>
#include <folly/synchronization/test/Barrier.h>
namespace folly {
template <typename RefCount>
void basicTest() {
constexpr size_t numIters = 100000;
constexpr size_t numThreads = 10;
size_t got0 = 0;
RefCount count;
folly::Baton<> b;
std::vector<std::thread> ts;
folly::Baton<> threadBatons[numThreads];
for (size_t t = 0; t < numThreads; ++t) {
ts.emplace_back([&count, &b, &got0, t, &threadBatons] {
for (size_t i = 0; i < numIters; ++i) {
auto ret = ++count;
EXPECT_TRUE(ret > 1);
if (i == 0) {
threadBatons[t].post();
}
}
if (t == 0) {
b.post();
}
for (size_t i = 0; i < numIters; ++i) {
auto ret = --count;
if (ret == 0) {
++got0;
EXPECT_EQ(numIters - 1, i);
}
}
});
}
for (size_t t = 0; t < numThreads; ++t) {
threadBatons[t].wait();
}
b.wait();
count.useGlobal();
if (--count == 0) {
++got0;
}
for (auto& t : ts) {
t.join();
}
EXPECT_EQ(1, got0);
EXPECT_EQ(0, ++count);
EXPECT_EQ(0, ++count);
}
template <typename RefCount>
void stressTest(size_t itersCount) {
for (size_t i = 0; i < itersCount; ++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(TLRefCount, Basic) {
basicTest<TLRefCount>();
}
TEST(TLRefCount, Stress) {
// This is absurdly slow, so we can't
// do it that many times.
stressTest<TLRefCount>(500);
}
TEST(TLRefCount, SafeToDeleteWhenReachingZero) {
// Tests that it is safe to delete a TLRefCount when it is already in global
// state and its ref count reaches zero. This is a reasonable assumption
// since data structures typically embed a TLRefCount object and delete
// themselves when the refcount reaches 0.
TLRefCount* count = new TLRefCount();
folly::test::Barrier done(2); // give time for TSAN to catch issues
folly::Baton<> batonUnref;
int times_deleted = 0;
std::thread t1([&] {
++*count;
batonUnref.post();
if (--*count == 0) {
times_deleted++;
delete count;
}
done.wait();
});
std::thread t2([&] {
// Make sure thread 1 already grabbed a reference first, otherwise we might
// destroy it before thread 1 had a chance.
batonUnref.wait();
count->useGlobal();
if (--*count == 0) {
times_deleted++;
delete count;
}
done.wait();
});
t1.join();
t2.join();
EXPECT_EQ(times_deleted, 1);
}
} // namespace folly