implement of the engine and rewrite factory

- allocator shared_ptr hold in Index and Dataset
- the allocator ptr always point to SafeAllocator
- allocator ptr is used inner index
- only Engine/Resource can create Allocator Object

Signed-off-by: LHT129 <[email protected]>

examples/cpp/CMakeLists.txt

@@ -7,6 +7,9 @@ target_link_libraries (example_fresh_hnsw vsag)
 add_executable (example_diskann example_diskann.cpp)
 target_link_libraries (example_diskann vsag_static)
 
+add_executable (example_engine example_engine.cpp)
+target_link_libraries (example_engine vsag_static)
+
 add_executable (simple_hnsw simple_hnsw.cpp)
 target_link_libraries (simple_hnsw vsag_static)
 

examples/cpp/example_engine.cpp

@@ -0,0 +1,134 @@
+
+// Copyright 2024-present the vsag project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <vsag/vsag.h>
+
+#include <iostream>
+
+class SimpleAllocator : public vsag::Allocator {
+public:
+    SimpleAllocator() = default;
+    ~SimpleAllocator() override = default;
+
+public:
+    std::string
+    Name() override {
+        return "DefaultAllocator";
+    }
+
+    void*
+    Allocate(size_t size) override {
+        auto ptr = malloc(size);
+        allocate_count_++;
+        return ptr;
+    }
+
+    void
+    Deallocate(void* p) override {
+        free(p);
+    }
+
+    void*
+    Reallocate(void* p, size_t size) override {
+        auto ptr = realloc(p, size);
+        return ptr;
+    }
+
+    void
+    Print() {
+        std::cout << "allocate times : " << allocate_count_ << std::endl;
+    }
+    uint64_t allocate_count_{0};
+};
+
+int
+main(int argc, char** argv) {
+    vsag::init();
+
+    int64_t num_vectors = 1000;
+    int64_t dim = 128;
+
+    // prepare ids and vectors
+    auto ids = new int64_t[num_vectors];
+    auto vectors = new float[dim * num_vectors];
+
+    std::mt19937 rng(47);
+    std::uniform_real_distribution<float> distrib_real;
+    for (int64_t i = 0; i < num_vectors; ++i) {
+        ids[i] = i;
+    }
+    for (int64_t i = 0; i < dim * num_vectors; ++i) {
+        vectors[i] = distrib_real(rng);
+    }
+
+    // create index
+    auto hnsw_build_paramesters = R"(
+    {
+        "dtype": "float32",
+        "metric_type": "l2",
+        "dim": 128,
+        "hnsw": {
+            "max_degree": 16,
+            "ef_construction": 100
+        }
+    }
+    )";
+
+    auto* allocator = new SimpleAllocator();
+
+    {
+        vsag::Resource resource(allocator);
+        vsag::Engine engine(&resource);
+        auto index = engine.CreateIndex("hnsw", hnsw_build_paramesters).value();
+        auto base = vsag::Dataset::Make();
+        base->NumElements(num_vectors)->Dim(dim)->Ids(ids)->Float32Vectors(vectors)->Owner(false);
+        index->Build(base);
+
+        // prepare a query vector
+        // memory will be released by query the dataset since owner is set to true when creating the query.
+        auto query_vector = new float[dim];
+        for (int64_t i = 0; i < dim; ++i) {
+            query_vector[i] = distrib_real(rng);
+        }
+
+        // search on the index
+        auto hnsw_search_parameters = R"(
+        {
+            "hnsw": {
+                "ef_search": 100
+            }
+        }
+        )";
+        int64_t topk = 10;
+        auto query = vsag::Dataset::Make();
+        query->NumElements(1)->Dim(dim)->Float32Vectors(query_vector)->Owner(true);
+        auto result = index->KnnSearch(query, topk, hnsw_search_parameters).value();
+
+        // print the results
+        std::cout << "results: " << std::endl;
+        for (int64_t i = 0; i < result->GetDim(); ++i) {
+            std::cout << result->GetIds()[i] << ": " << result->GetDistances()[i] << std::endl;
+        }
+        engine.Shutdown();
+        allocator->Print();
+    }
+
+    // free memory
+    delete[] ids;
+    delete[] vectors;
+    delete allocator;
+
+    return 0;
+}

include/vsag/engine.h

@@ -17,16 +17,19 @@
 
 #include <memory>
 
