milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
milvus/core/unittest/db/test_snapshot.cpp

1973 lines
71 KiB
C++
Raw Blame History

// 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 <fiu-control.h>
#include <fiu-local.h>
#include <gtest/gtest.h>
#include <string>
#include <set>
#include <algorithm>
#include "db/utils.h"
#include "db/snapshot/HandlerFactory.h"
Status
GetFirstCollectionID(ID_TYPE& result_id) {
std::vector<ID_TYPE> ids;
auto status = Snapshots::GetInstance().GetCollectionIds(ids);
if (status.ok()) {
result_id = ids.at(0);
}
return status;
}
TEST_F(SnapshotTest, ResourcesTest) {
int nprobe = 16;
milvus::json params = {{"nprobe", nprobe}};
ParamsField p_field(params);
ASSERT_EQ(params, p_field.GetParams());
auto nprobe_real = p_field.GetParams().at("nprobe").get<int>();
ASSERT_EQ(nprobe, nprobe_real);
}
TEST_F(SnapshotTest, ReferenceProxyTest) {
std::string status("raw");
const std::string CALLED = "CALLED";
auto callback = [&]() {
status = CALLED;
};
auto proxy = ReferenceProxy();
ASSERT_EQ(proxy.ref_count(), 0);
int refcnt = 3;
for (auto i = 0; i < refcnt; ++i) {
proxy.Ref();
}
ASSERT_EQ(proxy.ref_count(), refcnt);
proxy.RegisterOnNoRefCB(callback);
for (auto i = 0; i < refcnt; ++i) {
proxy.UnRef();
}
ASSERT_EQ(proxy.ref_count(), 0);
ASSERT_EQ(status, CALLED);
}
TEST_F(SnapshotTest, ScopedResourceTest) {
auto inner = std::make_shared<Collection>("c1");
ASSERT_EQ(inner->ref_count(), 0);
{
auto not_scoped = CollectionScopedT(inner, false);
ASSERT_EQ(not_scoped->ref_count(), 0);
not_scoped->Ref();
ASSERT_EQ(not_scoped->ref_count(), 1);
ASSERT_EQ(inner->ref_count(), 1);
auto not_scoped_2 = not_scoped;
ASSERT_EQ(not_scoped_2->ref_count(), 1);
ASSERT_EQ(not_scoped->ref_count(), 1);
ASSERT_EQ(inner->ref_count(), 1);
not_scoped_2->Ref();
ASSERT_EQ(not_scoped_2->ref_count(), 2);
ASSERT_EQ(not_scoped->ref_count(), 2);
ASSERT_EQ(inner->ref_count(), 2);
}
inner->UnRef();
ASSERT_EQ(inner->ref_count(), 1);
inner->UnRef();
ASSERT_EQ(inner->ref_count(), 0);
{
// Test scoped construct
auto scoped = CollectionScopedT(inner);
ASSERT_EQ(scoped->ref_count(), 1);
ASSERT_EQ(inner->ref_count(), 1);
{
// Test bool operator
CollectionScopedT other_scoped;
ASSERT_EQ(other_scoped, false);
// Test operator=
other_scoped = scoped;
ASSERT_EQ(other_scoped->ref_count(), 2);
ASSERT_EQ(scoped->ref_count(), 2);
ASSERT_EQ(inner->ref_count(), 2);
}
ASSERT_EQ(scoped->ref_count(), 1);
ASSERT_EQ(inner->ref_count(), 1);
{
// Test copy
auto other_scoped(scoped);
ASSERT_EQ(other_scoped->ref_count(), 2);
ASSERT_EQ(scoped->ref_count(), 2);
ASSERT_EQ(inner->ref_count(), 2);
}
ASSERT_EQ(scoped->ref_count(), 1);
ASSERT_EQ(inner->ref_count(), 1);
}
ASSERT_EQ(inner->ref_count(), 0);
}
TEST_F(SnapshotTest, ResourceHoldersTest) {
ID_TYPE collection_id;
ASSERT_TRUE(GetFirstCollectionID(collection_id).ok());
auto collection = CollectionsHolder::GetInstance().GetResource(collection_id, false);
auto prev_cnt = collection->ref_count();
{
auto collection_2 = CollectionsHolder::GetInstance().GetResource(collection_id, false);
ASSERT_EQ(collection->GetID(), collection_id);
ASSERT_EQ(collection->ref_count(), prev_cnt);
}
{
auto collection_3 = CollectionsHolder::GetInstance().GetResource(collection_id, true);
ASSERT_EQ(collection_3->GetID(), collection_id);
ASSERT_EQ(collection_3->ref_count(), 1+prev_cnt);
}
std::this_thread::sleep_for(std::chrono::milliseconds(80));
if (prev_cnt == 0) {
auto collection_4 = CollectionsHolder::GetInstance().GetResource(collection_id, false);
ASSERT_TRUE(!collection_4);
}
}
TEST_F(SnapshotTest, DeleteOperationTest) {
std::string collection_name = "test_c1";
LSN_TYPE lsn = 1;
auto ss = CreateCollection(collection_name, lsn);
ASSERT_TRUE(ss);
auto collection = CollectionsHolder::GetInstance().GetResource(ss->GetCollectionId());
ASSERT_EQ(collection->GetName(), collection_name);
{
auto soft_op = std::make_shared<SoftDeleteCollectionOperation>(collection->GetID());
auto status = soft_op->Push();
ASSERT_TRUE(status.ok());
CollectionPtr soft_deleted;
status = soft_op->GetResource(soft_deleted);
ASSERT_TRUE(status.ok());
ASSERT_TRUE(soft_deleted);
ASSERT_EQ(soft_deleted->GetID(), collection->GetID());
ASSERT_TRUE(soft_deleted->IsDeactive());
}
{
CollectionPtr loaded;
LoadOperationContext context;
context.id = collection->GetID();
auto load_op = std::make_shared<milvus::engine::snapshot::LoadOperation<Collection>>(context);
auto status = load_op->Push();
ASSERT_TRUE(status.ok());
status = load_op->GetResource(loaded);
ASSERT_TRUE(status.ok());
ASSERT_TRUE(loaded);
ASSERT_EQ(loaded->GetID(), collection->GetID());
ASSERT_TRUE(loaded->IsDeactive());
}
{
auto hard_op = std::make_shared<milvus::engine::snapshot::HardDeleteOperation<Collection>
>(collection->GetID());
auto status = hard_op->Push();
ASSERT_TRUE(status.ok());
}
{
CollectionPtr loaded;
LoadOperationContext context;
context.id = collection->GetID();
auto load_op = std::make_shared<milvus::engine::snapshot::LoadOperation<Collection>>(context);
auto status = load_op->Push();
if (!status.ok()) {
std::cout << status.ToString() << std::endl;
}
ASSERT_FALSE(status.ok());
}
{
auto soft_op = std::make_shared<SoftDeleteCollectionOperation>(collection->GetID());
auto status = soft_op->Push();
if (!status.ok()) {
std::cout << status.ToString() << std::endl;
}
ASSERT_FALSE(status.ok());
}
}
TEST_F(SnapshotTest, CreateCollectionOperationTest) {
ScopedSnapshotT expect_null;
auto status = Snapshots::GetInstance().GetSnapshot(expect_null, 100000);
ASSERT_FALSE(expect_null);
std::string collection_name = "test_c1";
LSN_TYPE lsn = 1;
auto ss = CreateCollection(collection_name, lsn);
ASSERT_TRUE(ss);
ScopedSnapshotT latest_ss;
status = Snapshots::GetInstance().GetSnapshot(latest_ss, "xxxx");
ASSERT_FALSE(status.ok());
status = Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
ASSERT_TRUE(latest_ss);
ASSERT_TRUE(latest_ss->GetName() == collection_name);
IDS_TYPE ids;
status = Snapshots::GetInstance().GetCollectionIds(ids);
ASSERT_EQ(ids.size(), 6);
ASSERT_EQ(ids.back(), latest_ss->GetCollectionId());
OperationContext sd_op_ctx;
sd_op_ctx.collection = latest_ss->GetCollection();
sd_op_ctx.lsn = latest_ss->GetMaxLsn() + 1;
ASSERT_TRUE(sd_op_ctx.collection->IsActive());
auto sd_op = std::make_shared<DropCollectionOperation>(sd_op_ctx, latest_ss);
status = sd_op->Push();
ASSERT_TRUE(status.ok());
ASSERT_TRUE(sd_op->GetStatus().ok());
ASSERT_FALSE(sd_op_ctx.collection->IsActive());
ASSERT_FALSE(latest_ss->GetCollection()->IsActive());
Snapshots::GetInstance().Reset();
}
TEST_F(SnapshotTest, DropCollectionTest) {
std::string collection_name = "test_c1";
LSN_TYPE lsn = 1;
auto ss = CreateCollection(collection_name, lsn);
ASSERT_TRUE(ss);
ScopedSnapshotT lss;
auto status = Snapshots::GetInstance().GetSnapshot(lss, collection_name);
ASSERT_TRUE(status.ok());
ASSERT_TRUE(lss);
ASSERT_EQ(ss->GetID(), lss->GetID());
auto prev_ss_id = ss->GetID();
auto prev_c_id = ss->GetCollection()->GetID();
lsn = ss->GetMaxLsn() + 1;
status = Snapshots::GetInstance().DropCollection(collection_name, lsn);
ASSERT_TRUE(status.ok());
status = Snapshots::GetInstance().GetSnapshot(lss, collection_name);
ASSERT_FALSE(status.ok());
auto ss_2 = CreateCollection(collection_name, ++lsn);
status = Snapshots::GetInstance().GetSnapshot(lss, collection_name);
// EXPECT_DEATH({assert(1 == 2);}, "nullptr")
ASSERT_TRUE(status.ok());
ASSERT_EQ(ss_2->GetID(), lss->GetID());
ASSERT_NE(prev_ss_id, ss_2->GetID());
ASSERT_NE(prev_c_id, ss_2->GetCollection()->GetID());
status = Snapshots::GetInstance().DropCollection(collection_name, ++lsn);
ASSERT_TRUE(status.ok());
status = Snapshots::GetInstance().DropCollection(collection_name, ++lsn);
ASSERT_FALSE(status.ok());
}
TEST_F(SnapshotTest, ConCurrentCollectionOperation) {
std::string collection_name("c1");
LSN_TYPE lsn = 1;
ID_TYPE stale_ss_id;
auto worker1 = [&]() {
Status status;
auto ss = CreateCollection(collection_name, ++lsn);
ASSERT_TRUE(ss);
ASSERT_EQ(ss->GetName(), collection_name);
stale_ss_id = ss->GetID();
ScopedSnapshotT a_ss;
status = Snapshots::GetInstance().GetSnapshot(a_ss, collection_name);
ASSERT_TRUE(status.ok());
std::this_thread::sleep_for(std::chrono::milliseconds(80));
ASSERT_FALSE(ss->GetCollection()->IsActive());
status = Snapshots::GetInstance().GetSnapshot(a_ss, collection_name);
ASSERT_FALSE(status.ok());
auto c_c = CollectionCommitsHolder::GetInstance().GetResource(stale_ss_id, false);
ASSERT_TRUE(c_c);
ASSERT_EQ(c_c->GetID(), stale_ss_id);
};
auto worker2 = [&] {
std::this_thread::sleep_for(std::chrono::milliseconds(50));
auto status = Snapshots::GetInstance().DropCollection(collection_name, ++lsn);
ASSERT_TRUE(status.ok());
ScopedSnapshotT a_ss;
status = Snapshots::GetInstance().GetSnapshot(a_ss, collection_name);
ASSERT_FALSE(status.ok());
};
auto worker3 = [&] {
std::this_thread::sleep_for(std::chrono::milliseconds(20));
auto ss = CreateCollection(collection_name, ++lsn);
ASSERT_FALSE(ss);
std::this_thread::sleep_for(std::chrono::milliseconds(80));
ss = CreateCollection(collection_name, ++lsn);
ASSERT_TRUE(ss);
ASSERT_EQ(ss->GetName(), collection_name);
};
std::thread t1 = std::thread(worker1);
std::thread t2 = std::thread(worker2);
std::thread t3 = std::thread(worker3);
t1.join();
t2.join();
t3.join();
auto c_c = CollectionCommitsHolder::GetInstance().GetResource(stale_ss_id, false);
ASSERT_FALSE(c_c);
}
TEST_F(SnapshotTest, PartitionTest) {
std::string collection_name("c1");
LSN_TYPE lsn = 1;
auto ss = CreateCollection(collection_name, ++lsn);
ASSERT_TRUE(ss);
ASSERT_EQ(ss->GetName(), collection_name);
ASSERT_EQ(ss->NumberOfPartitions(), 1);
auto partition_iterator = std::make_shared<PartitionCollector>(ss);
partition_iterator->Iterate();
ASSERT_TRUE(partition_iterator->GetStatus().ok());
ASSERT_EQ(partition_iterator->partition_names_.size(), 1);
OperationContext context;
context.lsn = ++lsn;
auto op = std::make_shared<CreatePartitionOperation>(context, ss);
std::string partition_name("p1");
PartitionContext p_ctx;
p_ctx.name = partition_name;
PartitionPtr partition;
auto status = op->CommitNewPartition(p_ctx, partition);
ASSERT_TRUE(status.ok());
ASSERT_TRUE(partition);
ASSERT_EQ(partition->GetName(), partition_name);
ASSERT_FALSE(partition->IsActive());
ASSERT_TRUE(partition->HasAssigned());
status = op->Push();
ASSERT_TRUE(status.ok());
ScopedSnapshotT curr_ss;
status = op->GetSnapshot(curr_ss);
ASSERT_TRUE(status.ok());
ASSERT_TRUE(curr_ss);
ASSERT_EQ(curr_ss->GetName(), ss->GetName());
ASSERT_GT(curr_ss->GetID(), ss->GetID());
ASSERT_EQ(curr_ss->NumberOfPartitions(), 2);
partition_iterator = std::make_shared<PartitionCollector>(curr_ss);
partition_iterator->Iterate();
ASSERT_TRUE(partition_iterator->GetStatus().ok());
ASSERT_EQ(partition_iterator->partition_names_.size(), 2);
p_ctx.lsn = ++lsn;
auto drop_op = std::make_shared<DropPartitionOperation>(p_ctx, curr_ss);
status = drop_op->Push();
ASSERT_TRUE(status.ok());
ScopedSnapshotT latest_ss;
status = drop_op->GetSnapshot(latest_ss);
ASSERT_TRUE(status.ok());
ASSERT_TRUE(latest_ss);
ASSERT_EQ(latest_ss->GetName(), ss->GetName());
ASSERT_TRUE(latest_ss->GetID() > curr_ss->GetID());
ASSERT_EQ(latest_ss->NumberOfPartitions(), 1);
p_ctx.lsn = ++lsn;
drop_op = std::make_shared<DropPartitionOperation>(p_ctx, latest_ss);
status = drop_op->Push();
std::cout << status.ToString() << std::endl;
ASSERT_FALSE(status.ok());
// TODO: Keep LSN in order
PartitionContext pp_ctx;
/* pp_ctx.name = "p2"; */
/* curr_ss = CreatePartition(collection_name, pp_ctx, lsn-1); */
/* ASSERT_TRUE(curr_ss); */
std::stringstream p_name_stream;
auto num = RandomInt(20, 30);
for (auto i = 0; i < num; ++i) {
p_name_stream.str("");
p_name_stream << "partition_" << i;
pp_ctx.name = p_name_stream.str();
curr_ss = CreatePartition(collection_name, pp_ctx, ++lsn);
ASSERT_TRUE(curr_ss);
ASSERT_EQ(curr_ss->NumberOfPartitions(), 2 + i);
}
auto total_partition_num = curr_ss->NumberOfPartitions();
ID_TYPE partition_id;
for (auto i = 0; i < num; ++i) {
p_name_stream.str("");
p_name_stream << "partition_" << i;
status = curr_ss->GetPartitionId(p_name_stream.str(), partition_id);
ASSERT_TRUE(status.ok());
status = Snapshots::GetInstance().DropPartition(curr_ss->GetCollectionId(), partition_id, ++lsn);
ASSERT_TRUE(status.ok());
status = Snapshots::GetInstance().GetSnapshot(curr_ss, curr_ss->GetCollectionId());
ASSERT_TRUE(status.ok());
ASSERT_EQ(curr_ss->NumberOfPartitions(), total_partition_num - i - 1);
}
}
TEST_F(SnapshotTest, PartitionTest2) {
std::string collection_name("c1");
LSN_TYPE lsn = 1;
milvus::Status status;
auto ss = CreateCollection(collection_name, ++lsn);
ASSERT_TRUE(ss);
ASSERT_EQ(lsn, ss->GetMaxLsn());
OperationContext context;
context.lsn = lsn;
auto cp_op = std::make_shared<CreatePartitionOperation>(context, ss);
std::string partition_name("p1");
PartitionContext p_ctx;
p_ctx.name = partition_name;
PartitionPtr partition;
status = cp_op->CommitNewPartition(p_ctx, partition);
ASSERT_TRUE(status.ok());
ASSERT_TRUE(partition);
ASSERT_EQ(partition->GetName(), partition_name);
ASSERT_FALSE(partition->IsActive());
ASSERT_TRUE(partition->HasAssigned());
status = cp_op->Push();
ASSERT_FALSE(status.ok());
}
TEST_F(SnapshotTest, IndexTest) {
LSN_TYPE lsn = 0;
auto next_lsn = [&]() -> decltype(lsn) {
return ++lsn;
};
std::vector<ID_TYPE> ids;
auto status = Snapshots::GetInstance().GetCollectionIds(ids);
ASSERT_TRUE(status.ok()) << status.message();
auto collection_id = ids.at(0);
ScopedSnapshotT ss;
status = Snapshots::GetInstance().GetSnapshot(ss, collection_id);
ASSERT_TRUE(status.ok()) << status.message();
SegmentFileContext sf_context;
SFContextBuilder(sf_context, ss);
OperationContext context;
context.lsn = next_lsn();
context.prev_partition = ss->GetResource<Partition>(sf_context.partition_id);
auto build_op = std::make_shared<ChangeSegmentFileOperation>(context, ss);
SegmentFilePtr seg_file;
status = build_op->CommitNewSegmentFile(sf_context, seg_file);
ASSERT_TRUE(status.ok());
ASSERT_TRUE(seg_file);
auto op_ctx = build_op->GetContext();
ASSERT_EQ(seg_file, op_ctx.new_segment_files[0]);
build_op->Push();
status = build_op->GetSnapshot(ss);
ASSERT_TRUE(status.ok()) << status.message();
auto filter = [&](SegmentFile::Ptr segment_file) -> bool {
return segment_file->GetSegmentId() == seg_file->GetSegmentId();
};
auto filter2 = [&](SegmentFile::Ptr segment_file) -> bool {
return true;
};
auto sf_collector = std::make_shared<SegmentFileCollector>(ss, filter);
sf_collector->Iterate();
auto it_found = sf_collector->segment_files_.find(seg_file->GetID());
ASSERT_NE(it_found, sf_collector->segment_files_.end());
status = Snapshots::GetInstance().GetSnapshot(ss, collection_id);
ASSERT_TRUE(status.ok()) << status.ToString();
OperationContext drop_ctx;
drop_ctx.lsn = next_lsn();
drop_ctx.stale_segment_files.push_back(seg_file);
auto drop_op = std::make_shared<DropIndexOperation>(drop_ctx, ss);
status = drop_op->Push();
ASSERT_TRUE(status.ok());
status = drop_op->GetSnapshot(ss);
ASSERT_TRUE(status.ok());
sf_collector = std::make_shared<SegmentFileCollector>(ss, filter);
sf_collector->Iterate();
it_found = sf_collector->segment_files_.find(seg_file->GetID());
ASSERT_EQ(it_found, sf_collector->segment_files_.end());
PartitionContext pp_ctx;
std::stringstream p_name_stream;
auto num = RandomInt(3, 5);
for (auto i = 0; i < num; ++i) {
p_name_stream.str("");
p_name_stream << "partition_" << i;
pp_ctx.name = p_name_stream.str();
ss = CreatePartition(ss->GetName(), pp_ctx, next_lsn());
ASSERT_TRUE(ss);
}
ASSERT_EQ(ss->NumberOfPartitions(), num + 1);
sf_collector = std::make_shared<SegmentFileCollector>(ss, filter2);
sf_collector->Iterate();
auto prev_total = sf_collector->segment_files_.size();
decltype(sf_context) osf_context;
SFContextBuilder(osf_context, ss, {sf_context.field_element_name});
std::vector<SegmentFileContext> sfs_context = {sf_context, osf_context};
auto new_total = 0;
auto partitions = ss->GetResources<Partition>();
for (auto& kv : partitions) {
num = RandomInt(2, 5);
auto row_cnt = 1024;
for (auto i = 0; i < num; ++i) {
ASSERT_TRUE(CreateSegment(ss, kv.first, next_lsn(), sfs_context, row_cnt).ok());
}
new_total += num * sfs_context.size();
}
status = Snapshots::GetInstance().GetSnapshot(ss, ss->GetName());
ASSERT_TRUE(status.ok());
std::cout << ss->ToString() << std::endl;
sf_collector = std::make_shared<SegmentFileCollector>(ss, filter2);
sf_collector->Iterate();
auto total = sf_collector->segment_files_.size();
ASSERT_EQ(total, prev_total + new_total);
std::set<ID_TYPE> fe_ids;
for (auto ctx : sfs_context) {
auto field_element_id = ss->GetFieldElementId(ctx.field_name,
ctx.field_element_name);
ASSERT_NE(field_element_id, 0);
fe_ids.insert(field_element_id);
}
auto filter3 = [&](SegmentFile::Ptr segment_file) -> bool {
return fe_ids.find(segment_file->GetFieldElementId()) != fe_ids.end();
};
sf_collector = std::make_shared<SegmentFileCollector>(ss, filter3);
sf_collector->Iterate();
auto specified_segment_files_cnt = sf_collector->segment_files_.size();
OperationContext d_a_i_ctx;
d_a_i_ctx.lsn = next_lsn();
for (auto fe_id : fe_ids) {
d_a_i_ctx.stale_field_elements.push_back(ss->GetResource<FieldElement>(fe_id));
}
auto drop_all_index_op = std::make_shared<DropAllIndexOperation>(d_a_i_ctx, ss);
status = drop_all_index_op->Push();
ASSERT_TRUE(status.ok());
status = drop_all_index_op->GetSnapshot(ss);
ASSERT_TRUE(status.ok());
/* { */
/* auto& fields = ss->GetResources<Field>(); */
/* for (auto field_kv : fields) { */
/* auto field = field_kv.second; */
/* std::cout << "field " << field->GetID() << " " << field->GetName() << std::endl; */
/* auto& field_elements = ss->GetResources<FieldElement>(); */
/* for (auto element_kv : field_elements) { */
/* auto element = element_kv.second; */
/* if (element->GetFieldId() != field->GetID()) { */
/* continue; */
/* } */
/* std::cout << "\tfield_element " << element->GetID() << " " << element->GetName() << std::endl; */
/* } */
/* } */
/* } */
std::cout << ss->ToString() << std::endl;
sf_collector = std::make_shared<SegmentFileCollector>(ss, filter2);
sf_collector->Iterate();
ASSERT_EQ(sf_collector->segment_files_.size(), total - specified_segment_files_cnt);
std::cout << "sf size = " << sf_collector->segment_files_.size() << std::endl;
std::cout << "total = " << total << std::endl;
std::cout << "specified_segment_files_cnt = " << specified_segment_files_cnt << std::endl;
{
auto& field_elements = ss->GetResources<FieldElement>();
for (auto& kv : field_elements) {
ASSERT_EQ(fe_ids.find(kv.second->GetID()), fe_ids.end());
}
}
}
TEST_F(SnapshotTest, OperationTest) {
std::string to_string;
LSN_TYPE lsn;
Status status;
/* ID_TYPE collection_id; */
/* ASSERT_TRUE(GetFirstCollectionID(collection_id).ok()); */
std::string collection_name("c1");
auto ss = CreateCollection(collection_name, ++lsn);
ASSERT_TRUE(ss);
SegmentFileContext sf_context;
SFContextBuilder(sf_context, ss);
auto& partitions = ss->GetResources<Partition>();
auto total_row_cnt = 0;
for (auto& kv : partitions) {
auto num = RandomInt(2, 5);
for (auto i = 0; i < num; ++i) {
auto row_cnt = RandomInt(100, 200);
ASSERT_TRUE(CreateSegment(ss, kv.first, ++lsn, sf_context, row_cnt).ok());
total_row_cnt += row_cnt;
}
}
status = Snapshots::GetInstance().GetSnapshot(ss, collection_name);
ASSERT_TRUE(status.ok());
ASSERT_EQ(total_row_cnt, ss->GetCollectionCommit()->GetRowCount());
auto total_size = ss->GetCollectionCommit()->GetSize();
auto ss_id = ss->GetID();
lsn = ss->GetMaxLsn() + 1;
// Check snapshot
{
auto collection_commit = CollectionCommitsHolder::GetInstance().GetResource(ss_id, false);
/* snapshot::SegmentCommitsHolder::GetInstance().GetResource(prev_segment_commit->GetID()); */
ASSERT_TRUE(collection_commit);
std::cout << collection_commit->ToString() << std::endl;
}
OperationContext merge_ctx;
auto merge_segment_row_cnt = 0;
std::set<ID_TYPE> stale_segment_commit_ids;
SFContextBuilder(sf_context, ss);
std::cout << ss->ToString() << std::endl;
ID_TYPE new_seg_id;
ScopedSnapshotT new_ss;
// Check build operation correctness
{
OperationContext context;
context.lsn = ++lsn;
auto build_op = std::make_shared<ChangeSegmentFileOperation>(context, ss);
SegmentFilePtr seg_file;
status = build_op->CommitNewSegmentFile(sf_context, seg_file);
std::cout << status.ToString() << std::endl;
ASSERT_TRUE(status.ok());
ASSERT_TRUE(seg_file);
auto prev_segment_commit = ss->GetSegmentCommitBySegmentId(seg_file->GetSegmentId());
auto prev_segment_commit_mappings = prev_segment_commit->GetMappings();
ASSERT_FALSE(prev_segment_commit->ToString().empty());
auto new_size = RandomInt(1000, 20000);
seg_file->SetSize(new_size);
total_size += new_size;
auto delta = prev_segment_commit->GetRowCount() / 2;
build_op->CommitRowCountDelta(delta);
total_row_cnt -= delta;
build_op->Push();
status = build_op->GetSnapshot(ss);
ASSERT_TRUE(status.ok());
ASSERT_GT(ss->GetID(), ss_id);
ASSERT_EQ(ss->GetCollectionCommit()->GetRowCount(), total_row_cnt);
ASSERT_EQ(ss->GetCollectionCommit()->GetSize(), total_size);
std::cout << ss->ToString() << std::endl;
auto segment_commit = ss->GetSegmentCommitBySegmentId(seg_file->GetSegmentId());
auto segment_commit_mappings = segment_commit->GetMappings();
MappingT expected_mappings = prev_segment_commit_mappings;
expected_mappings.insert(seg_file->GetID());
ASSERT_EQ(expected_mappings, segment_commit_mappings);
auto seg = ss->GetResource<Segment>(seg_file->GetSegmentId());
ASSERT_TRUE(seg);
merge_ctx.stale_segments.push_back(seg);
merge_segment_row_cnt += ss->GetSegmentCommitBySegmentId(seg->GetID())->GetRowCount();
stale_segment_commit_ids.insert(segment_commit->GetID());
}
// std::this_thread::sleep_for(std::chrono::milliseconds(100));
// // Check stale snapshot has been deleted from store
// {
// auto collection_commit = CollectionCommitsHolder::GetInstance().GetResource(ss_id, false);
// ASSERT_FALSE(collection_commit);
// }
ss_id = ss->GetID();
ID_TYPE partition_id;
{
std::vector<PartitionPtr> partitions;
auto executor = [&](const PartitionPtr& partition,
PartitionIterator* itr) -> Status {
if (partition->GetCollectionId() != ss->GetCollectionId()) {
return Status::OK();
}
partitions.push_back(partition);
return Status::OK();
};
auto iterator = std::make_shared<PartitionIterator>(ss, executor);
iterator->Iterate();
OperationContext context;
context.lsn = ++lsn;
context.prev_partition = ss->GetResource<Partition>(partitions[0]->GetID());
auto op = std::make_shared<NewSegmentOperation>(context, ss);
SegmentPtr new_seg;
status = op->CommitNewSegment(new_seg);
ASSERT_TRUE(status.ok());
ASSERT_FALSE(new_seg->ToString().empty());
SegmentFilePtr seg_file;
status = op->CommitNewSegmentFile(sf_context, seg_file);
ASSERT_TRUE(status.ok());
auto new_segment_row_cnt = RandomInt(100, 200);
status = op->CommitRowCount(new_segment_row_cnt);
ASSERT_TRUE(status.ok());
total_row_cnt += new_segment_row_cnt;
auto new_size = new_segment_row_cnt * 5;
seg_file->SetSize(new_size);
total_size += new_size;
status = op->Push();
ASSERT_TRUE(status.ok());
status = op->GetSnapshot(ss);
ASSERT_GT(ss->GetID(), ss_id);
ASSERT_TRUE(status.ok());
ASSERT_EQ(ss->GetCollectionCommit()->GetRowCount(), total_row_cnt);
ASSERT_EQ(ss->GetCollectionCommit()->GetSize(), total_size);
auto segment_commit = ss->GetSegmentCommitBySegmentId(seg_file->GetSegmentId());
auto segment_commit_mappings = segment_commit->GetMappings();
MappingT expected_segment_mappings;
expected_segment_mappings.insert(seg_file->GetID());
ASSERT_EQ(expected_segment_mappings, segment_commit_mappings);
merge_ctx.stale_segments.push_back(new_seg);
merge_segment_row_cnt += ss->GetSegmentCommitBySegmentId(new_seg->GetID())->GetRowCount();
partition_id = segment_commit->GetPartitionId();
stale_segment_commit_ids.insert(segment_commit->GetID());
auto partition = ss->GetResource<Partition>(partition_id);
merge_ctx.prev_partition = partition;
new_seg_id = seg_file->GetSegmentId();
new_ss = ss;
}
SegmentPtr merge_seg;
ss_id = ss->GetID();
{
auto prev_partition_commit = ss->GetPartitionCommitByPartitionId(partition_id);
auto expect_null = ss->GetPartitionCommitByPartitionId(11111111);
ASSERT_FALSE(expect_null);
ASSERT_FALSE(prev_partition_commit->ToString().empty());
merge_ctx.lsn = ++lsn;
auto op = std::make_shared<MergeOperation>(merge_ctx, ss);
SegmentPtr new_seg;
status = op->CommitNewSegment(new_seg);
sf_context.segment_id = new_seg->GetID();
SegmentFilePtr seg_file;
status = op->CommitNewSegmentFile(sf_context, seg_file);
ASSERT_TRUE(status.ok());
auto new_size = RandomInt(1000, 20000);
seg_file->SetSize(new_size);
auto stale_size = 0;
for (auto& stale_seg : merge_ctx.stale_segments) {
stale_size += ss->GetSegmentCommitBySegmentId(stale_seg->GetID())->GetSize();
}
total_size += new_size;
total_size -= stale_size;
status = op->Push();
ASSERT_TRUE(status.ok()) << status.ToString();
std::cout << op->ToString() << std::endl;
status = op->GetSnapshot(ss);
ASSERT_GT(ss->GetID(), ss_id);
ASSERT_TRUE(status.ok());
ASSERT_EQ(total_size, ss->GetCollectionCommit()->GetSize());
auto segment_commit = ss->GetSegmentCommitBySegmentId(new_seg->GetID());
ASSERT_EQ(segment_commit->GetRowCount(), merge_segment_row_cnt);
ASSERT_EQ(segment_commit->GetSize(), new_size);
ASSERT_EQ(ss->GetCollectionCommit()->GetRowCount(), total_row_cnt);
auto new_partition_commit = ss->GetPartitionCommitByPartitionId(partition_id);
auto new_mappings = new_partition_commit->GetMappings();
auto prev_mappings = prev_partition_commit->GetMappings();
auto expected_mappings = prev_mappings;
for (auto id : stale_segment_commit_ids) {
expected_mappings.erase(id);
}
expected_mappings.insert(segment_commit->GetID());
ASSERT_EQ(expected_mappings, new_mappings);
CollectionCommitsHolder::GetInstance().Dump();
merge_seg = new_seg;
}
// 1. New seg1, seg2
// 2. Build seg1 start
// 3. Merge seg1, seg2 to seg3
// 4. Commit new seg file of build operation -> Stale Segment Found Here!
{
OperationContext context;
context.lsn = ++lsn;
auto build_op = std::make_shared<ChangeSegmentFileOperation>(context, new_ss);
SegmentFilePtr seg_file;
auto new_sf_context = sf_context;
new_sf_context.segment_id = new_seg_id;
status = build_op->CommitNewSegmentFile(new_sf_context, seg_file);
ASSERT_FALSE(status.ok());
}
{
OperationContext context;
context.lsn = ++lsn;
auto op = std::make_shared<ChangeSegmentFileOperation>(context, ss);
SegmentFilePtr seg_file;
auto new_sf_context = sf_context;
new_sf_context.segment_id = merge_seg->GetID();
status = op->CommitNewSegmentFile(new_sf_context, seg_file);
ASSERT_TRUE(status.ok());
auto prev_sc = ss->GetSegmentCommitBySegmentId(merge_seg->GetID());
ASSERT_TRUE(prev_sc);
auto delta = prev_sc->GetRowCount() + 1;
op->CommitRowCountDelta(delta);
status = op->Push();
std::cout << status.ToString() << std::endl;
ASSERT_FALSE(status.ok());
}
std::string new_fe_name = "fe_index";
{
status = Snapshots::GetInstance().GetSnapshot(ss, collection_name);
ASSERT_TRUE(status.ok());
auto field = ss->GetField(sf_context.field_name);
ASSERT_TRUE(field);
auto new_fe = std::make_shared<FieldElement>(ss->GetCollectionId(),
field->GetID(), new_fe_name, milvus::engine::FieldElementType::FET_INDEX);
OperationContext context;
context.lsn = ++lsn;
context.new_field_elements.push_back(new_fe);
auto op = std::make_shared<AddFieldElementOperation>(context, ss);
status = op->Push();
ASSERT_TRUE(status.ok());
status = op->GetSnapshot(ss);
ASSERT_TRUE(status.ok());
std::cout << ss->ToString() << std::endl;
}
{
auto snapshot_id = ss->GetID();
auto field = ss->GetField(sf_context.field_name);
ASSERT_TRUE(field);
FieldElementPtr new_fe;
status = ss->GetFieldElement(sf_context.field_name, new_fe_name, new_fe);
ASSERT_TRUE(status.ok());
ASSERT_TRUE(new_fe);
OperationContext context;
context.lsn = ++lsn;
context.new_field_elements.push_back(new_fe);
auto op = std::make_shared<AddFieldElementOperation>(context, ss);
status = op->Push();
ASSERT_FALSE(status.ok());
status = Snapshots::GetInstance().GetSnapshot(ss, collection_name);
ASSERT_TRUE(status.ok());
ASSERT_EQ(snapshot_id, ss->GetID());
}
// 1. Build start
// 2. Commit new seg file of build operation
// 3. Drop collection
// 4. Commit build operation -> Stale Segment Found Here!
{
OperationContext context;
context.lsn = ++lsn;
auto build_op = std::make_shared<ChangeSegmentFileOperation>(context, ss);
SegmentFilePtr seg_file;
auto new_sf_context = sf_context;
new_sf_context.segment_id = merge_seg->GetID();
status = build_op->CommitNewSegmentFile(new_sf_context, seg_file);
ASSERT_TRUE(status.ok());
std::cout << build_op->ToString() << std::endl;
auto status = Snapshots::GetInstance().DropCollection(ss->GetName(), ++lsn);
ASSERT_TRUE(status.ok());
status = build_op->Push();
std::cout << status.ToString() << std::endl;
ASSERT_FALSE(status.ok());
ASSERT_FALSE(build_op->GetStatus().ok());
std::cout << build_op->ToString() << std::endl;
}
Snapshots::GetInstance().Reset();
}
TEST_F(SnapshotTest, OperationTest2) {
LSN_TYPE lsn = 0;
std::string collection_name("c1");
auto ss = CreateCollection(collection_name, ++lsn);
ASSERT_TRUE(ss);
SegmentFileContext sf_context;
SFContextBuilder(sf_context, ss);
auto& partitions = ss->GetResources<Partition>();
auto total_row_cnt = 0;
for (auto& kv : partitions) {
auto num = RandomInt(2, 5);
for (auto i = 0; i < num; ++i) {
auto row_cnt = RandomInt(100, 200);
ASSERT_TRUE(CreateSegment(ss, kv.first, ++lsn, sf_context, row_cnt).ok());
total_row_cnt += row_cnt;
}
}
auto status = Snapshots::GetInstance().GetSnapshot(ss, collection_name);
ASSERT_TRUE(status.ok());
ASSERT_EQ(total_row_cnt, ss->GetCollectionCommit()->GetRowCount());
auto total_size = ss->GetCollectionCommit()->GetSize();
auto target_segment = ss->GetResources<Segment>().begin()->second.Get();
auto sf_ids = ss->GetSegmentFileIds(target_segment->GetID());
ASSERT_GT(sf_ids.size(), 0);
auto stale_sf = ss->GetResource<SegmentFile>(*(sf_ids.begin()));
ASSERT_TRUE(stale_sf);
{
OperationContext context;
context.lsn = ++lsn;
context.stale_segment_files.push_back(stale_sf);
auto op = std::make_shared<ChangeSegmentFileOperation>(context, ss);
SegmentFilePtr seg_file;
sf_context.field_name = "vector";
sf_context.field_element_name = "_raw";
sf_context.segment_id = stale_sf->GetSegmentId();
sf_context.partition_id = stale_sf->GetPartitionId();
sf_context.collection_id = stale_sf->GetCollectionId();
status = op->CommitNewSegmentFile(sf_context, seg_file);
/* std::cout << status.ToString() << std::endl; */
ASSERT_TRUE(status.ok());
ASSERT_TRUE(seg_file);
auto prev_segment_commit = ss->GetSegmentCommitBySegmentId(seg_file->GetSegmentId());
auto prev_segment_commit_mappings = prev_segment_commit->GetMappings();
ASSERT_FALSE(prev_segment_commit->ToString().empty());
auto new_size = RandomInt(1000, 20000);
seg_file->SetSize(new_size);
total_size += new_size;
total_size -= stale_sf->GetSize();
auto delta = prev_segment_commit->GetRowCount() / 2;
op->CommitRowCountDelta(delta);
total_row_cnt -= delta;
status = op->Push();
ASSERT_TRUE(status.ok()) << status.message();
status = op->GetSnapshot(ss);
ASSERT_TRUE(status.ok());
ASSERT_EQ(ss->GetCollectionCommit()->GetRowCount(), total_row_cnt);
ASSERT_EQ(ss->GetCollectionCommit()->GetSize(), total_size);
std::cout << ss->ToString() << std::endl;
}
{
OperationContext context;
context.lsn = ++lsn;
SegmentFile::VecT stale_sfs;
std::set<ID_TYPE> modified_segments;
auto executor = [&](const SegmentFilePtr& sf,
SegmentFileIterator* itr) -> Status {
context.stale_segment_files.push_back(sf);
modified_segments.insert(sf->GetSegmentId());
return Status::OK();
};
auto iterator = std::make_shared<SegmentFileIterator>(ss, executor);
iterator->Iterate();
ASSERT_TRUE(iterator->GetStatus().ok());
{
auto op = std::make_shared<CompoundSegmentsOperation>(context, ss);
for (auto seg_id : modified_segments) {
auto sc = ss->GetSegmentCommitBySegmentId(seg_id);
ASSERT_TRUE(sc);
status = op->CommitRowCountDelta(seg_id, sc->GetRowCount() + 1, true);
ASSERT_TRUE(status.ok()) << status.ToString();
}
status = op->Push();
ASSERT_FALSE(status.ok()) << status.ToString();
}
{
auto op = std::make_shared<CompoundSegmentsOperation>(context, ss);
status = op->CommitRowCountDelta(100000000, 1, false);
ASSERT_FALSE(status.ok()) << status.ToString();
}
{
auto segment_cnt = ss->GetResources<Segment>().size();
context.lsn++;
auto op = std::make_shared<CompoundSegmentsOperation>(context, ss);
for (auto seg_id : modified_segments) {
auto sc = ss->GetSegmentCommitBySegmentId(seg_id);
ASSERT_TRUE(sc);
bool is_sub = (RandomInt(1, 10) >= 5);
auto delta = sc->GetRowCount() / 2;
status = op->CommitRowCountDelta(seg_id, delta, is_sub);
ASSERT_TRUE(status.ok()) << status.ToString();
if (is_sub) {
total_row_cnt -= delta;
} else {
total_row_cnt += delta;
}
}
auto new_size = 0;
for (auto stale_sf : context.stale_segment_files) {
SegmentFilePtr seg_file;
sf_context.field_name = "vector";
sf_context.field_element_name = "_raw";
sf_context.segment_id = stale_sf->GetSegmentId();
sf_context.partition_id = stale_sf->GetPartitionId();
sf_context.collection_id = stale_sf->GetCollectionId();
status = op->CommitNewSegmentFile(sf_context, seg_file);
auto size_delta = RandomInt(100, 1000);
seg_file->SetSize(size_delta);
new_size += size_delta;
ASSERT_TRUE(status.ok()) << status.ToString();
}
SegmentPtr new_seg;
OperationContext new_seg_ctx;
new_seg_ctx.prev_partition = ss->GetResource<Partition>(
context.stale_segment_files[0]->GetPartitionId());
status = op->CommitNewSegment(new_seg_ctx, new_seg);
ASSERT_TRUE(status.ok()) << status.ToString();
ASSERT_TRUE(new_seg);
{
SegmentFilePtr seg_file;
sf_context.field_name = "vector";
sf_context.field_element_name = "_raw";
sf_context.segment_id = new_seg->GetID();
sf_context.partition_id = new_seg->GetPartitionId();
sf_context.collection_id = new_seg->GetCollectionId();
status = op->CommitNewSegmentFile(sf_context, seg_file);
ASSERT_TRUE(status.ok()) << status.ToString();
auto size_delta = RandomInt(100, 1000);
seg_file->SetSize(size_delta);
new_size += size_delta;
ASSERT_TRUE(status.ok()) << status.ToString();
}
auto delta_row = RandomInt(100, 1000);
status = op->CommitRowCountDelta(new_seg->GetID(), delta_row, false);
ASSERT_TRUE(status.ok());
total_row_cnt += delta_row;
status = op->Push();
ASSERT_TRUE(status.ok()) << status.ToString();
status = op->GetSnapshot(ss);
ASSERT_TRUE(status.ok());
std::cout << ss->ToString() << std::endl;
ASSERT_EQ(segment_cnt + 1, ss->GetResources<Segment>().size());
auto expect_size = ss->GetCollectionCommit()->GetSize();
ASSERT_EQ(expect_size, new_size) << status.ToString();
auto expect_row_cnt = ss->GetCollectionCommit()->GetRowCount();
ASSERT_EQ(expect_row_cnt, total_row_cnt) << status.ToString();
}
{
modified_segments.clear();
context.stale_segment_files.clear();
iterator = std::make_shared<SegmentFileIterator>(ss, executor);
iterator->Iterate();
ASSERT_TRUE(iterator->GetStatus().ok());
context.lsn++;
auto op = std::make_shared<CompoundSegmentsOperation>(context, ss);
for (auto seg_id : modified_segments) {
auto sc = ss->GetSegmentCommitBySegmentId(seg_id);
ASSERT_TRUE(sc);
status = op->CommitRowCountDelta(seg_id, sc->GetRowCount(), true);
ASSERT_TRUE(status.ok()) << status.ToString();
}
status = op->Push();
ASSERT_TRUE(status.ok()) << status.ToString();
status = op->GetSnapshot(ss);
ASSERT_TRUE(status.ok());
std::cout << ss->ToString() << std::endl;
ASSERT_EQ(ss->GetCollectionCommit()->GetRowCount(), 0);
ASSERT_EQ(ss->GetCollectionCommit()->GetSize(), 0);
}
}
}
TEST_F(SnapshotTest, CompoundTest1) {
Status status;
std::atomic<LSN_TYPE> lsn = 0;
auto next_lsn = [&]() -> decltype(lsn) {
return ++lsn;
};
std::atomic<LSN_TYPE> pid = 0;
auto next_pid = [&]() -> decltype(pid) {
return ++pid;
};
std::string collection_name("c1");
auto ss = CreateCollection(collection_name, next_lsn());
ASSERT_TRUE(ss);
ASSERT_EQ(lsn, ss->GetMaxLsn());
Queue merge_queue;
Queue build_queue;
std::set<ID_TYPE> all_segments;
std::set<ID_TYPE> segment_in_building;
std::map<ID_TYPE, std::set<ID_TYPE>> merged_segs_history;
std::set<ID_TYPE> merged_segs;
std::set<ID_TYPE> built_segs;
std::set<ID_TYPE> build_stale_segs;
std::mutex all_mtx;
std::mutex merge_mtx;
std::mutex built_mtx;
std::mutex partition_mtx;
WaitableObj merge_waiter;
WaitableObj build_waiter;
SegmentFileContext sf_context;
sf_context.field_name = "vector";
sf_context.field_element_name = "ivfsq8";
sf_context.segment_id = 1;
sf_context.partition_id = 1;
IDS_TYPE partitions = {ss->GetResources<Partition>().begin()->second->GetID()};
auto do_build = [&] (const ID_TYPE& seg_id) {
ScopedSnapshotT latest_ss;
auto status = Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
ASSERT_TRUE(status.ok());
auto build_sf_context = sf_context;
OperationContext context;
context.lsn = next_lsn();
auto build_op = std::make_shared<ChangeSegmentFileOperation>(context, latest_ss);
SegmentFilePtr seg_file;
build_sf_context.segment_id = seg_id;
status = build_op->CommitNewSegmentFile(build_sf_context, seg_file);
if (!status.ok()) {
std::cout << status.ToString() << std::endl;
std::unique_lock<std::mutex> lock(merge_mtx);
auto it = merged_segs.find(seg_id);
ASSERT_NE(it, merged_segs.end());
return;
}
std::unique_lock<std::mutex> lock(built_mtx);
status = build_op->Push();
if (!status.ok()) {
std::cout << status.ToString() << std::endl;
std::unique_lock<std::mutex> lock(merge_mtx);
auto it = merged_segs.find(seg_id);
ASSERT_NE(it, merged_segs.end());
return;
}
ASSERT_TRUE(status.ok());
status = build_op->GetSnapshot(latest_ss);
ASSERT_TRUE(status.ok());
built_segs.insert(seg_id);
};
auto do_merge = [&] (std::set<ID_TYPE>& seg_ids, ID_TYPE& new_seg_id) {
if (seg_ids.size() == 0) {
return;
}
ScopedSnapshotT latest_ss;
auto status = Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
ASSERT_TRUE(status.ok());
PartitionPtr partition;
OperationContext context;
for (auto& id : seg_ids) {
auto seg = latest_ss->GetResource<Segment>(id);
if (!seg) {
std::cout << "Error seg=" << id << std::endl;
ASSERT_TRUE(seg);
}
if (!partition) {
partition = latest_ss->GetResource<Partition>(seg->GetPartitionId());
ASSERT_TRUE(partition);
} else {
ASSERT_EQ(seg->GetPartitionId(), partition->GetID());
}
context.stale_segments.push_back(seg);
if (!context.prev_partition) {
context.prev_partition = latest_ss->GetResource<Partition>(
seg->GetPartitionId());
}
}
context.lsn = next_lsn();
auto op = std::make_shared<MergeOperation>(context, latest_ss);
SegmentPtr new_seg;
status = op->CommitNewSegment(new_seg);
ASSERT_TRUE(status.ok());
sf_context.segment_id = new_seg->GetID();
SegmentFilePtr seg_file;
status = op->CommitNewSegmentFile(sf_context, seg_file);
ASSERT_TRUE(status.ok());
std::unique_lock<std::mutex> lock(merge_mtx);
status = op->Push();
if (!status.ok()) {
lock.unlock();
std::unique_lock<std::mutex> blk(built_mtx);
std::cout << status.ToString() << std::endl;
/* for (auto id : built_segs) { */
/* std::cout << "builted " << id << std::endl; */
/* } */
/* for (auto id : seg_ids) { */
/* std::cout << "to_merge " << id << std::endl; */
/* } */
bool stale_found = false;
for (auto& seg_id : seg_ids) {
auto it = built_segs.find(seg_id);
if (it != built_segs.end()) {
stale_found = true;
break;
}
}
ASSERT_TRUE(stale_found);
return;
}
ID_TYPE ss_id = latest_ss->GetID();
status = op->GetSnapshot(latest_ss);
ASSERT_TRUE(status.ok());
ASSERT_TRUE(latest_ss->GetID() > ss_id);
merged_segs_history[new_seg->GetID()] = seg_ids;
for (auto& seg_id : seg_ids) {
merged_segs.insert(seg_id);
}
new_seg_id = new_seg->GetID();
ASSERT_EQ(new_seg->GetPartitionId(), partition->GetID());
};
// TODO: If any Compound Operation find larger Snapshot. This Operation should be rollback to latest
auto handler_worker = [&] {
auto loop_cnt = RandomInt(10, 20);
ScopedSnapshotT latest_ss;
auto create_new_segment = [&]() {
ID_TYPE partition_id;
{
std::unique_lock<std::mutex> lock(partition_mtx);
auto idx = RandomInt(0, partitions.size() - 1);
partition_id = partitions[idx];
}
Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
OperationContext context;
context.prev_partition = latest_ss->GetResource<Partition>(partition_id);
context.lsn = next_lsn();
auto op = std::make_shared<NewSegmentOperation>(context, latest_ss);
SegmentPtr new_seg;
auto status = op->CommitNewSegment(new_seg);
if (!status.ok()) {
std::cout << status.ToString() << std::endl;
}
ASSERT_TRUE(status.ok());
SegmentFilePtr seg_file;
sf_context.segment_id = new_seg->GetID();
op->CommitNewSegmentFile(sf_context, seg_file);
op->Push();
status = op->GetSnapshot(latest_ss);
ASSERT_TRUE(status.ok());
{
std::unique_lock<std::mutex> lock(all_mtx);
all_segments.insert(new_seg->GetID());
}
if (RandomInt(0, 10) >= 7) {
build_queue.Put(new_seg->GetID());
}
merge_queue.Put(new_seg->GetID());
};
auto create_partition = [&]() {
std::stringstream ss;
ss << "fake_partition_" << next_pid();
PartitionContext context;
context.name = ss.str();
std::unique_lock<std::mutex> lock(partition_mtx);
auto latest_ss = CreatePartition(collection_name, context, next_lsn());
ASSERT_TRUE(latest_ss);
auto partition = latest_ss->GetPartition(ss.str());
partitions.push_back(partition->GetID());
if (latest_ss->NumberOfPartitions() != partitions.size()) {
for (auto& pid : partitions) {
std::cout << "PartitionId=" << pid << std::endl;
}
}
ASSERT_EQ(latest_ss->NumberOfPartitions(), partitions.size());
};
for (auto i = 0; i < loop_cnt; ++i) {
if (RandomInt(0, 10) > 7) {
create_partition();
}
create_new_segment();
}
};
std::map<ID_TYPE, std::set<ID_TYPE>> merge_segs;
auto merge_worker = [&] {
while (true) {
auto seg_id = merge_queue.Take();
if (seg_id == 0) {
std::cout << "Exiting Merge Worker" << std::endl;
break;
}
ScopedSnapshotT latest_ss;
auto status = Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
ASSERT_TRUE(status.ok());
auto seg = latest_ss->GetResource<Segment>(seg_id);
if (!seg) {
std::cout << "SegID=" << seg_id << std::endl;
std::cout << latest_ss->ToString() << std::endl;
ASSERT_TRUE(seg);
}
auto it_segs = merge_segs.find(seg->GetPartitionId());
if (it_segs == merge_segs.end()) {
merge_segs[seg->GetPartitionId()] = {seg->GetID()};
} else {
merge_segs[seg->GetPartitionId()].insert(seg->GetID());
}
auto& segs = merge_segs[seg->GetPartitionId()];
if ((segs.size() >= 2) && (RandomInt(0, 10) >= 2)) {
std::cout << "Merging partition " << seg->GetPartitionId() << " segs (";
for (auto seg : segs) {
std::cout << seg << ",";
}
std::cout << ")" << std::endl;
ID_TYPE new_seg_id = 0;
do_merge(segs, new_seg_id);
segs.clear();
if (new_seg_id == 0) {
continue;
}
if (RandomInt(0, 10) >= 6) {
build_queue.Put(new_seg_id);
}
} else {
continue;
}
}
merge_waiter.Notify();
build_queue.Put(0);
build_waiter.Wait();
};
auto build_worker = [&] {
while (true) {
auto seg_id = build_queue.Take();
if (seg_id == 0) {
std::cout << "Exiting Build Worker" << std::endl;
break;
}
std::cout << "Building " << seg_id << std::endl;
do_build(seg_id);
}
build_waiter.Notify();
};
std::vector<std::thread> handlers;
auto num_handlers = RandomInt(8, 9);
for (auto i = 0; i < num_handlers; ++i) {
handlers.emplace_back(handler_worker);
}
std::thread t3 = std::thread(merge_worker);
std::thread t4 = std::thread(build_worker);
for (auto& handler : handlers) {
handler.join();
}
merge_queue.Put(0);
t3.join();
t4.join();
/* for (auto& kv : merged_segs_history) { */
/* std::cout << "merged: ("; */
/* for (auto i : kv.second) { */
/* std::cout << i << ","; */
/* } */
/* std::cout << ") -> " << kv.first << std::endl; */
/* } */
/* for (auto& id : built_segs) { */
/* std::cout << "built: " << id << std::endl; */
/* } */
merge_waiter.Wait();
ScopedSnapshotT latest_ss;
status = Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
ASSERT_TRUE(status.ok());
auto expect_segments = all_segments;
for (auto& kv : merged_segs_history) {
expect_segments.insert(kv.first);
for (auto& id : kv.second) {
expect_segments.erase(id);
}
}
std::set<ID_TYPE> final_segments;
auto segments = latest_ss->GetResources<Segment>();
for (auto& kv : segments) {
final_segments.insert(kv.first);
}
ASSERT_EQ(final_segments, expect_segments);
auto final_segment_file_cnt = latest_ss->GetResources<SegmentFile>().size();
size_t expect_segment_file_cnt;
expect_segment_file_cnt = expect_segments.size();
expect_segment_file_cnt += built_segs.size();
std::cout << latest_ss->ToString() << std::endl;
std::vector<int> common_ids;
std::set_intersection(merged_segs.begin(), merged_segs.end(), built_segs.begin(), built_segs.end(),
std::back_inserter(common_ids));
expect_segment_file_cnt -= common_ids.size();
ASSERT_EQ(expect_segment_file_cnt, final_segment_file_cnt);
}
TEST_F(SnapshotTest, CompoundTest2) {
Status status;
LSN_TYPE lsn = 0;
auto next_lsn = [&]() -> LSN_TYPE& {
return ++lsn;
};
std::atomic<LSN_TYPE> pid = 0;
auto next_pid = [&]() -> LSN_TYPE {
return ++pid;
};
std::string collection_name("c1");
auto ss = CreateCollection(collection_name, next_lsn());
ASSERT_TRUE(ss);
ASSERT_EQ(lsn, ss->GetMaxLsn());
TQueue merge_queue;
TQueue build_queue;
std::set<ID_TYPE> all_segments;
std::set<ID_TYPE> segment_in_building;
std::map<ID_TYPE, std::set<ID_TYPE>> merged_segs_history;
std::set<ID_TYPE> merged_segs;
std::set<ID_TYPE> built_segs;
std::set<ID_TYPE> build_stale_segs;
std::mutex all_mtx;
std::mutex merge_mtx;
std::mutex built_mtx;
std::mutex partition_mtx;
std::mutex seg_p_mtx;
WaitableObj merge_waiter;
WaitableObj build_waiter;
SegmentFileContext sf_context;
sf_context.field_name = "vector";
sf_context.field_element_name = "ivfsq8";
sf_context.segment_id = 1;
sf_context.partition_id = 1;
IDS_TYPE partitions = {ss->GetResources<Partition>().begin()->second->GetID()};
std::set<ID_TYPE> stale_partitions;
std::map<ID_TYPE, ID_TYPE> seg_p_map;
auto do_build = [&] (const ID_TYPE& seg_id, const ID_TYPE& p_id) {
ScopedSnapshotT latest_ss;
auto status = Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
ASSERT_TRUE(status.ok());
auto build_sf_context = sf_context;
OperationContext context;
context.lsn = next_lsn();
auto build_op = std::make_shared<ChangeSegmentFileOperation>(context, latest_ss);
SegmentFilePtr seg_file;
build_sf_context.segment_id = seg_id;
status = build_op->CommitNewSegmentFile(build_sf_context, seg_file);
if (!status.ok()) {
std::cout << status.ToString() << std::endl;
{
std::unique_lock<std::mutex> lock(partition_mtx);
auto it = stale_partitions.find(p_id);
if (it != stale_partitions.end()) {
std::cout << "Segment " << seg_id << " host partition " << p_id << " is stale" << std::endl;
return;
}
}
std::unique_lock<std::mutex> lock(merge_mtx);
auto it = merged_segs.find(seg_id);
ASSERT_NE(it, merged_segs.end());
return;
}
std::unique_lock<std::mutex> lock(built_mtx);
status = build_op->Push();
if (!status.ok()) {
lock.unlock();
std::cout << status.ToString() << std::endl;
{
std::unique_lock<std::mutex> lock(partition_mtx);
auto it = stale_partitions.find(p_id);
if (it != stale_partitions.end()) {
std::cout << "Segment " << seg_id << " host partition " << p_id << " is stale" << std::endl;
return;
}
}
std::unique_lock<std::mutex> lock(merge_mtx);
auto it = merged_segs.find(seg_id);
ASSERT_NE(it, merged_segs.end());
return;
}
ASSERT_TRUE(status.ok());
status = build_op->GetSnapshot(latest_ss);
ASSERT_TRUE(status.ok());
built_segs.insert(seg_id);
};
auto do_merge = [&] (std::set<ID_TYPE>& seg_ids, ID_TYPE& new_seg_id, const ID_TYPE& p_id) {
if (seg_ids.size() == 0) {
return;
}
ScopedSnapshotT latest_ss;
auto status = Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
ASSERT_TRUE(status.ok());
PartitionPtr partition;
OperationContext context;
for (auto& id : seg_ids) {
auto seg = latest_ss->GetResource<Segment>(id);
if (!seg) {
std::cout << "Stale seg=" << id << std::endl;
std::unique_lock<std::mutex> lock(partition_mtx);
std::cout << ((stale_partitions.find(p_id) != stale_partitions.end()) ? " due stale partition"
: " unexpected") << std::endl;
ASSERT_TRUE(stale_partitions.find(p_id) != stale_partitions.end());
return;
}
if (!partition) {
partition = latest_ss->GetResource<Partition>(seg->GetPartitionId());
ASSERT_TRUE(partition);
} else {
ASSERT_EQ(seg->GetPartitionId(), partition->GetID());
}
context.stale_segments.push_back(seg);
if (!context.prev_partition) {
context.prev_partition = latest_ss->GetResource<Partition>(
seg->GetPartitionId());
}
}
context.lsn = next_lsn();
auto op = std::make_shared<MergeOperation>(context, latest_ss);
SegmentPtr new_seg;
status = op->CommitNewSegment(new_seg);
ASSERT_TRUE(status.ok());
sf_context.segment_id = new_seg->GetID();
SegmentFilePtr seg_file;
status = op->CommitNewSegmentFile(sf_context, seg_file);
ASSERT_TRUE(status.ok());
std::unique_lock<std::mutex> lock(merge_mtx);
status = op->Push();
if (!status.ok()) {
lock.unlock();
bool stale_found = false;
{
std::unique_lock<std::mutex> blk(built_mtx);
std::cout << status.ToString() << std::endl;
/* for (auto id : built_segs) { */
/* std::cout << "builted " << id << std::endl; */
/* } */
/* for (auto id : seg_ids) { */
/* std::cout << "to_merge " << id << std::endl; */
/* } */
for (auto& seg_id : seg_ids) {
auto it = built_segs.find(seg_id);
if (it != built_segs.end()) {
stale_found = true;
break;
}
}
}
if (!stale_found) {
std::unique_lock<std::mutex> lock(partition_mtx);
auto it = stale_partitions.find(p_id);
if (it != stale_partitions.end()) {
/* std::cout << "Segment " << seg_id << " host partition " << p_id << " is stale" << std::endl; */
stale_found = true;
}
}
ASSERT_TRUE(stale_found);
return;
}
ID_TYPE ss_id = latest_ss->GetID();
status = op->GetSnapshot(latest_ss);
ASSERT_TRUE(status.ok());
ASSERT_TRUE(latest_ss->GetID() > ss_id);
merged_segs_history[new_seg->GetID()] = seg_ids;
for (auto& seg_id : seg_ids) {
merged_segs.insert(seg_id);
}
new_seg_id = new_seg->GetID();
seg_p_map[new_seg_id] = new_seg->GetPartitionId();
ASSERT_EQ(new_seg->GetPartitionId(), partition->GetID());
};
auto handler_worker = [&] {
auto loop_cnt = RandomInt(30, 35);
ScopedSnapshotT latest_ss;
auto create_new_segment = [&]() {
ID_TYPE partition_id;
{
std::unique_lock<std::mutex> lock(partition_mtx);
auto idx = RandomInt(0, partitions.size() - 1);
partition_id = partitions[idx];
}
Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
OperationContext context;
context.prev_partition = latest_ss->GetResource<Partition>(partition_id);
context.lsn = next_lsn();
auto op = std::make_shared<NewSegmentOperation>(context, latest_ss);
SegmentPtr new_seg;
auto status = op->CommitNewSegment(new_seg);
if (!status.ok()) {
std::cout << status.ToString() << std::endl;
std::unique_lock<std::mutex> lock(partition_mtx);
ASSERT_TRUE(stale_partitions.find(partition_id) != stale_partitions.end());
return;
}
ASSERT_TRUE(status.ok());
SegmentFilePtr seg_file;
sf_context.segment_id = new_seg->GetID();
status = op->CommitNewSegmentFile(sf_context, seg_file);
if (!status.ok()) {
std::cout << status.ToString() << std::endl;
ASSERT_TRUE(status.ok());
}
status = op->Push();
if (!status.ok()) {
std::unique_lock<std::mutex> lock(partition_mtx);
/* if (stale_partitions.find(partition_id) == stale_partitions.end()) { */
/* for (auto p : stale_partitions) { */
/* std::cout << "stale p: " << p << std::endl; */
/* } */
/* ASSERT_TRUE(false); */
/* } */
ASSERT_TRUE(stale_partitions.find(partition_id) != stale_partitions.end());
return;
}
status = op->GetSnapshot(latest_ss);
if (!status.ok()) {
std::cout << status.ToString() << std::endl;
std::unique_lock<std::mutex> lock(partition_mtx);
ASSERT_TRUE(stale_partitions.find(partition_id) != stale_partitions.end());
return;
}
{
std::unique_lock<std::mutex> lock(all_mtx);
all_segments.insert(new_seg->GetID());
}
if (RandomInt(0, 10) >= 7) {
build_queue.Put({new_seg->GetID(), new_seg->GetPartitionId()});
}
seg_p_map[new_seg->GetID()] = new_seg->GetPartitionId();
merge_queue.Put({new_seg->GetID(), new_seg->GetPartitionId()});
};
auto create_partition = [&]() {
std::stringstream ss;
ss << "fake_partition_" << next_pid();
PartitionContext context;
context.name = ss.str();
std::unique_lock<std::mutex> lock(partition_mtx);
auto latest_ss = CreatePartition(collection_name, context, next_lsn());
ASSERT_TRUE(latest_ss);
auto partition = latest_ss->GetPartition(ss.str());
partitions.push_back(partition->GetID());
if (latest_ss->NumberOfPartitions() != partitions.size()) {
for (auto& pid : partitions) {
std::cout << "PartitionId=" << pid << std::endl;
}
}
ASSERT_EQ(latest_ss->NumberOfPartitions(), partitions.size());
};
auto drop_partition = [&]() {
std::unique_lock<std::mutex> lock(partition_mtx);
if (partitions.size() <= 2) {
return;
}
ScopedSnapshotT latest_ss;
Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
auto index = RandomInt(0, partitions.size() - 1);
auto pid = partitions[index];
std::cout << "Dropping partition " << pid << std::endl;
status = Snapshots::GetInstance().DropPartition(
latest_ss->GetCollectionId(), pid, next_lsn());
ASSERT_TRUE(status.ok());
stale_partitions.insert(pid);
partitions.erase(partitions.begin() + index);
};
for (auto i = 0; i < loop_cnt; ++i) {
if (RandomInt(0, 10) > 5) {
create_partition();
}
if (RandomInt(0, 10) > 7) {
drop_partition();
}
create_new_segment();
}
};
std::map<ID_TYPE, std::set<ID_TYPE>> merge_segs;
auto merge_worker = [&] {
while (true) {
auto ids = merge_queue.Take();
auto seg_id = std::get<0>(ids);
auto p_id = std::get<1>(ids);
if (seg_id == 0) {
std::cout << "Exiting Merge Worker" << std::endl;
break;
}
ScopedSnapshotT latest_ss;
auto status = Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
ASSERT_TRUE(status.ok());
auto seg = latest_ss->GetResource<Segment>(seg_id);
if (!seg) {
std::cout << "SegID=" << seg_id << std::endl;
std::cout << latest_ss->ToString() << std::endl;
std::unique_lock<std::mutex> lock(partition_mtx);
auto it = stale_partitions.find(p_id);
if (it != stale_partitions.end()) {
std::cout << "Segment " << seg_id << " host partition " << p_id << " is stale" << std::endl;
continue;
}
ASSERT_TRUE(seg);
}
auto it_segs = merge_segs.find(seg->GetPartitionId());
if (it_segs == merge_segs.end()) {
merge_segs[seg->GetPartitionId()] = {seg->GetID()};
} else {
merge_segs[seg->GetPartitionId()].insert(seg->GetID());
}
auto& segs = merge_segs[seg->GetPartitionId()];
if ((segs.size() >= 2) && (RandomInt(0, 10) >= 2)) {
std::cout << "Merging partition " << seg->GetPartitionId() << " segs (";
for (auto seg : segs) {
std::cout << seg << ",";
}
std::cout << ")" << std::endl;
ID_TYPE new_seg_id = 0;
ID_TYPE partition_id = 0;
do_merge(segs, new_seg_id, p_id);
segs.clear();
if (new_seg_id == 0) {
continue;
}
if (RandomInt(0, 10) >= 6) {
build_queue.Put({new_seg_id, p_id});
}
} else {
continue;
}
}
merge_waiter.Notify();
build_queue.Put({0, 0});
build_waiter.Wait();
};
auto build_worker = [&] {
while (true) {
auto ids = build_queue.Take();
auto seg_id = std::get<0>(ids);
auto p_id = std::get<1>(ids);
if (seg_id == 0) {
std::cout << "Exiting Build Worker" << std::endl;
break;
}
std::cout << "Building " << seg_id << std::endl;
do_build(seg_id, p_id);
}
build_waiter.Notify();
};
std::vector<std::thread> handlers;
auto num_handlers = RandomInt(6, 8);
for (auto i = 0; i < num_handlers; ++i) {
handlers.emplace_back(handler_worker);
}
std::thread t3 = std::thread(merge_worker);
std::thread t4 = std::thread(build_worker);
for (auto& handler : handlers) {
handler.join();
}
merge_queue.Put({0, 0});
t3.join();
t4.join();
/* for (auto& kv : merged_segs_history) { */
/* std::cout << "merged: ("; */
/* for (auto i : kv.second) { */
/* std::cout << i << ","; */
/* } */
/* std::cout << ") -> " << kv.first << std::endl; */
/* } */
/* for (auto& id : built_segs) { */
/* std::cout << "built: " << id << std::endl; */
/* } */
merge_waiter.Wait();
ScopedSnapshotT latest_ss;
status = Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
ASSERT_TRUE(status.ok());
auto expect_segments = all_segments;
for (auto& kv : merged_segs_history) {
expect_segments.insert(kv.first);
for (auto& id : kv.second) {
expect_segments.erase(id);
}
}
for (auto& seg_p : seg_p_map) {
auto it = stale_partitions.find(seg_p.second);
if (it == stale_partitions.end()) {
continue;
}
expect_segments.erase(seg_p.first);
}
std::set<ID_TYPE> final_segments;
auto segments = latest_ss->GetResources<Segment>();
for (auto& kv : segments) {
final_segments.insert(kv.first);
}
ASSERT_EQ(final_segments, expect_segments);
// TODO: Check Total Segment Files Cnt
}
struct GCSchedule {
static constexpr const char* Name = "GCSchedule";
};
struct FlushSchedule {
static constexpr const char* Name = "FlushSchedule";
};
using IEventHandler = milvus::engine::snapshot::IEventHandler;
/* struct SampleHandler : public IEventHandler { */
/* static constexpr const char* EventName = "SampleHandler"; */
/* const char* */
/* GetEventName() const override { */
/* return EventName; */
/* } */
/* }; */
REGISTER_HANDLER(GCSchedule, IEventHandler);
/* REGISTER_HANDLER(GCSchedule, SampleHandler); */
REGISTER_HANDLER(FlushSchedule, IEventHandler);
/* REGISTER_HANDLER(FlushSchedule, SampleHandler); */
using GCScheduleFactory = milvus::engine::snapshot::HandlerFactory<GCSchedule>;
using FlushScheduleFactory = milvus::engine::snapshot::HandlerFactory<GCSchedule>;
TEST_F(SnapshotTest, RegistryTest) {
{
auto& factory = GCScheduleFactory::GetInstance();
auto ihandler = factory.GetHandler(IEventHandler::EventName);
ASSERT_TRUE(ihandler);
/* auto sihandler = factory.GetHandler(SampleHandler::EventName); */
/* ASSERT_TRUE(sihandler); */
/* ASSERT_EQ(SampleHandler::EventName, sihandler->GetEventName()); */
}
{
/* auto& factory = FlushScheduleFactory::GetInstance(); */
/* auto ihandler = factory.GetHandler(IEventHandler::EventName); */
/* ASSERT_TRUE(ihandler); */
/* auto sihandler = factory.GetHandler(SampleHandler::EventName); */
/* ASSERT_TRUE(sihandler); */
/* ASSERT_EQ(SampleHandler::EventName, sihandler->GetEventName()); */
}
}
Reference in New Issue View Git Blame Copy Permalink