SERVER-24945 added support to ThreadPoolMock for post-initialization function on worker thread.

Author: benety

src/mongo/executor/network_interface_mock_test.cpp

@@ -47,7 +47,8 @@ namespace {
 
 class NetworkInterfaceMockTest : public mongo::unittest::Test {
 public:
-    NetworkInterfaceMockTest() : _net{}, _executor(&_net, 1), _tearDownCalled(false) {}
+    NetworkInterfaceMockTest()
+        : _net{}, _executor(&_net, 1, ThreadPoolMock::Options()), _tearDownCalled(false) {}
 
     NetworkInterfaceMock& net() {
         return _net;

src/mongo/executor/thread_pool_mock.cpp

@@ -38,8 +38,8 @@
 namespace mongo {
 namespace executor {
 
-ThreadPoolMock::ThreadPoolMock(NetworkInterfaceMock* net, int32_t prngSeed)
-    : _prng(prngSeed), _net(net) {}
+ThreadPoolMock::ThreadPoolMock(NetworkInterfaceMock* net, int32_t prngSeed, Options options)
+    : _options(std::move(options)), _prng(prngSeed), _net(net) {}
 
 ThreadPoolMock::~ThreadPoolMock() {
     stdx::unique_lock<stdx::mutex> lk(_mutex);
@@ -64,6 +64,7 @@ void ThreadPoolMock::startup() {
     invariant(!_worker.joinable());
     _started = true;
     _worker = stdx::thread([this] {
+        _options.onCreateThread();
         stdx::unique_lock<stdx::mutex> lk(_mutex);
         consumeTasks(&lk);
     });

src/mongo/executor/thread_pool_mock.h

@@ -49,11 +49,20 @@ class NetworkInterfaceMock;
  */
 class ThreadPoolMock final : public ThreadPoolInterface {
 public:
+    /**
+     * Structure used to configure an instance of ThreadPoolMock.
+     */
+    struct Options {
+        // This function is run before the worker thread begins consuming tasks.
+        using OnCreateThreadFn = stdx::function<void()>;
+        OnCreateThreadFn onCreateThread = []() {};
+    };
+
     /**
      * Create an instance that interlocks with "net". "prngSeed" seeds the pseudorandom number
      * generator that is used to determine which schedulable task runs next.
      */
-    ThreadPoolMock(NetworkInterfaceMock* net, int32_t prngSeed);
+    ThreadPoolMock(NetworkInterfaceMock* net, int32_t prngSeed, Options options);
     ~ThreadPoolMock();
 
     void startup() override;
@@ -64,6 +73,9 @@ class ThreadPoolMock final : public ThreadPoolInterface {
 private:
     void consumeTasks(stdx::unique_lock<stdx::mutex>* lk);
 
+    // These are the options with which the pool was configured at construction time.
+    const Options _options;
+
     stdx::mutex _mutex;
     stdx::thread _worker;
     std::vector<Task> _tasks;

src/mongo/executor/thread_pool_task_executor_test_fixture.cpp

@@ -37,15 +37,20 @@ namespace mongo {
 namespace executor {
 
 std::unique_ptr<ThreadPoolTaskExecutor> makeThreadPoolTestExecutor(
-    std::unique_ptr<NetworkInterfaceMock> net) {
+    std::unique_ptr<NetworkInterfaceMock> net, ThreadPoolMock::Options options) {
     auto netPtr = net.get();
-    return stdx::make_unique<ThreadPoolTaskExecutor>(stdx::make_unique<ThreadPoolMock>(netPtr, 1),
-                                                     std::move(net));
+    return stdx::make_unique<ThreadPoolTaskExecutor>(
+        stdx::make_unique<ThreadPoolMock>(netPtr, 1, std::move(options)), std::move(net));
+}
+
+ThreadPoolMock::Options ThreadPoolExecutorTest::makeThreadPoolMockOptions() const {
+    return ThreadPoolMock::Options();
 }
 
 std::unique_ptr<TaskExecutor> ThreadPoolExecutorTest::makeTaskExecutor(
     std::unique_ptr<NetworkInterfaceMock> net) {
-    return makeThreadPoolTestExecutor(std::move(net));
+    auto options = makeThreadPoolMockOptions();
+    return makeThreadPoolTestExecutor(std::move(net), std::move(options));
 }
 
 }  // namespace executor

src/mongo/executor/thread_pool_task_executor_test_fixture.h

@@ -32,6 +32,7 @@
 
 #include "mongo/executor/network_interface_mock.h"
 #include "mongo/executor/task_executor_test_fixture.h"
+#include "mongo/executor/thread_pool_mock.h"
 #include "mongo/executor/thread_pool_task_executor.h"
 
 namespace mongo {
@@ -41,13 +42,15 @@ namespace executor {
  * Makes a new ThreadPoolTaskExecutor for use in unit tests.
  */
 std::unique_ptr<ThreadPoolTaskExecutor> makeThreadPoolTestExecutor(
-    std::unique_ptr<NetworkInterfaceMock> net);
+    std::unique_ptr<NetworkInterfaceMock> net,
+    executor::ThreadPoolMock::Options options = executor::ThreadPoolMock::Options());
 
 /**
  * Useful fixture class for tests that use a ThreadPoolTaskExecutor.
  */
 class ThreadPoolExecutorTest : public TaskExecutorTest {
 private:
+    virtual ThreadPoolMock::Options makeThreadPoolMockOptions() const;
     std::unique_ptr<TaskExecutor> makeTaskExecutor(
         std::unique_ptr<NetworkInterfaceMock> net) override;
 };