From d803f1547ca0d58cc267c92d6f758bec9766a570 Mon Sep 17 00:00:00 2001
From: Jamie Reece Wilson <me@reece.sx>
Date: Fri, 12 Jul 2024 17:22:07 +0100
Subject: [PATCH] [+] Added AuAsync per thread io group singleton [*] Move
 static functions out of the thread pool interface. no idea how i had
 successful builds before

---
 Include/Aurora/Async/Async.hpp           |  5 +++
 Include/Aurora/Async/IThreadPool.hpp     | 10 +++--
 Include/Aurora/IO/IAsyncTransaction.hpp  | 13 ++++++
 Source/Async/Async.cpp                   | 52 ++++++++++++++++++++++++
 Source/Async/AsyncApp.cpp                |  5 +++
 Source/Async/AsyncApp.hpp                |  1 +
 Source/Async/AuThreadStateSingletons.cpp | 19 +++++++++
 Source/Async/AuThreadStateSingletons.hpp |  2 +
 Source/Async/ThreadPool.cpp              | 49 +++++-----------------
 Source/Async/ThreadPool.hpp              |  5 +--
 10 files changed, 114 insertions(+), 47 deletions(-)

diff --git a/Include/Aurora/Async/Async.hpp b/Include/Aurora/Async/Async.hpp
index 6fce5074..30cc915a 100644
--- a/Include/Aurora/Async/Async.hpp
+++ b/Include/Aurora/Async/Async.hpp
@@ -61,6 +61,11 @@ namespace Aurora::Async
     AUKN_SYM AuSPtr<IThreadPool> NewThreadPool();
 
     AUKN_SYM void SetMainThreadForSysPumpScheduling(WorkerPId_t pid);
+
+    AUKN_SYM AuSPtr<IO::IIOProcessor>                      GetSelfIOProcessor();
+    AUKN_SYM AuSPtr<IO::Net::INetInterface>                GetSelfIONetInterface();
+    AUKN_SYM AuSPtr<IO::Net::INetWorker>                   GetSelfIONetWorker();
+    AUKN_SYM AuSPtr<IO::CompletionGroup::ICompletionGroup> GetSelfIOGroup();
 }
 
 #if !defined(_CPPSHARP)
diff --git a/Include/Aurora/Async/IThreadPool.hpp b/Include/Aurora/Async/IThreadPool.hpp
index ac6552de..7c178df2 100644
--- a/Include/Aurora/Async/IThreadPool.hpp
+++ b/Include/Aurora/Async/IThreadPool.hpp
@@ -19,6 +19,11 @@ namespace Aurora::IO::Net
     struct INetWorker;
 }
 
+namespace Aurora::IO::CompletionGroup
+{
+    struct ICompletionGroup;
+}
+
 namespace Aurora::Async
 {
     struct IThreadPool
@@ -87,10 +92,7 @@ namespace Aurora::Async
         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;
-
-        static AuSPtr<IO::IIOProcessor>        GetSelfIOProcessor();
-        static AuSPtr<IO::Net::INetInterface>  GetSelfIONetInterface();
-        static AuSPtr<IO::Net::INetWorker>     GetSelfIONetWorker();
+        virtual AuSPtr<IO::CompletionGroup::ICompletionGroup> GetIOGroup(WorkerId_t id)   = 0;
 
         // Synchronization
         
diff --git a/Include/Aurora/IO/IAsyncTransaction.hpp b/Include/Aurora/IO/IAsyncTransaction.hpp
index 7b586f00..551788c7 100644
--- a/Include/Aurora/IO/IAsyncTransaction.hpp
+++ b/Include/Aurora/IO/IAsyncTransaction.hpp
@@ -17,6 +17,11 @@ namespace Aurora::IO::CompletionGroup
     struct ICompletionGroup;
 }
 
+namespace Aurora::Async
+{
+    AUKN_SYM AuSPtr<IO::CompletionGroup::ICompletionGroup> GetSelfIOGroup();
+}
+
 #include "CompletionGroup/ICompletionGroupWorkHandle.hpp"
 #include <Aurora/Debug/Debug.hpp>
 
@@ -132,6 +137,10 @@ namespace Aurora::IO
         {
             (void)pTransaction->TryAttachToCompletionGroup(pCompletionGroup.Value());
         }