+#include "dataset.h"
 #include "index.h"
 #include "resource.h"
 
 namespace vsag {
 class Engine {
 public:
-    explicit Engine(Resource* resource = nullptr);
+    explicit Engine();
+
+    explicit Engine(Resource* resource);
 
     void
-    Shutdown();
+    Shutdown();  // like ~Engine(), but will warn whether the resources are still reference outside
 
     tl::expected<std::shared_ptr<Index>, Error>
     CreateIndex(const std::string& name, const std::string& parameters);

include/vsag/resource.h

@@ -20,15 +20,19 @@
 #include "allocator.h"
 
 namespace vsag {
-struct Resource {
+class Resource {
 public:
     explicit Resource(Allocator* allocator = nullptr);
 
     virtual ~Resource() = default;
 
+    virtual std::shared_ptr<Allocator>
+    GetAllocator() {
+        return this->allocator;
+    }
+
 public:
     std::shared_ptr<Allocator> allocator;
 
-    // std::shared_ptr<> thread_pool;
 };
 }  // namespace vsag

include/vsag/vsag.h

@@ -41,6 +41,7 @@ init();
 #include "bitset.h"
 #include "constants.h"
 #include "dataset.h"
+#include "engine.h"
 #include "errors.h"
 #include "expected.hpp"
 #include "factory.h"

src/algorithm/hgraph.cpp

@@ -53,14 +53,14 @@ next_multiple_of_power_of_two(uint64_t x, uint64_t n) {
 HGraph::HGraph(const JsonType& index_param, const vsag::IndexCommonParam& common_param) noexcept
     : index_param_(index_param),
       common_param_(common_param),
-      label_lookup_(common_param.allocator_),
-      label_op_mutex_(MAX_LABEL_OPERATION_LOCKS, common_param_.allocator_),
-      neighbors_mutex_(0, common_param_.allocator_),
-      route_graphs_(common_param.allocator_),
-      labels_(common_param.allocator_) {
+      label_lookup_(common_param.allocator_.get()),
+      label_op_mutex_(MAX_LABEL_OPERATION_LOCKS, common_param.allocator_.get()),
+      neighbors_mutex_(0, common_param.allocator_.get()),
+      route_graphs_(common_param.allocator_.get()),
+      labels_(common_param.allocator_.get()) {
     this->dim_ = common_param.dim_;
     this->metric_ = common_param.metric_;
-    this->allocator_ = common_param.allocator_;
+    this->allocator_ = common_param.allocator_.get();
 }
 
 tl::expected<void, Error>

src/data_cell/flatten_datacell.h

@@ -130,7 +130,7 @@ template <typename QuantTmpl, typename IOTmpl>
 FlattenDataCell<QuantTmpl, IOTmpl>::FlattenDataCell(const JsonType& quantization_param,
                                                     const JsonType& io_param,
                                                     const IndexCommonParam& common_param)
-    : allocator_(common_param.allocator_) {
+    : allocator_(common_param.allocator_.get()) {
     this->quantizer_ = std::make_shared<QuantTmpl>(quantization_param, common_param);
     this->io_ = std::make_shared<IOTmpl>(io_param, common_param);
     this->code_size_ = quantizer_->GetCodeSize();

src/data_cell/flatten_datacell_test.cpp

@@ -16,13 +16,15 @@
 #include "flatten_datacell.h"
 
 #include <algorithm>
+#include <utility>
 
 #include "catch2/catch_template_test_macros.hpp"
 #include "default_allocator.h"
 #include "fixtures.h"
 #include "flatten_interface_test.h"
 #include "io/io_headers.h"
 #include "quantization/quantizer_headers.h"
+#include "safe_allocator.h"
 
 using namespace vsag;
 
@@ -36,7 +38,7 @@ TestFlattenDataCell(int dim,
     auto counts = {100, 1000};
     IndexCommonParam common;
     common.dim_ = dim;
-    common.allocator_ = allocator.get();
+    common.allocator_ = std::move(allocator);
     common.metric_ = metric;
     for (auto count : counts) {
         auto flatten =
@@ -67,7 +69,7 @@ TestFlattenDataCellFP32(int dim,
 }
 
 TEST_CASE("fp32", "[ut][flatten_data_cell]") {
-    auto allocator = std::make_shared<DefaultAllocator>();
+    auto allocator = SafeAllocator::FactoryDefaultAllocator();
     auto fp32_param = JsonType::parse("{}");
     auto io_param = JsonType::parse("{}");
     auto dims = {8, 64, 512};
@@ -96,7 +98,7 @@ TestFlattenDataCellSQ8(int dim,
 }
 
 TEST_CASE("sq8", "[ut][flatten_data_cell]") {
-    auto allocator = std::make_shared<DefaultAllocator>();
+    auto allocator = SafeAllocator::FactoryDefaultAllocator();
     auto sq8_param = JsonType::parse("{}");
     auto io_param = JsonType::parse("{}");
     auto dims = {32, 64, 512};

src/data_cell/graph_datacell_test.cpp

@@ -19,6 +19,7 @@
 #include "fmt/format-inl.h"
 #include "graph_interface_test.h"
 #include "io/io_headers.h"
+#include "safe_allocator.h"
 using namespace vsag;
 
 template <typename IOTemp>
@@ -37,7 +38,7 @@ TestGraphDataCell(const JsonType& graph_param,
 }
 
 TEST_CASE("graph basic test", "[ut][graph_datacell]") {
-    auto allocator = std::make_shared<DefaultAllocator>();
+    auto allocator = SafeAllocator::FactoryDefaultAllocator();
     auto dims = {32, 64};
     auto max_degrees = {5, 12, 24, 32, 64, 128};
     auto max_capacities = {1, 100, 10000, 10'000'000, 32'179'837};
@@ -58,7 +59,7 @@ TEST_CASE("graph basic test", "[ut][graph_datacell]") {
     for (auto dim : dims) {
         IndexCommonParam param;
         param.dim_ = dim;
-        param.allocator_ = allocator.get();
+        param.allocator_ = allocator;
         for (auto& gp : graph_params) {
             TestGraphDataCell<MemoryIO>(gp, io_param, param);
             TestGraphDataCell<MemoryBlockIO>(gp, io_param, param);

src/data_cell/graph_interface_test.cpp

@@ -20,12 +20,13 @@
 #include "catch2/catch_test_macros.hpp"
 #include "default_allocator.h"
 #include "fixtures.h"
+#include "safe_allocator.h"
 
 using namespace vsag;
 
 void
 GraphInterfaceTest::BasicTest(uint64_t max_id, uint64_t count, GraphInterfacePtr other) {
-    auto allocator = std::make_shared<DefaultAllocator>();
+    auto allocator = SafeAllocator::FactoryDefaultAllocator();
     auto max_degree = this->graph_->MaximumDegree();
     auto old_count = this->graph_->TotalCount();
     UnorderedMap<InnerIdType, std::shared_ptr<Vector<InnerIdType>>> maps(allocator.get());

src/data_cell/sparse_graph_datacell.cpp

@@ -19,7 +19,7 @@ namespace vsag {
 
 SparseGraphDataCell::SparseGraphDataCell(const JsonType& graph_param,
                                          const IndexCommonParam& common_param)
-    : allocator_(common_param.allocator_), neighbors_(common_param.allocator_) {
+    : allocator_(common_param.allocator_.get()), neighbors_(common_param.allocator_.get()) {
     if (graph_param.contains(GRAPH_PARAM_MAX_DEGREE)) {
         this->maximum_degree_ = graph_param[GRAPH_PARAM_MAX_DEGREE];
     }

src/data_cell/sparse_graph_datacell_test.cpp

@@ -19,6 +19,7 @@
 #include "default_allocator.h"
 #include "fmt/format-inl.h"
 #include "graph_interface_test.h"
+#include "safe_allocator.h"
 
 using namespace vsag;
 
@@ -35,7 +36,7 @@ TestSparseGraphDataCell(const JsonType& graph_param, const IndexCommonParam& par
 }
 
 TEST_CASE("graph basic test", "[ut][sparse_graph_datacell]") {
-    auto allocator = std::make_shared<DefaultAllocator>();
+    auto allocator = SafeAllocator::FactoryDefaultAllocator();
     auto dims = {32, 64};
     auto max_degrees = {5, 12, 24, 32, 64, 128};
     auto max_capacities = {1, 100, 10000, 10'000'000, 32'179'837};
@@ -55,7 +56,7 @@ TEST_CASE("graph basic test", "[ut][sparse_graph_datacell]") {
     for (auto dim : dims) {
         IndexCommonParam param;
         param.dim_ = dim;
-        param.allocator_ = allocator.get();
+        param.allocator_ = allocator;
         for (auto& gp : graph_params) {
             TestSparseGraphDataCell(gp, param);
         }

src/default_allocator.h

@@ -25,16 +25,6 @@
 namespace vsag {
 
 class DefaultAllocator : public Allocator {
-public:
-    static std::shared_ptr<Allocator>
-    Instance() {
-        static std::shared_ptr<Allocator> s_instance;
-        if (s_instance == nullptr) {
-            s_instance = std::make_shared<DefaultAllocator>();
-        }
-        return s_instance;
-    }
-
 public:
     DefaultAllocator() = default;
     ~DefaultAllocator() override {