1 files changed, 531 insertions, 0 deletions
diff --git a/src/mongo/util/producer_consumer_queue_test.cpp b/src/mongo/util/producer_consumer_queue_test.cpp
new file mode 100644
index 00000000000..f824b5a63de
--- /dev/null
+++ b/src/mongo/util/producer_consumer_queue_test.cpp
@@ -0,0 +1,531 @@
+/**
+ * Copyright (C) 2018 MongoDB Inc.
+ *
+ * This program is free software: you can redistribute it and/or  modify
+ * it under the terms of the GNU Affero General Public License, version 3,
+ * as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ * As a special exception, the copyright holders give permission to link the
+ * code of portions of this program with the OpenSSL library under certain
+ * conditions as described in each individual source file and distribute
+ * linked combinations including the program with the OpenSSL library. You
+ * must comply with the GNU Affero General Public License in all respects
+ * for all of the code used other than as permitted herein. If you modify
+ * file(s) with this exception, you may extend this exception to your
+ * version of the file(s), but you are not obligated to do so. If you do not
+ * wish to do so, delete this exception statement from your version. If you
+ * delete this exception statement from all source files in the program,
+ * then also delete it in the license file.
+ */
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/unittest/unittest.h"
+
+#include "mongo/util/producer_consumer_queue.h"
+
+#include "mongo/db/service_context_noop.h"
+#include "mongo/stdx/condition_variable.h"
+#include "mongo/stdx/memory.h"
+#include "mongo/stdx/mutex.h"
+#include "mongo/stdx/thread.h"
+#include "mongo/util/assert_util.h"
+
+namespace mongo {
+
+namespace {
+
+class ProducerConsumerQueueTest : public unittest::Test {
+public:
+    ProducerConsumerQueueTest() : _serviceCtx(stdx::make_unique<ServiceContextNoop>()) {}
+
+    template <typename Callback>
+    stdx::thread runThread(StringData name, Callback&& cb) {
+        return stdx::thread([this, name, cb] { cb(); });
+    }
+
+private:
+    std::unique_ptr<ServiceContext> _serviceCtx;
+};
+
+class MoveOnly {
+public:
+    struct CostFunc {
+        CostFunc() = default;
+        explicit CostFunc(size_t val) : val(val) {}
+
+        size_t operator()(const MoveOnly& mo) const {
+            return val + *mo._val;
+        }
+
+        const size_t val = 0;
+    };
+
+    explicit MoveOnly(int i) : _val(i) {}
+
+    MoveOnly(const MoveOnly&) = delete;
+    MoveOnly& operator=(const MoveOnly&) = delete;
+
+    MoveOnly(MoveOnly&& other) : _val(other._val) {
+        other._val.reset();
+    }
+
+    MoveOnly& operator=(MoveOnly&& other) {
+        if (&other == this) {
+            return *this;
+        }
+
+        _val = other._val;
+        other._val.reset();
+
+        return *this;
+    }
+
+    bool movedFrom() const {
+        return !_val;
+    }
+
+    friend bool operator==(const MoveOnly& lhs, const MoveOnly& rhs) {
+        return *lhs._val == *rhs._val;
+    }
+
+    friend bool operator!=(const MoveOnly& lhs, const MoveOnly& rhs) {
+        return !(lhs == rhs);
+    }
+
+    friend std::ostream& operator<<(std::ostream& os, const MoveOnly& mo) {
+        return (os << "MoveOnly(" << *mo._val << ")");
+    }
+
+private:
+    boost::optional<int> _val;
+};
+
+TEST_F(ProducerConsumerQueueTest, basicPushPop) {
+    ProducerConsumerQueue<MoveOnly> pcq{};
+
+    runThread("Producer", [&]() { pcq.push(MoveOnly(1)); }).join();
+
+    ASSERT_EQUALS(pcq.sizeForTest(), 1ul);
+
+    runThread("Consumer", [&]() { ASSERT_EQUALS(pcq.pop(), MoveOnly(1)); }).join();
+
+    ASSERT_TRUE(pcq.emptyForTest());
+}
+
+TEST_F(ProducerConsumerQueueTest, closeConsumerEnd) {
+    ProducerConsumerQueue<MoveOnly> pcq{1};
+
+    pcq.push(MoveOnly(1));
+
+    auto producer = runThread("Producer", [&]() {
+        ASSERT_THROWS_CODE(
+            pcq.push(MoveOnly(2)), DBException, ErrorCodes::ProducerConsumerQueueEndClosed);
+    });
+
+    ASSERT_EQUALS(pcq.sizeForTest(), 1ul);
+
+    pcq.closeConsumerEnd();
+
+    ASSERT_THROWS_CODE(pcq.pop(), DBException, ErrorCodes::ProducerConsumerQueueEndClosed);
+
+    producer.join();
+}
+
+TEST_F(ProducerConsumerQueueTest, closeProducerEndImmediate) {
+    ProducerConsumerQueue<MoveOnly> pcq{};
+
+    pcq.push(MoveOnly(1));
+    pcq.closeProducerEnd();
+
+    runThread("Consumer", [&]() {
+        ASSERT_EQUALS(pcq.pop(), MoveOnly(1));
+
+        ASSERT_THROWS_CODE(pcq.pop(), DBException, ErrorCodes::ProducerConsumerQueueEndClosed);
+    }).join();
+}
+
+TEST_F(ProducerConsumerQueueTest, closeProducerEndBlocking) {
+    ProducerConsumerQueue<MoveOnly> pcq{};
+
+    auto consumer = runThread("Consumer", [&]() {
+        ASSERT_THROWS_CODE(pcq.pop(), DBException, ErrorCodes::ProducerConsumerQueueEndClosed);
+    });
+
+    pcq.closeProducerEnd();
+
+    consumer.join();
+}
+
+TEST_F(ProducerConsumerQueueTest, popsWithTimeout) {
+    ProducerConsumerQueue<MoveOnly> pcq{};
+
+    runThread("Consumer", [&]() {
+        ASSERT_THROWS_CODE(pcq.pop(Milliseconds(100)), DBException, ErrorCodes::ExceededTimeLimit);
+
+        std::vector<MoveOnly> vec;
+        ASSERT_THROWS_CODE(pcq.popMany(std::back_inserter(vec), Milliseconds(100)),
+                           DBException,
+                           ErrorCodes::ExceededTimeLimit);
+
+        ASSERT_THROWS_CODE(pcq.popManyUpTo(1000, std::back_inserter(vec), Milliseconds(100)),
+                           DBException,
+                           ErrorCodes::ExceededTimeLimit);
+    }).join();
+
+    ASSERT_EQUALS(pcq.sizeForTest(), 0ul);
+}
+
+TEST_F(ProducerConsumerQueueTest, pushesWithTimeout) {
+    ProducerConsumerQueue<MoveOnly> pcq{1};
+
+    {
+        MoveOnly mo(1);
+        pcq.push(std::move(mo));
+        ASSERT(mo.movedFrom());
+    }
+
+    runThread("Consumer", [&]() {
+        {
+            MoveOnly mo(2);
+            ASSERT_THROWS_CODE(pcq.push(std::move(mo), Milliseconds(100)),
+                               DBException,
+                               ErrorCodes::ExceededTimeLimit);
+            ASSERT_EQUALS(pcq.sizeForTest(), 1ul);
+            ASSERT(!mo.movedFrom());
+            ASSERT_EQUALS(mo, MoveOnly(2));
+        }
+
+        {
+            std::vector<MoveOnly> vec;
+            vec.emplace_back(MoveOnly(2));
+
+            auto iter = begin(vec);
+            ASSERT_THROWS_CODE(pcq.pushMany(iter, end(vec), Milliseconds(100)),
+                               DBException,
+                               ErrorCodes::ExceededTimeLimit);
+            ASSERT_EQUALS(pcq.sizeForTest(), 1ul);
+            ASSERT(!vec[0].movedFrom());
+            ASSERT_EQUALS(vec[0], MoveOnly(2));
+        }
+    }).join();
+
+    ASSERT_EQUALS(pcq.sizeForTest(), 1ul);
+}
+
+TEST_F(ProducerConsumerQueueTest, basicPushPopWithBlocking) {
+    ProducerConsumerQueue<MoveOnly> pcq{};
+
+    auto consumer = runThread("Consumer", [&]() { ASSERT_EQUALS(pcq.pop(), MoveOnly(1)); });
+
+    auto producer = runThread("Producer", [&]() { pcq.push(MoveOnly(1)); });
+
+    consumer.join();
+    producer.join();
+
+    ASSERT_TRUE(pcq.emptyForTest());
+}
+
+TEST_F(ProducerConsumerQueueTest, multipleStepPushPopWithBlocking) {
+    ProducerConsumerQueue<MoveOnly> pcq{1};
+
+    auto consumer = runThread("Consumer", [&]() {
+        for (int i = 0; i < 10; ++i) {
+            ASSERT_EQUALS(pcq.pop(), MoveOnly(i));
+        }
+    });
+
+    auto producer = runThread("Producer", [&]() {
+        for (int i = 0; i < 10; ++i) {
+            pcq.push(MoveOnly(i));
+        }
+    });
+
+    consumer.join();
+    producer.join();
+
+    ASSERT_TRUE(pcq.emptyForTest());
+}
+
+
+TEST_F(ProducerConsumerQueueTest, pushTooLarge) {
+    {
+        ProducerConsumerQueue<MoveOnly, MoveOnly::CostFunc> pcq{1};
+
+        runThread("Producer", [&]() {
+            ASSERT_THROWS_CODE(
+                pcq.push(MoveOnly(2)), DBException, ErrorCodes::ProducerConsumerQueueBatchTooLarge);
+        }).join();
+    }
+
+    {
+        ProducerConsumerQueue<MoveOnly, MoveOnly::CostFunc> pcq{4};
+
+        std::vector<MoveOnly> vec;
+        vec.push_back(MoveOnly(3));
+        vec.push_back(MoveOnly(3));
+
+        runThread("Producer", [&]() {
+            ASSERT_THROWS_CODE(pcq.pushMany(begin(vec), end(vec)),
+                               DBException,
+                               ErrorCodes::ProducerConsumerQueueBatchTooLarge);
+        }).join();
+    }
+}
+
+TEST_F(ProducerConsumerQueueTest, pushManyPopWithoutBlocking) {
+    ProducerConsumerQueue<MoveOnly> pcq{};
+
+    runThread("Producer", [&]() {
+        std::vector<MoveOnly> vec;
+        for (int i = 0; i < 10; ++i) {
+            vec.emplace_back(MoveOnly(i));
+        }
+
+        pcq.pushMany(begin(vec), end(vec));
+    }).join();
+
+    runThread("Consumer", [&]() {
+        for (int i = 0; i < 10; ++i) {
+            ASSERT_EQUALS(pcq.pop(), MoveOnly(i));
+        }
+    }).join();
+
+    ASSERT_TRUE(pcq.emptyForTest());
+}
+
+TEST_F(ProducerConsumerQueueTest, popManyPopWithBlocking) {
+    ProducerConsumerQueue<MoveOnly> pcq{2};
+
+    auto consumer = runThread("Consumer", [&]() {
+        for (int i = 0; i < 10; i = i + 2) {
+            std::vector<MoveOnly> out;
+
+            pcq.popMany(std::back_inserter(out));
+
+            ASSERT_EQUALS(out.size(), 2ul);
+            ASSERT_EQUALS(out[0], MoveOnly(i));
+            ASSERT_EQUALS(out[1], MoveOnly(i + 1));
+        }
+    });
+
+    auto producer = runThread("Producer", [&]() {
+        std::vector<MoveOnly> vec;
+        for (int i = 0; i < 10; ++i) {
+            vec.emplace_back(MoveOnly(i));
+        }
+
+        for (auto iter = begin(vec); iter != end(vec); iter += 2) {
+            pcq.pushMany(iter, iter + 2);
+        }
+    });
+
+    consumer.join();
+    producer.join();
+
+    ASSERT_TRUE(pcq.emptyForTest());
+}
+
+TEST_F(ProducerConsumerQueueTest, popManyUpToPopWithBlocking) {
+    ProducerConsumerQueue<MoveOnly> pcq{4};
+
+    auto consumer = runThread("Consumer", [&]() {
+        for (int i = 0; i < 10; i = i + 2) {
+            std::vector<MoveOnly> out;
+
+            size_t spent;
+            std::tie(spent, std::ignore) = pcq.popManyUpTo(2, std::back_inserter(out));
+
+            ASSERT_EQUALS(spent, 2ul);
+            ASSERT_EQUALS(out.size(), 2ul);
+            ASSERT_EQUALS(out[0], MoveOnly(i));
+            ASSERT_EQUALS(out[1], MoveOnly(i + 1));
+        }
+    });
+
+    auto producer = runThread("Producer", [&]() {
+        std::vector<MoveOnly> vec;
+        for (int i = 0; i < 10; ++i) {
+            vec.emplace_back(MoveOnly(i));
+        }
+
+        for (auto iter = begin(vec); iter != end(vec); iter += 2) {
+            pcq.pushMany(iter, iter + 2);
+        }
+    });
+
+    consumer.join();
+    producer.join();
+
+    ASSERT_TRUE(pcq.emptyForTest());
+}
+
+TEST_F(ProducerConsumerQueueTest, popManyUpToPopWithBlockingWithSpecialCost) {
+    ProducerConsumerQueue<MoveOnly, MoveOnly::CostFunc> pcq{};
+
+    auto consumer = runThread("Consumer", [&]() {
+        {
+            std::vector<MoveOnly> out;
+            size_t spent;
+            std::tie(spent, std::ignore) = pcq.popManyUpTo(5, std::back_inserter(out));
+
+            ASSERT_EQUALS(spent, 6ul);
+            ASSERT_EQUALS(out.size(), 3ul);
+            ASSERT_EQUALS(out[0], MoveOnly(1));
+            ASSERT_EQUALS(out[1], MoveOnly(2));
+            ASSERT_EQUALS(out[2], MoveOnly(3));
+        }
+
+        {
+            std::vector<MoveOnly> out;
+            size_t spent;
+            std::tie(spent, std::ignore) = pcq.popManyUpTo(15, std::back_inserter(out));
+
+            ASSERT_EQUALS(spent, 9ul);
+            ASSERT_EQUALS(out.size(), 2ul);
+            ASSERT_EQUALS(out[0], MoveOnly(4));
+            ASSERT_EQUALS(out[1], MoveOnly(5));
+        }
+    });
+
+    auto producer = runThread("Producer", [&]() {
+        std::vector<MoveOnly> vec;
+        for (int i = 1; i < 6; ++i) {
+            vec.emplace_back(MoveOnly(i));
+        }
+
+        pcq.pushMany(begin(vec), end(vec));
+    });
+
+    consumer.join();
+    producer.join();
+
+    ASSERT_TRUE(pcq.emptyForTest());
+}
+
+TEST_F(ProducerConsumerQueueTest, singleProducerMultiConsumer) {
+    ProducerConsumerQueue<MoveOnly> pcq{};
+
+    stdx::mutex mutex;
+    size_t success = 0;
+    size_t failure = 0;
+
+    std::array<stdx::thread, 3> threads;
+    for (auto& thread : threads) {
+        thread = runThread("Consumer", [&]() {
+            {
+                try {
+                    pcq.pop();
+                    stdx::lock_guard<stdx::mutex> lk(mutex);
+                    success++;
+                } catch (const DBException& exception) {
+                    ASSERT_EQUALS(exception.getCode(), ErrorCodes::ProducerConsumerQueueEndClosed);
+                    stdx::lock_guard<stdx::mutex> lk(mutex);
+                    failure++;
+                }
+            }
+        });
+    }
+
+    pcq.push(MoveOnly(1));
+    pcq.push(MoveOnly(2));
+
+    pcq.closeProducerEnd();
+
+    for (auto& thread : threads) {
+        thread.join();
+    }
+
+    ASSERT_EQUALS(success, 2ul);
+    ASSERT_EQUALS(failure, 1ul);
+
+    ASSERT_TRUE(pcq.emptyForTest());
+}
+
+TEST_F(ProducerConsumerQueueTest, basicTryPop) {
+    ProducerConsumerQueue<MoveOnly> pcq{};
+
+    ASSERT_FALSE(pcq.tryPop());
+    ASSERT_TRUE(pcq.tryPush(MoveOnly(1)));
+    ASSERT_EQUALS(pcq.sizeForTest(), 1ul);
+
+    auto val = pcq.tryPop();
+
+    ASSERT_FALSE(pcq.tryPop());
+    ASSERT_TRUE(val);
+    ASSERT_EQUALS(*val, MoveOnly(1));
+
+    ASSERT_TRUE(pcq.emptyForTest());
+}
+
+TEST_F(ProducerConsumerQueueTest, basicTryPush) {
+    ProducerConsumerQueue<MoveOnly> pcq{1};
+
+    ASSERT_TRUE(pcq.tryPush(MoveOnly(1)));
+    ASSERT_FALSE(pcq.tryPush(MoveOnly(2)));
+
+    ASSERT_EQUALS(pcq.sizeForTest(), 1ul);
+
+    auto val = pcq.tryPop();
+    ASSERT_FALSE(pcq.tryPop());
+    ASSERT_TRUE(val);
+    ASSERT_EQUALS(*val, MoveOnly(1));
+
+    ASSERT_TRUE(pcq.emptyForTest());
+}
+
+TEST_F(ProducerConsumerQueueTest, tryPushWithSpecialCost) {
+    ProducerConsumerQueue<MoveOnly, MoveOnly::CostFunc> pcq{5};
+
+    ASSERT_TRUE(pcq.tryPush(MoveOnly(1)));
+    ASSERT_TRUE(pcq.tryPush(MoveOnly(2)));
+    ASSERT_FALSE(pcq.tryPush(MoveOnly(3)));
+
+    ASSERT_EQUALS(pcq.sizeForTest(), 3ul);
+
+    auto val1 = pcq.tryPop();
+    ASSERT_EQUALS(pcq.sizeForTest(), 2ul);
+    auto val2 = pcq.tryPop();
+    ASSERT_EQUALS(pcq.sizeForTest(), 0ul);
+    ASSERT_FALSE(pcq.tryPop());
+    ASSERT_TRUE(val1);
+    ASSERT_TRUE(val2);
+    ASSERT_EQUALS(*val1, MoveOnly(1));
+    ASSERT_EQUALS(*val2, MoveOnly(2));
+
+    ASSERT_TRUE(pcq.emptyForTest());
+}
+
+TEST_F(ProducerConsumerQueueTest, tryPushWithSpecialStatefulCost) {
+    ProducerConsumerQueue<MoveOnly, MoveOnly::CostFunc> pcq{5, MoveOnly::CostFunc(1)};
+
+    ASSERT_TRUE(pcq.tryPush(MoveOnly(1)));
+    ASSERT_TRUE(pcq.tryPush(MoveOnly(2)));
+    ASSERT_FALSE(pcq.tryPush(MoveOnly(3)));
+
+    ASSERT_EQUALS(pcq.sizeForTest(), 5ul);
+
+    auto val1 = pcq.tryPop();
+    ASSERT_EQUALS(pcq.sizeForTest(), 3ul);
+    auto val2 = pcq.tryPop();
+    ASSERT_EQUALS(pcq.sizeForTest(), 0ul);
+    ASSERT_FALSE(pcq.tryPop());
+    ASSERT_TRUE(val1);
+    ASSERT_TRUE(val2);
+    ASSERT_EQUALS(*val1, MoveOnly(1));
+    ASSERT_EQUALS(*val2, MoveOnly(2));
+
+    ASSERT_TRUE(pcq.emptyForTest());
+}
+
+}  // namespace
+
+}  // namespace mongo