From a189151c598bfd1886ae7134e2fde26ad3fa89e2 Mon Sep 17 00:00:00 2001
From: Jamie Reece Wilson <me@reece.sx>
Date: Fri, 1 Dec 2023 09:22:51 +0000
Subject: [PATCH] [+] AuAsync::IThreadPool::GetIOProcessor [+]
 AuAsync::IThreadPool::GetIONetInterface [+]
 AuAsync::IThreadPool::GetIONetWorker [+]
 AuAsync::IWorkItem::SetSchedByLoopSource

---
 Include/Aurora/Async/IThreadPool.hpp      |  11 +++
 Include/Aurora/Async/IWorkItem.hpp        |   2 +
 Include/Aurora/Async/IWorkItemHandler.hpp |   1 +
 Source/Async/AsyncApp.cpp                 |  21 ++++-
 Source/Async/AsyncApp.hpp                 |   4 +-
 Source/Async/AuThreadState.hpp            |   3 +
 Source/Async/AuThreadStateSingletons.cpp  |  54 ++++++++++++
 Source/Async/AuThreadStateSingletons.hpp  |  16 ++++
 Source/Async/ThreadPool.cpp               |  32 ++++++-
 Source/Async/ThreadPool.hpp               |   4 +
 Source/Async/WorkItem.cpp                 | 101 ++++++++++++++++++++--
 Source/Async/WorkItem.hpp                 |   3 +
 12 files changed, 240 insertions(+), 12 deletions(-)

diff --git a/Include/Aurora/Async/IThreadPool.hpp b/Include/Aurora/Async/IThreadPool.hpp
index d2fc99c9..cf6242cd 100644
--- a/Include/Aurora/Async/IThreadPool.hpp
+++ b/Include/Aurora/Async/IThreadPool.hpp
@@ -13,6 +13,12 @@ namespace Aurora::IO::Loop
     struct ILoopSource;
 }
 
