Optionally give node_locations_t store a max number of bytes to use

Author: joto

src/node-locations.cpp

@@ -21,8 +21,12 @@
 #include <cassert>
 #include <limits>
 
-void node_locations_t::set(osmid_t id, osmium::Location location)
+bool node_locations_t::set(osmid_t id, osmium::Location location)
 {
+    if (used_memory() >= m_max_size && will_resize()) {
+        return false;
+    }
+
     if (first_entry_in_block()) {
         m_did.clear();
         m_dx.clear();
@@ -37,6 +41,8 @@ void node_locations_t::set(osmid_t id, osmium::Location location)
         &m_data, protozero::encode_zigzag64(m_dy.update(location.y())));
 
     ++m_count;
+
+    return true;
 }
 
 osmium::Location node_locations_t::get(osmid_t id) const

src/node-locations.hpp

@@ -17,6 +17,7 @@
 
 #include <array>
 #include <cstddef>
+#include <limits>
 #include <string>
 #include <utility>
 
@@ -35,12 +36,25 @@
 class node_locations_t
 {
 public:
+    /**
+     * Construct a node locations store. Takes a single optional argument
+     * which gives the maximum number of bytes this store should be allowed
+     * to use. If this is not specified, the size is only limited by available
+     * memory. The store will try to keep the memory used under what's
+     * specified here.
+     */
+    explicit node_locations_t(
+        std::size_t max_size = std::numeric_limits<std::size_t>::max())
+    : m_max_size(max_size)
+    {}
+
     /**
      * Store a node location.
      *
      * \pre id must be strictly larger than all ids stored before.
+     * \return True if the entry was added, false if the index is full.
      */
-    void set(osmid_t id, osmium::Location location);
+    bool set(osmid_t id, osmium::Location location);
 
     /**
      * Retrieve a node location. If the location wasn't stored before, an
@@ -69,6 +83,17 @@ class node_locations_t
         return m_count % block_size == 0;
     }
 
+    /// The maximum number of bytes an entry will need in storage.
+    constexpr static std::size_t max_bytes_per_entry() noexcept {
+        return 10U /*max varint length*/ * 3U /*id, x, y*/;
+    }
+
+    bool will_resize() const noexcept
+    {
+        return m_index.will_resize() ||
+               (m_data.size() + max_bytes_per_entry() >= m_data.capacity());
+    }
+
     /**
      * The block size used for internal blocks. The larger the block size
      * the less memory is consumed but the more expensive the access is.
@@ -78,6 +103,9 @@ class node_locations_t
     ordered_index_t m_index;
     std::string m_data;
 
+    /// Maximum size in bytes this object may allocate.
+    std::size_t m_max_size;
+
     /// The number of (id, location) pairs stored.
     std::size_t m_count = 0;
 

src/ordered-index.hpp

@@ -150,6 +150,9 @@ class ordered_index_t
         m_size = 0;
     }
 
+    /// Return true if adding an entry to the index will make it resize.
+    bool will_resize() const noexcept { return m_size + 1 >= m_capacity; }
+
 private:
     struct second_level_index_entry
     {

tests/test-node-locations.cpp

@@ -16,8 +16,8 @@ TEST_CASE("node locations basics", "[NoDB]")
     node_locations_t nl;
     REQUIRE(nl.size() == 0);
 
-    nl.set(3, {1.2, 3.4});
-    nl.set(5, {5.6, 7.8});
+    REQUIRE(nl.set(3, {1.2, 3.4}));
+    REQUIRE(nl.set(5, {5.6, 7.8}));
 
     REQUIRE(nl.size() == 2);
 
@@ -76,3 +76,14 @@ TEST_CASE("node locations in more than one block", "[NoDB]")
     }
 }
 
+TEST_CASE("full node locations store", "[NoDB]")
+{
+    node_locations_t nl{30};
+    REQUIRE(nl.size() == 0);
+
+    REQUIRE(nl.set(3, {1.2, 3.4}));
+    REQUIRE_FALSE(nl.set(5, {5.6, 7.8}));
+
+    REQUIRE(nl.size() == 1);
+}
+