ThreadPool.cpp

/*
 * Copyright(c) Sophist Solutions, Inc. 1990-2026.  All rights reserved
 */
#include "Stroika/Foundation/StroikaPreComp.h"
 
#include "Stroika/Foundation/Characters/Format.h"
#include "Stroika/Foundation/Characters/StringBuilder.h"
#include "Stroika/Foundation/Characters/ToString.h"
#include "Stroika/Foundation/Debug/Main.h"
#include "Stroika/Foundation/Debug/TimingTrace.h"
#include "Stroika/Foundation/Execution/Common.h"
#include "Stroika/Foundation/Execution/Finally.h"
#include "Stroika/Foundation/Execution/Sleep.h"
#include "Stroika/Foundation/Execution/Timeout.h"
#include "Stroika/Foundation/Memory/BlockAllocated.h"
 
#include "ThreadPool.h"
 
using namespace Stroika::Foundation;
using namespace Stroika::Foundation::Characters;
using namespace Stroika::Foundation::Containers;
using namespace Stroika::Foundation::Execution;
 
using Memory::MakeSharedPtr;
 
// Comment this in to turn on aggressive noisy DbgTrace in this module
//#define   USE_NOISY_TRACE_IN_THIS_MODULE_       1
 
namespace {
    constexpr bool kEmitDbgTraceOnThreadPoolEntryExceptions_ = qStroika_Foundation_Debug_AssertionsChecked;
}
 
class ThreadPool::MyRunnable_ {
public:
    MyRunnable_ (ThreadPool& threadPool)
        : fThreadPool{threadPool}
    {
    }
 
public:
    // whether this thread is running a task - and if so, the task (assumes the caller holds fThreadPool.fCriticalSection_)
    bool IsRunning () const
    {
        return fCurTask.fTask != nullptr;
    }
    bool IsRunning (TaskID task) const
    {
        return IsRunning () and fCurTask.fID == task;
    }
 
public:
    nonvirtual void Run ()
    {
        // For NOW - allow ThreadAbort to just kill this thread. In the future - we may want to implement some sort of restartability
 
        // Keep grabbing new tasks, and running them
        while (true) {
            {
                if constexpr (qStroika_Foundation_Debug_AssertionsChecked) {
                    [[maybe_unused]] lock_guard critSec{fThreadPool.fCriticalSection_};
                    Assert (not IsRunning ());
                }
                // Subtle point, but we must copy directly into fCurTask (WaitForNextTask_ call) so its filled in under lock
                // while being moved so task moves from pending to in-use without ever temporarily disappearing from known tasks lists
                fThreadPool.WaitForNextTask_ (&fCurTask); // This will block INDEFINITELY until ThreadAbort throws out or we have a new task to run
                if constexpr (qStroika_Foundation_Debug_AssertionsChecked) {
                    [[maybe_unused]] lock_guard critSec{fThreadPool.fCriticalSection_};
                    Assert (IsRunning ());
                }
            }
            [[maybe_unused]] auto&& cleanup = Finally ([this] () noexcept {
                Time::TimePointSeconds taskStartedAt;
                {
                    [[maybe_unused]] lock_guard critSec{fThreadPool.fCriticalSection_};
                    fCurTask      = PendingTaskInfo_{};
                    taskStartedAt = fCurTaskStartedAt;
                }
                if (fThreadPool.fCollectStatistics_) {
                    [[maybe_unused]] lock_guard critSec{fThreadPool.fCriticalSection_};
                    ++fThreadPool.fCollectedTaskStats_.fNumberOfTasksCompleted;
                    ++fThreadPool.fCollectedTaskStats_.fNumberOfTasksReporting;
                    fThreadPool.fCollectedTaskStats_.fTotalTimeConsumed += Time::GetTickCount () - taskStartedAt;
                }
            });
            try {
                // Use lock to access fCurTask, but don't hold the lock during run, so others can call getcurrenttask
                ThreadPool::TaskType task2Run;
                {
                    [[maybe_unused]] lock_guard critSec{fThreadPool.fCriticalSection_};
                    task2Run          = fCurTask.fTask;
                    fCurTaskStartedAt = Time::GetTickCount ();
                }
                task2Run ();
            }
            catch (const Thread::AbortException&) {
                throw; // cancel this thread
            }
            catch (...) {
                // other exceptions WARNING WITH DEBUG MESSAGE - but otherwise - EAT/IGNORE
                if constexpr (kEmitDbgTraceOnThreadPoolEntryExceptions_ and qStroika_Foundation_Debug_DefaultTracingOn) {
                    DbgTrace ("in threadpool, ignoring exception running task: {}"_f, current_exception ());
                }
            }
        }
    }
 
public:
    ThreadPool& fThreadPool;
    // fThreadPool.fCriticalSection_ protect access to fCurTask/fCurTaskStartedAt - very short duration
    PendingTaskInfo_       fCurTask;              // the task running, its ID and name; fTask null if none
    Time::TimePointSeconds fCurTaskStartedAt{0s}; // meaningless if not IsRunning ()
};
 
