From dfc21453ef5377dedc93cd42eaf6841556d6353f Mon Sep 17 00:00:00 2001
From: Reece <code@reece.sx>
Date: Thu, 5 May 2022 11:24:44 +0100
Subject: [PATCH] [+] Add IAsyncTransaction::Reset [*] Bug fixes: Async.NT.cpp

---
 Include/Aurora/IO/FS/IAsyncTransaction.hpp |  5 ++
 Source/IO/FS/Async.NT.cpp                  | 88 ++++++++++++++++------
 Source/IO/FS/Async.NT.hpp                  |  7 +-
 3 files changed, 78 insertions(+), 22 deletions(-)

diff --git a/Include/Aurora/IO/FS/IAsyncTransaction.hpp b/Include/Aurora/IO/FS/IAsyncTransaction.hpp
index 9b10baaf..b6d82f30 100644
--- a/Include/Aurora/IO/FS/IAsyncTransaction.hpp
+++ b/Include/Aurora/IO/FS/IAsyncTransaction.hpp
@@ -45,5 +45,10 @@ namespace Aurora::IO::FS
          *        Polling the transaction may result in the execution of the callback.
         */ 
         virtual AuSPtr<Loop::ILoopSource> NewLoopSource() = 0;
+
+        /**
+         * @brief Resets ::Complete() and NewLoopSource()->IsSignaled() to false
+        */
+        virtual void Reset() = 0;
     };
 }
\ No newline at end of file
diff --git a/Source/IO/FS/Async.NT.cpp b/Source/IO/FS/Async.NT.cpp
index 2efb6369..964c19d7 100644
--- a/Source/IO/FS/Async.NT.cpp
+++ b/Source/IO/FS/Async.NT.cpp
@@ -29,7 +29,14 @@ namespace Aurora::IO::FS
 
     bool NtAsyncFileTransactionLoopSource::OnTrigger(AuUInt handle)
     {
-        return IsSignaled();
+        auto lock = caller_.lock();
+
+        if (lock)
+        {
+            return lock->Complete();
+        }
+
+        return true;
     }
 
     Loop::ELoopSource    NtAsyncFileTransactionLoopSource::GetType()
@@ -40,8 +47,17 @@ namespace Aurora::IO::FS
     bool NtAsyncFileTransactionLoopSource::IsSignaled()
     {
         auto lock = caller_.lock();
-        if (!lock) return LSHandle::IsSignaled();
-        return lock->Complete();
+        
+        //if (LSHandle::IsSignaled())
+        {
+            if (lock)
+            {
+                return lock->Complete();
+            }
+
+            return LSHandle::IsSignaled();
+        }
+        return false;
     }
 
     NtAsyncFileTransaction::~NtAsyncFileTransaction()
@@ -209,22 +225,17 @@ namespace Aurora::IO::FS
         this->overlap_.hEvent = this->event_ = CreateEventW(nullptr, true, false, nullptr);
         return this->overlap_.hEvent != INVALID_HANDLE_VALUE;
     }
-
     static bool TranslateNtStatus(NtAsyncFileTransaction *that, BOOL val)
     {
         if ((val) || 
-            (!val && GetLastError() == ERROR_IO_PENDING))
+            (!val && (GetLastError() == ERROR_IO_PENDING)))
         {
-            if (val)
-            {
-                SetEvent(that->event_);
-                that->DispatchCb(that->lastAbstractStat_);
-            }
             return true;
         }
         else
         {
             that->pin_.reset();
+            that->Reset();
             SysPushErrorFIO("Async FIO error: {} {}", that->GetFileHandle()->path, GetLastError());
             return false;
         }
