From d49d4ad4fcdbbf1c2c38a8e5b64a2842575267f5 Mon Sep 17 00:00:00 2001
From: Jamie Reece Wilson <me@reece.sx>
Date: Fri, 15 Sep 2023 16:53:18 +0100
Subject: [PATCH] [*] addd4080 cont

---
 Source/Async/WorkItem.cpp | 39 ++++++++++++++++++++++++++++++++++-----
 1 file changed, 34 insertions(+), 5 deletions(-)

diff --git a/Source/Async/WorkItem.cpp b/Source/Async/WorkItem.cpp
index 5d574ec1..ae6c8eec 100644
--- a/Source/Async/WorkItem.cpp
+++ b/Source/Async/WorkItem.cpp
@@ -478,7 +478,11 @@ namespace Aurora::Async
         }
 
         auto pWorker = GetCurrentWorkerPId().pool;
-        if (!pWorker) AuUnsafeRaiiToShared(gAsyncApp);
+        if (!pWorker)
+        {
+            pWorker = AuUnsafeRaiiToShared(gAsyncApp);
+        }
+
         return AuMakeShared<WorkItem>(ToInternal(pWorker).get(), WorkerPId_t { pWorker , worker }, task);
     }
 
@@ -492,7 +496,18 @@ namespace Aurora::Async
             return {};
         }
 
-        return AuMakeSharedThrow<FuncWorker>(ToInternal(worker.pool).get(), worker, AuMove(func));
+        auto pWorker = worker.pool;
+        if (!pWorker)
+        {
+            pWorker = GetCurrentWorkerPId().pool; 
+        }
+
+        if (!pWorker)
+        {
+            pWorker = AuUnsafeRaiiToShared(gAsyncApp);
+        }
+
+        return AuMakeSharedThrow<FuncWorker>(ToInternal(pWorker).get(), worker, AuMove(func));
     }
 
     AUKN_SYM AuSPtr<IWorkItem> NewWorkItem(const WorkerPId_t &worker, const AuSPtr<IWorkItemHandler> &task)
@@ -504,14 +519,28 @@ namespace Aurora::Async
             SysPushErrorArg("WorkItem has null task. Running out of memory?");
             return {};
         }
-        
-        return AuMakeSharedThrow<WorkItem>(ToInternal(worker.pool).get(), worker, task);
+
+        auto pWorker = worker.pool;
+        if (!pWorker)
+        {
+            pWorker = GetCurrentWorkerPId().pool;
+        }
+
+        if (!pWorker)
+        {
+            pWorker = AuUnsafeRaiiToShared(gAsyncApp);
+        }
+
+        return AuMakeSharedThrow<WorkItem>(ToInternal(pWorker).get(), worker, task);
     }
 
     AUKN_SYM AuSPtr<IWorkItem> NewFence()
     {
         auto pWorker = GetCurrentWorkerPId().pool;
-        if (!pWorker) AuUnsafeRaiiToShared(gAsyncApp);
+        if (!pWorker)
+        {
+            pWorker = AuUnsafeRaiiToShared(gAsyncApp);
+        }
         return AuMakeShared<WorkItem>((IThreadPoolInternal *)ToInternal(pWorker).get(), WorkerPId_t {}, AuSPtr<IWorkItemHandler>{});
     }