+        else if (auto pGroup = Aurora::Async::GetSelfIOGroup())
+        {
+            (void)pTransaction->TryAttachToCompletionGroup(pGroup);
+        }
 
         if (callback)
         {
@@ -157,6 +166,10 @@ namespace Aurora::IO
         {
             (void)pTransaction->TryAttachToCompletionGroup(pCompletionGroup.Value());
         }
+        else if (auto pGroup = Aurora::Async::GetSelfIOGroup())
+        {
+            (void)pTransaction->TryAttachToCompletionGroup(pGroup);
+        }
 
         if (callback)
         {
diff --git a/Source/Async/Async.cpp b/Source/Async/Async.cpp
index 6fdd798f..658ee5b9 100644
--- a/Source/Async/Async.cpp
+++ b/Source/Async/Async.cpp
@@ -23,4 +23,56 @@ namespace Aurora::Async
         DeinitSched();
         ReleaseApp();
     }
+
+    AuSPtr<IO::IIOProcessor> GetSelfIOProcessor()
+    {
+        auto pid = GetCurrentWorkerPId();
+        if (pid)
+        {
+            return pid.GetPool()->GetIOProcessor(pid);
+        }
+        else
+        {
+            return {};
+        }
+    }
+
+    AuSPtr<IO::Net::INetInterface> GetSelfIONetInterface()
+    {
+        auto pid = GetCurrentWorkerPId();
+        if (pid)
+        {
+            return pid.GetPool()->GetIONetInterface(pid);
+        }
+        else
+        {
+            return {};
+        }
+    }
+
+    AuSPtr<IO::Net::INetWorker> GetSelfIONetWorker()
+    {
+        auto pid = GetCurrentWorkerPId();
+        if (pid)
+        {
+            return pid.GetPool()->GetIONetWorker(pid);
+        }
+        else
+        {
+            return {};
+        }
+    }
+
+    AuSPtr<IO::CompletionGroup::ICompletionGroup> GetSelfIOGroup()
+    {
+        auto pid = GetCurrentWorkerPId();
+        if (pid)
+        {
+            return pid.GetPool()->GetIOGroup(pid);
+        }
+        else
+        {
+            return {};
+        }
+    }
 }
\ No newline at end of file
diff --git a/Source/Async/AsyncApp.cpp b/Source/Async/AsyncApp.cpp
index 885b2757..9fd5b40c 100644
--- a/Source/Async/AsyncApp.cpp
+++ b/Source/Async/AsyncApp.cpp
@@ -248,6 +248,11 @@ namespace Aurora::Async
         return ThreadPool::GetIONetWorker(pid);
     }
 
+    AuSPtr<IO::CompletionGroup::ICompletionGroup> AsyncApp::GetIOGroup(WorkerId_t pid)
+    {
+        return ThreadPool::GetIOGroup(pid);
+    }
+
     AuSPtr<AuLoop::ILoopQueue> AsyncApp::ToKernelWorkQueue()
     {
         return ThreadPool::ToKernelWorkQueue();
diff --git a/Source/Async/AsyncApp.hpp b/Source/Async/AsyncApp.hpp
index 30c587ed..5aeaa835 100644
--- a/Source/Async/AsyncApp.hpp
+++ b/Source/Async/AsyncApp.hpp
@@ -48,6 +48,7 @@ namespace Aurora::Async
         AuSPtr<AuIO::Net::INetWorker> GetIONetWorker(WorkerId_t pid)            override;
         AuSPtr<AuLoop::ILoopQueue> ToKernelWorkQueue()                          override;
         AuSPtr<AuLoop::ILoopQueue> ToKernelWorkQueue(WorkerId_t workerId)       override;
+        AuSPtr<IO::CompletionGroup::ICompletionGroup> GetIOGroup(WorkerId_t id) override;
 
         // Main thread logic
         void Start()                                                            override;
diff --git a/Source/Async/AuThreadStateSingletons.cpp b/Source/Async/AuThreadStateSingletons.cpp
index f01c46d3..6950c6f1 100644
--- a/Source/Async/AuThreadStateSingletons.cpp
+++ b/Source/Async/AuThreadStateSingletons.cpp
@@ -29,6 +29,25 @@ namespace Aurora::Async
         }
     }
 
