Inside Bitcoin

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
15struct ExpectedException : std::runtime_error {
16    explicit ExpectedException(const std::string& msg) : std::runtime_error(msg) {}
17};
18
19struct ThrowTask {
20    void operator()() const { throw ExpectedException("fail"); }
21};
22
23struct 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.
30static 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.
45ThreadPool g_pool{"fuzz"};
46// Global to verify we always have the same number of threads.
47size_t g_num_workers = 3;
48
49static void StartPoolIfNeeded()
50{
51    if (g_pool.WorkersCount() == g_num_workers) return;
52    g_pool.Start(g_num_workers);
53}
54
55static void setup_threadpool_test()
56{
57    // Disable logging entirely. It seems to cause memory leaks.
58    LogInstance().DisableLogging();
59}
60
61FUZZ_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