/*
 ********************************************************************************
 ************************ Execution::ThreadPool::Statistics *********************
 ********************************************************************************
 */
Characters::String ThreadPool::Statistics::ToString () const
{
    StringBuilder sb;
    sb << "{"sv;
    sb << "numberOfTasksAdded: "sv << fNumberOfTasksAdded;
    sb << ", numberOfTasksCompleted: "sv << fNumberOfTasksCompleted;
    sb << ", totalTimeConsumed: "sv << fTotalTimeConsumed;
    sb << "}"sv;
    return sb;
}
 
/*
 ********************************************************************************
 ****************************** Execution::ThreadPool ***************************
 ********************************************************************************
 */
ThreadPool::ThreadPool (const Options& options)
    : fCollectStatistics_{options.fCollectStatistics}
    , fDefaultQMax_{options.fQMax}
    , fThreadPoolName_{options.fThreadPoolName}
{
    Require (Debug::AppearsDuringMainLifetime ());
    SetPoolSize (options.fThreadCount);
}
 
ThreadPool::~ThreadPool ()
{
    Require (Debug::AppearsDuringMainLifetime ());
    AbortAndWaitForDone_ ();
}
 
unsigned int ThreadPool::GetPoolSize () const
{
    [[maybe_unused]] lock_guard critSec{fCriticalSection_};
    return static_cast<unsigned int> (fThreads_.size ());
}
 
void ThreadPool::SetPoolSize (unsigned int poolSize)
{
    Debug::TraceContextBumper ctx{"ThreadPool::SetPoolSize", "newPoolSize={}, *this={}"_f, poolSize, ToString ()};
    Require (not fAborted_);
    [[maybe_unused]] lock_guard critSec{fCriticalSection_};
    DbgTrace ("fThreads_.size ()={}"_f, fThreads_.size ());
    while (poolSize > fThreads_.size ()) {
        fThreads_.Add (mkThread_ ());
    }
 
    // Still quite weak implementation of REMOVAL
    while (poolSize < fThreads_.size ()) {
        // iterate over threads if any not busy, remove that them first
        bool anyFoundToKill = false;
        for (Iterator<TPInfo_> i = fThreads_.begin (); i != fThreads_.end (); ++i) {
            if (not i->fRunnable->IsRunning ()) {
                // since we have fCriticalSection_ - we can safely remove this thread
                fThreads_.Remove (i);
                anyFoundToKill = true;
                break;
            }
        }
        if (not anyFoundToKill) {
            // @todo - fix this better/eventually - either throw or wait...
            DbgTrace ("Failed to lower the loop size - cuz all threads busy - giving up"_f);
            return;
        }
    }
}
 
