From 0eff88b310ad179eef4cbb6e2d76f6beadd07e16 Mon Sep 17 00:00:00 2001
From: Yedidya Feldblum <yfeldblum@fb.com>
Date: Fri, 5 May 2017 11:52:53 -0700
Subject: [PATCH] Control the number of threads in TestExecutor

Summary:
[Folly] Control the number of threads in `TestExecutor`.

Let the caller control the number of threads to use in each given case.

Reviewed By: spacedentist

Differential Revision: D4999699

fbshipit-source-id: 4acf68cf17fbca14f0779daf0268d54c5606e4a8
---
 folly/futures/test/RetryingTest.cpp     | 4 +++-
 folly/futures/test/TestExecutor.cpp     | 8 ++++----
 folly/futures/test/TestExecutor.h       | 4 ++--
 folly/futures/test/TestExecutorTest.cpp | 3 ++-
 4 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/folly/futures/test/RetryingTest.cpp b/folly/futures/test/RetryingTest.cpp
index 2de5c8a4..43a15ea2 100644
--- a/folly/futures/test/RetryingTest.cpp
+++ b/folly/futures/test/RetryingTest.cpp
@@ -151,7 +151,7 @@ TEST(RetryingTest, large_retries) {
     PCHECK(setrlimit(RLIMIT_AS, &oldMemLimit) == 0);
   };
 
-  TestExecutor executor;
+  TestExecutor executor(4);
   // size of implicit promise is at least the size of the return.
   using LargeReturn = array<uint64_t, 16000>;
   auto func = [&executor](size_t retryNum) -> Future<LargeReturn> {
@@ -172,6 +172,8 @@ TEST(RetryingTest, large_retries) {
         func));
   }
 
+  // 40 * 10,000 * 16,000B > 1GB; we should avoid OOM
+
   for (auto& f : futures) {
     f.wait();
     EXPECT_TRUE(f.hasValue());
diff --git a/folly/futures/test/TestExecutor.cpp b/folly/futures/test/TestExecutor.cpp
index 16f8d26d..fc57e6de 100644
--- a/folly/futures/test/TestExecutor.cpp
+++ b/folly/futures/test/TestExecutor.cpp
@@ -20,9 +20,9 @@ using namespace std;
 
 namespace folly {
 
-TestExecutor::TestExecutor() {
-  const auto kWorkers = std::max(1U, thread::hardware_concurrency());
-  for (auto idx = 0U; idx < kWorkers; ++idx) {
+TestExecutor::TestExecutor(size_t numThreads) {
+  const auto kWorkers = std::max(size_t(1), numThreads);
+  for (auto idx = 0u; idx < kWorkers; ++idx) {
     workers_.emplace_back([this] {
       while (true) {
         Func work;
@@ -57,7 +57,7 @@ void TestExecutor::add(Func f) {
   }
 }
 
-uint32_t TestExecutor::numThreads() const {
+size_t TestExecutor::numThreads() const {
   return workers_.size();
 }
 
diff --git a/folly/futures/test/TestExecutor.h b/folly/futures/test/TestExecutor.h
index c09713be..f022bc03 100644
--- a/folly/futures/test/TestExecutor.h
+++ b/folly/futures/test/TestExecutor.h
@@ -29,13 +29,13 @@ namespace folly {
  */
 class TestExecutor : public Executor {
  public:
-  TestExecutor();
+  explicit TestExecutor(size_t numThreads);
 
   ~TestExecutor() override;
 
   void add(Func f) override;
 
-  uint32_t numThreads() const;
+  size_t numThreads() const;
 
  private:
   void addImpl(Func f);
diff --git a/folly/futures/test/TestExecutorTest.cpp b/folly/futures/test/TestExecutorTest.cpp
index 750e904b..5af5893b 100644
--- a/folly/futures/test/TestExecutorTest.cpp
+++ b/folly/futures/test/TestExecutorTest.cpp
@@ -24,8 +24,9 @@ using namespace folly;
 TEST(TestExecutor, parallel_run) {
   mutex m;
   set<thread::id> ids;
-  auto executor = std::make_unique<TestExecutor>();
+  auto executor = std::make_unique<TestExecutor>(4);
   const auto numThreads = executor->numThreads();
+  EXPECT_EQ(4, numThreads);
   for (auto idx = 0U; idx < numThreads * 10; ++idx) {
     executor->add([&m, &ids]() mutable {
       /* sleep override */ this_thread::sleep_for(milliseconds(100));
-- 
2.34.1