+    AuSPtr<AuIO::CompletionGroup::ICompletionGroup> ThreadStateSingletons::GetIOGroup()
+    {
+        if (this->pGroup)
+        {
+            return this->pGroup;
+        }
+
+        {
+            AU_LOCK_GUARD(this->mutex);
+
+            if (this->pGroup)
+            {
+                return this->pGroup;
+            }
+
+            return this->pGroup = AuIO::CompletionGroup::NewCompletionGroup();
+        }
+    }
+
     void ThreadStateSingletons::TryInitNet(AuWorkerPId_t pid)
     {
         if (this->pNetInterface)
diff --git a/Source/Async/AuThreadStateSingletons.hpp b/Source/Async/AuThreadStateSingletons.hpp
index 3c085087..9ccbe20d 100644
--- a/Source/Async/AuThreadStateSingletons.hpp
+++ b/Source/Async/AuThreadStateSingletons.hpp
@@ -20,7 +20,9 @@ namespace Aurora::Async
         AuSPtr<AuIO::IIOProcessor> pIOProcessor;
         AuSPtr<AuIO::Net::INetInterface> pNetInterface;
         AuSPtr<AuIO::Net::INetWorker> pNetWorker;
+        AuSPtr<AuIO::CompletionGroup::ICompletionGroup> pGroup;
 
+        AuSPtr<AuIO::CompletionGroup::ICompletionGroup> GetIOGroup();
         AuSPtr<AuIO::IIOProcessor> GetIOProcessor(AuWorkerPId_t pid);
         AuSPtr<AuIO::Net::INetInterface> GetIONetInterface(AuWorkerPId_t pid);
         AuSPtr<AuIO::Net::INetWorker> GetIONetWorker(AuWorkerPId_t pid);
diff --git a/Source/Async/ThreadPool.cpp b/Source/Async/ThreadPool.cpp
index 4f233278..edee1937 100644
--- a/Source/Async/ThreadPool.cpp
+++ b/Source/Async/ThreadPool.cpp
@@ -811,6 +811,16 @@ namespace Aurora::Async
         return {};
     }
 
+    AuSPtr<IO::CompletionGroup::ICompletionGroup> ThreadPool::GetIOGroup(WorkerId_t id)
+    {
+        if (auto pState = this->GetThreadHandle(id))
+        {
+            return pState->singletons.GetIOGroup();
+        }
+
+        return {};
+    }
+
     AuSPtr<AuIO::Net::INetInterface> ThreadPool::GetIONetInterface(WorkerId_t pid)
     {
         if (auto pState = this->GetThreadHandle(pid))
@@ -830,45 +840,6 @@ namespace Aurora::Async
 
         return {};
     }
-
-    AuSPtr<IO::IIOProcessor>       ThreadPool::GetSelfIOProcessor()
-    {
-        auto pid = GetCurrentWorkerPId();
-        if (pid)
-        {
-            return pid.GetPool()->GetIOProcessor(pid);
-        }
-        else
-        {
-            return {};
-        }
-    }
-
-    AuSPtr<IO::Net::INetInterface> ThreadPool::GetSelfIONetInterface()
-    {
-        auto pid = GetCurrentWorkerPId();
-        if (pid)
-        {
-            return pid.GetPool()->GetIONetInterface(pid);
-        }
-        else
-        {
-            return {};
-        }
-    }
-
-    AuSPtr<IO::Net::INetWorker>    ThreadPool::GetSelfIONetWorker()
-    {
-        auto pid = GetCurrentWorkerPId();
-        if (pid)
-        {
-            return pid.GetPool()->GetIONetWorker(pid);
-        }
-        else
-        {
-            return {};
-        }
-    }
   
     bool ThreadPool::Sync(WorkerId_t workerId, AuUInt32 timeoutMs, bool requireSignal)
     {
diff --git a/Source/Async/ThreadPool.hpp b/Source/Async/ThreadPool.hpp
index e6bff9ac..4b3a3c26 100644
--- a/Source/Async/ThreadPool.hpp
+++ b/Source/Async/ThreadPool.hpp
@@ -63,10 +63,7 @@ namespace Aurora::Async
         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;
-
-        AuSPtr<IO::IIOProcessor>       GetSelfIOProcessor();
-        AuSPtr<IO::Net::INetInterface> GetSelfIONetInterface();
-        AuSPtr<IO::Net::INetWorker>    GetSelfIONetWorker();
+        virtual AuSPtr<AuIO::CompletionGroup::ICompletionGroup> GetIOGroup(WorkerId_t id)              override;
 
         virtual bool Sync(WorkerId_t workerId, AuUInt32 timeoutMs, bool requireSignal)                 override;
         virtual void Signal(WorkerId_t workerId)                                                       override;