Fixed job handles

wopss · May 3, 2024 · 8badaf3 · 8badaf3
1 parent bb2909f
commit 8badaf3
Show file tree

Hide file tree

Showing 3 changed files with 101 additions and 12 deletions.
diff --git a/include/RED4ext/Detail/AddressHashes.hpp b/include/RED4ext/Detail/AddressHashes.hpp
@@ -130,11 +130,14 @@ constexpr std::uint32_t JobDispatcher_DispatchJob = 0x9C441E82;
 #pragma endregion
 
 #pragma region JobHandle
-constexpr std::uint32_t JobHandle_ctor = 0x6EFF1BD2;
 constexpr std::uint32_t JobHandle_dtor = 0x9E480A2E;
 constexpr std::uint32_t JobHandle_Join = 0x9C9C097C;
 #pragma endregion
 
+#pragma region JobInternalHandle
+constexpr std::uint32_t JobInternalHandle_Acquire = 0x6EFF1BD2;
+#pragma endregion
+
 #pragma region JobQueue
 constexpr std::uint32_t JobQueue_ctor_FromGroup = 0xE750D4B;
 constexpr std::uint32_t JobQueue_ctor_FromParams = 0x82BD14F8;

diff --git a/include/RED4ext/JobQueue-inl.hpp b/include/RED4ext/JobQueue-inl.hpp
@@ -4,6 +4,11 @@
 #include <RED4ext/JobQueue.hpp>
 #endif
 
+#include <cstdint>
+#include <memory>
+
+#include <Windows.h>
+
 #include <RED4ext/Detail/AddressHashes.hpp>
 #include <RED4ext/Hashing/FNV1a.hpp>
 #include <RED4ext/Relocation.hpp>
@@ -41,18 +46,44 @@ RED4EXT_INLINE RED4ext::JobParamSet::JobParamSet() noexcept
 
 RED4EXT_INLINE RED4ext::JobHandle::JobHandle(uintptr_t aUnk)
 {
-    using func_t = void* (*)(void*, uintptr_t);
-    static UniversalRelocFunc<func_t> func(Detail::AddressHashes::JobHandle_ctor);
+    AcquireInternalHandle(aUnk);
+}
 
-    unk00 = func(nullptr, aUnk);
+RED4EXT_INLINE RED4ext::JobHandle::JobHandle(const JobHandle& aOther)
+{
+    CopyInternalHandle(aOther);
 }
 
-RED4EXT_INLINE RED4ext::JobHandle::~JobHandle()
+RED4EXT_INLINE RED4ext::JobHandle::JobHandle(JobHandle&& aOther) noexcept
 {
-    using func_t = void (*)(JobHandle*);
-    static UniversalRelocFunc<func_t> func(Detail::AddressHashes::JobHandle_dtor);
+    MoveInternalHandle(aOther);
+}
 
-    func(this);
+RED4EXT_INLINE RED4ext::JobHandle& RED4ext::JobHandle::operator=(const JobHandle& aOther)
+{
+    if (this != std::addressof(aOther))
+    {
+        ReleaseInternalHandle();
+        CopyInternalHandle(aOther);
+    }
+
+    return *this;
+}
+
+RED4EXT_INLINE RED4ext::JobHandle& RED4ext::JobHandle::operator=(JobHandle&& aOther) noexcept
+{
+    if (this != std::addressof(aOther))
+    {
+        ReleaseInternalHandle();
+        MoveInternalHandle(aOther);
+    }
+
+    return *this;
+}
+
+RED4EXT_INLINE RED4ext::JobHandle::~JobHandle()
+{
+    ReleaseInternalHandle();
 }
 
 RED4EXT_INLINE void RED4ext::JobHandle::Join(const JobHandle& aOther)
@@ -63,6 +94,37 @@ RED4EXT_INLINE void RED4ext::JobHandle::Join(const JobHandle& aOther)
     func(this, aOther);
 }
 