+namespace Aurora::IO::Net
+{
+    struct INetInterface;
+    struct INetWorker;
+}
+
 namespace Aurora::Async
 {
     struct IThreadPool
@@ -77,6 +83,11 @@ namespace Aurora::Async
         //
         virtual WorkerId_t GetCurrentThread()                                             = 0;
 
+        //
+        virtual AuSPtr<IO::IIOProcessor>       GetIOProcessor(WorkerId_t id)              = 0;
+        virtual AuSPtr<IO::Net::INetInterface> GetIONetInterface(WorkerId_t id)           = 0;
+        virtual AuSPtr<IO::Net::INetWorker>    GetIONetWorker(WorkerId_t id)              = 0;
+
         // Synchronization
         // Note: syncing to yourself will nullify requireSignal to prevent deadlock conditions
         virtual bool Sync(WorkerId_t workerId,
diff --git a/Include/Aurora/Async/IWorkItem.hpp b/Include/Aurora/Async/IWorkItem.hpp
index 1b10de9d..b664e4a7 100644
--- a/Include/Aurora/Async/IWorkItem.hpp
+++ b/Include/Aurora/Async/IWorkItem.hpp
@@ -22,6 +22,8 @@ namespace Aurora::Async
         virtual AuSPtr<IWorkItem> WaitFor(const AuSPtr<IWorkItem> &pWorkItem)   = 0;
         virtual AuSPtr<IWorkItem> WaitFor(const AuList<AuSPtr<IWorkItem>> &workItems) = 0;
         
+        virtual AuSPtr<IWorkItem> SetSchedByLoopSource(const AuSPtr<IO::Loop::ILoopSource> &pLoopSource) = 0;
+
         // ms = time relative to the current time
         virtual AuSPtr<IWorkItem> SetSchedTime(AuUInt32 ms)                     = 0;
 
diff --git a/Include/Aurora/Async/IWorkItemHandler.hpp b/Include/Aurora/Async/IWorkItemHandler.hpp
index 01ee03f6..2bd9fe8e 100644
--- a/Include/Aurora/Async/IWorkItemHandler.hpp
+++ b/Include/Aurora/Async/IWorkItemHandler.hpp
@@ -33,6 +33,7 @@ namespace Aurora::Async
             AuUInt32 reschedClockAbsMs {};
             AuUInt64 reschedClockAbsNs {};
             AuUInt64 reschedSteadyClockAbsNs {};
+            AuSPtr<IO::Loop::ILoopSource> pLoopSource {};
             // @hideinitializer
             IThreadPool *pool;
         };
diff --git a/Source/Async/AsyncApp.cpp b/Source/Async/AsyncApp.cpp
index ab9a4f6c..770b2618 100644
--- a/Source/Async/AsyncApp.cpp
+++ b/Source/Async/AsyncApp.cpp
@@ -35,13 +35,13 @@ namespace Aurora::Async
     void AsyncApp::Start()
     {
         ThreadPool::SetRunningMode(true);
-        SysAssert(ThreadPool::Create({0, 0}));
+        SysAssert(ThreadPool::Create({ 0, 0 }));
         StartSched(); // this is now an init once function
     }
 
     void AsyncApp::Main()
     {
-        ThreadPool::Entrypoint({0, 0});
+        ThreadPool::Entrypoint({ 0, 0 });
     }
 
     void AsyncApp::SetConsoleCommandDispatcher(WorkerId_t id)
@@ -152,7 +152,7 @@ namespace Aurora::Async
     {
         return ThreadPool::RunAllPending();
     }
-   
+
     void AsyncApp::Shutdown()
     {
         ThreadPool::Shutdown();
@@ -228,6 +228,21 @@ namespace Aurora::Async
         ThreadPool::AssertWorker(id);
     }
 
+    AuSPtr<AuIO::IIOProcessor> AsyncApp::GetIOProcessor(WorkerId_t pid)
+    {
+        return ThreadPool::GetIOProcessor(pid);
+    }
+
+    AuSPtr<AuIO::Net::INetInterface> AsyncApp::GetIONetInterface(WorkerId_t pid)
+    {
+        return ThreadPool::GetIONetInterface(pid);
+    }
+
+    AuSPtr<AuIO::Net::INetWorker> AsyncApp::GetIONetWorker(WorkerId_t pid)
+    {
+        return ThreadPool::GetIONetWorker(pid);
+    }
+
     AuSPtr<AuLoop::ILoopQueue> AsyncApp::ToKernelWorkQueue()
     {
         return ThreadPool::ToKernelWorkQueue();
diff --git a/Source/Async/AsyncApp.hpp b/Source/Async/AsyncApp.hpp
index f5249160..8f4c070a 100644
--- a/Source/Async/AsyncApp.hpp
+++ b/Source/Async/AsyncApp.hpp
@@ -42,7 +42,9 @@ namespace Aurora::Async
         void AddFeature(WorkerId_t id, AuSPtr<Threading::Threads::IThreadFeature> feature, bool async)  override;
         void AssertInThreadGroup(ThreadGroup_t group)                           override;
         void AssertWorker(WorkerId_t id)                                        override;
-
+        AuSPtr<AuIO::IIOProcessor> GetIOProcessor(WorkerId_t pid)               override;
+        AuSPtr<AuIO::Net::INetInterface> GetIONetInterface(WorkerId_t pid)      override;
+        AuSPtr<AuIO::Net::INetWorker> GetIONetWorker(WorkerId_t pid)            override;
         AuSPtr<AuLoop::ILoopQueue> ToKernelWorkQueue()                          override;
         AuSPtr<AuLoop::ILoopQueue> ToKernelWorkQueue(WorkerId_t workerId)       override;
 
diff --git a/Source/Async/AuThreadState.hpp b/Source/Async/AuThreadState.hpp
index eb891c0e..2fb0aa7d 100644
--- a/Source/Async/AuThreadState.hpp
+++ b/Source/Async/AuThreadState.hpp
@@ -7,6 +7,8 @@
 ***/
 #pragma once
 
+#include "AuThreadStateSingletons.hpp"
+
 namespace Aurora::Async
 {
     struct AsyncLoop;
@@ -68,6 +70,7 @@ namespace Aurora::Async
         AuSPtr<AsyncLoop> asyncLoop;
         ThreadStateStack stackState;
         ThreadStateFeatureCallbacks tlsFeatures;
+        ThreadStateSingletons singletons;
 
         bool Init();
     };
diff --git a/Source/Async/AuThreadStateSingletons.cpp b/Source/Async/AuThreadStateSingletons.cpp
index 035ac357..f01c46d3 100644
--- a/Source/Async/AuThreadStateSingletons.cpp
+++ b/Source/Async/AuThreadStateSingletons.cpp
@@ -10,5 +10,59 @@
 
 namespace Aurora::Async
 {
+    AuSPtr<AuIO::IIOProcessor> ThreadStateSingletons::GetIOProcessor(AuWorkerPId_t pid)
+    {
+        if (this->pIOProcessors)
+        {
+            return this->pIOProcessors;
+        }
 
+        {
+            AU_LOCK_GUARD(this->mutex);
+
+            if (this->pIOProcessors)
+            {
+                return this->pIOProcessors;
+            }
+
+            return this->pIOProcessors = AuIO::NewIOProcessorOnThread(false, pid);
+        }
+    }
+
+    void ThreadStateSingletons::TryInitNet(AuWorkerPId_t pid)
+    {
+        if (this->pNetInterface)
+        {
+            return;
+        }
+
+        AU_LOCK_GUARD(this->mutex);
+
+        auto pNetProcessor = AuNet::NewNetworkInterface();
+        if (!pNetProcessor)
+        {
+            return;
+        }
+
+        auto pNetWorker = pNetProcessor->GetWorkersService()->Attach(this->GetIOProcessor(pid));
+        if (!pNetWorker)
+        {
+            return;
+        }
+
+        this->pNetWorker = pNetWorker;
+        this->pNetInterface = pNetProcessor;
+    }
+
+    AuSPtr<AuIO::Net::INetInterface> ThreadStateSingletons::GetIONetInterface(AuWorkerPId_t pid)
+    {
+        this->TryInitNet(pid);
+        return this->pNetInterface;
+    }
+
+    AuSPtr<AuIO::Net::INetWorker> ThreadStateSingletons::GetIONetWorker(AuWorkerPId_t pid)
+    {
+        this->TryInitNet(pid);
+        return this->pNetWorker;
+    }
 }
\ No newline at end of file
diff --git a/Source/Async/AuThreadStateSingletons.hpp b/Source/Async/AuThreadStateSingletons.hpp
index ce583d18..3c085087 100644
--- a/Source/Async/AuThreadStateSingletons.hpp
+++ b/Source/Async/AuThreadStateSingletons.hpp
@@ -7,7 +7,23 @@
 ***/
 #pragma once
 
+#include <Aurora/IO/Net/NetExperimental.hpp>
+
 namespace Aurora::Async
 {
+    struct ThreadStateSingletons
+    {
+        AuCriticalSection mutex;
 
+        AuSPtr<AuIO::IIOProcessor> pIOProcessors;
+
+        AuSPtr<AuIO::IIOProcessor> pIOProcessor;
+        AuSPtr<AuIO::Net::INetInterface> pNetInterface;
+        AuSPtr<AuIO::Net::INetWorker> pNetWorker;
+
+        AuSPtr<AuIO::IIOProcessor> GetIOProcessor(AuWorkerPId_t pid);
+        AuSPtr<AuIO::Net::INetInterface> GetIONetInterface(AuWorkerPId_t pid);
+        AuSPtr<AuIO::Net::INetWorker> GetIONetWorker(AuWorkerPId_t pid);
+        void TryInitNet(AuWorkerPId_t pid);
+    };
 }
\ No newline at end of file
diff --git a/Source/Async/ThreadPool.cpp b/Source/Async/ThreadPool.cpp
index c5876adc..1c0b286e 100644
--- a/Source/Async/ThreadPool.cpp
+++ b/Source/Async/ThreadPool.cpp
@@ -738,7 +738,37 @@ namespace Aurora::Async
     {
         return tlsWorkerId;
     }
-    
+
+    AuSPtr<AuIO::IIOProcessor> ThreadPool::GetIOProcessor(WorkerId_t pid)
+    {
+        if (auto pState = this->GetThreadHandle(pid))
+        {
+            return pState->singletons.GetIOProcessor({ this->SharedFromThis(), pid });
+        }
+
+        return {};
+    }
+
+    AuSPtr<AuIO::Net::INetInterface> ThreadPool::GetIONetInterface(WorkerId_t pid)
+    {
+        if (auto pState = this->GetThreadHandle(pid))
+        {
+            return pState->singletons.GetIONetInterface({ this->SharedFromThis(), pid });
+        }
+
+        return {};
+    }
+
+    AuSPtr<AuIO::Net::INetWorker> ThreadPool::GetIONetWorker(WorkerId_t pid)
+    {
+        if (auto pState = this->GetThreadHandle(pid))
+        {
+            return pState->singletons.GetIONetWorker({ this->SharedFromThis(), pid });
+        }
+
+        return {};
+    }
+
     bool ThreadPool::Sync(WorkerId_t workerId, AuUInt32 timeoutMs, bool requireSignal)
     {
         AU_LOCK_GUARD(this->pRWReadView);
diff --git a/Source/Async/ThreadPool.hpp b/Source/Async/ThreadPool.hpp
index 5177ed4d..977a10e5 100644
--- a/Source/Async/ThreadPool.hpp
+++ b/Source/Async/ThreadPool.hpp
@@ -60,6 +60,10 @@ namespace Aurora::Async
                                                                                                        
         virtual WorkerId_t GetCurrentThread()                                                          override;
 
+        virtual AuSPtr<AuIO::IIOProcessor>       GetIOProcessor(WorkerId_t id)                         override;
+        virtual AuSPtr<AuIO::Net::INetInterface> GetIONetInterface(WorkerId_t pid)                     override;
+        virtual AuSPtr<AuIO::Net::INetWorker>    GetIONetWorker(WorkerId_t id)                         override;
+
         virtual bool Sync(WorkerId_t workerId, AuUInt32 timeoutMs, bool requireSignal)                 override;
         virtual void Signal(WorkerId_t workerId)                                                       override;
         virtual AuSPtr<AuLoop::ILoopSource> WorkerToLoopSource(WorkerId_t id)                          override;
diff --git a/Source/Async/WorkItem.cpp b/Source/Async/WorkItem.cpp
index 76da1608..271c07f2 100644
--- a/Source/Async/WorkItem.cpp
+++ b/Source/Async/WorkItem.cpp
@@ -43,8 +43,10 @@ namespace Aurora::Async
 
     WorkItem::~WorkItem()
     {
-
-        //Fail();
+        if (auto pIOWatch = AuExchange(this->pIOWatch, {}))
+        {
+            pIOWatch->StopWatch();
+        }
     }
 
     AuSPtr<IWorkItem> WorkItem::WaitFor(const AuSPtr<IWorkItem> &workItem)
@@ -172,6 +174,12 @@ namespace Aurora::Async
         return AU_SHARED_FROM_THIS;
     }
 
+    AuSPtr<IWorkItem> WorkItem::SetSchedByLoopSource(const AuSPtr<IO::Loop::ILoopSource> &pLoopSource)
+    {
+        this->pIOWatchLS = pLoopSource;
+        return AU_SHARED_FROM_THIS;
+    }
+
     AuSPtr<IWorkItem> WorkItem::SetSchedTime(AuUInt32 ms)
     {
         this->dispatchTimeNs_  = Time::SteadyClockNS() + AuMSToNS<AuUInt64>(ms);
@@ -232,6 +240,22 @@ namespace Aurora::Async
 
         this->dispatchPending_ = true;
 
+        if (this->pIOWatchLS)
+        {
+            if (!this->pIOWatchLS->IsSignaled())
+            {
+                if (!Schedule())
+                {
+                    this->Fail();
+                }
+                return;
+            }
+            else
+            {
+                AuResetMember(this->pIOWatchLS);
+            }
+        }
+
         if (Time::SteadyClockNS() < this->dispatchTimeNs_)
         {
             if (!Schedule())
@@ -324,11 +348,14 @@ namespace Aurora::Async
 
     AuSPtr<ThreadState> WorkItem::GetState()
     {
-        if (!this->worker_.HasValue())
+        if (this->worker_.HasValue())
+        {
+            return this->owner_->GetThreadHandle(this->worker_.value());
+        }
+        else
         {
             return {};
         }
-        return this->owner_->GetThreadHandle(this->worker_.value());
     }
 
     void WorkItem::RunAsyncLocked2()
@@ -398,12 +425,15 @@ namespace Aurora::Async
             {
                 SetSchedSteadyTimeNsAbs(info.reschedSteadyClockAbsNs);
             }
+            else if (info.pLoopSource)
+            {
+                SetSchedByLoopSource(info.pLoopSource);
+            }
 
             if (!WaitForLocked(info.waitFor))
             {
                 this->Fail();
             }
-
         }
         [[fallthrough]];
         case ETickType::eRerun:
@@ -419,21 +449,38 @@ namespace Aurora::Async
         }
 
         this->finished = true;
+
         if (this->finishedEvent_)
         {
             this->finishedEvent_->Set();
         }
 
-        for (auto &waiter : this->waiters_)
+        for (auto &waiter : AuExchange(this->waiters_, {}))
         {
             AuReinterpretCast<WorkItem>(waiter)->DispatchExLocked(true);
         }
+
+        this->waitOn_.clear();
+
+        if (auto pIOWatch = AuExchange(this->pIOWatch, {}))
+        {
+            pIOWatch->StopWatch();
+        }
+
+        AuResetMember(this->pIOWatchLS);
     }
 
     void WorkItem::Fail()
     {
         failed = true;
 
+        if (auto pIOWatch = AuExchange(this->pIOWatch, {}))
+        {
+            pIOWatch->StopWatch();
+        }
+
+        AuResetMember(this->pIOWatchLS);
+
         if (auto task_ = AuExchange(this->task_, {}))
         {
             task_->OnFailure();
@@ -531,7 +578,47 @@ namespace Aurora::Async
     
     bool WorkItem::Schedule()
     {
-        return Async::Schedule(this->dispatchTimeNs_, this->owner_, this->worker_.value(), AuSharedFromThis());
+        if (auto pLoopSource = this->pIOWatchLS)
+        {
+            if (this->pIOWatch)
+            {
+                return true;
+            }
+
+            auto pState = this->GetState();
+            if (!pState)
+            {
+                return false;
+            }
+
+            auto pIOProcessor = pState->singletons.GetIOProcessor(this->worker_.value());
+            if (!pIOProcessor)
+            {
+                return false;
+            }
+
+            this->pIOWatch = pIOProcessor->StartSimpleLSWatch(pLoopSource, AuMakeSharedThrow<AuIO::IIOSimpleEventListenerFunctional>([=]()
+            {
+                this->Dispatch();
+            }, [=]()
+            {
+                this->Dispatch();
+            }, [=]()
+            {
+                this->Dispatch();
+            }));
+
+            if (!this->pIOWatch)
+            {
+                return false;
+            }
+
+            return true;
+        }
+        else
+        {
+            return Async::Schedule(this->dispatchTimeNs_, this->owner_, this->worker_.value(), AuSharedFromThis());
+        }
     }
 
     void WorkItem::SendOff()
diff --git a/Source/Async/WorkItem.hpp b/Source/Async/WorkItem.hpp
index 141c2a9f..3a422f99 100644
--- a/Source/Async/WorkItem.hpp
+++ b/Source/Async/WorkItem.hpp
@@ -32,6 +32,7 @@ namespace Aurora::Async
         AuSPtr<IWorkItem> SetSchedTimeAbs(AuUInt32 ms) override;
         AuSPtr<IWorkItem> SetSchedTimeNsAbs(AuUInt64 ns) override;
         AuSPtr<IWorkItem> SetSchedSteadyTimeNsAbs(AuUInt64 ns) override;
+        AuSPtr<IWorkItem> SetSchedByLoopSource(const AuSPtr<IO::Loop::ILoopSource> &pLoopSource) override;
         AuSPtr<IWorkItem> Then(const AuSPtr<IWorkItem> &next) override;
 
         AuSPtr<IWorkItem> Dispatch() override;
@@ -80,6 +81,8 @@ namespace Aurora::Async
         AuThreadPrimitives::Event finishedEvent_;
         AuUInt32 uShutdownCookie {};
         AuOptionalEx<AuUInt32> optOtherCookie {};
+        AuSPtr<AuIO::IIOProcessorItem> pIOWatch;
+        AuSPtr<IO::Loop::ILoopSource> pIOWatchLS;
 
         bool finished {};
         bool failed {};