threadpool.cpp raw

   1  // Copyright (c) The Bitcoin Core developers
   2  // Distributed under the MIT software license, see the accompanying
   3  // file COPYING or http://www.opensource.org/licenses/mit-license.php.
   4  
   5  #include <logging.h>
   6  #include <util/threadpool.h>
   7  
   8  #include <test/fuzz/FuzzedDataProvider.h>
   9  #include <test/fuzz/fuzz.h>
  10  
  11  #include <atomic>
  12  #include <future>
  13  #include <queue>
  14  
  15  struct ExpectedException : std::runtime_error {
  16      explicit ExpectedException(const std::string& msg) : std::runtime_error(msg) {}
  17  };
  18  
  19  struct ThrowTask {
  20      void operator()() const { throw ExpectedException("fail"); }
  21  };
  22  
  23  struct CounterTask {
  24      std::atomic_uint32_t& m_counter;
  25      explicit CounterTask(std::atomic_uint32_t& counter) : m_counter{counter} {}
  26      void operator()() const { m_counter.fetch_add(1, std::memory_order_relaxed); }
  27  };
  28  
  29  // Waits for a future to complete. Increments 'fail_counter' if the expected exception is thrown.
  30  static void GetFuture(std::future<void>& future, uint32_t& fail_counter)
  31  {
  32      try {
  33          future.get();
  34      } catch (const ExpectedException&) {
  35          fail_counter++;
  36      } catch (...) {
  37          assert(false && "Unexpected exception type");
  38      }
  39  }
  40  
  41  // Global thread pool for fuzzing. Persisting it across iterations prevents
  42  // the excessive thread creation/destruction overhead that can lead to
  43  // instability in the fuzzing environment.
  44  // This is also how we use it in the app's lifecycle.
  45  ThreadPool g_pool{"fuzz"};
  46  // Global to verify we always have the same number of threads.
  47  size_t g_num_workers = 3;
  48  
  49  static void StartPoolIfNeeded()
  50  {
  51      if (g_pool.WorkersCount() == g_num_workers) return;
  52      g_pool.Start(g_num_workers);
  53  }
  54  
  55  static void setup_threadpool_test()
  56  {
  57      // Disable logging entirely. It seems to cause memory leaks.
  58      LogInstance().DisableLogging();
  59  }
  60  
  61  FUZZ_TARGET(threadpool, .init = setup_threadpool_test)
  62  {
  63      // Because LibAFL calls fork() after calling the init setup function,
  64      // the child processes end up having one thread active and no workers.
  65      // To work around this limitation, start thread pool inside the first runner.
  66      StartPoolIfNeeded();
  67  
  68      FuzzedDataProvider fuzzed_data_provider(buffer.data(), buffer.size());
  69  
  70      const uint32_t num_tasks = fuzzed_data_provider.ConsumeIntegralInRange<uint32_t>(0, 1024);
  71      assert(g_pool.WorkersCount() == g_num_workers);
  72      assert(g_pool.WorkQueueSize() == 0);
  73  
  74      // Counters
  75      std::atomic_uint32_t task_counter{0};
  76      uint32_t fail_counter{0};
  77      uint32_t expected_task_counter{0};
  78      uint32_t expected_fail_tasks{0};
  79  
  80      std::queue<std::future<void>> futures;
  81      for (uint32_t i = 0; i < num_tasks; ++i) {
  82          const bool will_throw = fuzzed_data_provider.ConsumeBool();
  83          const bool wait_immediately = fuzzed_data_provider.ConsumeBool();
  84  
  85          std::future<void> fut;
  86          if (will_throw) {
  87              expected_fail_tasks++;
  88              fut = *Assert(g_pool.Submit(ThrowTask{}));
  89          } else {
  90              expected_task_counter++;
  91              fut = *Assert(g_pool.Submit(CounterTask{task_counter}));
  92          }
  93  
  94          // If caller wants to wait immediately, consume the future here (safe).
  95          if (wait_immediately) {
  96              // Waits for this task to complete immediately; prior queued tasks may also complete
  97              // as they were queued earlier.
  98              GetFuture(fut, fail_counter);
  99          } else {
 100              // Store task for a posterior check
 101              futures.emplace(std::move(fut));
 102          }
 103      }
 104  
 105      // Drain remaining futures
 106      while (!futures.empty()) {
 107          auto fut = std::move(futures.front());
 108          futures.pop();
 109          GetFuture(fut, fail_counter);
 110      }
 111  
 112      assert(g_pool.WorkQueueSize() == 0);
 113      assert(task_counter.load() == expected_task_counter);
 114      assert(fail_counter == expected_fail_tasks);
 115  }
 116