milvus/core/unittest/db/test_meta.cpp

306 lines
14 KiB
C++

// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// 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 "db/meta/MetaNames.h"
#include "db/meta/backend/MetaContext.h"
#include "db/snapshot/ResourceContext.h"
#include "db/utils.h"
#include "utils/Json.h"
template<typename T>
using ResourceContext = milvus::engine::snapshot::ResourceContext<T>;
template<typename T>
using ResourceContextBuilder = milvus::engine::snapshot::ResourceContextBuilder<T>;
using FType = milvus::engine::DataType;
using FEType = milvus::engine::FieldElementType;
using Op = milvus::engine::meta::MetaContextOp;
using State = milvus::engine::snapshot::State;
TEST_F(MetaTest, ApplyTest) {
ID_TYPE result_id;
auto collection = std::make_shared<Collection>("meta_test_c1");
auto c_ctx = ResourceContextBuilder<Collection>().SetResource(collection).CreatePtr();
auto status = meta_->Execute<Collection>(c_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
collection->SetID(result_id);
collection->Activate();
auto c2_ctx = ResourceContextBuilder<Collection>().SetResource(collection)
.SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
status = meta_->Execute<Collection>(c2_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
ASSERT_EQ(result_id, collection->GetID());
auto c3_ctx = ResourceContextBuilder<Collection>().SetID(result_id).SetOp(Op::oDelete).CreatePtr();
status = meta_->Execute<Collection>(c3_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
ASSERT_EQ(result_id, collection->GetID());
}
TEST_F(MetaTest, SessionTest) {
ID_TYPE result_id;
auto collection = std::make_shared<Collection>("meta_test_c1");
auto c_ctx = ResourceContextBuilder<Collection>().SetResource(collection).CreatePtr();
auto status = meta_->Execute<Collection>(c_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
collection->SetID(result_id);
auto partition = std::make_shared<Partition>("meta_test_p1", result_id);
auto p_ctx = ResourceContextBuilder<Partition>().SetResource(partition).CreatePtr();
status = meta_->Execute<Partition>(p_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
partition->SetID(result_id);
auto field = std::make_shared<Field>("meta_test_f1", 1, FType::INT64);
auto f_ctx = ResourceContextBuilder<Field>().SetResource(field).CreatePtr();
status = meta_->Execute<Field>(f_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
field->SetID(result_id);
auto field_element = std::make_shared<FieldElement>(collection->GetID(), field->GetID(),
"meta_test_f1_fe1", FEType::FET_RAW);
auto fe_ctx = ResourceContextBuilder<FieldElement>().SetResource(field_element).CreatePtr();
status = meta_->Execute<FieldElement>(fe_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
field_element->SetID(result_id);
auto session = meta_->CreateSession();
ASSERT_TRUE(collection->Activate());
auto c2_ctx = ResourceContextBuilder<Collection>().SetResource(collection)
.SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
status = session->Apply<Collection>(c2_ctx);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_TRUE(partition->Activate());
auto p2_ctx = ResourceContextBuilder<Partition>().SetResource(partition)
.SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
status = session->Apply<Partition>(p2_ctx);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_TRUE(field->Activate());
auto f2_ctx = ResourceContextBuilder<Field>().SetResource(field)
.SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
status = session->Apply<Field>(f2_ctx);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_TRUE(field_element->Activate());
auto fe2_ctx = ResourceContextBuilder<FieldElement>().SetResource(field_element)
.SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
status = session->Apply<FieldElement>(fe2_ctx);
ASSERT_TRUE(status.ok()) << status.ToString();
std::vector<ID_TYPE> result_ids;
status = session->Commit(result_ids);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_EQ(result_ids.size(), 4);
ASSERT_EQ(result_ids.at(0), collection->GetID());
ASSERT_EQ(result_ids.at(1), partition->GetID());
ASSERT_EQ(result_ids.at(2), field->GetID());
ASSERT_EQ(result_ids.at(3), field_element->GetID());
}
TEST_F(MetaTest, SelectTest) {
ID_TYPE result_id;
auto collection = std::make_shared<Collection>("meta_test_c1");
ASSERT_TRUE(collection->Activate());
auto c_ctx = ResourceContextBuilder<Collection>().SetResource(collection).CreatePtr();
auto status = meta_->Execute<Collection>(c_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
collection->SetID(result_id);
Collection::Ptr return_collection;
status = meta_->Select<Collection>(collection->GetID(), return_collection);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_EQ(collection->GetID(), return_collection->GetID());
ASSERT_EQ(collection->GetName(), return_collection->GetName());
auto collection2 = std::make_shared<Collection>("meta_test_c2");
ASSERT_TRUE(collection2->Activate());
auto c2_ctx = ResourceContextBuilder<Collection>().SetResource(collection2).CreatePtr();
status = meta_->Execute<Collection>(c2_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
collection2->SetID(result_id);
ASSERT_GT(collection2->GetID(), collection->GetID());
std::vector<Collection::Ptr> return_collections;
status = meta_->SelectBy<Collection, ID_TYPE>(milvus::engine::meta::F_ID,
{collection2->GetID()}, return_collections);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_EQ(return_collections.size(), 1);
ASSERT_EQ(return_collections.at(0)->GetID(), collection2->GetID());
ASSERT_EQ(return_collections.at(0)->GetName(), collection2->GetName());
return_collections.clear();
status = meta_->SelectBy<Collection, State>(milvus::engine::meta::F_STATE, {State::ACTIVE}, return_collections);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_EQ(return_collections.size(), 2);
std::vector<ID_TYPE> ids;
status = meta_->SelectResourceIDs<Collection, std::string>(ids, "", {""});
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_EQ(ids.size(), 2);
ids.clear();
status = meta_->SelectResourceIDs<Collection, std::string>(ids, milvus::engine::meta::F_NAME,
{collection->GetName()});
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_EQ(ids.size(), 1);
ASSERT_EQ(ids.at(0), collection->GetID());
}
TEST_F(MetaTest, TruncateTest) {
ID_TYPE result_id;
auto collection = std::make_shared<Collection>("meta_test_c1");
ASSERT_TRUE(collection->Activate());
auto c_ctx = ResourceContextBuilder<Collection>().SetResource(collection).CreatePtr();
auto status = meta_->Execute<Collection>(c_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
collection->SetID(result_id);
status = meta_->TruncateAll();
ASSERT_TRUE(status.ok()) << status.ToString();
Collection::Ptr return_collection;
status = meta_->Select<Collection>(collection->GetID(), return_collection);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_EQ(return_collection, nullptr);
}
TEST_F(MetaTest, MultiThreadRequestTest) {
auto request_worker = [&](size_t i) {
std::string collection_name_prefix = "meta_test_collection_" + std::to_string(i) + "_";
int64_t result_id;
for (size_t ii = 0; ii < 30; ii++) {
std::string collection_name = collection_name_prefix + std::to_string(ii);
auto collection = std::make_shared<Collection>(collection_name);
auto c_ctx = ResourceContextBuilder<Collection>().SetResource(collection).CreatePtr();
auto status = meta_->Execute<Collection>(c_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
collection->SetID(result_id);
collection->Activate();
auto c_ctx2 = ResourceContextBuilder<Collection>().SetResource(collection)
.SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
status = meta_->Execute<Collection>(c_ctx2, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
CollectionPtr collection2;
status = meta_->Select<Collection>(result_id, collection2);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_EQ(collection2->GetID(), result_id);
ASSERT_EQ(collection2->GetState(), State::ACTIVE);
ASSERT_EQ(collection2->GetName(), collection_name);
collection->Deactivate();
auto c_ctx3 = ResourceContextBuilder<Collection>().SetResource(collection)
.SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
status = meta_->Execute<Collection>(c_ctx3, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_EQ(result_id, collection->GetID());
auto c_ctx4 = ResourceContextBuilder<Collection>().SetID(result_id)
.SetOp(Op::oDelete).SetTable(Collection::Name).CreatePtr();
status = meta_->Execute<Collection>(c_ctx4, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
CollectionPtr collection3;
status = meta_->Select<Collection>(result_id, collection3);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_EQ(collection3, nullptr);
}
};
auto cc_task = [&](size_t j) {
std::string collection_name_prefix = "meta_test_collection_cc_" + std::to_string(j) + "_";
int64_t result_id;
Status status;
for (size_t jj = 0; jj < 20; jj ++) {
std::string collection_name = collection_name_prefix + std::to_string(jj);
milvus::json cj{{"segment_row_count", 1024}};
auto collection = std::make_shared<Collection>(collection_name, cj);
auto c_ctx = ResourceContextBuilder<Collection>().SetResource(collection).SetOp(Op::oAdd).CreatePtr();
status = meta_->Execute<Collection>(c_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
collection->SetID(result_id);
std::string partition_name = collection_name + "_p_" + std::to_string(jj);
auto partition = std::make_shared<Partition>(partition_name, collection->GetID());
auto p_ctx = ResourceContextBuilder<Partition>().SetResource(partition).SetOp(Op::oAdd).CreatePtr();
status = meta_->Execute<Partition>(p_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
partition->SetID(result_id);
std::string segment_name = partition_name + "_s_" + std::to_string(jj);
auto segment = std::make_shared<Segment>(collection->GetID(), partition->GetID());
auto s_ctx = ResourceContextBuilder<Segment>().SetResource(segment).SetOp(Op::oAdd).CreatePtr();
status = meta_->Execute<Segment>(s_ctx, result_id);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_GT(result_id, 0);
segment->SetID(result_id);
auto session = meta_->CreateSession();
collection->Activate();
auto c_ctx2 = ResourceContextBuilder<Collection>().SetResource(collection)
.SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
ASSERT_TRUE(session->Apply<Collection>(c_ctx2).ok());
partition->Activate();
auto p_ctx2 = ResourceContextBuilder<Partition>().SetResource(partition)
.SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
ASSERT_TRUE(session->Apply<Partition>(p_ctx2).ok());
segment->Activate();
auto s_ctx2 = ResourceContextBuilder<Segment>().SetResource(segment)
.SetOp(Op::oUpdate).AddAttr(milvus::engine::meta::F_STATE).CreatePtr();
ASSERT_TRUE(session->Apply<Segment>(s_ctx2).ok());
std::vector<int64_t> ids;
status = session->Commit(ids);
ASSERT_TRUE(status.ok()) << status.ToString();
}
};
unsigned int thread_hint = std::thread::hardware_concurrency();
std::vector<std::thread> request_threads;
for (size_t i = 0; i < 3 * thread_hint; i++) {
request_threads.emplace_back(request_worker, i);
}
std::vector<std::thread> cc_threads;
for (size_t j = 0; j < 3 * thread_hint; j++) {
cc_threads.emplace_back(cc_task, j);
}
for (auto& t : request_threads) {
t.join();
}
for (auto& t : cc_threads) {
t.join();
}
}