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"

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

        bool Init(const AuString &path, bool readOnly, bool directIO);

        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, std::enable_shared_from_this<NtAsyncFileTransaction>
    {
    public:
        ~NtAsyncFileTransaction();

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

        bool StartRead(AuUInt64 offset, void *, AuUInt32 length) override;
        bool StartWrite(AuUInt64 offset, const void *, AuUInt32 length) override;
        
        bool Complete() override;
        AuUInt32 GetLastPacketLength() override;
        
        void SetCallback(const AuSPtr<IAsyncFinishedSubscriber> &sub) override;
        
        bool Wait(AuUInt32 timeout) override;
    
        void DispatchCb();
        HANDLE GetHandle();
        AuSPtr<FileHandle> GetFileHandle();

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

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

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

    bool FileHandle::Init(const AuString &path, bool readOnly, bool directIO)
    {
        HANDLE fileHandle;

        auto pathex    = NormalizePathRet(path);
        auto win32Path = Locale::ConvertFromUTF8(pathex);

        auto flags = FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OVERLAPPED;

        if (directIO)
        {
            flags |= FILE_FLAG_NO_BUFFERING;
        }

        if (readOnly)
        {
            fileHandle = CreateFileW(win32Path.c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, flags, NULL);
        }
        else
        {
            CreateDirectories(pathex, true);
            fileHandle = CreateFileW(win32Path.c_str(), GENERIC_WRITE, FILE_SHARE_WRITE, NULL, CREATE_ALWAYS, flags, NULL);
        }

        if (fileHandle == INVALID_HANDLE_VALUE)
        {
            LogWarn("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, void *buffer, AuUInt32 length)
    {
        if (std::exchange(this->pin_, AuSharedFromThis()))
        {
            return {};
        }
        this->lastAbstractStat_ = length;
        this->lastAbstractOffset_ = offset;
        this->overlap_.Offset = offset & 0xFFFFFFFF;
        this->overlap_.OffsetHigh = (offset >> 32) & 0xFFFFFFFF;
        auto ret = ::ReadFileEx(this->handle_->handle, buffer, length, &overlap_, GenericCompletionRoutine);
        return TranslateNtStatus(this, ret);
    }

    bool NtAsyncFileTransaction::StartWrite(AuUInt64 offset, const void *buffer, AuUInt32 length)
    {
        if (std::exchange(this->pin_, AuSharedFromThis()))
        {
            return {};
        }
        this->lastAbstractStat_ = length;
        this->lastAbstractOffset_ = offset;
        this->overlap_.Offset = offset & 0xFFFFFFFF;
        this->overlap_.OffsetHigh = (offset >> 32) & 0xFFFFFFFF;
        auto ret = ::WriteFileEx(this->handle_->handle, buffer, length,  &overlap_, GenericCompletionRoutine);
        return TranslateNtStatus(this, ret);
    }

    void NtAsyncFileTransaction::DispatchCb()
    {
        if (std::exchange(this->latch_, true))
        {
            return;
        }
        
        auto hold = std::exchange(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 = WaitForSingleObject(this->event_, 0);
        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 = WaitForSingleObject(this->event_, timeout ? timeout : INFINITE);
        if (ret == WAIT_OBJECT_0)
        {
            DispatchCb();
            return true;
        }
        return false;
    }

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

    AUKN_SYM IAsyncFileStream *OpenAsyncNew(const AuString &path, bool readOnly, bool directIO)
    {
        auto fileHandle = AuMakeShared<FileHandle>();
        if (!fileHandle->Init(path, readOnly, directIO))
        {
            return {};
        }

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

        stream->Init(fileHandle);
        return stream;
    }

    AUKN_SYM void OpenAsyncRelease(IAsyncFileStream *handle)
    {
        SafeDelete<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(std::static_pointer_cast<NtAsyncFileTransaction>(file)->GetHandle());
        }

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

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

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