intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

5 files changed, 436 insertions, 0 deletions

diff --git a/library/cpp/threading/task_scheduler/task_scheduler.cpp b/library/cpp/threading/task_scheduler/task_scheduler.cpp
new file mode 100644
index 0000000000..174dde4bf7
--- /dev/null
+++ b/library/cpp/threading/task_scheduler/task_scheduler.cpp
@@ -0,0 +1,246 @@
+#include "task_scheduler.h"
+
+#include <util/system/thread.h>
+#include <util/string/cast.h>
+#include <util/stream/output.h>
+
+TTaskScheduler::ITask::~ITask() {}
+TTaskScheduler::IRepeatedTask::~IRepeatedTask() {}
+
+
+
+class TTaskScheduler::TWorkerThread
+    : public ISimpleThread
+{
+public:
+    TWorkerThread(TTaskScheduler& state)
+        : Scheduler_(state)
+    {
+    }
+
+    TString DebugState = "?";
+    TString DebugId = "";
+private:
+    void* ThreadProc() noexcept override {
+        Scheduler_.WorkerFunc(this);
+        return nullptr;
+    }
+private:
+    TTaskScheduler& Scheduler_;
+};
+
+
+
+TTaskScheduler::TTaskScheduler(size_t threadCount, size_t maxTaskCount)
+    : MaxTaskCount_(maxTaskCount)
+{
+    for (size_t i = 0; i < threadCount; ++i) {
+        Workers_.push_back(new TWorkerThread(*this));
+        Workers_.back()->DebugId = ToString(i);
+    }
+}
+
+TTaskScheduler::~TTaskScheduler() {
+    try {
+        Stop();
+    } catch (...) {
+        Cdbg << "task scheduled destruction error: " << CurrentExceptionMessage();
+    }
+}
+
+void TTaskScheduler::Start() {
+    for (auto& w : Workers_) {
+        w->Start();
+    }
+}
+
+void TTaskScheduler::Stop() {
+    with_lock (Lock_) {
+        IsStopped_ = true;
+        CondVar_.BroadCast();
+    }
+
+    for (auto& w: Workers_) {
+        w->Join();
+    }
+
+    Workers_.clear();
+    Queue_.clear();
+}
+
+size_t TTaskScheduler::GetTaskCount() const {
+    return static_cast<size_t>(AtomicGet(TaskCounter_));
+}
+
+namespace {
+    class TTaskWrapper
+        : public TTaskScheduler::ITask
+        , TNonCopyable
+    {
+    public:
+        TTaskWrapper(TTaskScheduler::ITaskRef task, TAtomic& counter)
+            : Task_(task)
+            , Counter_(counter)
+        {
+            AtomicIncrement(Counter_);
+        }
+
+        ~TTaskWrapper() override {
+            AtomicDecrement(Counter_);
+        }
+    private:
+        TInstant Process() override {
+            return Task_->Process();
+        }
+    private:
+        TTaskScheduler::ITaskRef Task_;
+        TAtomic& Counter_;
+    };
+}
+
+bool TTaskScheduler::Add(ITaskRef task, TInstant expire) {
+    with_lock (Lock_) {
+        if (!IsStopped_ && Workers_.size() > 0 && GetTaskCount() + 1 <= MaxTaskCount_) {
+            ITaskRef newTask = new TTaskWrapper(task, TaskCounter_);
+            Queue_.insert(std::make_pair(expire, TTaskHolder(newTask)));
+
+            if (!Queue_.begin()->second.WaitingWorker) {
+                CondVar_.Signal();
+            }
+            return true;
+        }
+    }
+
+    return false;
+}
+
+namespace {
+    class TRepeatedTask
+        : public TTaskScheduler::ITask
+    {
+    public:
+        TRepeatedTask(TTaskScheduler::IRepeatedTaskRef task, TDuration period, TInstant deadline)
+            : Task_(task)
+            , Period_(period)
+            , Deadline_(deadline)
+        {
+        }
+    private:
+        TInstant Process() final {
+            Deadline_ += Period_;
+            if (Task_->Process()) {
+                return Deadline_;
+            } else {
+                return TInstant::Max();
+            }
+        }
+    private:
+        TTaskScheduler::IRepeatedTaskRef Task_;
+        TDuration Period_;
+        TInstant Deadline_;
+    };
+}
+
+bool TTaskScheduler::Add(IRepeatedTaskRef task, TDuration period) {
+    const TInstant deadline = Now() + period;
+    ITaskRef t = new TRepeatedTask(task, period, deadline);
+    return Add(t, deadline);
+}
+
+
+const bool debugOutput = false;
+
+void TTaskScheduler::ChangeDebugState(TWorkerThread* thread, const TString& state) {
+    if (!debugOutput) {
+        Y_UNUSED(thread);
+        Y_UNUSED(state);
+        return;
+    }
+
+    thread->DebugState = state;
+
+    TStringStream ss;
+    ss << Now() << " " << thread->DebugId << ":\t";
+    for (auto& w : Workers_) {
+        ss << w->DebugState << " ";
+    }
+    ss << " [" << Queue_.size() << "] [" << TaskCounter_ << "]" <<  Endl;
+    Cerr << ss.Str();
+}
+
+bool TTaskScheduler::Wait(TWorkerThread* thread, TQueueIterator& toWait) {
+    ChangeDebugState(thread, "w");
+    toWait->second.WaitingWorker = thread;
+    return !CondVar_.WaitD(Lock_, toWait->first);
+}
+
+void TTaskScheduler::ChooseFromQueue(TQueueIterator& toWait) {
+    for (TQueueIterator it = Queue_.begin(); it != Queue_.end(); ++it) {
+        if (!it->second.WaitingWorker) {
+            if (toWait == Queue_.end()) {
+                toWait = it;
+            } else if (it->first < toWait->first) {
+                toWait->second.WaitingWorker = nullptr;
+                toWait = it;
+            }
+            break;
+        }
+    }
+}
+
+void TTaskScheduler::WorkerFunc(TWorkerThread* thread) {
+    TThread::SetCurrentThreadName("TaskSchedWorker");
+
+    TQueueIterator toWait = Queue_.end();
+    ITaskRef toDo;
+
+    for (;;) {
+        TInstant repeat = TInstant::Max();
+
+        if (!!toDo) {
+            try {
+                repeat = toDo->Process();
+            } catch (...) {
+               Cdbg << "task scheduler error: " << CurrentExceptionMessage();
+            }
+        }
+
+
+        with_lock (Lock_) {
+            ChangeDebugState(thread, "f");
+
+            if (IsStopped_) {
+                ChangeDebugState(thread, "s");
+                return ;
+            }
+
+            if (!!toDo) {
+                if (repeat < TInstant::Max()) {
+                    Queue_.insert(std::make_pair(repeat, TTaskHolder(toDo)));
+                }
+            }
+
+            toDo = nullptr;
+
+            ChooseFromQueue(toWait);
+
+            if (toWait != Queue_.end()) {
+                if (toWait->first <= Now() || Wait(thread, toWait)) {
+
+                    toDo = toWait->second.Task;
+                    Queue_.erase(toWait);
+                    toWait = Queue_.end();
+
+                    if (!Queue_.empty() && !Queue_.begin()->second.WaitingWorker && Workers_.size() > 1) {
+                        CondVar_.Signal();
+                    }
+
+                    ChangeDebugState(thread, "p");
+                }
+            } else {
+                ChangeDebugState(thread, "e");
+                CondVar_.WaitI(Lock_);
+            }
+        }
+    }
+}
diff --git a/library/cpp/threading/task_scheduler/task_scheduler.h b/library/cpp/threading/task_scheduler/task_scheduler.h
new file mode 100644
index 0000000000..df4da941a8
--- /dev/null
+++ b/library/cpp/threading/task_scheduler/task_scheduler.h
@@ -0,0 +1,86 @@
+#pragma once
+
+#include <util/generic/vector.h>
+#include <util/generic/ptr.h>
+#include <util/generic/map.h>
+
+#include <util/datetime/base.h>
+
+#include <util/system/condvar.h>
+#include <util/system/mutex.h>
+
+class TTaskScheduler {
+public:
+    class ITask;
+    using ITaskRef = TIntrusivePtr<ITask>;
+
+    class IRepeatedTask;
+    using IRepeatedTaskRef = TIntrusivePtr<IRepeatedTask>;
+public:
+    explicit TTaskScheduler(size_t threadCount = 1, size_t maxTaskCount = Max<size_t>());
+    ~TTaskScheduler();
+
+    void Start();
+    void Stop();
+
+    bool Add(ITaskRef task, TInstant expire);
+    bool Add(IRepeatedTaskRef task, TDuration period);
+
+    size_t GetTaskCount() const;
+private:
+    class TWorkerThread;
+
+    struct TTaskHolder {
+        explicit TTaskHolder(ITaskRef& task)
+            : Task(task)
+        {
+        }
+    public:
+        ITaskRef Task;
+        TWorkerThread* WaitingWorker = nullptr;
+    };
+
+    using TQueueType = TMultiMap<TInstant, TTaskHolder>;
+    using TQueueIterator = TQueueType::iterator;
+private:
+    void ChangeDebugState(TWorkerThread* thread, const TString& state);
+    void ChooseFromQueue(TQueueIterator& toWait);
+    bool Wait(TWorkerThread* thread, TQueueIterator& toWait);
+
+    void WorkerFunc(TWorkerThread* thread);
+private:
+    bool IsStopped_ = false;
+
+    TAtomic TaskCounter_ = 0;
+    TQueueType Queue_;
+
+    TCondVar CondVar_;
+    TMutex Lock_;
+
+    TVector<TAutoPtr<TWorkerThread>> Workers_;
+
+    const size_t MaxTaskCount_;
+};
+
+class TTaskScheduler::ITask
+    : public TAtomicRefCount<ITask>
+{
+public:
+    virtual ~ITask();
+
+    virtual TInstant Process() {//returns time to repeat this task
+        return TInstant::Max();
+    }
+};
+
+class TTaskScheduler::IRepeatedTask
+    : public TAtomicRefCount<IRepeatedTask>
+{
+public:
+    virtual ~IRepeatedTask();
+
+    virtual bool Process() {//returns if to repeat task again
+        return false;
+    }
+};
+
diff --git a/library/cpp/threading/task_scheduler/task_scheduler_ut.cpp b/library/cpp/threading/task_scheduler/task_scheduler_ut.cpp
new file mode 100644
index 0000000000..3b5203194a
--- /dev/null
+++ b/library/cpp/threading/task_scheduler/task_scheduler_ut.cpp
@@ -0,0 +1,86 @@
+#include <algorithm>
+#include <library/cpp/testing/unittest/registar.h>
+
+#include <util/stream/output.h>
+#include <util/system/atomic.h>
+#include <util/generic/vector.h>
+
+#include "task_scheduler.h"
+
+class TTaskSchedulerTest: public TTestBase {
+        UNIT_TEST_SUITE(TTaskSchedulerTest);
+            UNIT_TEST(Test);
+        UNIT_TEST_SUITE_END();
+
+        class TCheckTask: public TTaskScheduler::IRepeatedTask {
+            public:
+                TCheckTask(const TDuration& delay)
+                    : Start_(Now())
+                    , Delay_(delay)
+                {
+                    AtomicIncrement(ScheduledTaskCounter_);
+                }
+
+                ~TCheckTask() override {
+                }
+
+                bool Process() override {
+                    const TDuration delay = Now() - Start_;
+
+                    if (delay < Delay_) {
+                        AtomicIncrement(BadTimeoutCounter_);
+                    }
+
+                    AtomicIncrement(ExecutedTaskCounter_);
+
+                    return false;
+                }
+
+                static bool AllTaskExecuted() {
+                    return AtomicGet(ScheduledTaskCounter_) == AtomicGet(ExecutedTaskCounter_);
+                }
+
+                static size_t BadTimeoutCount() {
+                    return AtomicGet(BadTimeoutCounter_);
+                }
+
+            private:
+                TInstant Start_;
+                TDuration Delay_;
+                static TAtomic BadTimeoutCounter_;
+                static TAtomic ScheduledTaskCounter_;
+                static TAtomic ExecutedTaskCounter_;
+        };
+
+    public:
+        inline void Test() {
+            ScheduleCheckTask(200);
+            ScheduleCheckTask(100);
+            ScheduleCheckTask(1000);
+            ScheduleCheckTask(10000);
+            ScheduleCheckTask(5000);
+
+            Scheduler_.Start();
+
+            usleep(1000000);
+
+            UNIT_ASSERT_EQUAL(TCheckTask::BadTimeoutCount(), 0);
+            UNIT_ASSERT(TCheckTask::AllTaskExecuted());
+        }
+
+    private:
+        void ScheduleCheckTask(size_t delay) {
+            TDuration d = TDuration::MicroSeconds(delay);
+
+            Scheduler_.Add(new TCheckTask(d), d);
+        }
+
+    private:
+        TTaskScheduler Scheduler_;
+};
+
+TAtomic TTaskSchedulerTest::TCheckTask::BadTimeoutCounter_ = 0;
+TAtomic TTaskSchedulerTest::TCheckTask::ScheduledTaskCounter_ = 0;
+TAtomic TTaskSchedulerTest::TCheckTask::ExecutedTaskCounter_ = 0;
+
+UNIT_TEST_SUITE_REGISTRATION(TTaskSchedulerTest);
diff --git a/library/cpp/threading/task_scheduler/ut/ya.make b/library/cpp/threading/task_scheduler/ut/ya.make
new file mode 100644
index 0000000000..07ee8b0877
--- /dev/null
+++ b/library/cpp/threading/task_scheduler/ut/ya.make
@@ -0,0 +1,9 @@
+UNITTEST_FOR(library/cpp/threading/task_scheduler)
+
+OWNER(g:middle)
+
+SRCS(
+    task_scheduler_ut.cpp
+)
+
+END()
diff --git a/library/cpp/threading/task_scheduler/ya.make b/library/cpp/threading/task_scheduler/ya.make
new file mode 100644
index 0000000000..5b14c0aa63
--- /dev/null
+++ b/library/cpp/threading/task_scheduler/ya.make
@@ -0,0 +1,9 @@
+LIBRARY()
+
+OWNER(g:middle)
+
+SRCS(
+    task_scheduler.cpp
+)
+
+END()