Add memory tracking to Enumeration loading

This adds a new MemoryTracker type for tracking Enumeration data usage.
TileDB-Inc · Oct 3, 2023 · a390ad0 · a390ad0
1 parent 75caf3c
commit a390ad0
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 11 deletions.
diff --git a/test/src/unit-enumerations.cc b/test/src/unit-enumerations.cc
@@ -696,7 +696,9 @@ TEST_CASE_METHOD(
 
   auto enmr_path = schema->get_enumeration_path_name(enmr_name.value());
 
-  auto loaded = ad->load_enumerations_from_paths({enmr_path}, enc_key_);
+  MemoryTracker tracker;
+  auto loaded =
+      ad->load_enumerations_from_paths({enmr_path}, enc_key_, tracker);
   REQUIRE(loaded.size() == 1);
 
   auto enmr = loaded[0];
@@ -718,11 +720,44 @@ TEST_CASE_METHOD(
 
   auto schema = get_array_schema_latest();
   auto ad = get_array_directory();
+  MemoryTracker tracker;
 
   // Check that this function throws an exception when attempting to load
   // an unknown enumeration
-  REQUIRE_THROWS(
-      ad->load_enumerations_from_paths({"not_an_enumeration"}, enc_key_));
+  auto matcher =
+      Catch::Matchers::ContainsSubstring("No such file or directory");
+  REQUIRE_THROWS_WITH(
+      ad->load_enumerations_from_paths({"unknown_enmr"}, enc_key_, tracker),
+      matcher);
+}
+
+TEST_CASE_METHOD(
+    EnumerationFx,
+    "ArrayDirectory - Load Enumeration Memory Limit Exceeded",
+    "[enumeration][array-directory][error]") {
+  create_array();
+
+  auto schema = get_array_schema_latest();
+  auto ad = get_array_directory();
+
+  auto enmr_name = schema->attribute("attr1")->get_enumeration_name();
+  auto enmr_path = schema->get_enumeration_path_name(enmr_name.value());
+
+  MemoryTracker tracker;
+  tracker.set_budget(1);
+
+  // Check that this function throws an exception when attempting to load
+  // an enumeration that exceeds the memory budget.
+  auto matcher = Catch::Matchers::ContainsSubstring(
+      "Cannot load enumeration; Insufficient memory budget;");
+  REQUIRE_THROWS_WITH(
+      ad->load_enumerations_from_paths({enmr_path}, enc_key_, tracker),
+      matcher);
+
+  // Check that the fix is to increase the memory budget.
+  tracker.set_budget(std::numeric_limits<uint64_t>::max());
+  REQUIRE_NOTHROW(
+      ad->load_enumerations_from_paths({enmr_path}, enc_key_, tracker));
 }
 
 /* ********************************* */

diff --git a/tiledb/common/memory_tracker.h b/tiledb/common/memory_tracker.h
@@ -41,7 +41,13 @@ namespace sm {
 
 class MemoryTracker {
  public:
-  enum class MemoryType { RTREE, FOOTER, TILE_OFFSETS, MIN_MAX_SUM_NULL_COUNT };
+  enum class MemoryType {
+    RTREE,
+    FOOTER,
+    TILE_OFFSETS,
+    MIN_MAX_SUM_NULL_COUNT,
+    ENUMERATION
+  };
 
   /** Constructor. */
   MemoryTracker() {
@@ -150,4 +156,4 @@ class MemoryTracker {
 }  // namespace sm
 }  // namespace tiledb
 
-#endif  // TILEDB_OPEN_ARRAY_MEMORY_TRACKER_H
+#endif  // TILEDB_OPEN_ARRAY_MEMORY_TRACKER_H
diff --git a/tiledb/sm/array/array.cc b/tiledb/sm/array/array.cc
@@ -640,7 +640,7 @@ std::vector<shared_ptr<const Enumeration>> Array::get_enumerations(
 
       // Load the enumerations from storage
       loaded = array_dir_.load_enumerations_from_paths(
-          paths_to_load, get_encryption_key());
+          paths_to_load, get_encryption_key(), memory_tracker_);
     }
 
     // Store the loaded enumerations in the schema

diff --git a/tiledb/sm/array/array_directory.cc b/tiledb/sm/array/array_directory.cc
@@ -191,7 +191,8 @@ ArrayDirectory::load_all_array_schemas(
 std::vector<shared_ptr<const Enumeration>>
 ArrayDirectory::load_enumerations_from_paths(
     const std::vector<std::string>& enumeration_paths,
-    const EncryptionKey& encryption_key) const {
+    const EncryptionKey& encryption_key,
+    MemoryTracker& memory_tracker) const {
   // This should never be called with an empty list of enumeration paths, but
   // there's no reason to not check an early return case here given that code
   // changes.
@@ -202,7 +203,8 @@ ArrayDirectory::load_enumerations_from_paths(
   std::vector<shared_ptr<const Enumeration>> ret(enumeration_paths.size());
   auto& tp = resources_.get().io_tp();
   throw_if_not_ok(parallel_for(&tp, 0, enumeration_paths.size(), [&](size_t i) {
-    ret[i] = load_enumeration(enumeration_paths[i], encryption_key);
+    ret[i] =
+        load_enumeration(enumeration_paths[i], encryption_key, memory_tracker);
     return Status::Ok();
   }));
   return ret;
@@ -1313,13 +1315,26 @@ bool ArrayDirectory::consolidation_with_timestamps_supported(
 
 shared_ptr<const Enumeration> ArrayDirectory::load_enumeration(
     const std::string& enumeration_path,
-    const EncryptionKey& encryption_key) const {
+    const EncryptionKey& encryption_key,
+    MemoryTracker& memory_tracker) const {
   auto timer_se = resources_.get().stats().start_timer("sm_load_enumeration");
 
   auto enmr_uri = uri_.join_path(constants::array_schema_dir_name)
                       .join_path(constants::array_enumerations_dir_name)
                       .join_path(enumeration_path);
 
+  uint64_t size = 0;
+  throw_if_not_ok(resources_.get().vfs().file_size(enmr_uri, &size));
+
+  if (!memory_tracker.take_memory(
+          size, MemoryTracker::MemoryType::ENUMERATION)) {
+    throw ArrayDirectoryException(
+        "Cannot load enumeration; Insufficient memory budget; Needed " +
+        std::to_string(size) + " but only had " +
+        std::to_string(memory_tracker.get_memory_available()) +
+        " from budget " + std::to_string(memory_tracker.get_memory_budget()));
+  }
+
   auto&& tile = GenericTileIO::load(resources_, enmr_uri, 0, encryption_key);
   resources_.get().stats().add_counter("read_enumeration_size", tile.size());
 

diff --git a/tiledb/sm/array/array_directory.h b/tiledb/sm/array/array_directory.h
@@ -33,6 +33,7 @@
 #ifndef TILEDB_ARRAY_DIRECTORY_H
 #define TILEDB_ARRAY_DIRECTORY_H
 
+#include "tiledb/common/memory_tracker.h"
 #include "tiledb/common/status.h"
 #include "tiledb/common/thread_pool.h"
 #include "tiledb/sm/array_schema/array_schema.h"
@@ -394,7 +395,8 @@ class ArrayDirectory {
    */
   std::vector<shared_ptr<const Enumeration>> load_enumerations_from_paths(
       const std::vector<std::string>& enumeration_paths,
-      const EncryptionKey& encryption_key) const;
+      const EncryptionKey& encryption_key,
+      MemoryTracker& memory_tracker) const;
 
   /** Returns the array URI. */
   const URI& uri() const;
@@ -826,7 +828,8 @@ class ArrayDirectory {
    */
   shared_ptr<const Enumeration> load_enumeration(
       const std::string& enumeration_path,
-      const EncryptionKey& encryption_key) const;
+      const EncryptionKey& encryption_key,
+      MemoryTracker& memory_tracker) const;
 };
 
 }  // namespace sm