auto ThreadPool::AddTask (const TaskType& task, QMax qmax, const optional<Characters::String>& name) -> TaskID
{
    // also INTENTIONALLY dont hold lock long enuf to make this work 100% reliably cuz these magic numbers dont need to be precise, just approximate
    // @todo rewrite this with condition variables, so more efficient and wakes up/times out appropriately...
    Time::TimePointSeconds timeoutAt = Time::GetTickCount () + qmax.fAddBlockTimeout;
#if qStroika_Foundation_Debug_AssertionsChecked
    bool blockedAtLeastOnce = false;
#endif
    while (true) {
        if (GetPendingTasksCount () >= qmax.fLength) [[unlikely]] {
#if qStroika_Foundation_Debug_AssertionsChecked
            if (not blockedAtLeastOnce) {
                DbgTrace ("Blocking inside ThreadPool::AddTask due to excessive pending task count"_f);
                blockedAtLeastOnce = true;
            }
#endif
            ThrowTimeoutExceptionAfter (timeoutAt);
            Execution::Sleep (500ms); // @todo fix and use condition variable - but good news is can only happen if fAddBlockTimeout != 0s
        }
        else {
            return AddTask_ (task, name);
        }
    }
}
 
auto ThreadPool::AddTask_ (const TaskType& task, const optional<Characters::String>& name) -> TaskID
{
#if USE_NOISY_TRACE_IN_THIS_MODULE_
    Debug::TraceContextBumper ctx{"ThreadPool::AddTask_"};
#endif
    Require (not fAborted_);
    const TaskID id = TaskID::New ();
    {
        [[maybe_unused]] lock_guard critSec{fCriticalSection_};
        fPendingTasks_.push_back (PendingTaskInfo_{.fID = id, .fTask = task, .fName = name});
#if USE_NOISY_TRACE_IN_THIS_MODULE_
        DbgTrace ("fPendingTasks.size () now = {}"_f, (int)fPendingTasks_.size ());
#endif
        ++fCollectedTaskStats_.fNumberOfTasksAdded;
    }
 
    // Notify any waiting threads to wakeup and claim the next task
    fTasksMaybeAdded_.Set ();
    // this would be a race - if aborting and adding tasks at the same time
    Require (not fAborted_);
    return id;
}
 
void ThreadPool::AbortTask (TaskID task, Time::DurationSeconds timeout)
{
    Debug::TraceContextBumper ctx{"ThreadPool::AbortTask"};
    {
        // First see if its in the Q
        [[maybe_unused]] lock_guard critSec{fCriticalSection_};
        for (auto i = fPendingTasks_.begin (); i != fPendingTasks_.end (); ++i) {
            if (i->fID == task) {
                fPendingTasks_.erase (i);
                return;
            }
        }
    }
 
    // If we got here - its NOT in the task Q, so maybe its already running.
    //
    //
 
    // TODO:
    //      We walk the list of existing threads and ask each one if its (indirected - running task) is the given one and abort that task.
    //      But that requires we can RESTART an ABORTED thread (or that we remove it from the list - maybe thats better). THat COULD be OK
    //      actually since it involves on API changes and makes sense. The only slight issue is a peformace one but probably for something
    //      quite rare.
    //
    //      Anyhow SB OK for now to just not allow aborting a task which has already started....
    Thread::Ptr thread2Kill;
    {
        [[maybe_unused]] lock_guard critSec{fCriticalSection_};
        for (Iterator<TPInfo_> i = fThreads_.begin (); i != fThreads_.end (); ++i) {
            if (i->fRunnable->IsRunning (task)) {
                thread2Kill = i->fThread;
                fThreads_.Update (i, mkThread_ ());
                break;
            }
        }
    }
    if (thread2Kill != nullptr) {
        thread2Kill.AbortAndWaitForDone (timeout);
    }
}
 
