Add memory tracking for Domain. (#4758)

This adds a MemoryTracker to Domain and updates class members to use `tdb::pmr` containers. --- TYPE: NO_HISTORY DESC: Add memory tracking for Domain
TileDB-Inc · Mar 5, 2024 · faa42b9 · faa42b9
1 parent 9bd16ba
commit faa42b9
Show file tree

Hide file tree

Showing 35 changed files with 333 additions and 265 deletions.
diff --git a/test/src/unit-Reader.cc b/test/src/unit-Reader.cc
@@ -248,7 +248,7 @@ TEST_CASE_METHOD(
   auto d2{make_shared<Dimension>(HERE(), "d2", Datatype::INT32, tracker_)};
   CHECK(d2->set_domain(&domain_vec[2]).ok());
   CHECK(d2->set_tile_extent(&tile_extents_vec[1]).ok());
-  auto dom{make_shared<Domain>(HERE())};
+  auto dom{make_shared<Domain>(HERE(), tracker_)};
   CHECK(dom->add_dimension(d1).ok());
   CHECK(dom->add_dimension(d2).ok());
 

diff --git a/test/src/unit-capi-array.cc b/test/src/unit-capi-array.cc
@@ -79,6 +79,9 @@ using namespace tiledb::common;
 using namespace tiledb::sm;
 
 struct ArrayFx {
+  // The memory tracker
+  shared_ptr<tiledb::sm::MemoryTracker> memory_tracker_;
+
   // TileDB context
   tiledb_ctx_t* ctx_;
   tiledb_vfs_t* vfs_;
@@ -116,7 +119,8 @@ static const std::string test_ca_file =
     std::string(TILEDB_TEST_INPUTS_DIR) + "/test_certs/public.crt";
 
 ArrayFx::ArrayFx()
-    : fs_vec_(vfs_test_get_fs_vec()) {
+    : memory_tracker_(tiledb::test::create_test_memory_tracker())
+    , fs_vec_(vfs_test_get_fs_vec()) {
   // Initialize vfs test
   REQUIRE(vfs_test_init(fs_vec_, &ctx_, &vfs_).ok());
 }
@@ -2586,13 +2590,12 @@ TEST_CASE_METHOD(
       1,
       &buff);
   REQUIRE(rc == TILEDB_OK);
-  auto memory_tracker = ctx_->resources().create_memory_tracker();
   auto st = tiledb::sm::serialization::array_deserialize(
       array->array_.get(),
       tiledb::sm::SerializationType::CAPNP,
       buff->buffer(),
       ctx_->storage_manager(),
-      memory_tracker);
+      memory_tracker_);
   REQUIRE(st.ok());
 
   // 6. Server: Close array and clean up

diff --git a/test/src/unit-cppapi-schema-evolution.cc b/test/src/unit-cppapi-schema-evolution.cc
@@ -821,7 +821,8 @@ TEST_CASE(
   int range[2] = {0, 1000};
   throw_if_not_ok(dim->set_domain(range));
 
-  auto dom = make_shared<tiledb::sm::Domain>(HERE());
+  auto dom = make_shared<tiledb::sm::Domain>(
+      HERE(), tiledb::test::get_test_memory_tracker());
   throw_if_not_ok(dom->add_dimension(dim));
   throw_if_not_ok(schema->set_domain(dom));
 

diff --git a/test/src/unit-enumerations.cc b/test/src/unit-enumerations.cc
@@ -141,11 +141,11 @@ struct EnumerationFx {
 
   void rm_array();
 
+  shared_ptr<MemoryTracker> memory_tracker_;
   URI uri_;
   Config cfg_;
   Context ctx_;
   EncryptionKey enc_key_;
-  shared_ptr<MemoryTracker> memory_tracker_;
 };
 
 template <typename T>
@@ -2554,9 +2554,9 @@ struct TypeParams {
 };
 
 EnumerationFx::EnumerationFx()
-    : uri_("enumeration_test_array")
-    , ctx_(cfg_)
-    , memory_tracker_(tiledb::test::create_test_memory_tracker()) {
+    : memory_tracker_(tiledb::test::create_test_memory_tracker())
+    , uri_("enumeration_test_array")
+    , ctx_(cfg_) {
   rm_array();
   throw_if_not_ok(enc_key_.set_key(EncryptionType::NO_ENCRYPTION, nullptr, 0));
   memory_tracker_ = tiledb::test::create_test_memory_tracker();
@@ -2850,7 +2850,7 @@ shared_ptr<ArraySchema> EnumerationFx::create_schema() {
   int range[2] = {0, 1000};
   throw_if_not_ok(dim->set_domain(range));
 
-  auto dom = make_shared<Domain>(HERE());
+  auto dom = make_shared<Domain>(HERE(), memory_tracker_);
   throw_if_not_ok(dom->add_dimension(dim));
   throw_if_not_ok(schema->set_domain(dom));
 

diff --git a/test/src/unit-request-handlers.cc b/test/src/unit-request-handlers.cc
@@ -62,6 +62,7 @@ struct RequestHandlerFx {
 
   shared_ptr<Array> get_array(QueryType type);
 
+  shared_ptr<MemoryTracker> memory_tracker_;
   URI uri_;
   Config cfg_;
   Context ctx_;
@@ -96,16 +97,13 @@ struct HandleConsolidationPlanRequestFx : RequestHandlerFx {
   }
 
   virtual shared_ptr<ArraySchema> create_schema() override {
-    auto schema = make_shared<ArraySchema>(
-        HERE(), ArrayType::SPARSE, tiledb::test::create_test_memory_tracker());
+    auto schema =
+        make_shared<ArraySchema>(HERE(), ArrayType::SPARSE, memory_tracker_);
     auto dim = make_shared<Dimension>(
-        HERE(),
-        "dim1",
-        Datatype::INT32,
-        tiledb::test::get_test_memory_tracker());
+        HERE(), "dim1", Datatype::INT32, memory_tracker_);
     int range[2] = {0, 1000};
     throw_if_not_ok(dim->set_domain(range));
-    auto dom = make_shared<Domain>(HERE());
+    auto dom = make_shared<Domain>(HERE(), memory_tracker_);
     throw_if_not_ok(dom->add_dimension(dim));
     throw_if_not_ok(schema->set_domain(dom));
     return schema;
@@ -340,7 +338,8 @@ TEST_CASE_METHOD(
 /* ********************************* */
 
 RequestHandlerFx::RequestHandlerFx(const std::string uri)
-    : uri_(uri)
+    : memory_tracker_(tiledb::test::create_test_memory_tracker())
+    , uri_(uri)
     , ctx_(cfg_) {
   delete_array();
   throw_if_not_ok(enc_key_.set_key(EncryptionType::NO_ENCRYPTION, nullptr, 0));
@@ -402,14 +401,14 @@ HandleLoadArraySchemaRequestFx::create_string_enumeration(
 
 shared_ptr<ArraySchema> HandleLoadArraySchemaRequestFx::create_schema() {
   // Create a schema to serialize
-  auto schema = make_shared<ArraySchema>(
-      HERE(), ArrayType::SPARSE, tiledb::test::create_test_memory_tracker());
-  auto dim = make_shared<Dimension>(
-      HERE(), "dim1", Datatype::INT32, tiledb::test::get_test_memory_tracker());
+  auto schema =
+      make_shared<ArraySchema>(HERE(), ArrayType::SPARSE, memory_tracker_);
+  auto dim =
+      make_shared<Dimension>(HERE(), "dim1", Datatype::INT32, memory_tracker_);
   int range[2] = {0, 1000};
   throw_if_not_ok(dim->set_domain(range));
 
-  auto dom = make_shared<Domain>(HERE());
+  auto dom = make_shared<Domain>(HERE(), memory_tracker_);
   throw_if_not_ok(dom->add_dimension(dim));
   throw_if_not_ok(schema->set_domain(dom));
 
@@ -445,15 +444,13 @@ shared_ptr<ArraySchema> HandleLoadArraySchemaRequestFx::call_handler(
       resp_buf);
   REQUIRE(rval == TILEDB_OK);
 
-  auto memory_tracker =
-      ctx.ptr()->context().resources().create_memory_tracker();
   return serialization::deserialize_load_array_schema_response(
-      stype, resp_buf->buffer(), memory_tracker);
+      stype, resp_buf->buffer(), memory_tracker_);
 }
 
 shared_ptr<ArraySchema> HandleQueryPlanRequestFx::create_schema() {
-  auto schema = make_shared<ArraySchema>(
-      HERE(), ArrayType::DENSE, tiledb::test::create_test_memory_tracker());
+  auto schema =
+      make_shared<ArraySchema>(HERE(), ArrayType::DENSE, memory_tracker_);
   schema->set_capacity(10000);
   throw_if_not_ok(schema->set_cell_order(Layout::ROW_MAJOR));
   throw_if_not_ok(schema->set_tile_order(Layout::ROW_MAJOR));
@@ -466,7 +463,7 @@ shared_ptr<ArraySchema> HandleQueryPlanRequestFx::create_schema() {
       HERE(), "dim2", Datatype::INT32, tiledb::test::get_test_memory_tracker());
   throw_if_not_ok(dim2->set_domain(&dim_domain[2]));
 
-  auto dom = make_shared<Domain>(HERE());
+  auto dom = make_shared<Domain>(HERE(), memory_tracker_);
   throw_if_not_ok(dom->add_dimension(dim1));
   throw_if_not_ok(dom->add_dimension(dim2));
   throw_if_not_ok(schema->set_domain(dom));

diff --git a/tiledb/api/c_api/domain/domain_api.cc b/tiledb/api/c_api/domain/domain_api.cc
@@ -30,12 +30,16 @@
 #include "../dimension/dimension_api_internal.h"
 #include "domain_api_external.h"
 #include "domain_api_internal.h"
+#include "tiledb/common/memory_tracker.h"
 
 namespace tiledb::api {
 
-int32_t tiledb_domain_alloc(tiledb_domain_handle_t** domain) {
+int32_t tiledb_domain_alloc(
+    tiledb_ctx_t* ctx, tiledb_domain_handle_t** domain) {
   ensure_output_pointer_is_valid(domain);
-  *domain = tiledb_domain_handle_t::make_handle();
+  auto memory_tracker = ctx->resources().create_memory_tracker();
+  memory_tracker->set_type(tiledb::sm::MemoryTrackerType::ARRAY_CREATE);
+  *domain = tiledb_domain_handle_t::make_handle(memory_tracker);
   return TILEDB_OK;
 }
 
@@ -147,7 +151,8 @@ int32_t tiledb_domain_dump(const tiledb_domain_t* domain, FILE* out) {
 using tiledb::api::api_entry_context;
 
 CAPI_INTERFACE(domain_alloc, tiledb_ctx_t* ctx, tiledb_domain_t** domain) {
-  return api_entry_context<tiledb::api::tiledb_domain_alloc>(ctx, domain);
+  return tiledb::api::api_entry_with_context<tiledb::api::tiledb_domain_alloc>(
+      ctx, domain);
 }
 
 CAPI_INTERFACE_VOID(domain_free, tiledb_domain_t** domain) {

diff --git a/tiledb/api/c_api/domain/domain_api_internal.h b/tiledb/api/c_api/domain/domain_api_internal.h
@@ -35,6 +35,10 @@
 #include "tiledb/api/c_api_support/handle/handle.h"
 #include "tiledb/sm/array_schema/domain.h"
 
+namespace tiledb::sm {
+class MemoryTracker;
+}
+
 struct tiledb_domain_handle_t
     : public tiledb::api::CAPIHandle<tiledb_domain_handle_t> {
  private:
@@ -54,8 +58,9 @@ struct tiledb_domain_handle_t
    * `class Domain` is principally a container for `Dimension` objects. Domain
    * handles are first constructed as empty containers.
    */
-  tiledb_domain_handle_t()
-      : domain_{make_shared<tiledb::sm::Domain>(HERE())} {
+  explicit tiledb_domain_handle_t(
+      shared_ptr<tiledb::sm::MemoryTracker> memory_tracker)
+      : domain_{make_shared<tiledb::sm::Domain>(HERE(), memory_tracker)} {
   }
 
   /**

diff --git a/tiledb/api/c_api/domain/test/compile_capi_domain_stub_main.cc b/tiledb/api/c_api/domain/test/compile_capi_domain_stub_main.cc
@@ -29,7 +29,6 @@
 #include "../domain_api_internal.h"
 
 int main() {
-  // Domain is easy to deal with because it has a default constructor.
-  tiledb_domain_handle_t x{};
+  tiledb_domain_handle_t x{shared_ptr<tiledb::sm::Domain>()};
   return 0;
 }
diff --git a/tiledb/common/memory_tracker.cc b/tiledb/common/memory_tracker.cc
@@ -76,6 +76,8 @@ std::string memory_type_to_str(MemoryType type) {
       return "DimensionLabels";
     case MemoryType::DIMENSIONS:
       return "Dimensions";
+    case MemoryType::DOMAINS:
+      return "Domains";
     case MemoryType::TILE_SUMS:
       return "TileSums";
     case MemoryType::TILE_WRITER_DATA:

diff --git a/tiledb/common/memory_tracker.h b/tiledb/common/memory_tracker.h
@@ -115,6 +115,7 @@ enum class MemoryType {
   DIMENSION_LABELS,
   DIMENSIONS,
   METADATA,
+  DOMAINS,
   TILE_SUMS,
   TILE_WRITER_DATA
 };

diff --git a/tiledb/sm/array/test/unit_consistency.h b/tiledb/sm/array/test/unit_consistency.h
@@ -40,6 +40,7 @@
 
 #include "../array.h"
 #include "../consistency.h"
+#include "test/support/src/mem_helpers.h"
 #include "tiledb/sm/array_schema/dimension.h"
 #include "tiledb/sm/enums/array_type.h"
 #include "tiledb/sm/enums/encryption_type.h"
@@ -108,8 +109,8 @@ class WhiteboxConsistencyController : public ConsistencyController {
     throw_if_not_ok(dim->set_tile_extent(&tile_extent));
 
     std::vector<shared_ptr<Dimension>> dims = {dim};
-    shared_ptr<Domain> domain =
-        make_shared<Domain>(HERE(), Layout::ROW_MAJOR, dims, Layout::ROW_MAJOR);
+    shared_ptr<Domain> domain = make_shared<Domain>(
+        HERE(), Layout::ROW_MAJOR, dims, Layout::ROW_MAJOR, memory_tracker_);
 
     // Create the ArraySchema
     shared_ptr<ArraySchema> schema = make_shared<ArraySchema>(

diff --git a/tiledb/sm/array_schema/array_schema.h b/tiledb/sm/array_schema/array_schema.h
@@ -476,6 +476,7 @@ class ArraySchema {
    *
    * @param deserializer The deserializer to deserialize from.
    * @param uri The uri of the Array.
+   * @param memory_tracker The memory tracker to use.
    * @return A new ArraySchema.
    */
   static shared_ptr<ArraySchema> deserialize(

diff --git a/tiledb/sm/array_schema/dimension_label.cc b/tiledb/sm/array_schema/dimension_label.cc
@@ -28,6 +28,7 @@
 
 #include "tiledb/sm/array_schema/dimension_label.h"
 #include "tiledb/common/common.h"
+#include "tiledb/common/memory_tracker.h"
 #include "tiledb/sm/array_schema/array_schema.h"
 #include "tiledb/sm/array_schema/dimension.h"
 #include "tiledb/sm/array_schema/domain.h"
@@ -176,7 +177,11 @@ DimensionLabel::DimensionLabel(
   throw_if_not_ok(
       index_dims.back()->set_tile_extent(dim->tile_extent().data()));
   throw_if_not_ok(schema_->set_domain(make_shared<Domain>(
-      HERE(), Layout::ROW_MAJOR, index_dims, Layout::ROW_MAJOR)));
+      HERE(),
+      Layout::ROW_MAJOR,
+      index_dims,
+      Layout::ROW_MAJOR,
+      memory_tracker)));
 
   // Create and set dimension label attribute.
   auto label_attr = make_shared<Attribute>(