AuroraRuntime/Source/Threading/Threads/TLSView.cpp

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

    File: TLSView.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "TLSView.hpp"
#include "../Primitives/ThreadCookie.hpp"

namespace Aurora::Threading::Threads
{
    static AuUInt32 uAtomicCounter {};
    static thread_local AuUInt32 tlsThreadIndex { };
    static const auto kMaxMaxLinearThreads = AuNumericLimits<AuUInt16>::max() / 10;

    AuUInt32 GetThreadKey()
    {
        auto &uKey = tlsThreadIndex;

        if (!uKey)
        {
            uKey = AuAtomicAdd(&uAtomicCounter, 1u);

            if (uKey > kMaxMaxLinearThreads) [[unlikely]]
            {
                uKey = uAtomicCounter = kMaxMaxLinearThreads + 1;
            }
        }

        return uKey;
    }
    
    void SetThreadKey(AuUInt32 uThreadKey)
    {
        tlsThreadIndex = uThreadKey;
    }

    struct ViewEntry
    {
        ~ViewEntry();

        void *pVoid {};
        TLSView::TlsCb deinit;
    };

    struct ViewContext
    {
        // 256 bytes, less than some of microsofts stupid stl thread primitives
        // so, ill take this as acceptable
        /*void **/ ViewEntry table[64];

        AuThreadPrimitives::Mutex mutex;
        AuList<AuSPtr<ViewEntry>> ptrs;
        AuHashMap<AuUInt, AuSPtr<ViewEntry>> hashMap;

        ViewEntry &GetHandleReference(AuOptional<AuUInt32> optThreadTlsIdentityToken);
    };

    ViewEntry::~ViewEntry()
    {
        if (this->pVoid && this->deinit)
        {
            this->deinit(this->pVoid);
        }

        if (this->pVoid)
        {
            AuMemory::_Free(this->pVoid);
            this->pVoid = nullptr;
        }
    }

    ViewEntry &ViewContext::GetHandleReference(AuOptional<AuUInt32> optThreadTlsIdentityToken)
    {
        auto uKey = optThreadTlsIdentityToken.OrElse(&GetThreadKey).Value();

        if (uKey >= AuArraySize(this->table))
        {
            AU_LOCK_GUARD(this->mutex);

            if (uKey > kMaxMaxLinearThreads) [[unlikely]]
            {
                auto &sharedRef = this->hashMap[Primitives::GetThreadCookie()];
                if (!sharedRef)
                {
                    sharedRef = AuMakeSharedPanic<ViewEntry>();
                }

                return *sharedRef.get();
            }

            uKey -= AuArraySize(this->table);

            if (uKey >= this->ptrs.size())
            {
                this->ptrs.reserve(uKey * 2);
                this->ptrs.resize(uKey + 1);
            }

            auto &sharedRef = this->ptrs[uKey];
            if (!sharedRef)
            {
                sharedRef = AuMakeSharedPanic<ViewEntry>();
            }

            return *sharedRef.get();
        }
        else
        {
            return this->table[uKey];
        }
    }

    void *TLSView::InitTLS(TLSKey &key, AuMach length, AuOptional<AuUInt32> optThreadTlsIdentityToken, const TlsCb &init, const TlsCb &deinit)
    {
        ViewContext *pContext;

        if (!key.pContext)
        {
            AU_LOCK_GUARD(key.lock);
         
            if (key.pContext)
            {
                pContext = (ViewContext *)key.pContext;
            }
            else
            {
                key.pContext = pContext = new ViewContext();
            }
        }
        else
        {
            pContext = (ViewContext *)key.pContext;
        }

        auto &refEntry = pContext->GetHandleReference(optThreadTlsIdentityToken);
        if (refEntry.pVoid)
        {
            return refEntry.pVoid;
        }
        else
        {
            refEntry.pVoid = AuMemory::_FAlloc(length, 64);
            SysAssert(refEntry.pVoid, "TLS allocation failure");

            if (init)
            {
                init(refEntry.pVoid);
            }

            refEntry.deinit = deinit;
            return refEntry.pVoid;
        }
    }

    void  TLSView::Remove(TLSKey &key)
    {
        if (key.pContext)
        {
            delete (ViewContext *)key.pContext;
        }
    }
}