// Copyright (c) The Bitcoin Core developers

// Distributed under the MIT software license, see the accompanying

// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include <logging.h>

#include <util/threadpool.h>

#include <test/fuzz/FuzzedDataProvider.h>

#include <test/fuzz/fuzz.h>

#include <atomic>

#include <future>

#include <queue>

struct ExpectedException : std::runtime_error {

    explicit ExpectedException(const std::string& msg) : std::runtime_error(msg) {}

};

struct ThrowTask {

    void operator()() const { throw ExpectedException("fail"); }

};

struct CounterTask {

    std::atomic_uint32_t& m_counter;

    explicit CounterTask(std::atomic_uint32_t& counter) : m_counter{counter} {}

    void operator()() const { m_counter.fetch_add(1, std::memory_order_relaxed); }

};

// Waits for a future to complete. Increments 'fail_counter' if the expected exception is thrown.

static void GetFuture(std::future<void>& future, uint32_t& fail_counter)

{

    try {

        future.get();

    } catch (const ExpectedException&) {

        fail_counter++;

    } catch (...) {

        assert(false && "Unexpected exception type");

  Branch (37:9): [Folded - Ignored]
  Branch (37:9): [Folded - Ignored]
  Branch (37:9): [Folded - Ignored]

    }

}

// Global thread pool for fuzzing. Persisting it across iterations prevents

// the excessive thread creation/destruction overhead that can lead to

// instability in the fuzzing environment.

// This is also how we use it in the app's lifecycle.

ThreadPool g_pool{"fuzz"};

// Global to verify we always have the same number of threads.

size_t g_num_workers = 3;

static void StartPoolIfNeeded()

{

    if (g_pool.WorkersCount() == g_num_workers) return;

  Branch (51:9): [True: 133, False: 1]

    g_pool.Start(g_num_workers);

}

static void setup_threadpool_test()

{

    // Disable logging entirely. It seems to cause memory leaks.

    LogInstance().DisableLogging();

}

FUZZ_TARGET(threadpool, .init = setup_threadpool_test)

{

    // Because LibAFL calls fork() after calling the init setup function,

    // the child processes end up having one thread active and no workers.

    // To work around this limitation, start thread pool inside the first runner.

    StartPoolIfNeeded();

    FuzzedDataProvider fuzzed_data_provider(buffer.data(), buffer.size());

    const uint32_t num_tasks = fuzzed_data_provider.ConsumeIntegralInRange<uint32_t>(0, 1024);

    assert(g_pool.WorkersCount() == g_num_workers);

  Branch (71:5): [True: 134, False: 0]

    assert(g_pool.WorkQueueSize() == 0);

  Branch (72:5): [True: 134, False: 0]

    // Counters

    std::atomic_uint32_t task_counter{0};

    uint32_t fail_counter{0};

    uint32_t expected_task_counter{0};

    uint32_t expected_fail_tasks{0};

    std::queue<std::future<void>> futures;

    for (uint32_t i = 0; i < num_tasks; ++i) {

  Branch (81:26): [True: 81.1k, False: 134]

        const bool will_throw = fuzzed_data_provider.ConsumeBool();

        const bool wait_immediately = fuzzed_data_provider.ConsumeBool();

        std::future<void> fut;

        if (will_throw) {

  Branch (86:13): [True: 12.3k, False: 68.8k]

            expected_fail_tasks++;

            fut = *Assert(g_pool.Submit(ThrowTask{}));

        } else {

            expected_task_counter++;

            fut = *Assert(g_pool.Submit(CounterTask{task_counter}));

        }

        // If caller wants to wait immediately, consume the future here (safe).

        if (wait_immediately) {

  Branch (95:13): [True: 12.2k, False: 68.9k]

            // Waits for this task to complete immediately; prior queued tasks may also complete

            // as they were queued earlier.

            GetFuture(fut, fail_counter);

        } else {

            // Store task for a posterior check

            futures.emplace(std::move(fut));

        }

    }

    // Drain remaining futures

    while (!futures.empty()) {

  Branch (106:12): [True: 68.9k, False: 134]

        auto fut = std::move(futures.front());

        futures.pop();

        GetFuture(fut, fail_counter);

    }

    assert(g_pool.WorkQueueSize() == 0);

  Branch (112:5): [True: 134, False: 0]

    assert(task_counter.load() == expected_task_counter);

  Branch (113:5): [True: 134, False: 0]

    assert(fail_counter == expected_fail_tasks);

  Branch (114:5): [True: 134, False: 0]

}