folly/futures/test/TestExecutor.cpp

   1 /*
   2  * Copyright 2017 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 "TestExecutor.h"
  18
  19 using namespace std;
  20
  21 namespace folly {
  22
  23 TestExecutor::TestExecutor(size_t numThreads) {
  24   const auto kWorkers = std::max(size_t(1), numThreads);
  25   for (auto idx = 0u; idx < kWorkers; ++idx) {
  26     workers_.emplace_back([this] {
  27       while (true) {
  28         Func work;
  29         {
  30           unique_lock<mutex> lk(m_);
  31           cv_.wait(lk, [this] { return !workItems_.empty(); });
  32           work = std::move(workItems_.front());
  33           workItems_.pop();
  34         }
  35         if (!work) {
  36           break;
  37         }
  38         work();
  39       }
  40     });
  41   }
  42 }
  43
  44 TestExecutor::~TestExecutor() {
  45   for (auto& worker : workers_) {
  46     addImpl({});
  47   }
  48
  49   for (auto& worker : workers_) {
  50     worker.join();
  51   }
  52 }
  53
  54 void TestExecutor::add(Func f) {
  55   if (f) {
  56     addImpl(std::move(f));
  57   }
  58 }
  59
  60 size_t TestExecutor::numThreads() const {
  61   return workers_.size();
  62 }
  63
  64 void TestExecutor::addImpl(Func f) {
  65   {
  66     lock_guard<mutex> g(m_);
  67     workItems_.push(std::move(f));
  68   }
  69   cv_.notify_one();
  70 }
  71
  72 } // folly