partitioned.h

Go to the documentation of this file.

#pragma once

#include <algorithm>
#include <cstdlib>
#include <iterator>
#include <memory>
#include <numeric>
#include <utility>
#include <vector>

#include <boost/align/aligned_allocator.hpp>

#include "mongo/stdx/mutex.h"
#include "mongo/util/assert_util.h"
#include "mongo/util/with_alignment.h"

namespace mongo {

inline std::size_t partitionOf(const char x, const std::size_t nPartitions) {
    return static_cast<unsigned char>(x) % nPartitions;
}
inline std::size_t partitionOf(const unsigned char x, const std::size_t nPartitions) {
    return x % nPartitions;
}
inline std::size_t partitionOf(const signed char x, const std::size_t nPartitions) {
    return static_cast<unsigned char>(x) % nPartitions;
}
inline std::size_t partitionOf(const int x, const std::size_t nPartitions) {
    return static_cast<unsigned int>(x) % nPartitions;
}
inline std::size_t partitionOf(const unsigned int x, const std::size_t nPartitions) {
    return x % nPartitions;
}
inline std::size_t partitionOf(const short x, const std::size_t nPartitions) {
    return static_cast<unsigned short>(x) % nPartitions;
}
inline std::size_t partitionOf(const unsigned short x, const std::size_t nPartitions) {
    return x % nPartitions;
}
inline std::size_t partitionOf(const long x, const std::size_t nPartitions) {
    return static_cast<unsigned long>(x) % nPartitions;
}
inline std::size_t partitionOf(const unsigned long x, const std::size_t nPartitions) {
    return x % nPartitions;
}
inline std::size_t partitionOf(const long long x, const std::size_t nPartitions) {
    return static_cast<unsigned long long>(x) % nPartitions;
}
inline std::size_t partitionOf(const unsigned long long x, const std::size_t nPartitions) {
    return x % nPartitions;
}
inline std::size_t partitionOf(const wchar_t x, const std::size_t nPartitions) {
    return x % nPartitions;
}
inline std::size_t partitionOf(const char16_t x, const std::size_t nPartitions) {
    return x % nPartitions;
}
inline std::size_t partitionOf(const char32_t x, const std::size_t nPartitions) {
    return x % nPartitions;
}

template <typename T>
struct Partitioner {
    std::size_t operator()(const T& x, const std::size_t nPartitions) {
        return partitionOf(x, nPartitions);
    }
};

namespace partitioned_detail {

using CacheAlignedMutex = CacheAligned<stdx::mutex>;

template <typename Key, typename Value>
Key getKey(const std::pair<Key, Value>& pair) {
    return std::get<0>(pair);
}

template <typename Key>
Key getKey(const Key& key) {
    return key;
}

template <typename T>
inline std::vector<stdx::unique_lock<stdx::mutex>> lockAllPartitions(T& mutexes) {
    std::vector<stdx::unique_lock<stdx::mutex>> result;
    result.reserve(mutexes.size());
    std::transform(mutexes.begin(), mutexes.end(), std::back_inserter(result), [](auto&& mutex) {
        return stdx::unique_lock<stdx::mutex>{mutex};
    });
    return result;
}
}  // namespace partitioned_detail

template <typename AssociativeContainer,
          std::size_t nPartitions = 16,
          typename KeyPartitioner = Partitioner<typename AssociativeContainer::key_type>>
class Partitioned {
private:
    // Used to create an iterator representing the end of the partitioned structure.
    struct IteratorEndTag {};

public:
    static_assert(nPartitions > 0, "cannot create partitioned structure with 0 partitions");
    using value_type = typename AssociativeContainer::value_type;
    using key_type = typename AssociativeContainer::key_type;
    using PartitionId = std::size_t;

    class All {
    private:
        explicit All(Partitioned& partitionedContainer)
            : _lockGuards(partitioned_detail::lockAllPartitions(partitionedContainer._mutexes)),
              _partitionedContainer(&partitionedContainer) {}

    public:
        auto begin() & {
            return this->_partitionedContainer->_partitions.begin();
        }

        auto end() & {
            return this->_partitionedContainer->_partitions.end();
        }

        auto begin() const& {
            return this->_partitionedContainer->_partitions.begin();
        }
        void begin() && = delete;

        auto end() const& {
            return this->_partitionedContainer->_partitions.end();
        }
        void end() && = delete;