void ThreadPool::AbortTasks (Time::DurationSeconds timeout)
{
    Debug::TraceContextBumper ctx{"ThreadPool::AbortTasks"};
    auto                      tps = GetPoolSize ();
    {
        [[maybe_unused]] lock_guard critSec{fCriticalSection_};
        fPendingTasks_.clear ();
        for (const TPInfo_& ti : fThreads_) {
            ti.fThread.Abort ();
        }
    }
    {
        // @todo maybe fix unsafe - waiting here holding the critsec lock - seems deadlock waiting to happen - LGP 2023-11-05
        [[maybe_unused]] lock_guard critSec{fCriticalSection_};
        for (const TPInfo_& ti : fThreads_) {
            // @todo fix wrong timeout value here
            ti.fThread.AbortAndWaitForDone (timeout);
        }
        fThreads_.RemoveAll ();
    }
    // hack - not a good design or impl!!! - waste to recreate if not needed!
    SetPoolSize (tps);
}
 
bool ThreadPool::IsPresent (TaskID task) const
{
    {
        // First see if its in the Q
        [[maybe_unused]] lock_guard critSec{fCriticalSection_};
        for (auto i = fPendingTasks_.begin (); i != fPendingTasks_.end (); ++i) {
            if (i->fID == task) {
                return true;
            }
        }
        // then check if running
        for (auto i = fThreads_.begin (); i != fThreads_.end (); ++i) {
            if (i->fRunnable->IsRunning (task)) {
                return true;
            }
        }
    }
    return false;
}
 
bool ThreadPool::IsRunning (TaskID task) const
{
    [[maybe_unused]] lock_guard critSec{fCriticalSection_};
    for (auto i = fThreads_.begin (); i != fThreads_.end (); ++i) {
        if (i->fRunnable->IsRunning (task)) {
            return true;
        }
    }
    return false;
}
 
void ThreadPool::WaitForTask (TaskID task, Time::DurationSeconds timeout) const
{
    Debug::TraceContextBumper ctx{"ThreadPool::WaitForTask"};
    // Inefficient / VERY SLOPPY impl - @todo fix use WaitableEvent or condition variables...
    using Time::TimePointSeconds;
    TimePointSeconds timeoutAt = timeout + Time::GetTickCount ();
    while (true) {
        if (not IsPresent (task)) {
            return;
        }
        Time::DurationSeconds remaining = timeoutAt - Time::GetTickCount ();
        ThrowTimeoutExceptionAfter (timeoutAt);
        Execution::Sleep (min<Time::DurationSeconds> (remaining, 1.0s));
    }
}
 
auto ThreadPool::GetTasks () const -> Collection<TaskInfo>
{
    Collection<TaskInfo>        result;
    [[maybe_unused]] lock_guard critSec{fCriticalSection_};
    for (const auto& ti : fPendingTasks_) {
        result.Add (TaskInfo{.fID = ti.fID, .fTask = ti.fTask, .fName = ti.fName});
    }
    for (auto i = fThreads_.begin (); i != fThreads_.end (); ++i) {
        if (i->fRunnable->IsRunning ()) {
            const PendingTaskInfo_& cur = i->fRunnable->fCurTask;
            result.Add (TaskInfo{.fID = cur.fID, .fTask = cur.fTask, .fName = cur.fName, .fRunningSince = i->fRunnable->fCurTaskStartedAt});
        }
    }
    return result;
}
 
auto ThreadPool::GetRunningTasks () const -> Collection<TaskID>
{
    Collection<TaskID>          result;
    [[maybe_unused]] lock_guard critSec{fCriticalSection_};
    for (auto i = fThreads_.begin (); i != fThreads_.end (); ++i) {
        if (i->fRunnable->IsRunning ()) {
            result.Add (i->fRunnable->fCurTask.fID);
        }
    }
    return result;
}
 
