1 | /* |
2 | * Copyright 2012-present 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/experimental/EventCount.h> |
18 | |
19 | #include <algorithm> |
20 | #include <atomic> |
21 | #include <random> |
22 | #include <thread> |
23 | #include <vector> |
24 | |
25 | #include <glog/logging.h> |
26 | |
27 | #include <folly/Random.h> |
28 | #include <folly/portability/GTest.h> |
29 | |
30 | using namespace folly; |
31 | |
32 | namespace { |
33 | |
34 | class Semaphore { |
35 | public: |
36 | explicit Semaphore(int v = 0) : value_(v) {} |
37 | |
38 | void down() { |
39 | ec_.await([this] { return tryDown(); }); |
40 | } |
41 | |
42 | void up() { |
43 | ++value_; |
44 | ec_.notifyAll(); |
45 | } |
46 | |
47 | int value() const { |
48 | return value_; |
49 | } |
50 | |
51 | private: |
52 | bool tryDown() { |
53 | for (int v = value_; v != 0;) { |
54 | if (value_.compare_exchange_weak(v, v - 1)) { |
55 | return true; |
56 | } |
57 | } |
58 | return false; |
59 | } |
60 | |
61 | std::atomic<int> value_; |
62 | EventCount ec_; |
63 | }; |
64 | |
65 | template <class T, class Random> |
66 | void randomPartition( |
67 | Random& random, |
68 | T key, |
69 | int n, |
70 | std::vector<std::pair<T, int>>& out) { |
71 | while (n != 0) { |
72 | int m = std::min(n, 1000); |
73 | std::uniform_int_distribution<uint32_t> u(1, m); |
74 | int cut = u(random); |
75 | out.emplace_back(key, cut); |
76 | n -= cut; |
77 | } |
78 | } |
79 | |
80 | } // namespace |
81 | |
82 | TEST(EventCount, Simple) { |
83 | // We're basically testing for no deadlock. |
84 | static const size_t count = 300000; |
85 | |
86 | enum class Op { |
87 | UP, |
88 | DOWN, |
89 | }; |
90 | std::vector<std::pair<Op, int>> ops; |
91 | std::mt19937 rnd(randomNumberSeed()); |
92 | randomPartition(rnd, Op::UP, count, ops); |
93 | size_t uppers = ops.size(); |
94 | randomPartition(rnd, Op::DOWN, count, ops); |
95 | size_t downers = ops.size() - uppers; |
96 | VLOG(1) << "Using " << ops.size() << " threads: uppers=" << uppers |
97 | << " downers=" << downers << " sem_count=" << count; |
98 | |
99 | std::shuffle(ops.begin(), ops.end(), std::mt19937(std::random_device()())); |
100 | |
101 | std::vector<std::thread> threads; |
102 | threads.reserve(ops.size()); |
103 | |
104 | Semaphore sem; |
105 | for (auto& op : ops) { |
106 | int n = op.second; |
107 | if (op.first == Op::UP) { |
108 | auto fn = [&sem, n]() mutable { |
109 | while (n--) { |
110 | sem.up(); |
111 | } |
112 | }; |
113 | threads.push_back(std::thread(fn)); |
114 | } else { |
115 | auto fn = [&sem, n]() mutable { |
116 | while (n--) { |
117 | sem.down(); |
118 | } |
119 | }; |
120 | threads.push_back(std::thread(fn)); |
121 | } |
122 | } |
123 | |
124 | for (auto& thread : threads) { |
125 | thread.join(); |
126 | } |
127 | |
128 | EXPECT_EQ(0, sem.value()); |
129 | } |
130 | |