folly/futures/test/TestExecutor.cpp

   1 /*
   2  * Copyright 2017-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 "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     (void)worker;
  47     addImpl({});
  48   }
  49
  50   for (auto& worker : workers_) {
  51     worker.join();
  52   }
  53 }
  54
  55 void TestExecutor::add(Func f) {
  56   if (f) {
  57     addImpl(std::move(f));
  58   }
  59 }
  60
  61 size_t TestExecutor::numThreads() const {
  62   return workers_.size();
  63 }
  64
  65 void TestExecutor::addImpl(Func f) {
  66   {
  67     lock_guard<mutex> g(m_);
  68     workItems_.push(std::move(f));
  69   }
  70   cv_.notify_one();
  71 }
  72
  73 } // namespace folly