AttributeKeyedCache.h

#ifndef INCLUDED_FNGEOLIBUTIL_ATTRIBUTEKEYEDCACHE_H
#define INCLUDED_FNGEOLIBUTIL_ATTRIBUTEKEYEDCACHE_H

#include "ns.h"

#include <atomic>
#include <cstdlib>
#include <memory>
#include <mutex>
#include <string>
#ifdef ATTRIBUTEKEYEDCACHE_USE_BOOST
#include <boost/unordered_map.hpp>
#include <boost/unordered_set.hpp>
#else
#include <map>
#include <set>
#endif
#include <vector>

#ifndef ATTRIBUTEKEYEDCACHE_NO_TBB
#include <tbb/concurrent_hash_map.h>
#include <tbb/concurrent_queue.h>
#endif

#include <FnAttribute/FnAttribute.h>

FNGEOLIBUTIL_NAMESPACE_ENTER
{
    template <class T, class PointerT = typename std::shared_ptr<T>>
    class AttributeKeyedCacheMutex
    {
    public:
        typedef PointerT IMPLPtr;

#ifdef ATTRIBUTEKEYEDCACHE_USE_BOOST
        using KeyMapType = boost::unordered_map<uint64_t, IMPLPtr>;
        using KeySetType = boost::unorderd_set<uint64_t>;
#else
        using KeyMapType = std::map<uint64_t, IMPLPtr>;
        using KeySetType = std::set<uint64_t>;
#endif

        AttributeKeyedCacheMutex(
            const std::size_t maxNumEntries = 0xffffffff,
            const std::size_t maxNumInvalidKeys = 0xffffffff)
            : m_maxNumEntries(maxNumEntries),
              m_maxNumInvalidKeys(maxNumInvalidKeys)
        {
        }

        virtual ~AttributeKeyedCacheMutex() {}

        IMPLPtr getValue(const FnAttribute::Attribute& iAttr)
        {
            // For these purposes, good enough to use a simple 64-bit hash.
            // The rationale is that AttributeKeyedCache should not grow to
            // the point where 2**64th hash entries are insufficient.
            //
            // For example, with 65K worth of entries the odds are roughly
            // 1 in 10-billion.
            // For 6 million entries, the odds would still be 1 in a million
            // (If the num entries is greater than ~20 million, then all 128
            // bits of hash should be used)

            const uint64_t hash = iAttr.getHash().uint64();

            if (m_maxNumInvalidKeys > 0)
            {
                // lock for our search in the invalid set
                std::lock_guard<std::mutex> lock(m_invalidKeysMutex);

                if (m_invalidKeys.find(hash) != m_invalidKeys.end())
                {
                    return IMPLPtr();
                }
            }

            if (m_maxNumEntries > 0)
            {
                // lock for our search in the cache
                std::lock_guard<std::mutex> lock(m_entriesMutex);

                const auto it = m_entries.find(hash);
                if (it != m_entries.end())
                {
                    return it->second;
                }
            }

            IMPLPtr val = createValue(iAttr);
            if (val && m_maxNumEntries > 0)
            {
                std::lock_guard<std::mutex> lock(m_entriesMutex);
                if (!m_entries.count(hash))
                {
                    if (m_entries.size() == m_maxNumEntries)
                    {
                        const std::size_t idx =
                            std::rand() % m_entriesVector.size();
                        const uint64_t keyToDrop = m_entriesVector[idx];
                        m_entriesVector[idx] = hash;
                        m_entries.erase(keyToDrop);
                    }
                    else
                    {
                        m_entriesVector.push_back(hash);
                    }
                    m_entries.emplace(hash, val);
                }
            }
            else if (!val && m_maxNumInvalidKeys > 0)
            {
                std::lock_guard<std::mutex> lock(m_invalidKeysMutex);
                if (!m_invalidKeys.count(hash))
                {
                    if (m_invalidKeys.size() == m_maxNumInvalidKeys)
                    {
                        const std::size_t idx =
                            std::rand() % m_invalidKeysVector.size();
                        const uint64_t keyToDrop = m_invalidKeysVector[idx];
                        m_invalidKeysVector[idx] = hash;
                        m_invalidKeys.erase(keyToDrop);
                    }
                    else
                    {
                        m_invalidKeysVector.push_back(hash);
                    }
                    m_invalidKeys.emplace(hash);
                }
            }
            return val;
        }

        void clear()
        {
            {
                std::lock_guard<std::mutex> lock(m_entriesMutex);
                m_entries.clear();
                m_entriesVector.clear();
            }

            {
                std::lock_guard<std::mutex> lock(m_invalidKeysMutex);
                m_invalidKeys.clear();
                m_invalidKeysVector.clear();
            }
        }

    protected:
        virtual IMPLPtr createValue(const FnAttribute::Attribute& iAttr) = 0;

    private:
        const std::size_t m_maxNumEntries;
        const std::size_t m_maxNumInvalidKeys;

        KeyMapType m_entries;
        KeySetType m_invalidKeys;

        std::vector<uint64_t> m_entriesVector;
        std::vector<uint64_t> m_invalidKeysVector;

        std::mutex m_invalidKeysMutex;
        std::mutex m_entriesMutex;
    };

#ifndef ATTRIBUTEKEYEDCACHE_NO_TBB
    template <class T, class PointerT =
                           typename std::shared_ptr<T>>
    class AttributeKeyedCacheLockFree
    {
    public:
        typedef PointerT IMPLPtr;

        AttributeKeyedCacheLockFree(
            const std::size_t maxNumEntries = 0xffffffff,
            const std::size_t maxNumInvalidKeys = 0xffffffff)
            : m_maxNumEntries(maxNumEntries),
              m_maxNumInvalidKeys(maxNumInvalidKeys)
        {
        }

        virtual ~AttributeKeyedCacheLockFree() {}

        IMPLPtr getValue(const FnAttribute::Attribute& iAttr)
        {
            // For these purposes, good enough to use a simple 64-bit hash.
            // The rationale is that AttributeKeyedCache should not grow to
            // the point where 2**64th hash entries are insufficient.
            //
            // For example, with 65K worth of entries the odds are roughly
            // 1 in 10-billion.
            // For 6 million entries, the odds would still be 1 in a million
            // (If the num entries is greater than ~20 million, then all 128
            // bits of hash should be used)

            const uint64_t hash = iAttr.getHash().uint64();

            if (m_maxNumInvalidKeys > 0)
            {
                typename decltype(m_invalidKeys)::const_accessor acc;
                if (m_invalidKeys.find(acc, hash))
                {
                    return {};
                }
            }

            if (m_maxNumEntries > 0)
            {
                typename decltype(m_entries)::const_accessor acc;
                if (m_entries.find(acc, hash))
                {
                    return acc->second;
                }
            }

            IMPLPtr val = createValue(iAttr);
            if (val && m_maxNumEntries > 0)
            {
                insert<IMPLPtr, typename decltype(m_entries)::accessor>(
                    hash, val, m_maxNumEntries, m_numEntries, m_entries,
                    m_entriesQueue);
            }
            else if (!val && m_maxNumInvalidKeys > 0)
            {
                insert<bool, typename decltype(m_invalidKeys)::accessor>(
                    hash, false, m_maxNumInvalidKeys, m_numInvalidKeys,
                    m_invalidKeys, m_invalidKeysQueue);
            }
            return val;
        }

        void clear()
        {
            if (m_numEntries.exchange(0))
            {
                m_entries.clear();
                m_entriesQueue.clear();
            }

            if (m_numInvalidKeys.exchange(0))
            {
                m_invalidKeys.clear();
                m_invalidKeysQueue.clear();
            }
        }

    protected:
        virtual IMPLPtr createValue(const FnAttribute::Attribute& iAttr) = 0;

    private:
        template <typename ValueT, typename AccessorT>
        static void insert(const uint64_t key,
                           const ValueT& value,
                           const std::size_t& maxNumEntries,
                           std::atomic_size_t& numEntries,
                           tbb::concurrent_hash_map<uint64_t, ValueT>& entries,
                           tbb::concurrent_queue<uint64_t>& entriesQueue)
        {
            AccessorT acc;
            if (entries.insert(acc, key))
            {
                acc->second = value;
                acc.release();

                const bool dropKey = ++numEntries > maxNumEntries;
                if (dropKey)
                {
                    uint64_t keyToDrop = 0;
                    const bool retrieved = entriesQueue.try_pop(keyToDrop);
                    (void)retrieved;
                    assert(retrieved);
                    const bool erased = entries.erase(keyToDrop);
                    (void)erased;
                    assert(erased);
                    --numEntries;
                }

                entriesQueue.push(key);
            }
        }

    private:
        const std::size_t m_maxNumEntries;
        const std::size_t m_maxNumInvalidKeys;

        std::atomic_size_t m_numEntries{0};
        std::atomic_size_t m_numInvalidKeys{0};

        tbb::concurrent_hash_map<uint64_t, IMPLPtr> m_entries;
        tbb::concurrent_hash_map<uint64_t, bool /*unused*/> m_invalidKeys;

        tbb::concurrent_queue<uint64_t> m_entriesQueue;
        tbb::concurrent_queue<uint64_t> m_invalidKeysQueue;
    };
#endif

    template <class T,
              class PointerT =
                  typename std::shared_ptr<T>,
#ifdef ATTRIBUTEKEYEDCACHE_NO_TBB
              class Base = AttributeKeyedCacheMutex<T, PointerT> >
#else
              class Base = AttributeKeyedCacheLockFree<T, PointerT> >
#endif
    class AttributeKeyedCache : public Base
    {
    public:
        typedef PointerT IMPLPtr;

        AttributeKeyedCache(const std::size_t maxNumEntries = 0xffffffff,
                            const std::size_t maxNumInvalidKeys = 0xffffffff)
            : Base(maxNumEntries, maxNumInvalidKeys)
        {
        }

        virtual ~AttributeKeyedCache() {}

        IMPLPtr getValue(const FnAttribute::Attribute& iAttr)
        {
            return Base::getValue(iAttr);
        }

        void clear() { Base::clear(); }

    protected:
        virtual IMPLPtr createValue(const FnAttribute::Attribute& iAttr) = 0;
    };

}
FNGEOLIBUTIL_NAMESPACE_EXIT

#endif  // INCLUDED_FNGEOLIBUTIL_ATTRIBUTEKEYEDCACHE_H