@@ -237,14 +248,37 @@ namespace Aurora::IO::FS
     )
     {
         auto transaction = reinterpret_cast<NtAsyncFileTransaction *>(reinterpret_cast<AuUInt8*>(lpOverlapped) - offsetof(NtAsyncFileTransaction, overlap_));
-        transaction->Complete();
+        auto hold = AuExchange(transaction->pin_, {});
+        hold->Complete();
+    }
+
+    bool NtAsyncFileTransaction::IDontWannaUsePorts()
+    {
+        // TODO: maybe we could use a semaphore counter of read attempts
+        // to allow APC callbacks to drag behind
+
+        if (AuExchange(this->pin_, AuSharedFromThis()))
+        {
+            while (SleepEx(100, true) == WAIT_IO_COMPLETION)
+            {
+
+            }
+
+            if (AuExchange(this->pin_, AuSharedFromThis()))
+            {
+                SysPushErrorUnavailableError();
+                return {};
+            }
+        }
+
+        return true;
     }
 
     bool NtAsyncFileTransaction::StartRead(AuUInt64 offset, const AuSPtr<AuMemoryViewWrite> &memoryView)
     {
-        if (AuExchange(this->pin_, AuSharedFromThis()))
+        if (!IDontWannaUsePorts())
         {
-            return {};
+            return false;
         }
 
         this->latch_ = false;
@@ -264,9 +298,9 @@ namespace Aurora::IO::FS
 
     bool NtAsyncFileTransaction::StartWrite(AuUInt64 offset, const AuSPtr<AuMemoryViewRead> &memoryView)
     {
-        if (AuExchange(this->pin_, AuSharedFromThis()))
+        if (!IDontWannaUsePorts())
         {
-            return {};
+            return false;
         }
 
         this->latch_ = false;
@@ -290,28 +324,40 @@ namespace Aurora::IO::FS
             return;
         }
         
+        auto memoryHold = this->memoryHold_;
+        this->memoryHold_.reset();
+
         if (this->sub_)
         {
             this->sub_->OnAsyncFileOpFinished(this->lastAbstractOffset_, read);
         }
+    }
 
-        this->memoryHold_.reset();
+    void NtAsyncFileTransaction::Reset()
+    {
+        ::ResetEvent(this->event_);
+        this->lastAbstractStat_ = 0; // do not use latch
     }
 
     bool NtAsyncFileTransaction::Complete()
     {
         DWORD read;
-        
+
+        if (!this->lastAbstractStat_)
+        {
+            return false;
+        }
+
         if (this->latch_)
         {
-            return true;
+            return false;
         }
 
         if (GetOverlappedResult(this->handle_->handle, &this->overlap_, &read, false))
         {
-            auto hold = AuExchange(this->pin_, {});
-            hold->DispatchCb(read);
-            return true;
+            bool bLatched = this->latch_;
+            DispatchCb(read);
+            return !bLatched;
         }
 
         return false;
diff --git a/Source/IO/FS/Async.NT.hpp b/Source/IO/FS/Async.NT.hpp
index d5988d99..68617038 100644
--- a/Source/IO/FS/Async.NT.hpp
+++ b/Source/IO/FS/Async.NT.hpp
@@ -55,6 +55,10 @@ namespace Aurora::IO::FS
         bool Wait(AuUInt32 timeout) override;
         AuSPtr<Loop::ILoopSource> NewLoopSource() override;
 
+        void Reset();
+
+        bool IDontWannaUsePorts();
+
         void DispatchCb(AuUInt32 len);
         HANDLE GetHandle();
         AuSPtr<FileHandle> GetFileHandle();
@@ -64,10 +68,11 @@ namespace Aurora::IO::FS
         HANDLE event_ = INVALID_HANDLE_VALUE;
         AuSPtr<NtAsyncFileTransaction> pin_;
         AuUInt32 lastAbstractStat_ {}, lastAbstractOffset_ {};
+        bool latch_ {};
+
     private:
         AuSPtr<void> memoryHold_;
         AuSPtr<FileHandle> handle_;
-        bool latch_ {};
         AuSPtr<IAsyncFinishedSubscriber> sub_;
     };
 }
\ No newline at end of file