+RED4EXT_INLINE void RED4ext::JobHandle::AcquireInternalHandle(uintptr_t aUnk)
+{
+    using func_t = JobInternalHandle* (*)(void*, uintptr_t);
+    static UniversalRelocFunc<func_t> func(Detail::AddressHashes::JobInternalHandle_Acquire);
+
+    internal = func(nullptr, aUnk);
+}
+
+RED4EXT_INLINE void RED4ext::JobHandle::CopyInternalHandle(const JobHandle& aOther)
+{
+    internal = aOther.internal;
+    InterlockedExchangeAdd(&internal->unk1C, static_cast<uint32_t>(1));
+}
+
+RED4EXT_INLINE void RED4ext::JobHandle::MoveInternalHandle(JobHandle& aOther)
+{
+    internal = aOther.internal;
+    aOther.internal = nullptr;
+}
+
+RED4EXT_INLINE void RED4ext::JobHandle::ReleaseInternalHandle()
+{
+    if (internal)
+    {
+        using func_t = void (*)(JobHandle*);
+        static UniversalRelocFunc<func_t> func(Detail::AddressHashes::JobHandle_dtor);
+
+        func(this);
+    }
+}
+
 RED4EXT_INLINE RED4ext::JobQueue::JobQueue(const JobGroup& aGroup)
 {
     using func_t = JobQueue* (*)(JobQueue*, const JobGroup&);
@@ -105,11 +167,11 @@ RED4EXT_INLINE [[nodiscard]] RED4ext::JobHandle RED4ext::JobQueue::Capture()
 
 RED4EXT_INLINE void RED4ext::JobQueue::DispatchJob(const JobInstance& aJob)
 {
-    using func_t = uint32_t (*)(void*, const JobInstance&, uint8_t, JobHandle, JobHandle);
+    using func_t = uint32_t (*)(void*, const JobInstance&, uint8_t, JobInternalHandle*, JobInternalHandle*);
     static UniversalRelocFunc<func_t> func(Detail::AddressHashes::JobDispatcher_DispatchJob);
     static UniversalRelocPtr<void*> dispatcher(Detail::AddressHashes::JobDispatcher);
 
-    func(dispatcher, aJob, params.unk00, unk10, unk18);
+    func(dispatcher, aJob, params.unk00, unk10.internal, unk18.internal);
 }
 
 RED4EXT_INLINE void RED4ext::JobQueue::SyncWait()

diff --git a/include/RED4ext/JobQueue.hpp b/include/RED4ext/JobQueue.hpp
@@ -53,13 +53,31 @@ struct JobParamSet
 };
 RED4EXT_ASSERT_SIZE(JobParamSet, 0x3);
 
+struct JobInternalHandle
+{
+    uint64_t unk00;          // 00
+    uint64_t unk08;          // 08
+    uint64_t unk10;          // 10
+    volatile uint32_t unk18; // 18
+    volatile uint32_t unk1C; // 1C
+    uint8_t unk20;           // 20
+    JobParamSet params;      // 21
+};
+RED4EXT_ASSERT_SIZE(JobInternalHandle, 0x28);
+RED4EXT_ASSERT_OFFSET(JobInternalHandle, unk1C, 0x1C);
+
 /**
  * @brief A handle associated with a job queue or job group.
  * It's used to track job completion and sync job execution.
  */
-struct JobHandle
+class JobHandle
 {
+public:
     JobHandle(uintptr_t aUnk = 0);
+    JobHandle(const JobHandle&);
+    JobHandle(JobHandle&&) noexcept;
+    JobHandle& operator=(const JobHandle&);
+    JobHandle& operator=(JobHandle&&) noexcept;
     ~JobHandle();
 
     /**
@@ -70,7 +88,13 @@ struct JobHandle
      */
     void Join(const JobHandle& aOther);
 
-    void* unk00; // 00
+    JobInternalHandle* internal; // 00
+
+private:
+    void AcquireInternalHandle(uintptr_t aUnk);
+    void CopyInternalHandle(const JobHandle& aOther);
+    void MoveInternalHandle(JobHandle& aOther);
+    void ReleaseInternalHandle();
 };
 RED4EXT_ASSERT_SIZE(JobHandle, 0x8);