Merge pull request #10424 from vespa-engine/geirst/rewrite-enum-store-dictionary-3

Geirst/rewrite enum store dictionary 3

8 files changed, 183 insertions, 24 deletions

diff --git a/vespalib/CMakeLists.txt b/vespalib/CMakeLists.txt
index 041b2ccaba0..1a111ed49ef 100644
--- a/vespalib/CMakeLists.txt
+++ b/vespalib/CMakeLists.txt
@@ -42,6 +42,7 @@ vespa_define_module(
     src/tests/datastore/buffer_type
     src/tests/datastore/datastore
     src/tests/datastore/unique_store
+    src/tests/datastore/unique_store_dictionary
     src/tests/datastore/unique_store_string_allocator
     src/tests/delegatelist
     src/tests/dotproduct
diff --git a/vespalib/src/tests/datastore/unique_store_dictionary/CMakeLists.txt b/vespalib/src/tests/datastore/unique_store_dictionary/CMakeLists.txt
new file mode 100644
index 00000000000..b1478dea22c
--- /dev/null
+++ b/vespalib/src/tests/datastore/unique_store_dictionary/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(vespalib_unique_store_dictionary_test_app TEST
+    SOURCES
+    unique_store_dictionary_test.cpp
+    DEPENDS
+    vespalib
+    gtest
+)
+vespa_add_test(NAME vespalib_unique_store_dictionary_test_app COMMAND vespalib_unique_store_dictionary_test_app)
diff --git a/vespalib/src/tests/datastore/unique_store_dictionary/unique_store_dictionary_test.cpp b/vespalib/src/tests/datastore/unique_store_dictionary/unique_store_dictionary_test.cpp
new file mode 100644
index 00000000000..8b963dbf007
--- /dev/null
+++ b/vespalib/src/tests/datastore/unique_store_dictionary/unique_store_dictionary_test.cpp
@@ -0,0 +1,84 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/vespalib/datastore/unique_store.hpp>
+#include <vespa/vespalib/datastore/unique_store_dictionary.hpp>
+#include <vespa/vespalib/gtest/gtest.h>
+
+#include <vespa/log/log.h>
+LOG_SETUP("unique_store_dictionary_test");
+
+using namespace search::datastore;
+using namespace search::datastore::uniquestore;
+
+class Comparator : public EntryComparator {
+private:
+    EntryRef _to_find;
+
+    EntryRef resolve(EntryRef ref) const {
+        if (ref == EntryRef()) {
+            return _to_find;
+        }
+        return ref;
+    }
+
+public:
+    Comparator(uint32_t to_find)
+        : _to_find(to_find)
+    {}
+    bool operator()(const EntryRef lhs, const EntryRef rhs) const override {
+        return resolve(lhs).ref() < resolve(rhs).ref();
+    }
+};
+
+struct DictionaryReadTest : public ::testing::Test {
+    DefaultUniqueStoreDictionary dict;
+    UniqueStoreDictionaryBase::ReadSnapshot::UP snapshot;
+
+    DictionaryReadTest()
+        : dict(),
+          snapshot()
+    {
+    }
+    DictionaryReadTest& add(uint32_t value) {
+        auto result = dict.add(Comparator(value), [=]() { return EntryRef(value); });
+        assert(result.inserted());
+        return *this;
+    }
+    void take_snapshot() {
+        dict.freeze();
+        snapshot = dict.get_read_snapshot();
+    }
+};
+
+TEST_F(DictionaryReadTest, can_count_occurrences_of_a_key)
+{
+    add(3).add(5).take_snapshot();
+    EXPECT_EQ(0, snapshot->count(Comparator(2)));
+    EXPECT_EQ(1, snapshot->count(Comparator(3)));
+    EXPECT_EQ(0, snapshot->count(Comparator(4)));
+    EXPECT_EQ(1, snapshot->count(Comparator(5)));
+}
+
+TEST_F(DictionaryReadTest, can_count_occurrences_of_keys_in_a_range)
+{
+    add(3).add(5).add(7).add(9).take_snapshot();
+    EXPECT_EQ(1, snapshot->count_in_range(Comparator(3), Comparator(3)));
+    EXPECT_EQ(1, snapshot->count_in_range(Comparator(3), Comparator(4)));
+    EXPECT_EQ(2, snapshot->count_in_range(Comparator(3), Comparator(5)));
+    EXPECT_EQ(3, snapshot->count_in_range(Comparator(3), Comparator(7)));
+    EXPECT_EQ(4, snapshot->count_in_range(Comparator(3), Comparator(9)));
+    EXPECT_EQ(4, snapshot->count_in_range(Comparator(3), Comparator(10)));
+
+    EXPECT_EQ(0, snapshot->count_in_range(Comparator(5), Comparator(3)));
+}
+
+TEST_F(DictionaryReadTest, can_iterate_all_keys)
+{
+    using EntryRefVector = std::vector<EntryRef>;
+    add(3).add(5).add(7).take_snapshot();
+    EntryRefVector refs;
+    snapshot->foreach_key([&](EntryRef ref){ refs.emplace_back(ref); });
+    EXPECT_EQ(EntryRefVector({EntryRef(3), EntryRef(5), EntryRef(7)}), refs);
+}
+
+GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.h b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.h
index cbc681b96e3..cd13e88c77e 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.h
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.h
@@ -13,13 +13,26 @@ class EntryComparatorWrapper;
  * A dictionary for unique store. Mostly accessed via base class.
  */
 template <typename DictionaryT, typename ParentT = UniqueStoreDictionaryBase>
-class UniqueStoreDictionary : public ParentT
-{
+class UniqueStoreDictionary : public ParentT {
 protected:
     using DictionaryType = DictionaryT;
     using DataType = typename DictionaryType::DataType;
+    using FrozenView = typename DictionaryType::FrozenView;
+    using ReadSnapshot = typename ParentT::ReadSnapshot;
     using generation_t = typename ParentT::generation_t;
 
+    class ReadSnapshotImpl : public ReadSnapshot {
+    private:
+        FrozenView _frozen_view;
+
+    public:
+        ReadSnapshotImpl(FrozenView frozen_view);
+        EntryRef get_frozen_root() const override { return _frozen_view.getRoot(); }
+        size_t count(const EntryComparator& comp) const override;
+        size_t count_in_range(const EntryComparator& low, const EntryComparator& high) const override;
+        void foreach_key(std::function<void(EntryRef)> callback) const override;
+    };
+
     DictionaryType _dict;
 
 public:
@@ -35,8 +48,8 @@ public:
     uint32_t get_num_uniques() const override;
     vespalib::MemoryUsage get_memory_usage() const override;
     void build(const std::vector<EntryRef> &refs, const std::vector<uint32_t> &ref_counts, std::function<void(EntryRef)> hold) override;
+    std::unique_ptr<ReadSnapshot> get_read_snapshot() const override;
     EntryRef get_frozen_root() const override;
-    void foreach_key(EntryRef root, std::function<void(EntryRef)> callback) const override;
 };
 
 }
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.hpp b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.hpp
index 6e40b0b7c97..f3087bc5610 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.hpp
@@ -17,6 +17,47 @@
 namespace search::datastore {
 
 template <typename DictionaryT, typename ParentT>
+UniqueStoreDictionary<DictionaryT, ParentT>::
+ReadSnapshotImpl::ReadSnapshotImpl(FrozenView frozen_view)
+    : _frozen_view(frozen_view)
+{
+}
+
+template <typename DictionaryT, typename ParentT>
+size_t
+UniqueStoreDictionary<DictionaryT, ParentT>::
+ReadSnapshotImpl::count(const EntryComparator& comp) const
+{
+    auto itr = _frozen_view.lowerBound(EntryRef(), comp);
+    if (itr.valid() && !comp(EntryRef(), itr.getKey())) {
+        return 1u;
+    }
+    return 0u;
+}
+
+template <typename DictionaryT, typename ParentT>
+size_t
+UniqueStoreDictionary<DictionaryT, ParentT>::
+ReadSnapshotImpl::count_in_range(const EntryComparator& low,
+                                 const EntryComparator& high) const
+{
+    auto low_itr = _frozen_view.lowerBound(EntryRef(), low);
+    auto high_itr = low_itr;
+    if (high_itr.valid() && !high(EntryRef(), high_itr.getKey())) {
+        high_itr.seekPast(EntryRef(), high);
+    }
+    return high_itr - low_itr;
+}
+
+template <typename DictionaryT, typename ParentT>
+void
+UniqueStoreDictionary<DictionaryT, ParentT>::
+ReadSnapshotImpl::foreach_key(std::function<void(EntryRef)> callback) const
+{
+    _frozen_view.foreach_key(callback);
+}
+
+template <typename DictionaryT, typename ParentT>
 UniqueStoreDictionary<DictionaryT, ParentT>::UniqueStoreDictionary()
     : ParentT(),
       _dict()
@@ -135,18 +176,17 @@ UniqueStoreDictionary<DictionaryT, ParentT>::build(const std::vector<EntryRef> &
 }
 
 template <typename DictionaryT, typename ParentT>
-EntryRef
-UniqueStoreDictionary<DictionaryT, ParentT>::get_frozen_root() const
+std::unique_ptr<typename ParentT::ReadSnapshot>
+UniqueStoreDictionary<DictionaryT, ParentT>::get_read_snapshot() const
 {
-    return _dict.getFrozenView().getRoot();
+    return std::make_unique<ReadSnapshotImpl>(_dict.getFrozenView());
 }
 
 template <typename DictionaryT, typename ParentT>
-void
-UniqueStoreDictionary<DictionaryT, ParentT>::foreach_key(EntryRef root,
-                                                         std::function<void(EntryRef)> callback) const
+EntryRef
+UniqueStoreDictionary<DictionaryT, ParentT>::get_frozen_root() const
 {
-    _dict.getAllocator().getNodeStore().foreach_key(root, callback);
+    return _dict.getFrozenView().getRoot();
 }
 
 }
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary_base.h b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary_base.h
index 82c8687513a..e09fdc6093c 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary_base.h
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary_base.h
@@ -12,12 +12,27 @@ class EntryComparator;
 struct ICompactable;
 class UniqueStoreAddResult;
 
-/*
- * Interface class for unique store dictonary.
+/**
+ * Interface class for unique store dictionary.
  */
-class UniqueStoreDictionaryBase
-{
+class UniqueStoreDictionaryBase {
 public:
+    /**
+     * Class that provides a read snapshot of the dictionary.
+     *
+     * A generation guard that must be taken and held while the snapshot is considered valid.
+     */
+    class ReadSnapshot {
+    public:
+        using UP = std::unique_ptr<ReadSnapshot>;
+        virtual ~ReadSnapshot() = default;
+        // TODO: Remove when all relevant functions have been migrated to this API.
+        virtual EntryRef get_frozen_root() const = 0;
+        virtual size_t count(const EntryComparator& comp) const = 0;
+        virtual size_t count_in_range(const EntryComparator& low, const EntryComparator& high) const = 0;
+        virtual void foreach_key(std::function<void(EntryRef)> callback) const = 0;
+    };
+
     using generation_t = vespalib::GenerationHandler::generation_t;
     virtual ~UniqueStoreDictionaryBase() = default;
     virtual void freeze() = 0;
@@ -30,8 +45,8 @@ public:
     virtual uint32_t get_num_uniques() const = 0;
     virtual vespalib::MemoryUsage get_memory_usage() const = 0;
     virtual void build(const std::vector<EntryRef> &refs, const std::vector<uint32_t> &ref_counts, std::function<void(EntryRef)> hold) = 0;
+    virtual std::unique_ptr<ReadSnapshot> get_read_snapshot() const = 0;
     virtual EntryRef get_frozen_root() const = 0;
-    virtual void foreach_key(EntryRef root, std::function<void(EntryRef)> callback) const = 0;
 };
 
 }
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_enumerator.h b/vespalib/src/vespa/vespalib/datastore/unique_store_enumerator.h
index a2e626c6b20..3f260bfbc15 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_enumerator.h
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_enumerator.h
@@ -19,24 +19,22 @@ public:
     using EnumValues = std::vector<std::vector<uint32_t>>;
 
 private:
-    const UniqueStoreDictionaryBase &_dict;
-    EntryRef _frozen_root;
+    UniqueStoreDictionaryBase::ReadSnapshot::UP _dict_snapshot;
     const DataStoreBase &_store;
     EnumValues _enumValues;
     uint32_t _next_enum_val;
 public:
     UniqueStoreEnumerator(const UniqueStoreDictionaryBase &dict, const DataStoreBase &store);
     ~UniqueStoreEnumerator();
-    EntryRef get_frozen_root() const { return _frozen_root; }
+    EntryRef get_frozen_root() const { return _dict_snapshot->get_frozen_root(); }
     void enumerateValue(EntryRef ref);
     void enumerateValues();
     void clear();
 
     template <typename Function>
     void
-    foreach_key(Function &&func) const
-    {
-        _dict.foreach_key(_frozen_root, func);
+    foreach_key(Function &&func) const {
+        _dict_snapshot->foreach_key(func);
     }
 
     uint32_t mapEntryRefToEnumValue(EntryRef ref) const {
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_enumerator.hpp b/vespalib/src/vespa/vespalib/datastore/unique_store_enumerator.hpp
index a6053d4f64e..10ca0944519 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_enumerator.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_enumerator.hpp
@@ -8,8 +8,7 @@ namespace search::datastore {
 
 template <typename RefT>
 UniqueStoreEnumerator<RefT>::UniqueStoreEnumerator(const UniqueStoreDictionaryBase &dict, const DataStoreBase &store)
-    : _dict(dict),
-      _frozen_root(_dict.get_frozen_root()),
+    : _dict_snapshot(dict.get_read_snapshot()),
       _store(store),
       _enumValues(),
       _next_enum_val(1)
@@ -46,7 +45,7 @@ UniqueStoreEnumerator<RefT>::enumerateValues()
         }
     }
     _next_enum_val = 1;
-    _dict.foreach_key(_frozen_root, [this](EntryRef ref) { enumerateValue(ref); });
+    _dict_snapshot->foreach_key([this](EntryRef ref) { enumerateValue(ref); });
 }
 
 template <typename RefT>