size_t ThreadPool::GetTasksCount () const
{
    // First see if its in the Q
    [[maybe_unused]] lock_guard critSec{fCriticalSection_};
    size_t                      count = fPendingTasks_.size ();
    for (auto i = fThreads_.begin (); i != fThreads_.end (); ++i) {
        AssertNotNull (i->fRunnable);
        if (i->fRunnable->IsRunning ()) {
            ++count;
        }
    }
    return count;
}
 
auto ThreadPool::GetPendingTasks () const -> Collection<TaskID>
{
    Collection<TaskID>          result;
    [[maybe_unused]] lock_guard critSec{fCriticalSection_};
    for (const auto& i : fPendingTasks_) {
        result.Add (i.fID);
    }
    return result;
}
 
size_t ThreadPool::GetPendingTasksCount () const
{
    [[maybe_unused]] lock_guard critSec{fCriticalSection_};
    return fPendingTasks_.size ();
}
 
void ThreadPool::WaitForTasksDoneUntil (const Iterable<TaskID>& tasks, Time::TimePointSeconds timeoutAt) const
{
#if USE_NOISY_TRACE_IN_THIS_MODULE_
    Debug::TraceContextBumper ctx{Stroika_Foundation_Debug_OptionalizeTraceArgs (
        "ThreadPool::WaitForTasksDoneUntil", "*this={}, tasks={}, timeoutAt={}"_f, ToString (), ToString (tasks), timeoutAt)};
#endif
    Thread::CheckForInterruption ();
    for (const auto& task : tasks) {
        auto now = Time::GetTickCount ();
        ThrowTimeoutExceptionAfter (timeoutAt);
        WaitForTask (task, timeoutAt - now);
    }
}
 
void ThreadPool::WaitForTasksDoneUntil (Time::TimePointSeconds timeoutAt) const
{
#if USE_NOISY_TRACE_IN_THIS_MODULE_
    Debug::TraceContextBumper ctx{Stroika_Foundation_Debug_OptionalizeTraceArgs ("ThreadPool::WaitForTasksDoneUntil",
                                                                                 "*this={}, timeoutAt={}"_f, ToString (), timeoutAt)};
#endif
    Thread::CheckForInterruption ();
    // @todo - use WaitablEvent - this is a horrible implementation
    while (GetTasksCount () != 0) {
        ThrowTimeoutExceptionAfter (timeoutAt);
        Execution::Sleep (100ms);
    }
}
 
void ThreadPool::ResetStatistics ()
{
    [[maybe_unused]] lock_guard critSec{fCriticalSection_};
    Require (fCollectStatistics_);
    fCollectedTaskStats_ = {};
}
 
auto ThreadPool::GetCurrentStatistics () const -> Statistics
{
    [[maybe_unused]] lock_guard critSec{fCriticalSection_};
    Require (fCollectStatistics_);
    return fCollectedTaskStats_;
}
 
void ThreadPool::Abort_ () noexcept
{
    Thread::SuppressInterruptionInContext suppressCtx; // must cleanly shut down each of our sub-threads - even if our thread is aborting... don't be half-way aborted
    Debug::TraceContextBumper ctx{Stroika_Foundation_Debug_OptionalizeTraceArgs ("ThreadPool::Abort_", "*this={}"_f, ToString ())};
    Debug::TimingTrace        tt{1.0s};
    fAborted_ = true; // No race, because fAborted never 'unset'
                      // no need to set fTasksMaybeAdded_, since aborting each thread should be sufficient
    {
        // Clear the task Q and then abort each thread
        [[maybe_unused]] lock_guard critSec{fCriticalSection_};
        fPendingTasks_.clear ();
        for (const TPInfo_& ti : fThreads_) {
            ti.fThread.Abort ();
        }
    }
}
 
