AuroraRuntime/Source/IO/FS/Async.NT.cpp

/***
    Copyright (C) 2021 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: Async.NT.cpp
    Date: 2021-9-13
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "FS.hpp"
#include "Async.NT.hpp"
#include "FileAdvisory.NT.hpp"
#include <Source/Loop/Loop.hpp>
#include <Source/Loop/LSHandle.hpp>

namespace Aurora::IO::FS
{
    struct FileHandle
    {
        ~FileHandle();

        bool Init(const AuString &path, EFileOpenMode openMode, bool directIO, EFileAdvisoryLockLevel lock);

        HANDLE handle = INVALID_HANDLE_VALUE;
        AuString path;
        bool readOnly;
        bool directIO;
    };

    class NtAsyncFileStream : public IAsyncFileStream
    {
    public:
        AuSPtr<IAsyncTransaction> NewTransaction() override;

        void Init(const AuSPtr<FileHandle> &handle);

        AuSPtr<FileHandle> GetHandle();

    private:
        AuSPtr<FileHandle> handle_;
    };

    class NtAsyncFileTransaction : public IAsyncTransaction, AuEnableSharedFromThis<NtAsyncFileTransaction>
    {
    public:
        ~NtAsyncFileTransaction();

        bool Init(const AuSPtr<FileHandle> &handle);

        bool StartRead(AuUInt64 offset, const AuSPtr<AuMemoryViewWrite> &memoryView) override;
        bool StartWrite(AuUInt64 offset, const AuSPtr<AuMemoryViewRead> &memoryView) override;
        
        bool Complete() override;
        AuUInt32 GetLastPacketLength() override;
        
        void SetCallback(const AuSPtr<IAsyncFinishedSubscriber> &sub) override;
        
        bool Wait(AuUInt32 timeout) override;
        AuSPtr<Loop::ILoopSource> NewLoopSource() override;
    
        void DispatchCb();
        HANDLE GetHandle();
        AuSPtr<FileHandle> GetFileHandle();

        // Required for a very evil hack
        OVERLAPPED overlap_ {};
    private:
        AuSPtr<void> memoryHold_;
        AuSPtr<NtAsyncFileTransaction> pin_;
        AuSPtr<FileHandle> handle_;
        HANDLE event_ = INVALID_HANDLE_VALUE;
        AuUInt32 lastAbstractStat_ {}, lastAbstractOffset_ {};
        bool latch_ {};
        AuSPtr<IAsyncFinishedSubscriber> sub_;
    };

    struct NtAsyncFileTransactionLoopSource : public Loop::LSHandle
    {
        NtAsyncFileTransactionLoopSource(AuSPtr<NtAsyncFileTransaction> that) : caller_(that), Loop::LSHandle(AuUInt(that->GetFileHandle()->handle))
        {}

        virtual bool IsSignaled() override;
        virtual bool OnTrigger(AuUInt handle) override;
        virtual Loop::ELoopSource GetType() override;

    private:

        AuWPtr<NtAsyncFileTransaction> caller_;

    };

    bool NtAsyncFileTransactionLoopSource::OnTrigger(AuUInt handle)
    {
        return IsSignaled();
    }

    Loop::ELoopSource    NtAsyncFileTransactionLoopSource::GetType()
    {
        return Loop::ELoopSource::eSourceAIO;
    }

    bool NtAsyncFileTransactionLoopSource::IsSignaled()
    {
        auto lock = caller_.lock();
        if (!lock) return false;
        return lock->Complete();
    }

    NtAsyncFileTransaction::~NtAsyncFileTransaction()
    {
        AuWin32CloseHandle(this->event_);
    }

    FileHandle::~FileHandle()
    {
        AuWin32CloseHandle(this->handle);
    }

    bool FileHandle::Init(const AuString &path, EFileOpenMode openMode, bool directIO, EFileAdvisoryLockLevel lock)
    {
        HANDLE fileHandle;

        auto pathex    = NormalizePathRet(path);
        if (!pathex.empty())
        {
            return false;
        }

        auto win32Path = Locale::ConvertFromUTF8(pathex);
        if (!win32Path.empty())
        {
            return false;
        }

        auto flags = FILE_FLAG_OVERLAPPED;

        if (directIO)
        {
            flags |= FILE_FLAG_NO_BUFFERING;
        }

        fileHandle = INVALID_HANDLE_VALUE;

        switch (openMode)
        {
        case EFileOpenMode::eRead:
        {
            fileHandle = CreateFileW(win32Path.c_str(), GENERIC_READ, NtLockAdvisoryToShare(lock), NULL, OPEN_EXISTING, flags, NULL);
            break;
        }
        case EFileOpenMode::eReadWrite:
        case EFileOpenMode::eWrite:
        {
            CreateDirectories(pathex, true);
            fileHandle = CreateFileW(win32Path.c_str(), (openMode == EFileOpenMode::eReadWrite) ? (GENERIC_WRITE | GENERIC_READ) : GENERIC_WRITE, NtLockAdvisoryToShare(lock), NULL, CREATE_ALWAYS, flags, NULL);
            break;
        }
        }

        if (fileHandle == INVALID_HANDLE_VALUE)
        {
            AuLogWarn("Missing file: {}", path);
            SysPushErrorIO("Missing file: {}", path);
            return {};
        }

        this->directIO = directIO;
        this->handle = fileHandle;
        this->readOnly = readOnly;
        return true;
    }
    
    AuSPtr<FileHandle> NtAsyncFileStream::GetHandle()
    {
        return handle_;
    }

    void NtAsyncFileStream::Init(const AuSPtr<FileHandle> &handle)
    {
        this->handle_ = handle;
    }

    AuSPtr<IAsyncTransaction> NtAsyncFileStream::NewTransaction()
    {
        auto shared = AuMakeShared<NtAsyncFileTransaction>();
        if (!shared)
        {
            return {};
        }

        if (!shared->Init(this->handle_))
        {
            return {};
        }

        return shared;
    }

    bool NtAsyncFileTransaction::Init(const AuSPtr<FileHandle> &handle)
    {
        this->handle_ = handle;
        this->overlap_.hEvent = this->event_ = CreateEvent(NULL, true, 0, NULL);
        return this->overlap_.hEvent != INVALID_HANDLE_VALUE;
    }

    static bool TranslateNtStatus(NtAsyncFileTransaction *that, BOOL val)
    {
        if ((val) || 
            (!val && GetLastError() == ERROR_IO_PENDING))
        {
            if (val)
            {
                that->DispatchCb();
            }
            return true;
        }
        else
        {
            SysPushErrorFIO("Async FIO error: {}", that->GetFileHandle()->path);
            return false;
        }
    }

    static VOID WINAPI GenericCompletionRoutine(
        _In_    DWORD dwErrorCode,
        _In_    DWORD dwNumberOfBytesTransfered,
        _Inout_ LPOVERLAPPED lpOverlapped
    )
    {
        auto transaction = reinterpret_cast<NtAsyncFileTransaction *>(reinterpret_cast<AuUInt8*>(lpOverlapped) - offsetof(NtAsyncFileTransaction, overlap_));
        transaction->Complete();
    }

    bool NtAsyncFileTransaction::StartRead(AuUInt64 offset, const AuSPtr<AuMemoryViewWrite> &memoryView)
    {
        if (AuExchange(this->pin_, AuSharedFromThis()))
        {
            return {};
        }

        this->memoryHold_ = memoryView;
        this->lastAbstractStat_ = memoryView->length;
        this->lastAbstractOffset_ = offset;
        this->overlap_.Offset = offset & 0xFFFFFFFF;
        this->overlap_.OffsetHigh = (offset >> 32) & 0xFFFFFFFF;
        auto ret = ::ReadFileEx(this->handle_->handle, memoryView->ptr, memoryView->length, &overlap_, GenericCompletionRoutine);
        return TranslateNtStatus(this, ret);
    }

    bool NtAsyncFileTransaction::StartWrite(AuUInt64 offset, const AuSPtr<AuMemoryViewRead> &memoryView)
    {
        if (AuExchange(this->pin_, AuSharedFromThis()))
        {
            return {};
        }

        this->memoryHold_ = memoryView;
        this->lastAbstractStat_ = memoryView->length;
        this->lastAbstractOffset_ = offset;
        this->overlap_.Offset = offset & 0xFFFFFFFF;
        this->overlap_.OffsetHigh = (offset >> 32) & 0xFFFFFFFF;
        auto ret = ::WriteFileEx(this->handle_->handle, memoryView->ptr, memoryView->length,  &overlap_, GenericCompletionRoutine);
        return TranslateNtStatus(this, ret);
    }

    void NtAsyncFileTransaction::DispatchCb()
    {
        if (AuExchange(this->latch_, true))
        {
            return;
        }
        
        this->memoryHold_.reset();
        auto hold = AuExchange(this->pin_, {});

        if (hold->sub_)
        {
            DWORD read {};
            GetOverlappedResult(hold->handle_->handle, &hold->overlap_, &read, false);
            hold->sub_->OnAsyncFileOpFinished(hold->lastAbstractOffset_, read);
        }
    }

    bool NtAsyncFileTransaction::Complete()
    {
        auto ret = WaitForSingleObjectEx(this->event_, 0, true);
        if (ret == WAIT_OBJECT_0)
        {
            DispatchCb();
            return true;
        }
        return false;
    }

    AuUInt32 NtAsyncFileTransaction::GetLastPacketLength()
    {
        DWORD read {};
        GetOverlappedResult(this->handle_->handle, &this->overlap_, &read, false);
        return read;
    }
    
    void NtAsyncFileTransaction::SetCallback(const AuSPtr<IAsyncFinishedSubscriber> &sub)
    {
        this->sub_ = sub;
    }
    
    bool NtAsyncFileTransaction::Wait(AuUInt32 timeout) 
    {
        auto ret = WaitForSingleObjectEx(this->event_, timeout ? timeout : INFINITE, true);
        if (ret == WAIT_OBJECT_0)
        {
            DispatchCb();
            return true;
        }
        return false;
    }

    HANDLE NtAsyncFileTransaction::GetHandle()
    {
        return this->event_;
    }

    AuSPtr<FileHandle> NtAsyncFileTransaction::GetFileHandle()
    {
        return this->handle_;
    }

    AuSPtr<Loop::ILoopSource> NtAsyncFileTransaction::NewLoopSource()
    {
        return AuMakeShared<NtAsyncFileTransactionLoopSource>(AuSharedFromThis());
    }

    AUKN_SYM IAsyncFileStream *OpenAsyncNew(const AuString &path, EFileOpenMode openMode, bool directIO, EFileAdvisoryLockLevel lock)
    {
        AuSPtr<FileHandle> fileHandle;
        NtAsyncFileStream *stream;

        if (path.empty())
        {
            SysPushErrorParam("Empty path");
            return {};
        }

        if (!EFileOpenModeIsValid(openMode))
        {
            SysPushErrorParam("Invalid open mode");
            return {};
        }

        try
        {
            fileHandle = AuMakeShared<FileHandle>();
            if (!fileHandle->Init(path, openMode, directIO, lock))
            {
                return {};
            }
        }
        catch (...)
        {
            SysPushErrorCatch("Couldn't initialize FileHandle");
            return {};
        }

        try
        {
            stream = _new NtAsyncFileStream();
            if (!stream)
            {
                return {};
            }

            stream->Init(fileHandle);
        }
        catch (...)
        {
            if (stream) delete stream;
            SysPushErrorCatch("Couldn't initialize NtAsyncFileStream");
            return {};
        }

        return stream;
    }

    AUKN_SYM void OpenAsyncRelease(IAsyncFileStream *handle)
    {
        AuSafeDelete<NtAsyncFileStream *>(handle);
    }
    
    AUKN_SYM bool WaitMultiple(const AuList<AuSPtr<IAsyncTransaction>> &files, AuUInt32 timeout)
    {
        AuList<HANDLE> handles;
        
        if (files.empty())
        {
            return true;
        }

        for (const auto & file : files)
        {
            handles.push_back(AuStaticPointerCast<NtAsyncFileTransaction>(file)->GetHandle());
        }

        auto ret = WaitForMultipleObjectsEx(handles.size(), handles.data(), false, timeout ? timeout : INFINITE, TRUE);

    #if 0
        for (const auto &file : files)
        {
            AuStaticPointerCast<NtAsyncFileTransaction>(file)->Complete();
        }
    #endif

        return ret != WAIT_TIMEOUT && ret != WAIT_FAILED;
    }
}