        std::size_t size() const {
            return std::accumulate(
                this->_partitionedContainer->_partitions.begin(),
                this->_partitionedContainer->_partitions.end(),
                std::size_t{0},
                [](auto&& total, auto&& partition) { return total + partition.size(); });
        }

        bool empty() const {
            return std::all_of(this->_partitionedContainer->_partitions.begin(),
                               this->_partitionedContainer->_partitions.end(),
                               [](auto&& partition) { return partition.empty(); });
        }

        std::size_t count(const key_type& key) const {
            auto partitionId = KeyPartitioner()(key, nPartitions);
            return this->_partitionedContainer->_partitions[partitionId].count(key);
        }

        void clear() {
            for (auto&& partition : this->_partitionedContainer->_partitions) {
                partition.clear();
            }
        }

        void insert(value_type value) & {
            const auto partitionId =
                KeyPartitioner()(partitioned_detail::getKey(value), nPartitions);
            this->_partitionedContainer->_partitions[partitionId].insert(std::move(value));
        }
        void insert(value_type)&& = delete;

        std::size_t erase(const key_type& key) & {
            const auto partitionId = KeyPartitioner()(key, nPartitions);
            return this->_partitionedContainer->_partitions[partitionId].erase(key);
        }
        void erase(const key_type&) && = delete;

    private:
        friend class Partitioned;

        std::vector<stdx::unique_lock<stdx::mutex>> _lockGuards;
        Partitioned* _partitionedContainer;
    };

    class OnePartition {
    public:
        AssociativeContainer* operator->() const& {
            return &this->_partitioned->_partitions[_id];
        }
        void operator->() && = delete;

        AssociativeContainer& operator*() const& {
            return this->_partitioned->_partitions[_id];
        }
        void operator*() && = delete;

    private:
        friend class Partitioned;

        OnePartition(Partitioned& partitioned, PartitionId partitionId)
            : _partitionLock(partitioned._mutexes[partitionId]),
              _partitioned(&partitioned),
              _id(partitionId) {}

        stdx::unique_lock<stdx::mutex> _partitionLock;
        Partitioned* _partitioned;
        PartitionId _id;
    };

    Partitioned() : _mutexes(nPartitions), _partitions(nPartitions) {}

    Partitioned(const Partitioned&) = delete;
    Partitioned(Partitioned&&) = default;
    Partitioned& operator=(const Partitioned&) = delete;
    Partitioned& operator=(Partitioned&&) = default;
    ~Partitioned() = default;

    bool empty() const {
        const auto all = partitioned_detail::lockAllPartitions(this->_mutexes);
        return std::all_of(this->_partitions.begin(),
                           this->_partitions.end(),
                           [](auto&& partition) { return partition.empty(); });
    }

    std::size_t size() const {
        const auto all = partitioned_detail::lockAllPartitions(this->_mutexes);
        return std::accumulate(
            this->_partitions.begin(),
            this->_partitions.end(),
            std::size_t{0},
            [](auto&& total, auto&& partition) { return total + partition.size(); });
    }

    std::size_t count(const key_type& key) & {
        auto partition = this->lockOnePartition(key);
        return partition->count(key);
    }
    void count(const key_type&) && = delete;

    void clear() {
        auto all = this->lockAllPartitions();
        all.clear();
    }

    void insert(const value_type value) & {
        auto partition = this->lockOnePartitionById(
            KeyPartitioner()(partitioned_detail::getKey(value), nPartitions));
        partition->insert(std::move(value));
    }
    void insert(const value_type) && = delete;

    std::size_t erase(const key_type& key) & {
        auto partition = this->lockOnePartition(key);
        return partition->erase(key);
    }
    void erase(const key_type&) && = delete;

    All lockAllPartitions() & {
        return All{*this};
    }

    OnePartition lockOnePartition(const key_type key) & {
        return OnePartition{*this, KeyPartitioner()(key, nPartitions)};
    }

    OnePartition lockOnePartitionById(PartitionId id) & {
        return OnePartition{*this, id};
    }

private:
    using CacheAlignedAssociativeContainer = CacheAligned<AssociativeContainer>;

    template <typename T>
    using AlignedVector = std::vector<T, boost::alignment::aligned_allocator<T>>;

    // These two vectors parallel each other, but we keep them separate so that we can return an
    // iterator over `_partitions` from within All.
    mutable AlignedVector<partitioned_detail::CacheAlignedMutex> _mutexes;
    AlignedVector<CacheAlignedAssociativeContainer> _partitions;
};
}  // namespace mongo