Pull from FB rev 63ce89e2f2301e6bba44a111cc7d4218022156f6
[folly.git] / folly / test / ProducerConsumerQueueTest.cpp
1 /*
2  * Copyright 2012 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/ProducerConsumerQueue.h"
18
19 #include <gtest/gtest.h>
20 #include <vector>
21 #include <atomic>
22 #include <chrono>
23 #include <memory>
24 #include <thread>
25 #include <glog/logging.h>
26
27 //////////////////////////////////////////////////////////////////////
28
29 namespace {
30
31 template<class T> struct TestTraits {
32   T limit() const { return 1 << 24; }
33   T generate() const { return rand() % 26; }
34 };
35
36 template<> struct TestTraits<std::string> {
37   int limit() const { return 1 << 21; }
38   std::string generate() const { return std::string(12, ' '); }
39 };
40
41 template<class QueueType, size_t Size>
42 struct PerfTest {
43   typedef typename QueueType::value_type T;
44
45   explicit PerfTest() : queue_(Size), done_(false) {}
46
47   void operator()() {
48     using namespace std::chrono;
49     auto const startTime = system_clock::now();
50
51     std::thread producer([this] { this->producer(); });
52     std::thread consumer([this] { this->consumer(); });
53
54     producer.join();
55     done_ = true;
56     consumer.join();
57
58     auto duration = duration_cast<milliseconds>(
59       system_clock::now() - startTime);
60     LOG(INFO) << "     done: " << duration.count() << "ms";
61   }
62
63   void producer() {
64     for (int i = 0; i < traits_.limit(); ++i) {
65       while (!queue_.write(traits_.generate())) {
66       }
67     }
68   }
69
70   void consumer() {
71     while (!done_) {
72       T data;
73       queue_.read(data);
74     }
75   }
76
77   QueueType queue_;
78   std::atomic<bool> done_;
79   TestTraits<T> traits_;
80 };
81
82 template<class TestType> void doTest(const char* name) {
83   LOG(INFO) << "  testing: " << name;
84   std::unique_ptr<TestType> const t(new TestType());
85   (*t)();
86 }
87
88 template<class T> void perfTestType(const char* type) {
89   const size_t size = 0xfffe;
90
91   LOG(INFO) << "Type: " << type;
92   doTest<PerfTest<folly::ProducerConsumerQueue<T>,size> >(
93     "ProducerConsumerQueue");
94 }
95
96 template<class QueueType, size_t Size>
97 struct CorrectnessTest {
98   typedef typename QueueType::value_type T;
99
100   explicit CorrectnessTest()
101     : queue_(Size)
102     , done_(false)
103   {
104     const size_t testSize = traits_.limit();
105     testData_.reserve(testSize);
106     for (int i = 0; i < testSize; ++i) {
107       testData_.push_back(traits_.generate());
108     }
109   }
110
111   void operator()() {
112     std::thread producer([this] { this->producer(); });
113     std::thread consumer([this] { this->consumer(); });
114
115     producer.join();
116     done_ = true;
117     consumer.join();
118   }
119
120   void producer() {
121     for (auto& data : testData_) {
122       while (!queue_.write(data)) {
123       }
124     }
125   }
126
127   void consumer() {
128     for (auto& expect : testData_) {
129     again:
130       T data;
131       if (!queue_.read(data)) {
132         if (done_) {
133           // Try one more read; unless there's a bug in the queue class
134           // there should still be more data sitting in the queue even
135           // though the producer thread exited.
136           if (!queue_.read(data)) {
137             EXPECT_TRUE(0 && "Finished too early ...");
138             return;
139           }
140         } else {
141           goto again;
142         }
143       }
144       EXPECT_EQ(data, expect);
145     }
146   }
147
148   std::vector<T> testData_;
149   QueueType queue_;
150   TestTraits<T> traits_;
151   std::atomic<bool> done_;
152 };
153
154 template<class T> void correctnessTestType(const std::string& type) {
155   LOG(INFO) << "Type: " << type;
156   doTest<CorrectnessTest<folly::ProducerConsumerQueue<T>,0xfffe> >(
157     "ProducerConsumerQueue");
158 }
159
160 struct DtorChecker {
161   static int numInstances;
162   DtorChecker() { ++numInstances; }
163   DtorChecker(const DtorChecker& o) { ++numInstances; }
164   ~DtorChecker() { --numInstances; }
165 };
166
167 int DtorChecker::numInstances = 0;
168
169 }
170
171 //////////////////////////////////////////////////////////////////////
172
173 TEST(PCQ, QueueCorrectness) {
174   correctnessTestType<std::string>("string");
175   correctnessTestType<int>("int");
176   correctnessTestType<unsigned long long>("unsigned long long");
177 }
178
179 TEST(PCQ, PerfTest) {
180   perfTestType<std::string>("string");
181   perfTestType<int>("int");
182   perfTestType<unsigned long long>("unsigned long long");
183 }
184
185 TEST(PCQ, Destructor) {
186   // Test that orphaned elements in a ProducerConsumerQueue are
187   // destroyed.
188   {
189     folly::ProducerConsumerQueue<DtorChecker> queue(1024);
190     for (int i = 0; i < 10; ++i) {
191       EXPECT_TRUE(queue.write(DtorChecker()));
192     }
193
194     EXPECT_EQ(DtorChecker::numInstances, 10);
195
196     {
197       DtorChecker ignore;
198       EXPECT_TRUE(queue.read(ignore));
199       EXPECT_TRUE(queue.read(ignore));
200     }
201
202     EXPECT_EQ(DtorChecker::numInstances, 8);
203   }
204
205   EXPECT_EQ(DtorChecker::numInstances, 0);
206
207   // Test the same thing in the case that the queue write pointer has
208   // wrapped, but the read one hasn't.
209   {
210     folly::ProducerConsumerQueue<DtorChecker> queue(4);
211     for (int i = 0; i < 3; ++i) {
212       EXPECT_TRUE(queue.write(DtorChecker()));
213     }
214     EXPECT_EQ(DtorChecker::numInstances, 3);
215     {
216       DtorChecker ignore;
217       EXPECT_TRUE(queue.read(ignore));
218     }
219     EXPECT_EQ(DtorChecker::numInstances, 2);
220     EXPECT_TRUE(queue.write(DtorChecker()));
221     EXPECT_EQ(DtorChecker::numInstances, 3);
222   }
223   EXPECT_EQ(DtorChecker::numInstances, 0);
224 }