void ThreadPool::AbortAndWaitForDone_ () noexcept
{
    Thread::SuppressInterruptionInContext suppressCtx; // cuz we must shutdown owned threads
#if USE_NOISY_TRACE_IN_THIS_MODULE_
    Debug::TraceContextBumper ctx{Stroika_Foundation_Debug_OptionalizeTraceArgs ("ThreadPool::AbortAndWaitForDone_", "*this={}"_f, ToString ())};
    Debug::TimingTrace tt{1.0s};
#endif
    try {
        Abort_ (); // to get the rest of the threadpool abort stuff triggered - flag saying aborting
        Collection<Thread::Ptr> threadsToShutdown;
        {
            [[maybe_unused]] lock_guard critSec{fCriticalSection_};
            fThreads_.Apply ([&] (const TPInfo_& i) { threadsToShutdown.Add (i.fThread); });
        }
        Thread::AbortAndWaitForDone (threadsToShutdown);
        {
            [[maybe_unused]] lock_guard critSec{fCriticalSection_}; // they should all be shutdown now
            fThreads_.Apply ([&] (const TPInfo_& i) { Assert (i.fThread.IsDone ()); });
            fThreads_.RemoveAll ();
        }
    }
    catch (...) {
        DbgTrace ("ThreadPool::AbortAndWaitForDone_: serious bug/exception"_f);
        AssertNotReached (); // this should never happen due to the SuppressInterruptionInContext...
    }
    Ensure (fThreads_.empty ());
}
 
String ThreadPool::ToString () const
{
    StringBuilder sb;
    sb << "{"sv;
    {
        [[maybe_unused]] lock_guard critSec{fCriticalSection_};
        if (fThreadPoolName_) {
            sb << "pool-name: '{}'"_f(*fThreadPoolName_);
        }
    }
    if (fThreadPoolName_) {
        sb << ", "sv;
    }
    sb << "pending-task-count: {}"_f(GetPendingTasksCount ());
    sb << ", running-task-count: {}"_f(GetRunningTasks ().size ());
    {
        [[maybe_unused]] lock_guard critSec{fCriticalSection_};
        sb << ", pool-thread-count: {}"_f(fThreads_.size ());
    }
    sb << "}"sv;
    return sb;
}
 
// THIS is called NOT from 'this' - but from the context of an OWNED thread of the pool
void ThreadPool::WaitForNextTask_ (PendingTaskInfo_* result)
{
    RequireNotNull (result);
    Require (result->fTask == nullptr);
    while (true) {
        if (fAborted_) [[unlikely]] {
            Throw (Thread::AbortException::kThe);
        }
 
        {
            [[maybe_unused]] lock_guard critSec{fCriticalSection_};
            if (not fPendingTasks_.empty ()) {
                *result = fPendingTasks_.front ();
                fPendingTasks_.pop_front ();
                DbgTrace ("ThreadPool::WaitForNextTask_ () pulled a new task from 'pending-tasks' to run on this thread, leaving pending-task-list-size = {}"_f,
                          fPendingTasks_.size ());
                Ensure (result->fTask != nullptr);
                return;
            }
        }
 
        // Prevent spinwaiting... This event is SET when any new item arrives
        //DbgTrace ("ThreadPool::WaitForNextTask_ () - about to wait for added tasks");
        fTasksMaybeAdded_.WaitAndReset ();
        //DbgTrace ("ThreadPool::WaitForNextTask_ () - completed wait for added tasks");
    }
}
 
ThreadPool::TPInfo_ ThreadPool::mkThread_ ()
{
    shared_ptr<MyRunnable_> r{MakeSharedPtr<ThreadPool::MyRunnable_> (*this)};
    StringBuilder entryName = "TPE #{}"_f(fNextThreadEntryNumber_++); // make name so short cuz unix only shows first 15 chars - http://man7.org/linux/man-pages/man3/pthread_setname_np.3.html
    entryName += " {"sv + fThreadPoolName_.value_or ("anonymous-thread-pool"sv) + "}"sv;
    Thread::Ptr t = Thread::New ([r] () { r->Run (); }, Thread::eAutoStart, entryName); // race condition for updating this number, but who cares - its purely cosmetic...
    return TPInfo_{t, r};
}