furyhawk/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-07-21 10:15:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

enhance:refactor flush source (#50300)

Browse Source 
issue: #50425

Signed-off-by: luzhang <luzhang@zilliz.com>
Co-authored-by: luzhang <luzhang@zilliz.com>
This commit is contained in:
zhagnlu
2026-06-10 10:30:19 +08:00
committed by GitHub
co-authored by luzhang
parent a87bc5844c
commit f756c52a22
105 changed files with 7834 additions and 2727 deletions
Download Patch File Download Diff File Expand all files Collapse all files
configs/milvus.yaml
+1
View File
@@ -994,6 +994,7 @@ common:
enabled: true # enable split by average size policy in storage v2
threshold: 1024 # split by average size policy threshold(in bytes) in storage v2
useLoonFFI: false
enableGrowingSourceFlush: true # enable flushing growing segment payload from QueryNode growing source through StorageV3 manifest path
# Default value: auto
# Valid values: [auto, avx512, avx2, avx, sse4_2]
# This configuration is only used by querynode and indexnode, it selects CPU instruction set for Searching and Index-building.
internal/core/src/common/FieldData.cpp
+16 -10
View File
@@ -312,17 +312,23 @@ FieldDataImpl<Type, is_type_entire_row>::FillFieldData(
std::dynamic_pointer_cast<arrow::BinaryArray>(array);
AssertInfo(geometry_array != nullptr,
"null geometry arrow binary array");
std::vector<uint8_t> values(element_count);
for (size_t index = 0; index < element_count; ++index) {
values[index] = *geometry_array->GetValue(index, 0);
if constexpr (std::is_same_v<Type, std::string>) {
std::vector<std::string> values(element_count);
for (size_t index = 0; index < element_count; ++index) {
auto sv = geometry_array->GetView(index);
values[index].assign(sv.data(), sv.size());
}
if (nullable_) {
return FillFieldData(values.data(),
array->null_bitmap_data(),
element_count,
array->offset());
}
return FillFieldData(values.data(), element_count);
}
if (nullable_) {
return FillFieldData(values.data(),
array->null_bitmap_data(),
element_count,
array->offset());
}
return FillFieldData(values.data(), element_count);
ThrowInfo(DataTypeInvalid,
"GEOMETRY arrow data must be filled into string-backed "
"FieldData");
}
case DataType::ARRAY: {
auto array_array =
internal/core/src/segcore/ConcurrentVector.cpp
+1
View File
@@ -136,6 +136,7 @@ VectorBase::set_data_raw(ssize_t element_offset,
FIELD_DATA(data, timestamptz).data(),
element_count);
}
case DataType::STRING:
case DataType::VARCHAR:
case DataType::TEXT: {
auto& field_data = FIELD_DATA(data, string);
internal/core/src/segcore/InsertRecord.h
+1
View File
@@ -1650,6 +1650,7 @@ class InsertRecordGrowing {
field_id, size_per_chunk, scalar_mmap_descriptor);
return;
}
case DataType::STRING:
case DataType::VARCHAR:
case DataType::TEXT: {
this->append_data<std::string>(
internal/core/src/segcore/SegcoreConfig.h
+11
View File
@@ -72,6 +72,16 @@ class SegcoreConfig {
return enable_interim_segment_index_;
}
void
set_enable_growing_source_flush(bool enable_growing_source_flush) {
this->enable_growing_source_flush_ = enable_growing_source_flush;
}
bool
get_enable_growing_source_flush() const {
return enable_growing_source_flush_;
}
void
set_sub_dim(int64_t sub_dim) {
sub_dim_ = sub_dim;
@@ -205,6 +215,7 @@ class SegcoreConfig {
knowhere::IndexEnum::INDEX_FAISS_SCANN_DVR,
};
inline static bool enable_interim_segment_index_ = false;
inline static bool enable_growing_source_flush_ = false;
inline static int64_t chunk_rows_ = 32 * 1024;
inline static int64_t nlist_ = 100;
inline static int64_t nprobe_ = 4;
internal/core/src/segcore/SegmentGrowingImpl.cpp
+63 -35
View File
@@ -259,6 +259,17 @@ ExtractArrayLengths(const proto::schema::FieldData& field_data,
}
}
bool
SchemaHasTextField(const Schema& schema) {
for ([[maybe_unused]] const auto& [field_id, field_meta] :
schema.get_fields()) {
if (field_meta.get_data_type() == DataType::TEXT) {
return true;
}
}
return false;
}
} // anonymous namespace
void
@@ -312,6 +323,16 @@ SegmentGrowingImpl::mask_with_delete(BitsetTypeView& bitset,
void
SegmentGrowingImpl::try_remove_chunks(FieldId fieldId) {
if (SchemaHasTextField(*schema_) ||
segcore_config_.get_enable_growing_source_flush()) {
// Growing-source flush needs raw field chunks to persist a growing
// segment through milvus-storage. Interim indexes may also contain raw
// vector data, but FlushGrowingSegmentData reads directly from
// insert_record_ chunks. Keep raw chunks until the flush path can
// reliably export from indexes too.
return;
}
//remove the chunk data to reduce memory consumption
auto& field_meta = schema_->operator[](fieldId);
auto data_type = field_meta.get_data_type();
@@ -627,36 +648,36 @@ SegmentGrowingImpl::Insert(int64_t reserved_offset,
&insert_record_proto->fields_data(data_offset),
field_meta);
}
if (!indexing_record_.HasRawData(field_id)) {
// special handling for TEXT fields with spillover
if (field_meta.get_data_type() == DataType::TEXT) {
auto spillover = GetTextLobSpillover(field_id);
AssertInfo(spillover != nullptr,
"TEXT field must have spillover");
const auto& field_data =
insert_record_proto->fields_data(data_offset);
const auto& string_data = field_data.scalars().string_data();
// Growing-source flush reads raw data from insert_record_. Keep it
// populated even when the interim index can also serve raw vector data.
// Otherwise later inserts after index sync would be visible by row count
// but missing from field chunks during FlushGrowingSegmentData.
if (field_meta.get_data_type() == DataType::TEXT) {
auto spillover = GetTextLobSpillover(field_id);
AssertInfo(spillover != nullptr, "TEXT field must have spillover");
const auto& field_data =
insert_record_proto->fields_data(data_offset);
const auto& string_data = field_data.scalars().string_data();
std::vector<std::string> ref_strings(num_rows);
for (int64_t i = 0; i < num_rows; i++) {
const auto& text = string_data.data(i);
ref_strings[i] = spillover->WriteAndEncode(text);
}
auto* vec_base = insert_record_.get_data_base(field_id);
auto* string_vec =
dynamic_cast<ConcurrentVector<std::string>*>(vec_base);
AssertInfo(string_vec != nullptr,
"TEXT field must use ConcurrentVector<std::string>");
string_vec->set_data_raw(
reserved_offset, ref_strings.data(), num_rows);
} else {
insert_record_.get_data_base(field_id)->set_data_raw(
reserved_offset,
num_rows,
&insert_record_proto->fields_data(data_offset),
field_meta);
std::vector<std::string> ref_strings(num_rows);
for (int64_t i = 0; i < num_rows; i++) {
const auto& text = string_data.data(i);
ref_strings[i] = spillover->WriteAndEncode(text);
}
auto* vec_base = insert_record_.get_data_base(field_id);
auto* string_vec =
dynamic_cast<ConcurrentVector<std::string>*>(vec_base);
AssertInfo(string_vec != nullptr,
"TEXT field must use ConcurrentVector<std::string>");
string_vec->set_data_raw(
reserved_offset, ref_strings.data(), num_rows);
} else {
insert_record_.get_data_base(field_id)->set_data_raw(
reserved_offset,
num_rows,
&insert_record_proto->fields_data(data_offset),
field_meta);
}
//insert vector data into index
@@ -762,8 +783,12 @@ SegmentGrowingImpl::LoadFieldData(const LoadFieldDataInfo& infos,
milvus::OpContext* op_ctx) {
// Note: op_ctx is currently unused in growing segments but kept for interface consistency
(void)op_ctx;
AssertInfo(
!(infos.storage_version == STORAGE_V2 && SchemaHasTextField(*schema_)),
"TEXT growing segment cannot be loaded from StorageV2 binlogs; "
"StorageV3 manifest is required");
switch (infos.storage_version) {
case 2:
case STORAGE_V2:
load_column_group_data_internal(infos);
break;
default:
@@ -882,13 +907,13 @@ SegmentGrowingImpl::load_field_data_common(
auto field_meta = (*schema_)[field_id];
if (!indexing_record_.HasRawData(field_id)) {
if (insert_record_.is_valid_data_exist(field_id)) {
insert_record_.get_valid_data(field_id)->set_data_raw(field_data);
}
insert_record_.get_data_base(field_id)->set_data_raw(reserved_offset,
field_data);
// Growing-source flush reads raw data from insert_record_. Keep it
// populated even when an interim index can provide raw vector data.
if (insert_record_.is_valid_data_exist(field_id)) {
insert_record_.get_valid_data(field_id)->set_data_raw(field_data);
}
insert_record_.get_data_base(field_id)->set_data_raw(reserved_offset,
field_data);
if (segcore_config_.get_enable_interim_segment_index()) {
auto offset = reserved_offset;
for (auto& data : field_data) {
@@ -2258,6 +2283,9 @@ SegmentGrowingImpl::Load(milvus::tracer::TraceContext& trace_ctx,
field_data_info.load_priority = load_info_.priority();
auto manifest_path = load_info_.manifest_path();
AssertInfo(!(manifest_path.empty() && SchemaHasTextField(*schema_)),
"TEXT growing segment cannot be loaded without a StorageV3 "
"manifest");
if (manifest_path != "") {
LoadColumnsGroups(manifest_path);
return;
internal/core/src/segcore/Utils.cpp
+2
View File
@@ -648,6 +648,7 @@ CreateScalarDataArrayFrom(const void* data_raw,
obj->mutable_data()->Add(data, data + count);
break;
}
case DataType::STRING:
case DataType::VARCHAR:
case DataType::TEXT: {
auto data = reinterpret_cast<const std::string*>(data_raw);
@@ -1025,6 +1026,7 @@ MergeDataArray(std::vector<MergeBase>& merge_bases,
*(obj->mutable_data()->Add()) = data[src_offset];
break;
}
case DataType::STRING:
case DataType::VARCHAR:
case DataType::TEXT: {
auto* mutable_src = src_field_data->mutable_scalars()
internal/core/src/segcore/flush_growing_segment_test.cpp
+894
View File
@@ -451,6 +451,756 @@ TEST_F(FlushGrowingSegmentTest, FlushWithNullableFields) {
FreeFlushResult(&result);
}
TEST_F(FlushGrowingSegmentTest, FlushOrdinaryFieldSemanticsRoundTrip) {
auto schema = std::make_shared<Schema>();
auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
DefaultValueType default_value;
default_value.set_int_data(42);
auto default_fid = schema->AddDebugFieldWithDefaultValue(
"default_i32", DataType::INT32, default_value);
auto nullable_str_fid =
schema->AddDebugField("nullable_str", DataType::VARCHAR, true);
auto json_fid = schema->AddDebugField("json", DataType::JSON, true);
auto array_fid =
schema->AddDebugField("array", DataType::ARRAY, DataType::INT64);
auto geometry_fid =
schema->AddDebugField("geometry", DataType::GEOMETRY, true);
schema->set_primary_field_id(pk_fid);
auto segment = CreateGrowingSegment(schema, empty_index_meta);
ASSERT_NE(segment, nullptr);
constexpr int N = 3;
std::vector<int64_t> row_ids = {10, 11, 12};
std::vector<Timestamp> timestamps = {100, 101, 102};
std::vector<int64_t> pks = {1000, 1001, 1002};
std::vector<std::string> nullable_strings = {"alpha", "", "gamma"};
bool nullable_string_valid[N] = {true, false, true};
std::vector<std::string> json_values = {
R"({"k":1,"nested":{"v":"a"}})", "", R"({"k":3,"array":[1,2]})"};
bool json_valid[N] = {true, false, true};
std::vector<ScalarFieldProto> array_values(N);
array_values[0].mutable_long_data()->add_data(1);
array_values[0].mutable_long_data()->add_data(2);
array_values[1].mutable_long_data()->add_data(3);
array_values[2].mutable_long_data()->add_data(4);
array_values[2].mutable_long_data()->add_data(5);
array_values[2].mutable_long_data()->add_data(6);
auto geos_ctx = GEOS_init_r();
std::vector<std::string> geometry_values;
geometry_values.emplace_back(
Geometry(geos_ctx, "POINT (1 2)").to_wkb_string());
geometry_values.emplace_back("");
geometry_values.emplace_back(
Geometry(geos_ctx, "POINT (3 4)").to_wkb_string());
GEOS_finish_r(geos_ctx);
bool geometry_valid[N] = {true, false, true};
auto insert_data = std::make_unique<InsertRecordProto>();
insert_data->set_num_rows(N);
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(pks.data(), nullptr, N, (*schema)[pk_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(nullable_strings.data(),
nullable_string_valid,
N,
(*schema)[nullable_str_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
json_values.data(), json_valid, N, (*schema)[json_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
array_values.data(), nullptr, N, (*schema)[array_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
geometry_values.data(), geometry_valid, N, (*schema)[geometry_fid])
.release());
// default_i32 is intentionally absent from insert_data. Growing insert must
// fill it exactly as the DataNode write-buffer path does.
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
std::string segment_path = test_dir_ + "/segment_semantics";
config.segment_path = segment_path.c_str();
config.read_version = -1;
config.retry_limit = 3;
config.text_field_ids = nullptr;
config.text_lob_paths = nullptr;
config.num_text_columns = 0;
CFlushResult result;
auto status =
FlushGrowingSegmentData(segment.get(), 0, N, &config, &result);
ASSERT_EQ(status.error_code, Success) << status.error_msg;
ASSERT_EQ(result.num_rows, N);
auto default_datas = ReadFlushedFieldData(
segment_path, result, default_fid, DataType::INT32, true, 0);
ASSERT_EQ(default_datas.size(), 1);
ASSERT_EQ(default_datas[0]->get_num_rows(), N);
ASSERT_EQ(default_datas[0]->get_valid_rows(), N);
for (int i = 0; i < N; i++) {
ASSERT_TRUE(default_datas[0]->is_valid(i));
EXPECT_EQ(*static_cast<const int32_t*>(default_datas[0]->RawValue(i)),
42);
}
auto string_datas = ReadFlushedFieldData(
segment_path, result, nullable_str_fid, DataType::VARCHAR, true, 0);
ASSERT_EQ(string_datas.size(), 1);
ASSERT_TRUE(string_datas[0]->is_valid(0));
EXPECT_FALSE(string_datas[0]->is_valid(1));
ASSERT_TRUE(string_datas[0]->is_valid(2));
EXPECT_EQ(*static_cast<const std::string*>(string_datas[0]->RawValue(0)),
"alpha");
EXPECT_EQ(*static_cast<const std::string*>(string_datas[0]->RawValue(2)),
"gamma");
auto json_datas = ReadFlushedFieldData(
segment_path, result, json_fid, DataType::JSON, true, 0);
ASSERT_EQ(json_datas.size(), 1);
ASSERT_TRUE(json_datas[0]->is_valid(0));
EXPECT_FALSE(json_datas[0]->is_valid(1));
ASSERT_TRUE(json_datas[0]->is_valid(2));
EXPECT_EQ(std::string(
static_cast<const Json*>(json_datas[0]->RawValue(0))->data()),
json_values[0]);
EXPECT_EQ(std::string(
static_cast<const Json*>(json_datas[0]->RawValue(2))->data()),
json_values[2]);
auto array_datas = ReadFlushedFieldData(segment_path,
result,
array_fid,
DataType::ARRAY,
false,
0,
DataType::INT64);
ASSERT_EQ(array_datas.size(), 1);
auto row0_array =
static_cast<const Array*>(array_datas[0]->RawValue(0))->output_data();
auto row2_array =
static_cast<const Array*>(array_datas[0]->RawValue(2))->output_data();
EXPECT_EQ(row0_array.long_data().data().size(), 2);
EXPECT_EQ(row0_array.long_data().data(0), 1);
EXPECT_EQ(row0_array.long_data().data(1), 2);
EXPECT_EQ(row2_array.long_data().data().size(), 3);
EXPECT_EQ(row2_array.long_data().data(0), 4);
EXPECT_EQ(row2_array.long_data().data(2), 6);
auto geometry_datas = ReadFlushedFieldData(
segment_path, result, geometry_fid, DataType::GEOMETRY, true, 0);
ASSERT_EQ(geometry_datas.size(), 1);
ASSERT_TRUE(geometry_datas[0]->is_valid(0));
EXPECT_FALSE(geometry_datas[0]->is_valid(1));
ASSERT_TRUE(geometry_datas[0]->is_valid(2));
EXPECT_EQ(*static_cast<const std::string*>(geometry_datas[0]->RawValue(0)),
geometry_values[0]);
EXPECT_EQ(*static_cast<const std::string*>(geometry_datas[0]->RawValue(2)),
geometry_values[2]);
FreeFlushResult(&result);
}
TEST_F(FlushGrowingSegmentTest, FlushTimestamptzPartialRangeRoundTrip) {
auto schema = std::make_shared<Schema>();
auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
auto timestamptz_fid =
schema->AddDebugField("event_time", DataType::TIMESTAMPTZ, true);
schema->set_primary_field_id(pk_fid);
auto segment = CreateGrowingSegment(schema, empty_index_meta);
ASSERT_NE(segment, nullptr);
constexpr int N = 5;
std::vector<int64_t> row_ids = {0, 1, 2, 3, 4};
std::vector<Timestamp> timestamps = {100, 101, 102, 103, 104};
std::vector<int64_t> pks = {10, 11, 12, 13, 14};
std::vector<int64_t> event_times = {1700000000000,
1700000001000,
1700000002000,
1700000003000,
1700000004000};
bool event_time_valid[N] = {true, false, true, false, true};
auto insert_data = std::make_unique<InsertRecordProto>();
insert_data->set_num_rows(N);
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(pks.data(), nullptr, N, (*schema)[pk_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
event_times.data(), event_time_valid, N, (*schema)[timestamptz_fid])
.release());
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
std::string segment_path = test_dir_ + "/segment_timestamptz";
config.segment_path = segment_path.c_str();
config.read_version = -1;
config.retry_limit = 3;
config.text_field_ids = nullptr;
config.text_lob_paths = nullptr;
config.num_text_columns = 0;
CFlushResult result;
constexpr int64_t start = 1;
constexpr int64_t end = 4;
auto status =
FlushGrowingSegmentData(segment.get(), start, end, &config, &result);
ASSERT_EQ(status.error_code, Success) << status.error_msg;
ASSERT_EQ(result.num_rows, end - start);
auto field_datas = ReadFlushedFieldData(
segment_path, result, timestamptz_fid, DataType::TIMESTAMPTZ, true, 0);
ASSERT_EQ(field_datas.size(), 1);
ASSERT_EQ(field_datas[0]->get_num_rows(), end - start);
EXPECT_FALSE(field_datas[0]->is_valid(0));
ASSERT_TRUE(field_datas[0]->is_valid(1));
EXPECT_EQ(*static_cast<const int64_t*>(field_datas[0]->RawValue(1)),
event_times[2]);
EXPECT_FALSE(field_datas[0]->is_valid(2));
FreeFlushResult(&result);
}
TEST_F(FlushGrowingSegmentTest, FlushNullableScalarTypesPartialRangeRoundTrip) {
auto schema = std::make_shared<Schema>();
auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
auto bool_fid = schema->AddDebugField("bool_field", DataType::BOOL, true);
auto int8_fid = schema->AddDebugField("int8_field", DataType::INT8, true);
auto int16_fid =
schema->AddDebugField("int16_field", DataType::INT16, true);
auto int64_fid =
schema->AddDebugField("int64_field", DataType::INT64, true);
auto float_fid =
schema->AddDebugField("float_field", DataType::FLOAT, true);
auto double_fid =
schema->AddDebugField("double_field", DataType::DOUBLE, true);
schema->set_primary_field_id(pk_fid);
auto segment = CreateGrowingSegment(schema, empty_index_meta);
ASSERT_NE(segment, nullptr);
constexpr int N = 5;
std::vector<int64_t> row_ids = {0, 1, 2, 3, 4};
std::vector<Timestamp> timestamps = {100, 101, 102, 103, 104};
std::vector<int64_t> pks = {10, 11, 12, 13, 14};
bool bool_values[N] = {true, false, true, false, true};
std::vector<int8_t> int8_values = {-2, -1, 0, 1, 2};
std::vector<int16_t> int16_values = {-200, -100, 0, 100, 200};
std::vector<int64_t> int64_values = {1000, 1001, 1002, 1003, 1004};
std::vector<float> float_values = {1.25F, 2.5F, 3.75F, 4.5F, 5.25F};
std::vector<double> double_values = {10.25, 20.5, 30.75, 40.5, 50.25};
bool valid_data[N] = {true, false, true, false, true};
auto insert_data = std::make_unique<InsertRecordProto>();
insert_data->set_num_rows(N);
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(pks.data(), nullptr, N, (*schema)[pk_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(bool_values, valid_data, N, (*schema)[bool_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
int8_values.data(), valid_data, N, (*schema)[int8_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
int16_values.data(), valid_data, N, (*schema)[int16_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
int64_values.data(), valid_data, N, (*schema)[int64_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
float_values.data(), valid_data, N, (*schema)[float_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
double_values.data(), valid_data, N, (*schema)[double_fid])
.release());
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
std::string segment_path = test_dir_ + "/segment_nullable_scalars";
config.segment_path = segment_path.c_str();
config.read_version = -1;
config.retry_limit = 3;
config.text_field_ids = nullptr;
config.text_lob_paths = nullptr;
config.num_text_columns = 0;
CFlushResult result;
constexpr int64_t start = 1;
constexpr int64_t end = 5;
auto status =
FlushGrowingSegmentData(segment.get(), start, end, &config, &result);
ASSERT_EQ(status.error_code, Success) << status.error_msg;
ASSERT_EQ(result.num_rows, end - start);
auto bool_datas = ReadFlushedFieldData(
segment_path, result, bool_fid, DataType::BOOL, true, 0);
ASSERT_EQ(bool_datas.size(), 1);
EXPECT_FALSE(bool_datas[0]->is_valid(0));
ASSERT_TRUE(bool_datas[0]->is_valid(1));
EXPECT_EQ(*static_cast<const bool*>(bool_datas[0]->RawValue(1)),
bool_values[2]);
EXPECT_FALSE(bool_datas[0]->is_valid(2));
ASSERT_TRUE(bool_datas[0]->is_valid(3));
EXPECT_EQ(*static_cast<const bool*>(bool_datas[0]->RawValue(3)),
bool_values[4]);
auto int8_datas = ReadFlushedFieldData(
segment_path, result, int8_fid, DataType::INT8, true, 0);
ASSERT_EQ(int8_datas.size(), 1);
ASSERT_TRUE(int8_datas[0]->is_valid(1));
EXPECT_EQ(*static_cast<const int8_t*>(int8_datas[0]->RawValue(1)),
int8_values[2]);
ASSERT_TRUE(int8_datas[0]->is_valid(3));
EXPECT_EQ(*static_cast<const int8_t*>(int8_datas[0]->RawValue(3)),
int8_values[4]);
auto int16_datas = ReadFlushedFieldData(
segment_path, result, int16_fid, DataType::INT16, true, 0);
ASSERT_EQ(int16_datas.size(), 1);
ASSERT_TRUE(int16_datas[0]->is_valid(1));
EXPECT_EQ(*static_cast<const int16_t*>(int16_datas[0]->RawValue(1)),
int16_values[2]);
ASSERT_TRUE(int16_datas[0]->is_valid(3));
EXPECT_EQ(*static_cast<const int16_t*>(int16_datas[0]->RawValue(3)),
int16_values[4]);
auto int64_datas = ReadFlushedFieldData(
segment_path, result, int64_fid, DataType::INT64, true, 0);
ASSERT_EQ(int64_datas.size(), 1);
ASSERT_TRUE(int64_datas[0]->is_valid(1));
EXPECT_EQ(*static_cast<const int64_t*>(int64_datas[0]->RawValue(1)),
int64_values[2]);
ASSERT_TRUE(int64_datas[0]->is_valid(3));
EXPECT_EQ(*static_cast<const int64_t*>(int64_datas[0]->RawValue(3)),
int64_values[4]);
auto float_datas = ReadFlushedFieldData(
segment_path, result, float_fid, DataType::FLOAT, true, 0);
ASSERT_EQ(float_datas.size(), 1);
ASSERT_TRUE(float_datas[0]->is_valid(1));
EXPECT_FLOAT_EQ(*static_cast<const float*>(float_datas[0]->RawValue(1)),
float_values[2]);
ASSERT_TRUE(float_datas[0]->is_valid(3));
EXPECT_FLOAT_EQ(*static_cast<const float*>(float_datas[0]->RawValue(3)),
float_values[4]);
auto double_datas = ReadFlushedFieldData(
segment_path, result, double_fid, DataType::DOUBLE, true, 0);
ASSERT_EQ(double_datas.size(), 1);
ASSERT_TRUE(double_datas[0]->is_valid(1));
EXPECT_DOUBLE_EQ(*static_cast<const double*>(double_datas[0]->RawValue(1)),
double_values[2]);
ASSERT_TRUE(double_datas[0]->is_valid(3));
EXPECT_DOUBLE_EQ(*static_cast<const double*>(double_datas[0]->RawValue(3)),
double_values[4]);
FreeFlushResult(&result);
}
TEST_F(FlushGrowingSegmentTest, FlushVectorArrayRoundTrip) {
auto schema = std::make_shared<Schema>();
auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
auto vec_array_fid = schema->AddDebugVectorArrayField(
"emb_list", DataType::VECTOR_FLOAT, 4, knowhere::metric::L2);
schema->set_primary_field_id(pk_fid);
auto segment = CreateGrowingSegment(schema, empty_index_meta);
ASSERT_NE(segment, nullptr);
constexpr int N = 4;
constexpr int array_len = 3;
auto dataset = DataGen(schema, N, 42, 0, 1, array_len);
segment->PreInsert(N);
segment->Insert(0,
N,
dataset.row_ids_.data(),
dataset.timestamps_.data(),
dataset.raw_);
CFlushConfig config{};
std::string segment_path = test_dir_ + "/segment_vector_array";
config.segment_path = segment_path.c_str();
config.read_version = -1;
config.retry_limit = 3;
config.text_field_ids = nullptr;
config.text_lob_paths = nullptr;
config.num_text_columns = 0;
CFlushResult result;
auto status =
FlushGrowingSegmentData(segment.get(), 0, N, &config, &result);
ASSERT_EQ(status.error_code, Success) << status.error_msg;
ASSERT_EQ(result.num_rows, N);
auto field_datas = ReadFlushedFieldData(segment_path,
result,
vec_array_fid,
DataType::VECTOR_ARRAY,
false,
4,
DataType::VECTOR_FLOAT);
ASSERT_EQ(field_datas.size(), 1);
ASSERT_EQ(field_datas[0]->get_num_rows(), N);
auto expected = dataset.get_col<VectorFieldProto>(vec_array_fid);
for (int i = 0; i < N; i++) {
auto actual = static_cast<const milvus::VectorArray*>(
field_datas[0]->RawValue(i));
ASSERT_NE(actual, nullptr);
EXPECT_EQ(actual->length(), array_len);
EXPECT_EQ(actual->output_data().SerializeAsString(),
expected[i].SerializeAsString());
}
FreeFlushResult(&result);
}
TEST_F(FlushGrowingSegmentTest, FlushVectorArrayElementTypesRoundTrip) {
struct Case {
DataType element_type;
int64_t dim;
std::string metric_type;
std::string segment_suffix;
};
std::vector<Case> cases = {
{DataType::VECTOR_BINARY, 16, "JACCARD", "binary"},
{DataType::VECTOR_FLOAT16, 4, "L2", "fp16"},
{DataType::VECTOR_BFLOAT16, 4, "L2", "bf16"},
{DataType::VECTOR_INT8, 4, "L2", "int8"},
};
for (const auto& test_case : cases) {
auto schema = std::make_shared<Schema>();
auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
auto vec_array_fid =
schema->AddDebugVectorArrayField("emb_list",
test_case.element_type,
test_case.dim,
test_case.metric_type);
schema->set_primary_field_id(pk_fid);
auto segment = CreateGrowingSegment(schema, empty_index_meta);
ASSERT_NE(segment, nullptr);
constexpr int N = 4;
constexpr int array_len = 3;
auto dataset = DataGen(schema, N, 42, 0, 1, array_len);
segment->PreInsert(N);
segment->Insert(0,
N,
dataset.row_ids_.data(),
dataset.timestamps_.data(),
dataset.raw_);
CFlushConfig config{};
std::string segment_path =
test_dir_ + "/segment_vector_array_" + test_case.segment_suffix;
config.segment_path = segment_path.c_str();
config.read_version = -1;
config.retry_limit = 3;
config.text_field_ids = nullptr;
config.text_lob_paths = nullptr;
config.num_text_columns = 0;
CFlushResult result;
auto status =
FlushGrowingSegmentData(segment.get(), 0, N, &config, &result);
ASSERT_EQ(status.error_code, Success) << status.error_msg;
ASSERT_EQ(result.num_rows, N);
auto field_datas = ReadFlushedFieldData(segment_path,
result,
vec_array_fid,
DataType::VECTOR_ARRAY,
false,
test_case.dim,
test_case.element_type);
ASSERT_EQ(field_datas.size(), 1);
ASSERT_EQ(field_datas[0]->get_num_rows(), N);
auto expected = dataset.get_col<VectorFieldProto>(vec_array_fid);
for (int i = 0; i < N; i++) {
auto actual = static_cast<const milvus::VectorArray*>(
field_datas[0]->RawValue(i));
ASSERT_NE(actual, nullptr);
EXPECT_EQ(actual->length(), array_len);
EXPECT_EQ(actual->get_element_type(), test_case.element_type);
EXPECT_EQ(actual->output_data().SerializeAsString(),
expected[i].SerializeAsString());
}
FreeFlushResult(&result);
}
}
TEST_F(FlushGrowingSegmentTest, FlushStringAndTextRoundTrip) {
auto schema = std::make_shared<Schema>();
auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
auto string_fid = schema->AddDebugField("str", DataType::STRING, true);
auto text_fid = schema->AddDebugField("text", DataType::TEXT, true);
schema->set_primary_field_id(pk_fid);
auto segment = CreateGrowingSegment(schema, empty_index_meta);
ASSERT_NE(segment, nullptr);
constexpr int N = 3;
std::vector<int64_t> row_ids = {0, 1, 2};
std::vector<Timestamp> timestamps = {10, 11, 12};
std::vector<int64_t> pks = {100, 101, 102};
std::vector<std::string> strings = {"plain-string", "", "tail-string"};
bool string_valid[N] = {true, false, true};
std::vector<std::string> texts = {
"short text", "", std::string(512, 'x') + "-tail"};
bool text_valid[N] = {true, false, true};
auto insert_data = std::make_unique<InsertRecordProto>();
insert_data->set_num_rows(N);
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(pks.data(), nullptr, N, (*schema)[pk_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
strings.data(), string_valid, N, (*schema)[string_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(texts.data(), text_valid, N, (*schema)[text_fid])
.release());
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
std::string segment_path = test_dir_ + "/segment_string_text";
config.segment_path = segment_path.c_str();
config.read_version = -1;
config.retry_limit = 3;
config.text_field_ids = nullptr;
config.text_lob_paths = nullptr;
config.num_text_columns = 0;
CFlushResult result;
auto status =
FlushGrowingSegmentData(segment.get(), 0, N, &config, &result);
ASSERT_EQ(status.error_code, Success) << status.error_msg;
ASSERT_EQ(result.num_rows, N);
auto string_datas = ReadFlushedFieldData(
segment_path, result, string_fid, DataType::STRING, true, 0);
ASSERT_EQ(string_datas.size(), 1);
ASSERT_TRUE(string_datas[0]->is_valid(0));
EXPECT_FALSE(string_datas[0]->is_valid(1));
ASSERT_TRUE(string_datas[0]->is_valid(2));
EXPECT_EQ(*static_cast<const std::string*>(string_datas[0]->RawValue(0)),
strings[0]);
EXPECT_EQ(*static_cast<const std::string*>(string_datas[0]->RawValue(2)),
strings[2]);
auto text_datas = ReadFlushedFieldData(
segment_path, result, text_fid, DataType::TEXT, true, 0);
ASSERT_EQ(text_datas.size(), 1);
ASSERT_TRUE(text_datas[0]->is_valid(0));
EXPECT_FALSE(text_datas[0]->is_valid(1));
ASSERT_TRUE(text_datas[0]->is_valid(2));
EXPECT_EQ(*static_cast<const std::string*>(text_datas[0]->RawValue(0)),
texts[0]);
EXPECT_EQ(*static_cast<const std::string*>(text_datas[0]->RawValue(2)),
texts[2]);
FreeFlushResult(&result);
}
TEST_F(FlushGrowingSegmentTest, FlushArrayElementTypesRoundTrip) {
auto schema = std::make_shared<Schema>();
auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
auto bool_array_fid =
schema->AddDebugField("bool_array", DataType::ARRAY, DataType::BOOL);
auto int_array_fid =
schema->AddDebugField("int_array", DataType::ARRAY, DataType::INT32);
auto string_array_fid = schema->AddDebugField(
"string_array", DataType::ARRAY, DataType::VARCHAR);
auto double_array_fid = schema->AddDebugField(
"double_array", DataType::ARRAY, DataType::DOUBLE, true);
schema->set_primary_field_id(pk_fid);
auto segment = CreateGrowingSegment(schema, empty_index_meta);
ASSERT_NE(segment, nullptr);
constexpr int N = 3;
std::vector<int64_t> row_ids = {0, 1, 2};
std::vector<Timestamp> timestamps = {10, 11, 12};
std::vector<int64_t> pks = {100, 101, 102};
std::vector<ScalarFieldProto> bool_arrays(N);
bool_arrays[0].mutable_bool_data()->add_data(true);
bool_arrays[0].mutable_bool_data()->add_data(false);
bool_arrays[1].mutable_bool_data()->add_data(false);
bool_arrays[2].mutable_bool_data()->add_data(true);
bool_arrays[2].mutable_bool_data()->add_data(true);
std::vector<ScalarFieldProto> int_arrays(N);
int_arrays[0].mutable_int_data()->add_data(1);
int_arrays[0].mutable_int_data()->add_data(2);
int_arrays[1].mutable_int_data()->add_data(3);
int_arrays[2].mutable_int_data()->add_data(4);
int_arrays[2].mutable_int_data()->add_data(5);
std::vector<ScalarFieldProto> string_arrays(N);
string_arrays[0].mutable_string_data()->add_data("a");
string_arrays[0].mutable_string_data()->add_data("b");
string_arrays[1].mutable_string_data()->add_data("c");
string_arrays[2].mutable_string_data()->add_data("d");
string_arrays[2].mutable_string_data()->add_data("e");
std::vector<ScalarFieldProto> double_arrays(N);
double_arrays[0].mutable_double_data()->add_data(1.25);
double_arrays[0].mutable_double_data()->add_data(2.5);
double_arrays[1].mutable_double_data()->add_data(3.75);
double_arrays[2].mutable_double_data()->add_data(4.5);
bool double_valid[N] = {true, false, true};
auto insert_data = std::make_unique<InsertRecordProto>();
insert_data->set_num_rows(N);
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(pks.data(), nullptr, N, (*schema)[pk_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
bool_arrays.data(), nullptr, N, (*schema)[bool_array_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
int_arrays.data(), nullptr, N, (*schema)[int_array_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
string_arrays.data(), nullptr, N, (*schema)[string_array_fid])
.release());
insert_data->mutable_fields_data()->AddAllocated(
CreateDataArrayFrom(
double_arrays.data(), double_valid, N, (*schema)[double_array_fid])
.release());
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
std::string segment_path = test_dir_ + "/segment_arrays";
config.segment_path = segment_path.c_str();
config.read_version = -1;
config.retry_limit = 3;
config.text_field_ids = nullptr;
config.text_lob_paths = nullptr;
config.num_text_columns = 0;
CFlushResult result;
auto status =
FlushGrowingSegmentData(segment.get(), 0, N, &config, &result);
ASSERT_EQ(status.error_code, Success) << status.error_msg;
ASSERT_EQ(result.num_rows, N);
auto bool_datas = ReadFlushedFieldData(segment_path,
result,
bool_array_fid,
DataType::ARRAY,
false,
0,
DataType::BOOL);
ASSERT_EQ(bool_datas.size(), 1);
EXPECT_EQ(static_cast<const Array*>(bool_datas[0]->RawValue(0))
->output_data()
.SerializeAsString(),
bool_arrays[0].SerializeAsString());
EXPECT_EQ(static_cast<const Array*>(bool_datas[0]->RawValue(2))
->output_data()
.SerializeAsString(),
bool_arrays[2].SerializeAsString());
auto int_datas = ReadFlushedFieldData(segment_path,
result,
int_array_fid,
DataType::ARRAY,
false,
0,
DataType::INT32);
ASSERT_EQ(int_datas.size(), 1);
EXPECT_EQ(static_cast<const Array*>(int_datas[0]->RawValue(0))
->output_data()
.SerializeAsString(),
int_arrays[0].SerializeAsString());
EXPECT_EQ(static_cast<const Array*>(int_datas[0]->RawValue(2))
->output_data()
.SerializeAsString(),
int_arrays[2].SerializeAsString());
auto string_datas = ReadFlushedFieldData(segment_path,
result,
string_array_fid,
DataType::ARRAY,
false,
0,
DataType::VARCHAR);
ASSERT_EQ(string_datas.size(), 1);
EXPECT_EQ(static_cast<const Array*>(string_datas[0]->RawValue(0))
->output_data()
.SerializeAsString(),
string_arrays[0].SerializeAsString());
EXPECT_EQ(static_cast<const Array*>(string_datas[0]->RawValue(2))
->output_data()
.SerializeAsString(),
string_arrays[2].SerializeAsString());
auto double_datas = ReadFlushedFieldData(segment_path,
result,
double_array_fid,
DataType::ARRAY,
true,
0,
DataType::DOUBLE);
ASSERT_EQ(double_datas.size(), 1);
ASSERT_TRUE(double_datas[0]->is_valid(0));
EXPECT_FALSE(double_datas[0]->is_valid(1));
ASSERT_TRUE(double_datas[0]->is_valid(2));
EXPECT_EQ(static_cast<const Array*>(double_datas[0]->RawValue(0))
->output_data()
.SerializeAsString(),
double_arrays[0].SerializeAsString());
EXPECT_EQ(static_cast<const Array*>(double_datas[0]->RawValue(2))
->output_data()
.SerializeAsString(),
double_arrays[2].SerializeAsString());
FreeFlushResult(&result);
}
TEST_F(FlushGrowingSegmentTest, FlushNullableFloatVectorKeepsCompactMapping) {
auto schema = std::make_shared<Schema>();
auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
@@ -586,6 +1336,150 @@ TEST_F(FlushGrowingSegmentTest, FlushNullableInt8VectorKeepsCompactMapping) {
FreeFlushResult(&result);
}
TEST_F(FlushGrowingSegmentTest, FlushNullableFixedWidthVectorTypesRoundTrip) {
struct Case {
DataType data_type;
int64_t dim;
std::string metric_type;
std::string segment_suffix;
};
std::vector<Case> cases = {
{DataType::VECTOR_BINARY, 16, "JACCARD", "binary"},
{DataType::VECTOR_FLOAT16, 2, "L2", "fp16"},
{DataType::VECTOR_BFLOAT16, 2, "L2", "bf16"},
};
for (const auto& test_case : cases) {
auto schema = std::make_shared<Schema>();
auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
auto vec_fid = schema->AddDebugField("vec",
test_case.data_type,
test_case.dim,
test_case.metric_type,
true);
schema->set_primary_field_id(pk_fid);
auto segment = CreateGrowingSegment(schema, empty_index_meta);
ASSERT_NE(segment, nullptr);
constexpr int N = 3;
std::vector<int64_t> row_ids = {0, 1, 2};
std::vector<Timestamp> timestamps = {10, 11, 12};
std::vector<int64_t> pks = {100, 101, 102};
bool valid_data[N] = {true, false, true};
auto insert_data = std::make_unique<InsertRecordProto>();
insert_data->set_num_rows(N);
auto pk_array =
CreateDataArrayFrom(pks.data(), nullptr, N, (*schema)[pk_fid]);
insert_data->mutable_fields_data()->AddAllocated(pk_array.release());
std::vector<uint8_t> binary_vectors = {0x01, 0x02, 0xA0, 0xB0};
std::vector<float16> fp16_vectors = {1, 2, 3, 4};
std::vector<bfloat16> bf16_vectors = {5, 6, 7, 8};
const void* vector_data = nullptr;
switch (test_case.data_type) {
case DataType::VECTOR_BINARY:
vector_data = binary_vectors.data();
break;
case DataType::VECTOR_FLOAT16:
vector_data = fp16_vectors.data();
break;
case DataType::VECTOR_BFLOAT16:
vector_data = bf16_vectors.data();
break;
default:
FAIL() << "unexpected vector type";
}
auto vec_array = CreateVectorDataArrayFrom(
vector_data, valid_data, N, 2, (*schema)[vec_fid]);
insert_data->mutable_fields_data()->AddAllocated(vec_array.release());
segment->PreInsert(N);
segment->Insert(
0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
std::string segment_path =
test_dir_ + "/segment_nullable_" + test_case.segment_suffix;
config.segment_path = segment_path.c_str();
config.read_version = -1;
config.retry_limit = 3;
config.text_field_ids = nullptr;
config.text_lob_paths = nullptr;
config.num_text_columns = 0;
CFlushResult result;
auto status =
FlushGrowingSegmentData(segment.get(), 0, N, &config, &result);
ASSERT_EQ(status.error_code, Success) << status.error_msg;
ASSERT_EQ(result.num_rows, N);
auto field_datas = ReadFlushedFieldData(segment_path,
result,
vec_fid,
test_case.data_type,
true,
test_case.dim);
ASSERT_EQ(field_datas.size(), 1);
auto field_data = field_datas[0];
ASSERT_EQ(field_data->get_num_rows(), N);
ASSERT_EQ(field_data->get_valid_rows(), 2);
ASSERT_TRUE(field_data->is_valid(0));
EXPECT_FALSE(field_data->is_valid(1));
ASSERT_TRUE(field_data->is_valid(2));
switch (test_case.data_type) {
case DataType::VECTOR_BINARY: {
auto row0 =
static_cast<const uint8_t*>(field_data->RawValue(0));
auto row2 =
static_cast<const uint8_t*>(field_data->RawValue(2));
ASSERT_NE(row0, nullptr);
ASSERT_NE(row2, nullptr);
EXPECT_EQ(std::vector<uint8_t>(row0, row0 + 2),
std::vector<uint8_t>({0x01, 0x02}));
EXPECT_EQ(std::vector<uint8_t>(row2, row2 + 2),
std::vector<uint8_t>({0xA0, 0xB0}));
break;
}
case DataType::VECTOR_FLOAT16: {
auto row0 =
static_cast<const float16*>(field_data->RawValue(0));
auto row2 =
static_cast<const float16*>(field_data->RawValue(2));
ASSERT_NE(row0, nullptr);
ASSERT_NE(row2, nullptr);
EXPECT_EQ(std::vector<float16>(row0, row0 + 2),
std::vector<float16>({1, 2}));
EXPECT_EQ(std::vector<float16>(row2, row2 + 2),
std::vector<float16>({3, 4}));
break;
}
case DataType::VECTOR_BFLOAT16: {
auto row0 =
static_cast<const bfloat16*>(field_data->RawValue(0));
auto row2 =
static_cast<const bfloat16*>(field_data->RawValue(2));
ASSERT_NE(row0, nullptr);
ASSERT_NE(row2, nullptr);
EXPECT_EQ(std::vector<bfloat16>(row0, row0 + 2),
std::vector<bfloat16>({5, 6}));
EXPECT_EQ(std::vector<bfloat16>(row2, row2 + 2),
std::vector<bfloat16>({7, 8}));
break;
}
default:
FAIL() << "unexpected vector type";
}
FreeFlushResult(&result);
}
}
TEST_F(FlushGrowingSegmentTest, FlushNullableSparseVectorKeepsCompactMapping) {
auto schema = std::make_shared<Schema>();
auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
internal/core/src/segcore/segcore_init_c.cpp
+7
View File
@@ -50,6 +50,13 @@ SegcoreSetEnableInterminSegmentIndex(const bool value) {
config.set_enable_interim_segment_index(value);
}
extern "C" void
SegcoreSetEnableGrowingSourceFlush(const bool value) {
milvus::segcore::SegcoreConfig& config =
milvus::segcore::SegcoreConfig::default_config();
config.set_enable_growing_source_flush(value);
}
extern "C" void
SegcoreSetEnableGeometryCache(const bool value) {
milvus::segcore::SegcoreConfig& config =
internal/core/src/segcore/segcore_init_c.h
+3
View File
@@ -29,6 +29,9 @@ SegcoreSetChunkRows(const int64_t);
void
SegcoreSetEnableInterminSegmentIndex(const bool);
void
SegcoreSetEnableGrowingSourceFlush(const bool);
void
SegcoreSetEnableGeometryCache(const bool);
internal/core/src/segcore/segment_c.cpp
+84 -1
View File
@@ -937,6 +937,7 @@ struct FieldInfo {
milvus::FieldId field_id;
std::string field_name;
milvus::DataType data_type;
milvus::DataType element_type;
bool nullable;
int64_t dim; // for vector types
const milvus::segcore::VectorBase* vec_base;
@@ -958,6 +959,7 @@ GetElementByteWidth(milvus::DataType data_type, int64_t dim) {
case milvus::DataType::FLOAT:
return 4;
case milvus::DataType::INT64:
case milvus::DataType::TIMESTAMPTZ:
case milvus::DataType::DOUBLE:
return 8;
case milvus::DataType::VECTOR_FLOAT:
@@ -1275,6 +1277,47 @@ BuildTextArrayForChunkWithSpillover(
return builder.Finish();
}
arrow::Result<std::shared_ptr<arrow::Array>>
BuildVectorArrayForChunk(const FieldInfo& field_info,
int64_t start_offset,
int64_t num_rows) {
if (field_info.valid_data) {
return arrow::Status::Invalid(
"VECTOR_ARRAY does not support null rows");
}
auto vector_array_vec = dynamic_cast<
const milvus::segcore::ConcurrentVector<milvus::VectorArray>*>(
field_info.vec_base);
if (!vector_array_vec) {
return arrow::Status::Invalid("Expected ConcurrentVector<VectorArray>");
}
auto byte_width = milvus::vector_bytes_per_element(field_info.element_type,
field_info.dim);
auto value_builder = std::make_shared<arrow::FixedSizeBinaryBuilder>(
arrow::fixed_size_binary(byte_width));
arrow::ListBuilder builder(arrow::default_memory_pool(), value_builder);
ARROW_RETURN_NOT_OK(builder.Reserve(num_rows));
for (int64_t i = 0; i < num_rows; i++) {
const auto& vector_array = (*vector_array_vec)[start_offset + i];
if (vector_array.get_element_type() != field_info.element_type) {
return arrow::Status::Invalid("VECTOR_ARRAY element type mismatch");
}
if (vector_array.dim() != field_info.dim) {
return arrow::Status::Invalid("VECTOR_ARRAY dim mismatch");
}
ARROW_RETURN_NOT_OK(builder.Append());
ARROW_RETURN_NOT_OK(value_builder->AppendValues(
reinterpret_cast<const uint8_t*>(vector_array.data()),
vector_array.length()));
}
return builder.Finish();
}
// build boolean array for a chunk - booleans need special handling
arrow::Result<std::shared_ptr<arrow::Array>>
BuildBoolArrayForChunk(
@@ -1343,6 +1386,7 @@ BuildArrayForChunk(const FieldInfo& field_info,
global_offset);
case milvus::DataType::INT64:
case milvus::DataType::TIMESTAMPTZ:
return WrapChunkAsArrowArray<arrow::Int64Array>(
get_data_ptr(),
num_rows,
@@ -1450,6 +1494,29 @@ BuildArrayForChunk(const FieldInfo& field_info,
return builder.Finish();
}
case milvus::DataType::GEOMETRY: {
auto geometry_vec = dynamic_cast<
const milvus::segcore::ConcurrentVector<std::string>*>(
field_info.vec_base);
if (!geometry_vec) {
return arrow::Status::Invalid(
"Expected ConcurrentVector<std::string> for GEOMETRY");
}
arrow::BinaryBuilder builder;
ARROW_RETURN_NOT_OK(builder.Reserve(num_rows));
for (int64_t i = 0; i < num_rows; i++) {
int64_t offset = global_offset + i;
if (field_info.valid_data &&
!field_info.valid_data->is_valid(offset)) {
ARROW_RETURN_NOT_OK(builder.AppendNull());
} else {
auto wkb = geometry_vec->view_element(offset);
ARROW_RETURN_NOT_OK(builder.Append(wkb.data(), wkb.size()));
}
}
return builder.Finish();
}
case milvus::DataType::VECTOR_FLOAT:
case milvus::DataType::VECTOR_BINARY:
case milvus::DataType::VECTOR_FLOAT16:
@@ -1473,6 +1540,10 @@ BuildArrayForChunk(const FieldInfo& field_info,
return BuildSparseFloatVectorArrayForChunk(
field_info, global_offset, num_rows);
case milvus::DataType::VECTOR_ARRAY:
return BuildVectorArrayForChunk(
field_info, global_offset, num_rows);
default:
return arrow::Status::NotImplemented("Unsupported data type");
}
@@ -1537,6 +1608,7 @@ FlushGrowingSegmentData(CSegmentInterface c_segment,
info.field_id = RowFieldID;
info.field_name = field_meta.get_name().get();
info.data_type = field_meta.get_data_type();
info.element_type = milvus::DataType::NONE;
info.nullable = field_meta.is_nullable();
info.dim = 0;
info.vec_base = &insert_record.row_ids_;
@@ -1586,7 +1658,10 @@ FlushGrowingSegmentData(CSegmentInterface c_segment,
!milvus::IsSparseFloatVectorDataType(data_type)
? field_meta.get_dim()
: 0;
auto arrow_type = milvus::GetArrowDataType(data_type, dim);
auto arrow_type = data_type == milvus::DataType::VECTOR_ARRAY
? milvus::GetArrowDataTypeForVectorArray(
field_meta.get_element_type(), dim)
: milvus::GetArrowDataType(data_type, dim);
if (field_meta.is_nullable() &&
IsSupportedNullableVectorDataType(data_type)) {
arrow_type = arrow::binary();
@@ -1596,6 +1671,7 @@ FlushGrowingSegmentData(CSegmentInterface c_segment,
info.field_id = field_id;
info.field_name = field_meta.get_name().get();
info.data_type = data_type;
info.element_type = field_meta.get_element_type();
info.nullable = field_meta.is_nullable();
info.dim = dim;
info.vec_base = vec_base;
@@ -1624,6 +1700,13 @@ FlushGrowingSegmentData(CSegmentInterface c_segment,
metadata_keys.push_back(DIM_KEY);
metadata_values.push_back(std::to_string(dim));
}
if (data_type == milvus::DataType::VECTOR_ARRAY) {
metadata_keys.push_back(ELEMENT_TYPE_KEY_FOR_ARROW);
metadata_values.push_back(std::to_string(
static_cast<int>(field_meta.get_element_type())));
metadata_keys.push_back(DIM_KEY);
metadata_values.push_back(std::to_string(dim));
}
auto metadata =
arrow::KeyValueMetadata::Make(metadata_keys, metadata_values);
arrow_fields.push_back(arrow::field(std::to_string(field_id.get()),
internal/datacoord/services.go
+32 -3
View File
@@ -647,6 +647,10 @@ func (s *Server) SaveBinlogPaths(ctx context.Context, req *datapb.SaveBinlogPath
log.Warn("failed to get segment, the segment not healthy", zap.Error(err))
return merr.Status(err), nil
}
if err := s.validateTextSegmentStorage(req); err != nil {
log.Warn("invalid TEXT segment storage format", zap.Error(err))
return merr.Status(err), nil
}
// Set storage version
operators = append(operators, SetStorageVersion(req.GetSegmentID(), req.GetStorageVersion()))
@@ -739,6 +743,27 @@ func (s *Server) SaveBinlogPaths(ctx context.Context, req *datapb.SaveBinlogPath
return merr.Success(), nil
}
func (s *Server) validateTextSegmentStorage(req *datapb.SaveBinlogPathsRequest) error {
if req.GetSegLevel() == datapb.SegmentLevel_L0 || req.GetDropped() {
return nil
}
if !s.meta.collectionHasTextFields(req.GetCollectionID()) {
return nil
}
if req.GetStorageVersion() < storage.StorageV3 {
return merr.WrapErrParameterInvalidMsg(
"TEXT segment %d must be saved with StorageV3 manifest, got storage version %d",
req.GetSegmentID(),
req.GetStorageVersion())
}
if req.GetManifestPath() == "" {
return merr.WrapErrParameterInvalidMsg(
"TEXT segment %d requires non-empty StorageV3 manifest path",
req.GetSegmentID())
}
return nil
}
// DropVirtualChannel notifies vchannel dropped
// And contains the remaining data log & checkpoint to update
func (s *Server) DropVirtualChannel(ctx context.Context, req *datapb.DropVirtualChannelRequest) (*datapb.DropVirtualChannelResponse, error) {
@@ -1499,12 +1524,16 @@ func (s *Server) GetFlushState(ctx context.Context, req *datapb.GetFlushStateReq
for _, channel := range channels {
cp := s.meta.GetChannelCheckpoint(channel.GetName())
if cp == nil || cp.GetTimestamp() < req.GetFlushTs() {
cpTs := uint64(0)
if cp != nil {
cpTs = cp.GetTimestamp()
}
if cp == nil || cpTs < req.GetFlushTs() {
resp.Flushed = false
log.RatedInfo(10, "GetFlushState failed, channel unflushed", zap.String("channel", channel.GetName()),
zap.Time("CP", tsoutil.PhysicalTime(cp.GetTimestamp())),
zap.Duration("lag", tsoutil.PhysicalTime(req.GetFlushTs()).Sub(tsoutil.PhysicalTime(cp.GetTimestamp()))))
zap.Time("CP", tsoutil.PhysicalTime(cpTs)),
zap.Duration("lag", tsoutil.PhysicalTime(req.GetFlushTs()).Sub(tsoutil.PhysicalTime(cpTs))))
return resp, nil
}
}
internal/datacoord/services_test.go
+101 -1
View File
@@ -79,7 +79,7 @@ func (s *ServerSuite) SetupSuite() {
<-ctx.Done()
return ctx.Err()
})
b.EXPECT().GetLatestWALLocated(mock.Anything, mock.Anything).Return(0, true)
b.EXPECT().GetLatestWALLocated(mock.Anything, mock.Anything).Return(0, true).Maybe()
balance.Register(b)
}
@@ -149,6 +149,21 @@ func (s *ServerSuite) TestGetFlushState_ByFlushTs() {
}, resp)
}
func (s *ServerSuite) TestGetFlushState_ByFlushTsMissingCheckpoint() {
s.mockMixCoord.EXPECT().DescribeCollectionInternal(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
Status: merr.Success(),
CollectionID: 0,
VirtualChannelNames: []string{"missing-cp-channel"},
}, nil)
resp, err := s.testServer.GetFlushState(context.TODO(), &datapb.GetFlushStateRequest{FlushTs: 13})
s.NoError(err)
s.EqualValues(&milvuspb.GetFlushStateResponse{
Status: merr.Success(),
Flushed: false,
}, resp)
}
func (s *ServerSuite) TestGetFlushState_BySegment() {
s.mockMixCoord.EXPECT().DescribeCollectionInternal(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
@@ -327,6 +342,91 @@ func (s *ServerSuite) TestSaveBinlogPath_SaveDroppedSegment() {
}
}
func (s *ServerSuite) TestSaveBinlogPath_TextRequiresStorageV3Manifest() {
s.testServer.meta.AddCollection(&collectionInfo{
ID: 0,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, DataType: schemapb.DataType_Text},
},
},
})
err := s.testServer.meta.AddSegment(context.TODO(), NewSegmentInfo(&datapb.SegmentInfo{
ID: 10,
CollectionID: 0,
PartitionID: 1,
InsertChannel: "ch1",
State: commonpb.SegmentState_Sealed,
Level: datapb.SegmentLevel_L1,
NumOfRows: 1,
}))
s.Require().NoError(err)
resp, err := s.testServer.SaveBinlogPaths(context.Background(), &datapb.SaveBinlogPathsRequest{
Base: &commonpb.MsgBase{Timestamp: uint64(time.Now().Unix())},
SegmentID: 10,
CollectionID: 0,
PartitionID: 1,
Channel: "ch1",
SegLevel: datapb.SegmentLevel_L1,
Flushed: true,
StorageVersion: storage.StorageV2,
})
s.NoError(err)
s.ErrorIs(merr.Error(resp), merr.ErrParameterInvalid)
resp, err = s.testServer.SaveBinlogPaths(context.Background(), &datapb.SaveBinlogPathsRequest{
Base: &commonpb.MsgBase{Timestamp: uint64(time.Now().Unix())},
SegmentID: 10,
CollectionID: 0,
PartitionID: 1,
Channel: "ch1",
SegLevel: datapb.SegmentLevel_L1,
Flushed: true,
StorageVersion: storage.StorageV3,
})
s.NoError(err)
s.ErrorIs(merr.Error(resp), merr.ErrParameterInvalid)
resp, err = s.testServer.SaveBinlogPaths(context.Background(), &datapb.SaveBinlogPathsRequest{
Base: &commonpb.MsgBase{Timestamp: uint64(time.Now().Unix())},
SegmentID: 10,
CollectionID: 0,
PartitionID: 1,
Channel: "ch1",
SegLevel: datapb.SegmentLevel_L1,
Flushed: false,
StorageVersion: storage.StorageV3,
})
s.NoError(err)
s.ErrorIs(merr.Error(resp), merr.ErrParameterInvalid)
err = s.testServer.meta.AddSegment(context.TODO(), NewSegmentInfo(&datapb.SegmentInfo{
ID: 11,
CollectionID: 0,
PartitionID: 1,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
Level: datapb.SegmentLevel_L1,
NumOfRows: 1,
}))
s.Require().NoError(err)
resp, err = s.testServer.SaveBinlogPaths(context.Background(), &datapb.SaveBinlogPathsRequest{
Base: &commonpb.MsgBase{Timestamp: uint64(time.Now().Unix())},
SegmentID: 11,
CollectionID: 0,
PartitionID: 1,
Channel: "ch1",
SegLevel: datapb.SegmentLevel_L1,
Flushed: true,
StorageVersion: storage.StorageV2,
})
s.NoError(err)
s.True(merr.Ok(resp))
}
func (s *ServerSuite) TestSaveBinlogPath_L0Segment() {
s.testServer.meta.AddCollection(&collectionInfo{ID: 0})
internal/distributed/streaming/streaming.go
+5
View File
@@ -173,6 +173,11 @@ type WALAccesser interface {
// Local returns the local services.
Local() Local
// PrepareReleaseManualFlush prepares process-local release handoff.
// Returns false when the current process is not the local flush owner or
// the channel does not need growing-source retention.
PrepareReleaseManualFlush(ctx context.Context, collectionID int64, vchannel string, releaseSegmentIDs []int64) (bool, error)
// RawAppend writes a records to the log.
RawAppend(ctx context.Context, msgs message.MutableMessage, opts ...AppendOption) (*types.AppendResult, error)
internal/distributed/streaming/test_streaming.go
+4
View File
@@ -188,6 +188,10 @@ func (n *noopWALAccesser) Local() Local {
return &noopLocal{}
}
func (n *noopWALAccesser) PrepareReleaseManualFlush(ctx context.Context, collectionID int64, vchannel string, releaseSegmentIDs []int64) (bool, error) {
return false, getExpectErr()
}
func (n *noopWALAccesser) RawAppend(ctx context.Context, msgs message.MutableMessage, opts ...AppendOption) (*types.AppendResult, error) {
if err := getExpectErr(); err != nil {
return nil, err
internal/distributed/streaming/wal.go
+4
View File
@@ -83,6 +83,10 @@ func (w *walAccesserImpl) Local() Local {
return localServiceImpl{w}
}
func (w *walAccesserImpl) PrepareReleaseManualFlush(ctx context.Context, collectionID int64, vchannel string, releaseSegmentIDs []int64) (bool, error) {
return w.handlerClient.PrepareReleaseManualFlush(ctx, collectionID, vchannel, releaseSegmentIDs)
}
// ControlChannel returns the control channel name of the wal.
func (w *walAccesserImpl) ControlChannel() string {
last, err := w.streamingCoordClient.Assignment().GetLatestAssignments(context.Background())
internal/distributed/streaming/wal_test.go
+14
View File
@@ -192,6 +192,20 @@ func TestReleaseTimeout(t *testing.T) {
}, time.Second, 10*time.Millisecond)
}
func TestWALAccesserPrepareReleaseManualFlush(t *testing.T) {
ctx := context.Background()
w, _, _, handler := createMockWAL(t)
defer w.Close()
handler.EXPECT().
PrepareReleaseManualFlush(mock.Anything, int64(100), vChannel1, []int64{1001}).
Return(true, nil)
prepared, err := w.PrepareReleaseManualFlush(ctx, 100, vChannel1, []int64{1001})
assert.NoError(t, err)
assert.True(t, prepared)
}
func newInsertMessage(vChannel string) message.MutableMessage {
msg, err := message.NewInsertMessageBuilderV1().
WithVChannel(vChannel).
internal/flushcommon/metacache/actions.go
+20
View File
@@ -223,6 +223,14 @@ func StartSyncing(batchSize int64) SegmentAction {
}
}
func AbortSyncing(batchSize int64) SegmentAction {
return func(info *SegmentInfo) {
info.syncingRows -= batchSize
info.bufferRows += batchSize
info.syncingTasks--
}
}
func FinishSyncing(batchSize int64) SegmentAction {
return func(info *SegmentInfo) {
info.flushedRows += batchSize
@@ -249,6 +257,18 @@ func UpdateManifestPath(manifestPath string) SegmentAction {
}
}
// SetFlushSourceMode records which subsystem owns the segment's payload at
// flush time. The decision is sticky: once a non-Unknown mode is set, later
// calls with a different mode are no-ops, so the source for a given segment
// stays consistent across its lifetime.
func SetFlushSourceMode(mode FlushSourceMode) SegmentAction {
return func(info *SegmentInfo) {
if info.flushSourceMode == FlushSourceUnknown {
info.flushSourceMode = mode
}
}
}
// MergeSegmentAction is the util function to merge multiple SegmentActions into one.
func MergeSegmentAction(actions ...SegmentAction) SegmentAction {
return func(info *SegmentInfo) {
internal/flushcommon/metacache/flush_source.go
+61
View File
@@ -0,0 +1,61 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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.
package metacache
// FlushSourceMode indicates which subsystem is responsible for supplying the
// binlog payload of a segment when the WAL flush task executes.
//
// This is process-local runtime state, in the same family as bufferRows /
// syncingRows / syncingTasks. It is NEVER persisted to etcd or DataCoord. The
// mode is decided the first time a segment receives data and remains sticky
// for the segment's lifetime inside metacache. When the segment is removed
// from metacache (via RemoveSegments after flush/drop), the mode is naturally
// garbage-collected together with the SegmentInfo.
//
// On restart, sealed/flushed segments do not need this field (no growing
// phase remains). Growing segments re-derive the value through the regular
// write-buffer source decision path during WAL replay.
//
// Currently consumed only by the TEXT flush owner so that the writeBuffer
// can route between the live insert payload and an out-of-process growing
// source provider without keeping its own shadow map. The field is generic
// so future large-field offload paths (BM25, vector, etc.) can reuse the
// same machinery.
type FlushSourceMode int32
const (
// FlushSourceUnknown means the source has not yet been decided. New
// segments default to this value until the first insert arrives.
FlushSourceUnknown FlushSourceMode = 0
// FlushSourceWriteBuffer means the WriteBuffer's local insert payload is
// the source of truth for sync tasks of this segment.
FlushSourceWriteBuffer FlushSourceMode = 1
// FlushSourceGrowing means the data lives in an external growing source
// and the WriteBuffer only tracks progress.
FlushSourceGrowing FlushSourceMode = 2
)
func (m FlushSourceMode) String() string {
switch m {
case FlushSourceWriteBuffer:
return "WriteBuffer"
case FlushSourceGrowing:
return "Growing"
default:
return "Unknown"
}
}
internal/flushcommon/metacache/segment.go
+12
View File
@@ -49,6 +49,10 @@ type SegmentInfo struct {
bm25logs []*datapb.FieldBinlog
currentSplit []storagecommon.ColumnGroup
manifestPath string
// flushSourceMode is process-local runtime state; not persisted.
// See FlushSourceMode docs for lifecycle semantics.
flushSourceMode FlushSourceMode
}
func (s *SegmentInfo) SegmentID() int64 {
@@ -134,6 +138,13 @@ func (s *SegmentInfo) ManifestPath() string {
return s.manifestPath
}
// FlushSourceMode returns the sticky decision of which subsystem owns this
// segment's payload at flush time. The value is process-local and not
// persisted; see FlushSourceMode docs for details.
func (s *SegmentInfo) FlushSourceMode() FlushSourceMode {
return s.flushSourceMode
}
func (s *SegmentInfo) Clone() *SegmentInfo {
return &SegmentInfo{
segmentID: s.segmentID,
@@ -156,6 +167,7 @@ func (s *SegmentInfo) Clone() *SegmentInfo {
bm25logs: s.bm25logs,
currentSplit: s.currentSplit,
manifestPath: s.manifestPath,
flushSourceMode: s.flushSourceMode,
}
}
internal/flushcommon/syncmgr/growing_source.go
+621
View File
@@ -0,0 +1,621 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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.
package syncmgr
import (
"context"
"fmt"
"path"
"sort"
"sync"
"github.com/cockroachdb/errors"
"github.com/samber/lo"
"go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/flushcommon/metacache"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/storagev2/packed"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/log"
"github.com/milvus-io/milvus/pkg/v3/metrics"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/metautil"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/retry"
"github.com/milvus-io/milvus/pkg/v3/util/timerecord"
)
type GrowingFlushConfig struct {
SegmentBasePath string
PartitionBasePath string
CollectionID int64
PartitionID int64
TextFieldIDs []int64
TextLobPaths []string
ReadVersion int64
}
type GrowingFlushResult struct {
ManifestPath string
NumRows int64
}
type GrowingFlushSource interface {
CurrentOffset() int64
FlushGrowingData(ctx context.Context, startOffset, endOffset int64, config *GrowingFlushConfig) (*GrowingFlushResult, error)
Release()
}
type GrowingFlushSourceCommitter interface {
CommitGrowingFlush(targetOffset int64)
}
type GrowingSourceState int
const (
GrowingSourceUnavailable GrowingSourceState = iota
GrowingSourcePending
GrowingSourceUsable
)
type GrowingSourceProvider interface {
GetGrowingFlushSource(segmentID int64, targetOffset int64, endPos *msgpb.MsgPosition) (GrowingFlushSource, GrowingSourceState)
}
type GrowingSourceReleaseHandoffSegment struct {
SegmentID int64
TargetOffset int64
}
type GrowingSourceReleaseHandoffProvider interface {
PrepareGrowingSourceReleaseHandoff(ctx context.Context, fenceTs uint64, segments []GrowingSourceReleaseHandoffSegment) error
IsReleaseAllowed(segmentID int64, checkpointTs uint64) bool
IsReleasePrepared(segmentID int64, checkpointTs uint64) bool
MarkReleaseDetached(segmentID int64)
ClearReleasePrepared(segmentID int64)
ReleasePreparedSegments() []int64
}
type GrowingSourceRegistry struct {
mu sync.RWMutex
nextToken uint64
providers map[string]map[uint64]GrowingSourceProvider
}
type GrowingSourceRegistration struct {
registry *GrowingSourceRegistry
channel string
token uint64
}
func NewGrowingSourceRegistry() *GrowingSourceRegistry {
return &GrowingSourceRegistry{
providers: make(map[string]map[uint64]GrowingSourceProvider),
}
}
func (r *GrowingSourceRegistry) Register(channel string, provider GrowingSourceProvider) *GrowingSourceRegistration {
if provider == nil {
return nil
}
r.mu.Lock()
defer r.mu.Unlock()
r.nextToken++
token := r.nextToken
if _, ok := r.providers[channel]; !ok {
r.providers[channel] = make(map[uint64]GrowingSourceProvider)
}
r.providers[channel][token] = provider
return &GrowingSourceRegistration{
registry: r,
channel: channel,
token: token,
}
}
func (r *GrowingSourceRegistry) Unregister(registration *GrowingSourceRegistration) {
if registration == nil {
return
}
r.mu.Lock()
defer r.mu.Unlock()
providers, ok := r.providers[registration.channel]
if !ok {
return
}
delete(providers, registration.token)
if len(providers) == 0 {
delete(r.providers, registration.channel)
}
}
func (r *GrowingSourceRegistry) ProviderCount(channel string) int {
r.mu.RLock()
defer r.mu.RUnlock()
return len(r.providers[channel])
}
func (r *GrowingSourceRegistry) getProviders(channel string) []GrowingSourceProvider {
r.mu.RLock()
channelProviders := r.providers[channel]
tokens := make([]uint64, 0, len(channelProviders))
for token := range channelProviders {
tokens = append(tokens, token)
}
sort.Slice(tokens, func(i, j int) bool {
return tokens[i] < tokens[j]
})
providers := make([]GrowingSourceProvider, 0, len(tokens))
for _, token := range tokens {
providers = append(providers, channelProviders[token])
}
r.mu.RUnlock()
return providers
}
func (r *GrowingSourceRegistry) PrepareGrowingSourceReleaseHandoff(ctx context.Context, channel string, fenceTs uint64, segments []GrowingSourceReleaseHandoffSegment) error {
handoffProviders := make([]GrowingSourceReleaseHandoffProvider, 0)
for _, provider := range r.getProviders(channel) {
handoffProvider, ok := provider.(GrowingSourceReleaseHandoffProvider)
if !ok {
continue
}
handoffProviders = append(handoffProviders, handoffProvider)
}
if len(handoffProviders) == 0 {
return merr.WrapErrChannelNotAvailable(channel, "no local growing-source release handoff provider")
}
for _, handoffProvider := range handoffProviders {
if err := handoffProvider.PrepareGrowingSourceReleaseHandoff(ctx, fenceTs, segments); err != nil {
return err
}
}
for _, segment := range segments {
if segment.TargetOffset > 0 {
if !r.IsReleasePrepared(channel, segment.SegmentID, fenceTs) {
return merr.WrapErrSegmentNotFound(segment.SegmentID, "growing-source release handoff source is not prepared")
}
continue
}
if !r.IsReleaseAllowed(channel, segment.SegmentID, fenceTs) {
return merr.WrapErrSegmentNotFound(segment.SegmentID, "growing-source release handoff is not allowed")
}
}
return nil
}
func (r *GrowingSourceRegistry) IsReleaseAllowed(channel string, segmentID int64, checkpointTs uint64) bool {
for _, provider := range r.getProviders(channel) {
handoffProvider, ok := provider.(GrowingSourceReleaseHandoffProvider)
if !ok {
continue
}
if handoffProvider.IsReleaseAllowed(segmentID, checkpointTs) {
return true
}
}
return false
}
func (r *GrowingSourceRegistry) IsReleasePrepared(channel string, segmentID int64, checkpointTs uint64) bool {
for _, provider := range r.getProviders(channel) {
handoffProvider, ok := provider.(GrowingSourceReleaseHandoffProvider)
if !ok {
continue
}
if handoffProvider.IsReleasePrepared(segmentID, checkpointTs) {
return true
}
}
return false
}
func (r *GrowingSourceRegistry) ClearReleasePrepared(channel string, segmentID int64) {
for _, provider := range r.getProviders(channel) {
handoffProvider, ok := provider.(GrowingSourceReleaseHandoffProvider)
if !ok {
continue
}
handoffProvider.ClearReleasePrepared(segmentID)
}
}
func (r *GrowingSourceRegistry) MarkReleaseDetached(channel string, segmentID int64) {
for _, provider := range r.getProviders(channel) {
handoffProvider, ok := provider.(GrowingSourceReleaseHandoffProvider)
if !ok {
continue
}
handoffProvider.MarkReleaseDetached(segmentID)
}
}
func (r *GrowingSourceRegistry) ReleasePreparedSegments(channel string) []int64 {
var segments []int64
for _, provider := range r.getProviders(channel) {
handoffProvider, ok := provider.(GrowingSourceReleaseHandoffProvider)
if !ok {
continue
}
segments = append(segments, handoffProvider.ReleasePreparedSegments()...)
}
return lo.Uniq(segments)
}
func (r *GrowingSourceRegistry) Resolve(channel string, segmentID int64, targetOffset int64, endPos *msgpb.MsgPosition) (GrowingFlushSource, GrowingSourceState) {
hasPending := false
for _, provider := range r.getProviders(channel) {
if provider == nil {
continue
}
source, state := provider.GetGrowingFlushSource(segmentID, targetOffset, endPos)
if source == nil {
continue
}
switch state {
case GrowingSourceUsable:
return source, GrowingSourceUsable
case GrowingSourcePending:
hasPending = true
source.Release()
default:
source.Release()
}
}
if hasPending {
return nil, GrowingSourcePending
}
return nil, GrowingSourceUnavailable
}
var defaultGrowingSourceRegistry = NewGrowingSourceRegistry()
func DefaultGrowingSourceRegistry() *GrowingSourceRegistry {
return defaultGrowingSourceRegistry
}
type GrowingSourceSyncTask struct {
collectionID int64
partitionID int64
segmentID int64
channelName string
startPosition *msgpb.MsgPosition
checkpoint *msgpb.MsgPosition
batchRows int64
targetOffset int64
level datapb.SegmentLevel
isFlush bool
isDrop bool
metacache metacache.MetaCache
metaWriter MetaWriter
schema *schemapb.CollectionSchema
source GrowingFlushSource
chunkManager storage.ChunkManager
manifestPath string
flushedSize int64
writeRetryOpts []retry.Option
failureCallback func(error)
tr *timerecord.TimeRecorder
}
func NewGrowingSourceSyncTask() *GrowingSourceSyncTask {
return new(GrowingSourceSyncTask)
}
func (t *GrowingSourceSyncTask) WithCollectionID(collectionID int64) *GrowingSourceSyncTask {
t.collectionID = collectionID
return t
}
func (t *GrowingSourceSyncTask) WithPartitionID(partitionID int64) *GrowingSourceSyncTask {
t.partitionID = partitionID
return t
}
func (t *GrowingSourceSyncTask) WithSegmentID(segmentID int64) *GrowingSourceSyncTask {
t.segmentID = segmentID
return t
}
func (t *GrowingSourceSyncTask) WithChannelName(channelName string) *GrowingSourceSyncTask {
t.channelName = channelName
return t
}
func (t *GrowingSourceSyncTask) WithStartPosition(position *msgpb.MsgPosition) *GrowingSourceSyncTask {
t.startPosition = position
return t
}
func (t *GrowingSourceSyncTask) WithCheckpoint(position *msgpb.MsgPosition) *GrowingSourceSyncTask {
t.checkpoint = position
return t
}
func (t *GrowingSourceSyncTask) WithBatchRows(batchRows int64) *GrowingSourceSyncTask {
t.batchRows = batchRows
return t
}
func (t *GrowingSourceSyncTask) WithTargetOffset(targetOffset int64) *GrowingSourceSyncTask {
t.targetOffset = targetOffset
return t
}
func (t *GrowingSourceSyncTask) WithLevel(level datapb.SegmentLevel) *GrowingSourceSyncTask {
t.level = level
return t
}
func (t *GrowingSourceSyncTask) WithFlush() *GrowingSourceSyncTask {
t.isFlush = true
return t
}
func (t *GrowingSourceSyncTask) WithDrop() *GrowingSourceSyncTask {
t.isDrop = true
return t
}
func (t *GrowingSourceSyncTask) WithMetaCache(metacache metacache.MetaCache) *GrowingSourceSyncTask {
t.metacache = metacache
return t
}
func (t *GrowingSourceSyncTask) WithMetaWriter(metaWriter MetaWriter) *GrowingSourceSyncTask {
t.metaWriter = metaWriter
return t
}
func (t *GrowingSourceSyncTask) WithSchema(schema *schemapb.CollectionSchema) *GrowingSourceSyncTask {
t.schema = schema
return t
}
func (t *GrowingSourceSyncTask) WithSource(source GrowingFlushSource) *GrowingSourceSyncTask {
t.source = source
return t
}
func (t *GrowingSourceSyncTask) WithChunkManager(cm storage.ChunkManager) *GrowingSourceSyncTask {
t.chunkManager = cm
return t
}
func (t *GrowingSourceSyncTask) WithWriteRetryOptions(opts ...retry.Option) *GrowingSourceSyncTask {
t.writeRetryOpts = opts
return t
}
func (t *GrowingSourceSyncTask) WithFailureCallback(callback func(error)) *GrowingSourceSyncTask {
t.failureCallback = callback
return t
}
func (t *GrowingSourceSyncTask) SegmentID() int64 {
return t.segmentID
}
func (t *GrowingSourceSyncTask) Checkpoint() *msgpb.MsgPosition {
return t.checkpoint
}
func (t *GrowingSourceSyncTask) StartPosition() *msgpb.MsgPosition {
return t.startPosition
}
func (t *GrowingSourceSyncTask) ChannelName() string {
return t.channelName
}
func (t *GrowingSourceSyncTask) IsFlush() bool {
return t.isFlush
}
func (t *GrowingSourceSyncTask) IsDrop() bool {
return t.isDrop
}
func (t *GrowingSourceSyncTask) ManifestPath() string {
return t.manifestPath
}
func (t *GrowingSourceSyncTask) BatchRows() int64 {
return t.batchRows
}
func (t *GrowingSourceSyncTask) TargetOffset() int64 {
return t.targetOffset
}
func (t *GrowingSourceSyncTask) HandleError(err error) {
if t.failureCallback != nil {
t.failureCallback(err)
}
metrics.DataNodeFlushBufferCount.WithLabelValues(paramtable.GetStringNodeID(), metrics.FailLabel, t.level.String()).Inc()
}
func (t *GrowingSourceSyncTask) ReleaseSource() {
if t.source != nil {
t.source.Release()
t.source = nil
}
}
func (t *GrowingSourceSyncTask) Run(ctx context.Context) (err error) {
t.tr = timerecord.NewTimeRecorder("growingSourceSyncTask")
log := log.Ctx(ctx).With(
zap.Int64("collectionID", t.collectionID),
zap.Int64("partitionID", t.partitionID),
zap.Int64("segmentID", t.segmentID),
zap.String("channel", t.channelName),
)
commitSource := false
defer func() {
committer, shouldCommit := t.source.(GrowingFlushSourceCommitter)
t.ReleaseSource()
if commitSource && shouldCommit && (t.IsFlush() || t.IsDrop()) {
committer.CommitGrowingFlush(t.targetOffset)
}
if err != nil {
t.HandleError(err)
}
}()
segment, ok := t.metacache.GetSegmentByID(t.segmentID)
if !ok {
if t.isDrop {
log.Info("segment dropped, discard growing source sync task")
return nil
}
log.Warn("segment not found in metacache")
return nil
}
if t.source == nil {
return errors.New("growing flush source is nil")
}
if t.source.CurrentOffset() < t.targetOffset {
return errors.Errorf("growing flush source is behind target offset, current=%d target=%d", t.source.CurrentOffset(), t.targetOffset)
}
expectedRows := t.targetOffset - segment.FlushedRows()
if expectedRows < 0 {
return errors.Errorf("growing source target offset is behind flushed rows, flushedRows=%d targetOffset=%d segmentID=%d",
segment.FlushedRows(), t.targetOffset, t.segmentID)
}
if expectedRows == 0 {
t.manifestPath = segment.ManifestPath()
} else {
config := t.buildFlushConfig(segment)
result, err := t.source.FlushGrowingData(ctx, segment.FlushedRows(), t.targetOffset, config)
if err != nil {
return errors.Wrap(err, "flush growing source data")
}
if result == nil || result.ManifestPath == "" {
return errors.New("growing source flush returned empty manifest")
}
if result.NumRows != expectedRows {
return errors.Errorf("growing source flush row count mismatch, expected=%d actual=%d flushedRows=%d targetOffset=%d segmentID=%d",
expectedRows, result.NumRows, segment.FlushedRows(), t.targetOffset, t.segmentID)
}
t.manifestPath = result.ManifestPath
}
t.flushedSize = expectedRows
if t.metaWriter != nil {
if err := t.metaWriter.UpdateGrowingSourceSync(ctx, t); err != nil {
return err
}
}
actions := make([]metacache.SegmentAction, 0, 3)
if t.batchRows > 0 {
actions = append(actions, metacache.FinishSyncing(t.batchRows))
}
if t.manifestPath != "" {
actions = append(actions, metacache.UpdateManifestPath(t.manifestPath))
}
if t.IsFlush() {
actions = append(actions, metacache.UpdateState(commonpb.SegmentState_Flushed))
}
t.metacache.UpdateSegments(metacache.MergeSegmentAction(actions...), metacache.WithSegmentIDs(t.segmentID))
if t.isDrop {
t.metacache.RemoveSegments(metacache.WithSegmentIDs(t.segmentID))
log.Info("dropped growing source segment removed")
}
commitSource = true
metrics.DataNodeWriteDataCount.WithLabelValues(paramtable.GetStringNodeID(), metrics.StreamingDataSourceLabel, metrics.InsertLabel, fmt.Sprint(t.collectionID)).Add(float64(t.batchRows))
metrics.DataNodeFlushedRows.WithLabelValues(paramtable.GetStringNodeID(), metrics.StreamingDataSourceLabel).Add(float64(t.batchRows))
metrics.DataNodeFlushBufferCount.WithLabelValues(paramtable.GetStringNodeID(), metrics.SuccessLabel, t.level.String()).Inc()
log.Info("growing source sync task done",
zap.Int64("targetOffset", t.targetOffset),
zap.Int64("batchRows", t.batchRows),
zap.String("manifestPath", t.manifestPath),
zap.Duration("timeTaken", t.tr.ElapseSpan()))
return nil
}
func (t *GrowingSourceSyncTask) buildFlushConfig(segment *metacache.SegmentInfo) *GrowingFlushConfig {
segmentBasePath := path.Join(t.chunkManager.RootPath(), common.SegmentInsertLogPath,
metautil.JoinIDPath(t.collectionID, t.partitionID, t.segmentID))
partitionBasePath := path.Join(t.chunkManager.RootPath(), common.SegmentInsertLogPath,
metautil.JoinIDPath(t.collectionID, t.partitionID))
var textFieldIDs []int64
var textLobPaths []string
if t.schema != nil {
for _, field := range t.schema.GetFields() {
if field.GetDataType() == schemapb.DataType_Text {
fieldID := field.GetFieldID()
textFieldIDs = append(textFieldIDs, fieldID)
textLobPaths = append(textLobPaths, fmt.Sprintf("%s/lobs/%d", partitionBasePath, fieldID))
}
}
}
return &GrowingFlushConfig{
SegmentBasePath: segmentBasePath,
PartitionBasePath: partitionBasePath,
CollectionID: t.collectionID,
PartitionID: t.partitionID,
TextFieldIDs: textFieldIDs,
TextLobPaths: textLobPaths,
ReadVersion: manifestVersion(segment.ManifestPath()),
}
}
func manifestVersion(manifestPath string) int64 {
if manifestPath == "" {
return -1
}
if _, version, err := packedManifestVersion(manifestPath); err == nil {
return version
}
return -1
}
func packedManifestVersion(manifestPath string) (string, int64, error) {
return packed.UnmarshalManifestPath(manifestPath)
}
func (t *GrowingSourceSyncTask) startPositions() []*datapb.SegmentStartPosition {
startPos := lo.Map(t.metacache.GetSegmentsBy(
metacache.WithSegmentState(commonpb.SegmentState_Growing, commonpb.SegmentState_Sealed, commonpb.SegmentState_Flushing),
metacache.WithLevel(datapb.SegmentLevel_L1),
metacache.WithStartPosNotRecorded(),
), func(info *metacache.SegmentInfo, _ int) *datapb.SegmentStartPosition {
return &datapb.SegmentStartPosition{
SegmentID: info.SegmentID(),
StartPosition: info.StartPosition(),
}
})
if t.level == datapb.SegmentLevel_L0 {
startPos = append(startPos, &datapb.SegmentStartPosition{SegmentID: t.segmentID, StartPosition: t.startPosition})
}
return startPos
}
internal/flushcommon/syncmgr/growing_source_test.go
+91
View File
@@ -0,0 +1,91 @@
package syncmgr
import (
"context"
"testing"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus/internal/flushcommon/metacache"
"github.com/milvus-io/milvus/internal/flushcommon/metacache/pkoracle"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
)
type fakeCommitGrowingFlushSource struct {
commits []int64
}
func (s *fakeCommitGrowingFlushSource) CurrentOffset() int64 {
return 10
}
func (s *fakeCommitGrowingFlushSource) FlushGrowingData(context.Context, int64, int64, *GrowingFlushConfig) (*GrowingFlushResult, error) {
return &GrowingFlushResult{ManifestPath: "manifest", NumRows: 10}, nil
}
func (s *fakeCommitGrowingFlushSource) Release() {
}
func (s *fakeCommitGrowingFlushSource) CommitGrowingFlush(targetOffset int64) {
s.commits = append(s.commits, targetOffset)
}
func TestGrowingSourceSyncTaskCommitRetainedSourceOnlyOnFinalization(t *testing.T) {
run := func(t *testing.T, finalize func(*GrowingSourceSyncTask), expectCommit bool) {
segmentID := int64(1)
mc := metacache.NewMockMetaCache(t)
segment := metacache.NewSegmentInfo(&datapb.SegmentInfo{
ID: segmentID,
PartitionID: 2,
ManifestPath: "manifest",
}, pkoracle.NewBloomFilterSet(), nil)
metacache.UpdateNumOfRows(10)(segment)
source := &fakeCommitGrowingFlushSource{}
mc.EXPECT().GetSegmentByID(segmentID).Return(segment, true)
mc.EXPECT().UpdateSegments(mock.Anything, mock.Anything).Run(func(action metacache.SegmentAction, filters ...metacache.SegmentFilter) {
action(segment)
}).Return()
mc.EXPECT().RemoveSegments(mock.Anything).Return(nil).Maybe()
task := NewGrowingSourceSyncTask().
WithCollectionID(3).
WithPartitionID(2).
WithSegmentID(segmentID).
WithChannelName("ch").
WithStartPosition(&msgpb.MsgPosition{Timestamp: 100}).
WithCheckpoint(&msgpb.MsgPosition{Timestamp: 200}).
WithBatchRows(0).
WithTargetOffset(10).
WithMetaCache(mc).
WithSource(source)
if finalize != nil {
finalize(task)
}
require.NoError(t, task.Run(context.Background()))
if expectCommit {
require.Equal(t, []int64{10}, source.commits)
} else {
require.Empty(t, source.commits)
}
}
t.Run("non_final_sync_does_not_commit_retained_source", func(t *testing.T) {
run(t, nil, false)
})
t.Run("flush_finalization_commits_retained_source", func(t *testing.T) {
run(t, func(task *GrowingSourceSyncTask) {
task.WithFlush()
}, true)
})
t.Run("drop_finalization_commits_retained_source", func(t *testing.T) {
run(t, func(task *GrowingSourceSyncTask) {
task.WithDrop()
}, true)
})
}
internal/flushcommon/syncmgr/key_lock_dispatcher.go
+1
View File
@@ -21,6 +21,7 @@ type Task interface {
Run(context.Context) error
HandleError(error)
IsFlush() bool
IsDrop() bool
}
// pendingTask wraps a task queued for execution.
internal/flushcommon/syncmgr/meta_writer.go
+71
View File
@@ -21,6 +21,7 @@ import (
// MetaWriter is the interface for SyncManager to write segment sync meta.
type MetaWriter interface {
UpdateSync(context.Context, *SyncTask) error
UpdateGrowingSourceSync(context.Context, *GrowingSourceSyncTask) error
DropChannel(context.Context, string) error
}
@@ -162,6 +163,76 @@ func (b *brokerMetaWriter) UpdateSync(ctx context.Context, pack *SyncTask) error
return nil
}
func (b *brokerMetaWriter) UpdateGrowingSourceSync(ctx context.Context, task *GrowingSourceSyncTask) error {
segment, ok := task.metacache.GetSegmentByID(task.segmentID)
if !ok {
return merr.WrapErrSegmentNotFound(task.segmentID)
}
startPos := task.startPositions()
checkPoints := []*datapb.CheckPoint{{
SegmentID: task.segmentID,
NumOfRows: segment.FlushedRows() + task.batchRows,
Position: task.checkpoint,
}}
log.Info("SaveBinlogPath for growing source sync",
zap.Int64("SegmentID", task.segmentID),
zap.Int64("CollectionID", task.collectionID),
zap.Int64("ParitionID", task.partitionID),
zap.Any("startPos", startPos),
zap.Any("checkPoints", checkPoints),
zap.String("manifestPath", task.manifestPath),
zap.String("vChannelName", task.channelName),
)
req := &datapb.SaveBinlogPathsRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(0),
commonpbutil.WithMsgID(0),
commonpbutil.WithSourceID(b.serverID),
),
SegmentID: task.segmentID,
CollectionID: task.collectionID,
PartitionID: task.partitionID,
CheckPoints: checkPoints,
StartPositions: startPos,
Flushed: task.IsFlush(),
Dropped: task.IsDrop(),
Channel: task.channelName,
SegLevel: task.level,
StorageVersion: storage.StorageV3,
WithFullBinlogs: false,
ManifestPath: task.manifestPath,
}
err := retry.Handle(ctx, func() (bool, error) {
err := b.broker.SaveBinlogPaths(ctx, req)
if errors.IsAny(err, merr.ErrSegmentNotFound, merr.ErrChannelNotFound) {
log.Warn("meta error found, fail growing source sync",
zap.String("channel", task.channelName),
zap.Int64("segmentID", task.segmentID),
zap.Error(err))
return false, err
}
if err != nil {
return !merr.IsCanceledOrTimeout(err), err
}
return false, nil
}, b.opts...)
if err != nil {
log.Warn("failed to SaveBinlogPaths for growing source sync",
zap.Int64("segmentID", task.segmentID),
zap.Error(err))
return err
}
task.metacache.UpdateSegments(metacache.SetStartPosRecorded(true), metacache.WithSegmentIDs(lo.Map(startPos, func(pos *datapb.SegmentStartPosition, _ int) int64 {
return pos.GetSegmentID()
})...))
return nil
}
func (b *brokerMetaWriter) DropChannel(ctx context.Context, channelName string) error {
err := retry.Handle(ctx, func() (bool, error) {
status, err := b.broker.DropVirtualChannel(context.Background(), &datapb.DropVirtualChannelRequest{
internal/flushcommon/syncmgr/meta_writer_test.go
+51
View File
@@ -8,10 +8,12 @@ import (
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus/internal/flushcommon/broker"
"github.com/milvus-io/milvus/internal/flushcommon/metacache"
"github.com/milvus-io/milvus/internal/flushcommon/metacache/pkoracle"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/retry"
)
@@ -100,6 +102,55 @@ func (s *MetaWriterSuite) TestReturnError() {
s.Error(err)
}
func (s *MetaWriterSuite) TestGrowingSourceSyncMetaErrorsReturnError() {
testCases := []struct {
name string
err error
targetErr error
}{
{
name: "segment_not_found",
err: merr.WrapErrSegmentNotFound(1),
targetErr: merr.ErrSegmentNotFound,
},
{
name: "channel_not_found",
err: merr.WrapErrChannelNotFound("ch"),
targetErr: merr.ErrChannelNotFound,
},
}
for _, tc := range testCases {
s.Run(tc.name, func() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
bfs := pkoracle.NewBloomFilterSet()
seg := metacache.NewSegmentInfo(&datapb.SegmentInfo{
ID: 1,
PartitionID: 2,
}, bfs, nil)
s.metacache.EXPECT().GetSegmentByID(int64(1)).Return(seg, true).Once()
s.metacache.EXPECT().GetSegmentsBy(mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
s.broker.EXPECT().SaveBinlogPaths(mock.Anything, mock.Anything).Return(tc.err).Once()
task := NewGrowingSourceSyncTask().
WithCollectionID(3).
WithPartitionID(2).
WithSegmentID(1).
WithChannelName("ch").
WithCheckpoint(&msgpb.MsgPosition{Timestamp: 100}).
WithBatchRows(10).
WithTargetOffset(10).
WithMetaCache(s.metacache)
task.manifestPath = "manifest"
err := s.writer.UpdateGrowingSourceSync(ctx, task)
s.ErrorIs(err, tc.targetErr)
})
}
}
func TestMetaWriter(t *testing.T) {
suite.Run(t, new(MetaWriterSuite))
}
internal/flushcommon/syncmgr/mock_meta_writer.go
+47
View File
@@ -115,6 +115,53 @@ func (_c *MockMetaWriter_UpdateSync_Call) RunAndReturn(run func(context.Context,
return _c
}
// UpdateGrowingSourceSync provides a mock function with given fields: _a0, _a1
func (_m *MockMetaWriter) UpdateGrowingSourceSync(_a0 context.Context, _a1 *GrowingSourceSyncTask) error {
ret := _m.Called(_a0, _a1)
if len(ret) == 0 {
panic("no return value specified for UpdateGrowingSourceSync")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, *GrowingSourceSyncTask) error); ok {
r0 = rf(_a0, _a1)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockMetaWriter_UpdateGrowingSourceSync_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UpdateGrowingSourceSync'
type MockMetaWriter_UpdateGrowingSourceSync_Call struct {
*mock.Call
}
// UpdateGrowingSourceSync is a helper method to define mock.On call
// - _a0 context.Context
// - _a1 *GrowingSourceSyncTask
func (_e *MockMetaWriter_Expecter) UpdateGrowingSourceSync(_a0 interface{}, _a1 interface{}) *MockMetaWriter_UpdateGrowingSourceSync_Call {
return &MockMetaWriter_UpdateGrowingSourceSync_Call{Call: _e.mock.On("UpdateGrowingSourceSync", _a0, _a1)}
}
func (_c *MockMetaWriter_UpdateGrowingSourceSync_Call) Run(run func(_a0 context.Context, _a1 *GrowingSourceSyncTask)) *MockMetaWriter_UpdateGrowingSourceSync_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*GrowingSourceSyncTask))
})
return _c
}
func (_c *MockMetaWriter_UpdateGrowingSourceSync_Call) Return(_a0 error) *MockMetaWriter_UpdateGrowingSourceSync_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockMetaWriter_UpdateGrowingSourceSync_Call) RunAndReturn(run func(context.Context, *GrowingSourceSyncTask) error) *MockMetaWriter_UpdateGrowingSourceSync_Call {
_c.Call.Return(run)
return _c
}
// NewMockMetaWriter creates a new instance of MockMetaWriter. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewMockMetaWriter(t interface {
internal/flushcommon/syncmgr/mock_task.go
+45
View File
@@ -147,6 +147,51 @@ func (_c *MockTask_HandleError_Call) RunAndReturn(run func(error)) *MockTask_Han
return _c
}
// IsDrop provides a mock function with no fields
func (_m *MockTask) IsDrop() bool {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for IsDrop")
}
var r0 bool
if rf, ok := ret.Get(0).(func() bool); ok {
r0 = rf()
} else {
r0 = ret.Get(0).(bool)
}
return r0
}
// MockTask_IsDrop_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IsDrop'
type MockTask_IsDrop_Call struct {
*mock.Call
}
// IsDrop is a helper method to define mock.On call
func (_e *MockTask_Expecter) IsDrop() *MockTask_IsDrop_Call {
return &MockTask_IsDrop_Call{Call: _e.mock.On("IsDrop")}
}
func (_c *MockTask_IsDrop_Call) Run(run func()) *MockTask_IsDrop_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockTask_IsDrop_Call) Return(_a0 bool) *MockTask_IsDrop_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockTask_IsDrop_Call) RunAndReturn(run func() bool) *MockTask_IsDrop_Call {
_c.Call.Return(run)
return _c
}
// IsFlush provides a mock function with no fields
func (_m *MockTask) IsFlush() bool {
ret := _m.Called()
internal/flushcommon/syncmgr/sync_manager.go
+4
View File
@@ -121,6 +121,8 @@ func (mgr *syncManager) SyncData(ctx context.Context, task Task, callbacks ...fu
switch t := task.(type) {
case *SyncTask:
t.WithChunkManager(mgr.chunkManager)
case *GrowingSourceSyncTask:
t.WithChunkManager(mgr.chunkManager)
}
return mgr.safeSubmitTask(ctx, task, callbacks...), nil
@@ -134,6 +136,8 @@ func (mgr *syncManager) SyncDataWithChunkManager(ctx context.Context, task Task,
switch t := task.(type) {
case *SyncTask:
t.WithChunkManager(chunkManager)
case *GrowingSourceSyncTask:
t.WithChunkManager(chunkManager)
}
return mgr.safeSubmitTask(ctx, task, callbacks...), nil
internal/flushcommon/syncmgr/task.go
+4
View File
@@ -306,6 +306,10 @@ func (t *SyncTask) IsFlush() bool {
return t.pack.isFlush
}
func (t *SyncTask) IsDrop() bool {
return t.pack.isDrop
}
func (t *SyncTask) Binlogs() (map[int64]*datapb.FieldBinlog, map[int64]*datapb.FieldBinlog, *datapb.FieldBinlog, map[int64]*datapb.FieldBinlog) {
return t.insertBinlogs, t.statsBinlogs, t.deltaBinlog, t.bm25Binlogs
}
internal/flushcommon/writebuffer/l0_write_buffer.go
+13 -9
View File
@@ -63,17 +63,20 @@ func (wb *l0WriteBuffer) BufferData(insertData []*InsertData, deleteMsgs []*msgs
wb.mut.Lock()
defer wb.mut.Unlock()
// For TEXT collections, skip Insert buffer here.
// Insert data will be flushed by QueryNode's Growing Segment (via GrowingFlushManager).
// This avoids duplicate flush paths for TEXT data.
if !wb.hasTextFields {
// buffer insert data and add segment if not exists
for _, inData := range insertData {
err := wb.bufferInsert(inData, startPos, endPos, schemaVersion)
if err != nil {
return err
for _, inData := range insertData {
if wb.useGrowingSourceFlush {
targetOffset := wb.growingSourceTargetOffset(inData.segmentID, inData.rowNum)
decision := wb.decideGrowingFlushSource(inData.segmentID, targetOffset, endPos)
if decision.sourceType == metacache.FlushSourceGrowing {
wb.recordGrowingSourceProgress(inData, startPos, endPos, schemaVersion, targetOffset)
continue
}
}
err := wb.bufferInsert(inData, startPos, endPos, schemaVersion)
if err != nil {
return err
}
}
// In streaming service mode, flushed segments no longer maintain a bloom filter.
@@ -82,6 +85,7 @@ func (wb *l0WriteBuffer) BufferData(insertData []*InsertData, deleteMsgs []*msgs
wb.dispatchDeleteMsgsWithoutFilter(deleteMsgs, startPos, endPos)
// update buffer last checkpoint
wb.checkpoint = endPos
wb.updateProcessedTsLocked(endPos.GetTimestamp())
segmentsSync := wb.triggerSync()
for _, segment := range segmentsSync {
internal/flushcommon/writebuffer/l0_write_buffer_test.go
+1403
View File
File diff suppressed because it is too large Load Diff
internal/flushcommon/writebuffer/manager.go
+40 -4
View File
@@ -47,8 +47,11 @@ type BufferManager interface {
// NotifyCheckpointUpdated notify write buffer checkpoint updated to reset flushTs.
NotifyCheckpointUpdated(channel string, ts uint64)
// HasTextFields returns true if the collection on this channel has TEXT fields.
HasTextFields(channel string) bool
// UseGrowingSourceFlush returns true if the collection on this channel has growing-source fields.
UseGrowingSourceFlush(channel string) bool
// GetGrowingFlushProgress returns growing-source progress for the given channel.
// If segmentIDs is empty, all tracked growing-source segments are returned.
GetGrowingFlushProgress(ctx context.Context, channel string, segmentIDs []int64, fenceTs uint64) ([]GrowingFlushSegmentProgress, error)
// Start makes the background check start to work.
Start()
@@ -56,6 +59,10 @@ type BufferManager interface {
Stop()
}
type ReleaseManualFlushNeedChecker interface {
CheckReleaseManualFlushNeed(ctx context.Context, channel string, segmentIDs []int64) (bool, error)
}
// NewManager returns initialized manager as `Manager`
func NewManager(syncMgr syncmgr.SyncManager) BufferManager {
return &bufferManager{
@@ -241,12 +248,41 @@ func (m *bufferManager) BufferData(channel string, insertData []*InsertData, del
return buf.BufferData(insertData, deleteMsgs, startPos, endPos, schemaVersion)
}
func (m *bufferManager) HasTextFields(channel string) bool {
func (m *bufferManager) UseGrowingSourceFlush(channel string) bool {
buf, loaded := m.buffers.Get(channel)
if !loaded {
return false
}
return buf.HasTextFields()
return buf.UseGrowingSourceFlush()
}
func (m *bufferManager) CheckReleaseManualFlushNeed(ctx context.Context, channel string, segmentIDs []int64) (bool, error) {
buf, loaded := m.buffers.Get(channel)
if !loaded {
log.Ctx(ctx).Warn("write buffer not found when checking release manual flush",
zap.String("channel", channel),
zap.Int64s("segmentIDs", segmentIDs))
return true, merr.WrapErrChannelNotFound(channel)
}
checker, ok := buf.(interface {
CheckReleaseManualFlushNeed(segmentIDs []int64) bool
})
if !ok {
return true, nil
}
return checker.CheckReleaseManualFlushNeed(segmentIDs), nil
}
func (m *bufferManager) GetGrowingFlushProgress(ctx context.Context, channel string, segmentIDs []int64, fenceTs uint64) ([]GrowingFlushSegmentProgress, error) {
buf, loaded := m.buffers.Get(channel)
if !loaded {
log.Ctx(ctx).Warn("write buffer not found when get growing flush progress",
zap.String("channel", channel),
zap.Int64s("segmentIDs", segmentIDs),
zap.Uint64("fenceTs", fenceTs))
return nil, merr.WrapErrChannelNotFound(channel)
}
return buf.GetGrowingFlushProgress(ctx, segmentIDs, fenceTs)
}
// GetCheckpoint returns checkpoint for provided channel.
internal/flushcommon/writebuffer/mock_manager.go
+107 -46
View File
@@ -242,52 +242,6 @@ func (_c *MockBufferManager_FlushChannel_Call) RunAndReturn(run func(context.Con
return _c
}
// HasTextFields provides a mock function with given fields: channel
func (_m *MockBufferManager) HasTextFields(channel string) bool {
ret := _m.Called(channel)
if len(ret) == 0 {
panic("no return value specified for HasTextFields")
}
var r0 bool
if rf, ok := ret.Get(0).(func(string) bool); ok {
r0 = rf(channel)
} else {
r0 = ret.Get(0).(bool)
}
return r0
}
// MockBufferManager_HasTextFields_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'HasTextFields'
type MockBufferManager_HasTextFields_Call struct {
*mock.Call
}
// HasTextFields is a helper method to define mock.On call
// - channel string
func (_e *MockBufferManager_Expecter) HasTextFields(channel interface{}) *MockBufferManager_HasTextFields_Call {
return &MockBufferManager_HasTextFields_Call{Call: _e.mock.On("HasTextFields", channel)}
}
func (_c *MockBufferManager_HasTextFields_Call) Run(run func(channel string)) *MockBufferManager_HasTextFields_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(string))
})
return _c
}
func (_c *MockBufferManager_HasTextFields_Call) Return(_a0 bool) *MockBufferManager_HasTextFields_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockBufferManager_HasTextFields_Call) RunAndReturn(run func(string) bool) *MockBufferManager_HasTextFields_Call {
_c.Call.Return(run)
return _c
}
// GetCheckpoint provides a mock function with given fields: channel
func (_m *MockBufferManager) GetCheckpoint(channel string) (*msgpb.MsgPosition, bool, error) {
ret := _m.Called(channel)
@@ -353,6 +307,113 @@ func (_c *MockBufferManager_GetCheckpoint_Call) RunAndReturn(run func(string) (*
return _c
}
// GetGrowingFlushProgress provides a mock function with given fields: ctx, channel, segmentIDs, fenceTs
func (_m *MockBufferManager) GetGrowingFlushProgress(ctx context.Context, channel string, segmentIDs []int64, fenceTs uint64) ([]GrowingFlushSegmentProgress, error) {
ret := _m.Called(ctx, channel, segmentIDs, fenceTs)
if len(ret) == 0 {
panic("no return value specified for GetGrowingFlushProgress")
}
var r0 []GrowingFlushSegmentProgress
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, string, []int64, uint64) ([]GrowingFlushSegmentProgress, error)); ok {
return rf(ctx, channel, segmentIDs, fenceTs)
}
if rf, ok := ret.Get(0).(func(context.Context, string, []int64, uint64) []GrowingFlushSegmentProgress); ok {
r0 = rf(ctx, channel, segmentIDs, fenceTs)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]GrowingFlushSegmentProgress)
}
}
if rf, ok := ret.Get(1).(func(context.Context, string, []int64, uint64) error); ok {
r1 = rf(ctx, channel, segmentIDs, fenceTs)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockBufferManager_GetGrowingFlushProgress_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetGrowingFlushProgress'
type MockBufferManager_GetGrowingFlushProgress_Call struct {
*mock.Call
}
// GetGrowingFlushProgress is a helper method to define mock.On call
// - ctx context.Context
// - channel string
// - segmentIDs []int64
// - fenceTs uint64
func (_e *MockBufferManager_Expecter) GetGrowingFlushProgress(ctx interface{}, channel interface{}, segmentIDs interface{}, fenceTs interface{}) *MockBufferManager_GetGrowingFlushProgress_Call {
return &MockBufferManager_GetGrowingFlushProgress_Call{Call: _e.mock.On("GetGrowingFlushProgress", ctx, channel, segmentIDs, fenceTs)}
}
func (_c *MockBufferManager_GetGrowingFlushProgress_Call) Run(run func(ctx context.Context, channel string, segmentIDs []int64, fenceTs uint64)) *MockBufferManager_GetGrowingFlushProgress_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string), args[2].([]int64), args[3].(uint64))
})
return _c
}
func (_c *MockBufferManager_GetGrowingFlushProgress_Call) Return(_a0 []GrowingFlushSegmentProgress, _a1 error) *MockBufferManager_GetGrowingFlushProgress_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockBufferManager_GetGrowingFlushProgress_Call) RunAndReturn(run func(context.Context, string, []int64, uint64) ([]GrowingFlushSegmentProgress, error)) *MockBufferManager_GetGrowingFlushProgress_Call {
_c.Call.Return(run)
return _c
}
// UseGrowingSourceFlush provides a mock function with given fields: channel
func (_m *MockBufferManager) UseGrowingSourceFlush(channel string) bool {
ret := _m.Called(channel)
if len(ret) == 0 {
panic("no return value specified for UseGrowingSourceFlush")
}
var r0 bool
if rf, ok := ret.Get(0).(func(string) bool); ok {
r0 = rf(channel)
} else {
r0 = ret.Get(0).(bool)
}
return r0
}
// MockBufferManager_UseGrowingSourceFlush_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UseGrowingSourceFlush'
type MockBufferManager_UseGrowingSourceFlush_Call struct {
*mock.Call
}
// UseGrowingSourceFlush is a helper method to define mock.On call
// - channel string
func (_e *MockBufferManager_Expecter) UseGrowingSourceFlush(channel interface{}) *MockBufferManager_UseGrowingSourceFlush_Call {
return &MockBufferManager_UseGrowingSourceFlush_Call{Call: _e.mock.On("UseGrowingSourceFlush", channel)}
}
func (_c *MockBufferManager_UseGrowingSourceFlush_Call) Run(run func(channel string)) *MockBufferManager_UseGrowingSourceFlush_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(string))
})
return _c
}
func (_c *MockBufferManager_UseGrowingSourceFlush_Call) Return(_a0 bool) *MockBufferManager_UseGrowingSourceFlush_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockBufferManager_UseGrowingSourceFlush_Call) RunAndReturn(run func(string) bool) *MockBufferManager_UseGrowingSourceFlush_Call {
_c.Call.Return(run)
return _c
}
// NotifyCheckpointUpdated provides a mock function with given fields: channel, ts
func (_m *MockBufferManager) NotifyCheckpointUpdated(channel string, ts uint64) {
_m.Called(channel, ts)
internal/flushcommon/writebuffer/mock_write_buffer.go
+105 -45
View File
@@ -223,51 +223,6 @@ func (_c *MockWriteBuffer_EvictBuffer_Call) RunAndReturn(run func(...SyncPolicy)
return _c
}
// HasTextFields provides a mock function with no fields
func (_m *MockWriteBuffer) HasTextFields() bool {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for HasTextFields")
}
var r0 bool
if rf, ok := ret.Get(0).(func() bool); ok {
r0 = rf()
} else {
r0 = ret.Get(0).(bool)
}
return r0
}
// MockWriteBuffer_HasTextFields_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'HasTextFields'
type MockWriteBuffer_HasTextFields_Call struct {
*mock.Call
}
// HasTextFields is a helper method to define mock.On call
func (_e *MockWriteBuffer_Expecter) HasTextFields() *MockWriteBuffer_HasTextFields_Call {
return &MockWriteBuffer_HasTextFields_Call{Call: _e.mock.On("HasTextFields")}
}
func (_c *MockWriteBuffer_HasTextFields_Call) Run(run func()) *MockWriteBuffer_HasTextFields_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockWriteBuffer_HasTextFields_Call) Return(_a0 bool) *MockWriteBuffer_HasTextFields_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockWriteBuffer_HasTextFields_Call) RunAndReturn(run func() bool) *MockWriteBuffer_HasTextFields_Call {
_c.Call.Return(run)
return _c
}
// GetCheckpoint provides a mock function with no fields
func (_m *MockWriteBuffer) GetCheckpoint() *msgpb.MsgPosition {
ret := _m.Called()
@@ -360,6 +315,66 @@ func (_c *MockWriteBuffer_GetFlushTimestamp_Call) RunAndReturn(run func() uint64
return _c
}
// GetGrowingFlushProgress provides a mock function with given fields: ctx, segmentIDs, fenceTs
func (_m *MockWriteBuffer) GetGrowingFlushProgress(ctx context.Context, segmentIDs []int64, fenceTs uint64) ([]GrowingFlushSegmentProgress, error) {
ret := _m.Called(ctx, segmentIDs, fenceTs)
if len(ret) == 0 {
panic("no return value specified for GetGrowingFlushProgress")
}
var r0 []GrowingFlushSegmentProgress
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, []int64, uint64) ([]GrowingFlushSegmentProgress, error)); ok {
return rf(ctx, segmentIDs, fenceTs)
}
if rf, ok := ret.Get(0).(func(context.Context, []int64, uint64) []GrowingFlushSegmentProgress); ok {
r0 = rf(ctx, segmentIDs, fenceTs)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]GrowingFlushSegmentProgress)
}
}
if rf, ok := ret.Get(1).(func(context.Context, []int64, uint64) error); ok {
r1 = rf(ctx, segmentIDs, fenceTs)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockWriteBuffer_GetGrowingFlushProgress_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetGrowingFlushProgress'
type MockWriteBuffer_GetGrowingFlushProgress_Call struct {
*mock.Call
}
// GetGrowingFlushProgress is a helper method to define mock.On call
// - ctx context.Context
// - segmentIDs []int64
// - fenceTs uint64
func (_e *MockWriteBuffer_Expecter) GetGrowingFlushProgress(ctx interface{}, segmentIDs interface{}, fenceTs interface{}) *MockWriteBuffer_GetGrowingFlushProgress_Call {
return &MockWriteBuffer_GetGrowingFlushProgress_Call{Call: _e.mock.On("GetGrowingFlushProgress", ctx, segmentIDs, fenceTs)}
}
func (_c *MockWriteBuffer_GetGrowingFlushProgress_Call) Run(run func(ctx context.Context, segmentIDs []int64, fenceTs uint64)) *MockWriteBuffer_GetGrowingFlushProgress_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].([]int64), args[2].(uint64))
})
return _c
}
func (_c *MockWriteBuffer_GetGrowingFlushProgress_Call) Return(_a0 []GrowingFlushSegmentProgress, _a1 error) *MockWriteBuffer_GetGrowingFlushProgress_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockWriteBuffer_GetGrowingFlushProgress_Call) RunAndReturn(run func(context.Context, []int64, uint64) ([]GrowingFlushSegmentProgress, error)) *MockWriteBuffer_GetGrowingFlushProgress_Call {
_c.Call.Return(run)
return _c
}
// HasSegment provides a mock function with given fields: segmentID
func (_m *MockWriteBuffer) HasSegment(segmentID int64) bool {
ret := _m.Called(segmentID)
@@ -406,6 +421,51 @@ func (_c *MockWriteBuffer_HasSegment_Call) RunAndReturn(run func(int64) bool) *M
return _c
}
// UseGrowingSourceFlush provides a mock function with no fields
func (_m *MockWriteBuffer) UseGrowingSourceFlush() bool {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for UseGrowingSourceFlush")
}
var r0 bool
if rf, ok := ret.Get(0).(func() bool); ok {
r0 = rf()
} else {
r0 = ret.Get(0).(bool)
}
return r0
}
// MockWriteBuffer_UseGrowingSourceFlush_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UseGrowingSourceFlush'
type MockWriteBuffer_UseGrowingSourceFlush_Call struct {
*mock.Call
}
// UseGrowingSourceFlush is a helper method to define mock.On call
func (_e *MockWriteBuffer_Expecter) UseGrowingSourceFlush() *MockWriteBuffer_UseGrowingSourceFlush_Call {
return &MockWriteBuffer_UseGrowingSourceFlush_Call{Call: _e.mock.On("UseGrowingSourceFlush")}
}
func (_c *MockWriteBuffer_UseGrowingSourceFlush_Call) Run(run func()) *MockWriteBuffer_UseGrowingSourceFlush_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockWriteBuffer_UseGrowingSourceFlush_Call) Return(_a0 bool) *MockWriteBuffer_UseGrowingSourceFlush_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockWriteBuffer_UseGrowingSourceFlush_Call) RunAndReturn(run func() bool) *MockWriteBuffer_UseGrowingSourceFlush_Call {
_c.Call.Return(run)
return _c
}
// MemorySize provides a mock function with no fields
func (_m *MockWriteBuffer) MemorySize() int64 {
ret := _m.Called()
internal/flushcommon/writebuffer/options.go
+16
View File
@@ -3,6 +3,7 @@ package writebuffer
import (
"time"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus/internal/allocator"
"github.com/milvus-io/milvus/internal/flushcommon/metacache"
"github.com/milvus-io/milvus/internal/flushcommon/syncmgr"
@@ -13,6 +14,12 @@ type WriteBufferOption func(opt *writeBufferOption)
type TaskObserverCallback func(t syncmgr.Task, err error)
// GrowingSourceResolver resolves an optional in-memory growing segment source
// for growing-source flush. GrowingSourcePending means the growing source exists but has not
// caught up to targetOffset yet; WriteBuffer should only be used when the state
// is GrowingSourceUnavailable.
type GrowingSourceResolver func(segmentID int64, targetOffset int64, endPos *msgpb.MsgPosition) (syncmgr.GrowingFlushSource, syncmgr.GrowingSourceState)
type writeBufferOption struct {
idAllocator allocator.Interface
syncPolicies []SyncPolicy
@@ -22,6 +29,9 @@ type writeBufferOption struct {
errorHandler func(error)
taskObserverCallback TaskObserverCallback
storageVersion int64
growingSourceResolver GrowingSourceResolver
growingSourceRetryInterval time.Duration
}
func defaultWBOption(metacache metacache.MetaCache) *writeBufferOption {
@@ -76,3 +86,9 @@ func WithTaskObserverCallback(callback TaskObserverCallback) WriteBufferOption {
opt.taskObserverCallback = callback
}
}
func WithGrowingSourceResolver(resolver GrowingSourceResolver) WriteBufferOption {
return func(opt *writeBufferOption) {
opt.growingSourceResolver = resolver
}
}
internal/flushcommon/writebuffer/write_buffer.go
+782 -36
View File
File diff suppressed because it is too large Load Diff
internal/flushcommon/writebuffer/write_buffer_test.go
+190 -3
View File
@@ -17,12 +17,14 @@ import (
"github.com/milvus-io/milvus/internal/flushcommon/metacache"
"github.com/milvus-io/milvus/internal/flushcommon/metacache/pkoracle"
"github.com/milvus-io/milvus/internal/flushcommon/syncmgr"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/util/conc"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
)
type WriteBufferSuite struct {
@@ -74,6 +76,100 @@ func (s *WriteBufferSuite) TestHasSegment() {
s.True(s.wb.HasSegment(segmentID))
}
func (s *WriteBufferSuite) TestCreateNewGrowingSegmentStorageVersion() {
param := paramtable.Get()
param.Save(param.CommonCfg.UseLoonFFI.Key, "false")
defer param.Reset(param.CommonCfg.UseLoonFFI.Key)
param.Save(param.CommonCfg.EnableGrowingSourceFlush.Key, "false")
defer param.Reset(param.CommonCfg.EnableGrowingSourceFlush.Key)
s.Run("non_text_uses_v2_when_ffi_disabled", func() {
s.wb.useGrowingSourceFlush = false
s.False(s.wb.UseGrowingSourceFlush())
s.metacache.EXPECT().GetSegmentByID(int64(2001)).Return(nil, false).Once()
s.metacache.EXPECT().AddSegment(mock.MatchedBy(func(info *datapb.SegmentInfo) bool {
return info.GetStorageVersion() == storage.StorageV2 &&
info.GetManifestPath() == "" &&
info.GetSchemaVersion() == 11
}), mock.Anything, mock.Anything, mock.Anything).Return().Once()
s.wb.CreateNewGrowingSegment(10, 2001, nil, 11)
})
s.Run("growing_source_does_not_force_v3_manifest_when_ffi_disabled", func() {
s.wb.useGrowingSourceFlush = true
s.metacache.EXPECT().GetSegmentByID(int64(2002)).Return(nil, false).Once()
s.metacache.EXPECT().AddSegment(mock.MatchedBy(func(info *datapb.SegmentInfo) bool {
return info.GetStorageVersion() == storage.StorageV2 &&
info.GetManifestPath() == "" &&
info.GetSchemaVersion() == 12
}), mock.Anything, mock.Anything, mock.Anything).Return().Once()
s.wb.CreateNewGrowingSegment(10, 2002, nil, 12)
})
s.Run("text_schema_does_not_enable_growing_source_when_ffi_disabled", func() {
textSchema := &schemapb.CollectionSchema{
Name: "wb_text_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_Int64, IsPrimaryKey: true, Name: "pk"},
{FieldID: 101, DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "128"},
}},
{FieldID: 102, DataType: schemapb.DataType_Text, Name: "text"},
},
}
mc := metacache.NewMockMetaCache(s.T())
mc.EXPECT().GetSchema(mock.Anything).Return(textSchema).Maybe()
mc.EXPECT().Collection().Return(s.collID).Maybe()
wb, err := newWriteBufferBase(s.channelName, mc, s.syncMgr, &writeBufferOption{})
s.Require().NoError(err)
s.False(wb.UseGrowingSourceFlush())
mc.EXPECT().GetSegmentByID(int64(2003)).Return(nil, false).Once()
mc.EXPECT().AddSegment(mock.MatchedBy(func(info *datapb.SegmentInfo) bool {
return info.GetStorageVersion() == storage.StorageV2 &&
info.GetManifestPath() == "" &&
info.GetSchemaVersion() == 13
}), mock.Anything, mock.Anything, mock.Anything).Return().Once()
wb.CreateNewGrowingSegment(10, 2003, nil, 13)
})
s.Run("text_schema_uses_v3_manifest_when_ffi_enabled", func() {
param.Save(param.CommonCfg.UseLoonFFI.Key, "true")
defer param.Save(param.CommonCfg.UseLoonFFI.Key, "false")
textSchema := &schemapb.CollectionSchema{
Name: "wb_text_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_Int64, IsPrimaryKey: true, Name: "pk"},
{FieldID: 101, DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "128"},
}},
{FieldID: 102, DataType: schemapb.DataType_Text, Name: "text"},
},
}
mc := metacache.NewMockMetaCache(s.T())
mc.EXPECT().GetSchema(mock.Anything).Return(textSchema).Maybe()
mc.EXPECT().Collection().Return(s.collID).Maybe()
wb, err := newWriteBufferBase(s.channelName, mc, s.syncMgr, &writeBufferOption{})
s.Require().NoError(err)
s.True(wb.UseGrowingSourceFlush())
mc.EXPECT().GetSegmentByID(int64(2004)).Return(nil, false).Once()
mc.EXPECT().AddSegment(mock.MatchedBy(func(info *datapb.SegmentInfo) bool {
return info.GetStorageVersion() == storage.StorageV3 &&
info.GetManifestPath() != "" &&
info.GetSchemaVersion() == 14
}), mock.Anything, mock.Anything, mock.Anything).Return().Once()
wb.CreateNewGrowingSegment(10, 2004, nil, 14)
})
}
func (s *WriteBufferSuite) TestFlushSegments() {
segmentID := int64(1001)
@@ -90,7 +186,7 @@ func (s *WriteBufferSuite) TestSealSegmentsMissingSegment() {
segmentID := int64(1001)
s.Run("non_text_returns_error", func() {
s.wb.hasTextFields = false
s.wb.useGrowingSourceFlush = false
s.metacache.EXPECT().GetSegmentByID(segmentID).Return(nil, false).Once()
err := s.wb.SealSegments(context.Background(), []int64{segmentID})
@@ -98,9 +194,9 @@ func (s *WriteBufferSuite) TestSealSegmentsMissingSegment() {
})
s.Run("text_skips_missing_segment", func() {
s.wb.hasTextFields = true
s.wb.useGrowingSourceFlush = true
defer func() {
s.wb.hasTextFields = false
s.wb.useGrowingSourceFlush = false
}()
s.metacache.EXPECT().GetSegmentByID(segmentID).Return(nil, false).Once()
@@ -436,6 +532,97 @@ func (s *WriteBufferSuite) TestEvictBuffer() {
})
}
func (s *WriteBufferSuite) TestGrowingSourceProgressSelectedByPolicy() {
paramtable.Get().Save(paramtable.Get().DataNodeCfg.SyncPeriod.Key, "1")
defer paramtable.Get().Reset(paramtable.Get().DataNodeCfg.SyncPeriod.Key)
paramtable.Get().Save(paramtable.Get().DataNodeCfg.FlushInsertBufferSize.Key, "100")
defer paramtable.Get().Reset(paramtable.Get().DataNodeCfg.FlushInsertBufferSize.Key)
originalEstSize := s.wb.estSizePerRecord
s.wb.estSizePerRecord = 10
defer func() {
s.wb.estSizePerRecord = originalEstSize
}()
now := time.Now()
recentTs := tsoutil.ComposeTSByTime(now.Add(500*time.Millisecond), 0)
staleTs := tsoutil.ComposeTSByTime(now.Add(2*time.Second), 0)
startTs := tsoutil.ComposeTSByTime(now, 0)
s.Run("pending_flush", func() {
selected := s.wb.growingSourceProgressSelectedByPolicy(recentTs, 1001, &growingSourceProgress{
segmentID: 1001,
pendingFlush: true,
})
s.True(selected)
})
s.Run("sealed_segment", func() {
segment := metacache.NewSegmentInfo(&datapb.SegmentInfo{
ID: 1002,
State: commonpb.SegmentState_Sealed,
}, nil, nil)
s.metacache.EXPECT().GetSegmentByID(int64(1002)).Return(segment, true).Once()
selected := s.wb.growingSourceProgressSelectedByPolicy(recentTs, 1002, &growingSourceProgress{
segmentID: 1002,
})
s.True(selected)
})
s.Run("recent_progress", func() {
s.metacache.EXPECT().GetSegmentByID(int64(1003)).Return(nil, false).Once()
selected := s.wb.growingSourceProgressSelectedByPolicy(recentTs, 1003, &growingSourceProgress{
segmentID: 1003,
batches: []growingSourceProgressBatch{
{startPosition: &msgpb.MsgPosition{Timestamp: startTs}},
},
})
s.False(selected)
})
s.Run("below_row_threshold", func() {
segment := metacache.NewSegmentInfo(&datapb.SegmentInfo{
ID: 1004,
}, nil, nil)
s.metacache.EXPECT().GetSegmentByID(int64(1004)).Return(segment, true).Once()
selected := s.wb.growingSourceProgressSelectedByPolicy(recentTs, 1004, &growingSourceProgress{
segmentID: 1004,
targetOffset: 9,
batches: []growingSourceProgressBatch{
{startPosition: &msgpb.MsgPosition{Timestamp: startTs}},
},
})
s.False(selected)
})
s.Run("row_threshold", func() {
segment := metacache.NewSegmentInfo(&datapb.SegmentInfo{
ID: 1005,
}, nil, nil)
s.metacache.EXPECT().GetSegmentByID(int64(1005)).Return(segment, true).Once()
selected := s.wb.growingSourceProgressSelectedByPolicy(recentTs, 1005, &growingSourceProgress{
segmentID: 1005,
targetOffset: 10,
})
s.True(selected)
})
s.Run("stale_progress", func() {
s.metacache.EXPECT().GetSegmentByID(int64(1006)).Return(nil, false).Once()
selected := s.wb.growingSourceProgressSelectedByPolicy(staleTs, 1006, &growingSourceProgress{
segmentID: 1006,
batches: []growingSourceProgressBatch{
{startPosition: &msgpb.MsgPosition{Timestamp: startTs}},
},
})
s.True(selected)
})
}
func (s *WriteBufferSuite) TestDropPartitions() {
wb, err := newWriteBufferBase(s.channelName, s.metacache, s.syncMgr, &writeBufferOption{
pkStatsFactory: func(vchannel *datapb.SegmentInfo) pkoracle.PkStat {
internal/mocks/distributed/mock_streaming/mock_WALAccesser.go
+65 -6
View File
@@ -27,7 +27,7 @@ func (_m *MockWALAccesser) EXPECT() *MockWALAccesser_Expecter {
}
// AppendMessages provides a mock function with given fields: ctx, msgs
func (_m *MockWALAccesser) AppendMessages(ctx context.Context, msgs ...message.MutableMessage) streaming.AppendResponses {
func (_m *MockWALAccesser) AppendMessages(ctx context.Context, msgs ...message.MutableMessage) types.AppendResponses {
_va := make([]interface{}, len(msgs))
for _i := range msgs {
_va[_i] = msgs[_i]
@@ -41,11 +41,11 @@ func (_m *MockWALAccesser) AppendMessages(ctx context.Context, msgs ...message.M
panic("no return value specified for AppendMessages")
}
var r0 streaming.AppendResponses
if rf, ok := ret.Get(0).(func(context.Context, ...message.MutableMessage) streaming.AppendResponses); ok {
var r0 types.AppendResponses
if rf, ok := ret.Get(0).(func(context.Context, ...message.MutableMessage) types.AppendResponses); ok {
r0 = rf(ctx, msgs...)
} else {
r0 = ret.Get(0).(streaming.AppendResponses)
r0 = ret.Get(0).(types.AppendResponses)
}
return r0
@@ -77,12 +77,12 @@ func (_c *MockWALAccesser_AppendMessages_Call) Run(run func(ctx context.Context,
return _c
}
func (_c *MockWALAccesser_AppendMessages_Call) Return(_a0 streaming.AppendResponses) *MockWALAccesser_AppendMessages_Call {
func (_c *MockWALAccesser_AppendMessages_Call) Return(_a0 types.AppendResponses) *MockWALAccesser_AppendMessages_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockWALAccesser_AppendMessages_Call) RunAndReturn(run func(context.Context, ...message.MutableMessage) streaming.AppendResponses) *MockWALAccesser_AppendMessages_Call {
func (_c *MockWALAccesser_AppendMessages_Call) RunAndReturn(run func(context.Context, ...message.MutableMessage) types.AppendResponses) *MockWALAccesser_AppendMessages_Call {
_c.Call.Return(run)
return _c
}
@@ -320,6 +320,65 @@ func (_c *MockWALAccesser_Local_Call) RunAndReturn(run func() streaming.Local) *
return _c
}
// PrepareReleaseManualFlush provides a mock function with given fields: ctx, collectionID, vchannel, releaseSegmentIDs
func (_m *MockWALAccesser) PrepareReleaseManualFlush(ctx context.Context, collectionID int64, vchannel string, releaseSegmentIDs []int64) (bool, error) {
ret := _m.Called(ctx, collectionID, vchannel, releaseSegmentIDs)
if len(ret) == 0 {
panic("no return value specified for PrepareReleaseManualFlush")
}
var r0 bool
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, int64, string, []int64) (bool, error)); ok {
return rf(ctx, collectionID, vchannel, releaseSegmentIDs)
}
if rf, ok := ret.Get(0).(func(context.Context, int64, string, []int64) bool); ok {
r0 = rf(ctx, collectionID, vchannel, releaseSegmentIDs)
} else {
r0 = ret.Get(0).(bool)
}
if rf, ok := ret.Get(1).(func(context.Context, int64, string, []int64) error); ok {
r1 = rf(ctx, collectionID, vchannel, releaseSegmentIDs)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockWALAccesser_PrepareReleaseManualFlush_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'PrepareReleaseManualFlush'
type MockWALAccesser_PrepareReleaseManualFlush_Call struct {
*mock.Call
}
// PrepareReleaseManualFlush is a helper method to define mock.On call
// - ctx context.Context
// - collectionID int64
// - vchannel string
// - releaseSegmentIDs []int64
func (_e *MockWALAccesser_Expecter) PrepareReleaseManualFlush(ctx interface{}, collectionID interface{}, vchannel interface{}, releaseSegmentIDs interface{}) *MockWALAccesser_PrepareReleaseManualFlush_Call {
return &MockWALAccesser_PrepareReleaseManualFlush_Call{Call: _e.mock.On("PrepareReleaseManualFlush", ctx, collectionID, vchannel, releaseSegmentIDs)}
}
func (_c *MockWALAccesser_PrepareReleaseManualFlush_Call) Run(run func(ctx context.Context, collectionID int64, vchannel string, releaseSegmentIDs []int64)) *MockWALAccesser_PrepareReleaseManualFlush_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(int64), args[2].(string), args[3].([]int64))
})
return _c
}
func (_c *MockWALAccesser_PrepareReleaseManualFlush_Call) Return(_a0 bool, _a1 error) *MockWALAccesser_PrepareReleaseManualFlush_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockWALAccesser_PrepareReleaseManualFlush_Call) RunAndReturn(run func(context.Context, int64, string, []int64) (bool, error)) *MockWALAccesser_PrepareReleaseManualFlush_Call {
_c.Call.Return(run)
return _c
}
// RawAppend provides a mock function with given fields: ctx, msgs, opts
func (_m *MockWALAccesser) RawAppend(ctx context.Context, msgs message.MutableMessage, opts ...streaming.AppendOption) (*types.AppendResult, error) {
_va := make([]interface{}, len(opts))
internal/mocks/streamingnode/client/mock_handler/mock_HandlerClient.go
+93 -29
View File
@@ -5,12 +5,17 @@ package mock_handler
import (
context "context"
consumer "github.com/milvus-io/milvus/internal/streamingnode/client/handler/consumer"
handler "github.com/milvus-io/milvus/internal/streamingnode/client/handler"
mock "github.com/stretchr/testify/mock"
producer "github.com/milvus-io/milvus/internal/streamingnode/client/handler/producer"
types "github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
wal "github.com/milvus-io/milvus/internal/streamingnode/server/wal"
utility "github.com/milvus-io/milvus/internal/streamingnode/server/wal/utility"
)
// MockHandlerClient is an autogenerated mock type for the HandlerClient type
@@ -59,23 +64,23 @@ func (_c *MockHandlerClient_Close_Call) RunAndReturn(run func()) *MockHandlerCli
}
// CreateConsumer provides a mock function with given fields: ctx, opts
func (_m *MockHandlerClient) CreateConsumer(ctx context.Context, opts *handler.ConsumerOptions) (handler.Consumer, error) {
func (_m *MockHandlerClient) CreateConsumer(ctx context.Context, opts *handler.ConsumerOptions) (consumer.Consumer, error) {
ret := _m.Called(ctx, opts)
if len(ret) == 0 {
panic("no return value specified for CreateConsumer")
}
var r0 handler.Consumer
var r0 consumer.Consumer
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *handler.ConsumerOptions) (handler.Consumer, error)); ok {
if rf, ok := ret.Get(0).(func(context.Context, *handler.ConsumerOptions) (consumer.Consumer, error)); ok {
return rf(ctx, opts)
}
if rf, ok := ret.Get(0).(func(context.Context, *handler.ConsumerOptions) handler.Consumer); ok {
if rf, ok := ret.Get(0).(func(context.Context, *handler.ConsumerOptions) consumer.Consumer); ok {
r0 = rf(ctx, opts)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(handler.Consumer)
r0 = ret.Get(0).(consumer.Consumer)
}
}
@@ -107,34 +112,34 @@ func (_c *MockHandlerClient_CreateConsumer_Call) Run(run func(ctx context.Contex
return _c
}
func (_c *MockHandlerClient_CreateConsumer_Call) Return(_a0 handler.Consumer, _a1 error) *MockHandlerClient_CreateConsumer_Call {
func (_c *MockHandlerClient_CreateConsumer_Call) Return(_a0 consumer.Consumer, _a1 error) *MockHandlerClient_CreateConsumer_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockHandlerClient_CreateConsumer_Call) RunAndReturn(run func(context.Context, *handler.ConsumerOptions) (handler.Consumer, error)) *MockHandlerClient_CreateConsumer_Call {
func (_c *MockHandlerClient_CreateConsumer_Call) RunAndReturn(run func(context.Context, *handler.ConsumerOptions) (consumer.Consumer, error)) *MockHandlerClient_CreateConsumer_Call {
_c.Call.Return(run)
return _c
}
// CreateProducer provides a mock function with given fields: ctx, opts
func (_m *MockHandlerClient) CreateProducer(ctx context.Context, opts *handler.ProducerOptions) (handler.Producer, error) {
func (_m *MockHandlerClient) CreateProducer(ctx context.Context, opts *handler.ProducerOptions) (producer.Producer, error) {
ret := _m.Called(ctx, opts)
if len(ret) == 0 {
panic("no return value specified for CreateProducer")
}
var r0 handler.Producer
var r0 producer.Producer
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *handler.ProducerOptions) (handler.Producer, error)); ok {
if rf, ok := ret.Get(0).(func(context.Context, *handler.ProducerOptions) (producer.Producer, error)); ok {
return rf(ctx, opts)
}
if rf, ok := ret.Get(0).(func(context.Context, *handler.ProducerOptions) handler.Producer); ok {
if rf, ok := ret.Get(0).(func(context.Context, *handler.ProducerOptions) producer.Producer); ok {
r0 = rf(ctx, opts)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(handler.Producer)
r0 = ret.Get(0).(producer.Producer)
}
}
@@ -166,12 +171,12 @@ func (_c *MockHandlerClient_CreateProducer_Call) Run(run func(ctx context.Contex
return _c
}
func (_c *MockHandlerClient_CreateProducer_Call) Return(_a0 handler.Producer, _a1 error) *MockHandlerClient_CreateProducer_Call {
func (_c *MockHandlerClient_CreateProducer_Call) Return(_a0 producer.Producer, _a1 error) *MockHandlerClient_CreateProducer_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockHandlerClient_CreateProducer_Call) RunAndReturn(run func(context.Context, *handler.ProducerOptions) (handler.Producer, error)) *MockHandlerClient_CreateProducer_Call {
func (_c *MockHandlerClient_CreateProducer_Call) RunAndReturn(run func(context.Context, *handler.ProducerOptions) (producer.Producer, error)) *MockHandlerClient_CreateProducer_Call {
_c.Call.Return(run)
return _c
}
@@ -234,23 +239,23 @@ func (_c *MockHandlerClient_GetLatestMVCCTimestampIfLocal_Call) RunAndReturn(run
}
// GetReplicateCheckpoint provides a mock function with given fields: ctx, channelName
func (_m *MockHandlerClient) GetReplicateCheckpoint(ctx context.Context, channelName string) (*wal.ReplicateCheckpoint, error) {
func (_m *MockHandlerClient) GetReplicateCheckpoint(ctx context.Context, channelName string) (*utility.ReplicateCheckpoint, error) {
ret := _m.Called(ctx, channelName)
if len(ret) == 0 {
panic("no return value specified for GetReplicateCheckpoint")
}
var r0 *wal.ReplicateCheckpoint
var r0 *utility.ReplicateCheckpoint
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, string) (*wal.ReplicateCheckpoint, error)); ok {
if rf, ok := ret.Get(0).(func(context.Context, string) (*utility.ReplicateCheckpoint, error)); ok {
return rf(ctx, channelName)
}
if rf, ok := ret.Get(0).(func(context.Context, string) *wal.ReplicateCheckpoint); ok {
if rf, ok := ret.Get(0).(func(context.Context, string) *utility.ReplicateCheckpoint); ok {
r0 = rf(ctx, channelName)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*wal.ReplicateCheckpoint)
r0 = ret.Get(0).(*utility.ReplicateCheckpoint)
}
}
@@ -282,34 +287,34 @@ func (_c *MockHandlerClient_GetReplicateCheckpoint_Call) Run(run func(ctx contex
return _c
}
func (_c *MockHandlerClient_GetReplicateCheckpoint_Call) Return(_a0 *wal.ReplicateCheckpoint, _a1 error) *MockHandlerClient_GetReplicateCheckpoint_Call {
func (_c *MockHandlerClient_GetReplicateCheckpoint_Call) Return(_a0 *utility.ReplicateCheckpoint, _a1 error) *MockHandlerClient_GetReplicateCheckpoint_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockHandlerClient_GetReplicateCheckpoint_Call) RunAndReturn(run func(context.Context, string) (*wal.ReplicateCheckpoint, error)) *MockHandlerClient_GetReplicateCheckpoint_Call {
func (_c *MockHandlerClient_GetReplicateCheckpoint_Call) RunAndReturn(run func(context.Context, string) (*utility.ReplicateCheckpoint, error)) *MockHandlerClient_GetReplicateCheckpoint_Call {
_c.Call.Return(run)
return _c
}
// GetSalvageCheckpoint provides a mock function with given fields: ctx, channelName
func (_m *MockHandlerClient) GetSalvageCheckpoint(ctx context.Context, channelName string) ([]*wal.ReplicateCheckpoint, error) {
func (_m *MockHandlerClient) GetSalvageCheckpoint(ctx context.Context, channelName string) ([]*utility.ReplicateCheckpoint, error) {
ret := _m.Called(ctx, channelName)
if len(ret) == 0 {
panic("no return value specified for GetSalvageCheckpoint")
}
var r0 []*wal.ReplicateCheckpoint
var r0 []*utility.ReplicateCheckpoint
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, string) ([]*wal.ReplicateCheckpoint, error)); ok {
if rf, ok := ret.Get(0).(func(context.Context, string) ([]*utility.ReplicateCheckpoint, error)); ok {
return rf(ctx, channelName)
}
if rf, ok := ret.Get(0).(func(context.Context, string) []*wal.ReplicateCheckpoint); ok {
if rf, ok := ret.Get(0).(func(context.Context, string) []*utility.ReplicateCheckpoint); ok {
r0 = rf(ctx, channelName)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*wal.ReplicateCheckpoint)
r0 = ret.Get(0).([]*utility.ReplicateCheckpoint)
}
}
@@ -341,12 +346,12 @@ func (_c *MockHandlerClient_GetSalvageCheckpoint_Call) Run(run func(ctx context.
return _c
}
func (_c *MockHandlerClient_GetSalvageCheckpoint_Call) Return(_a0 []*wal.ReplicateCheckpoint, _a1 error) *MockHandlerClient_GetSalvageCheckpoint_Call {
func (_c *MockHandlerClient_GetSalvageCheckpoint_Call) Return(_a0 []*utility.ReplicateCheckpoint, _a1 error) *MockHandlerClient_GetSalvageCheckpoint_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockHandlerClient_GetSalvageCheckpoint_Call) RunAndReturn(run func(context.Context, string) ([]*wal.ReplicateCheckpoint, error)) *MockHandlerClient_GetSalvageCheckpoint_Call {
func (_c *MockHandlerClient_GetSalvageCheckpoint_Call) RunAndReturn(run func(context.Context, string) ([]*utility.ReplicateCheckpoint, error)) *MockHandlerClient_GetSalvageCheckpoint_Call {
_c.Call.Return(run)
return _c
}
@@ -409,6 +414,65 @@ func (_c *MockHandlerClient_GetWALMetricsIfLocal_Call) RunAndReturn(run func(con
return _c
}
// PrepareReleaseManualFlush provides a mock function with given fields: ctx, collectionID, vchannel, releaseSegmentIDs
func (_m *MockHandlerClient) PrepareReleaseManualFlush(ctx context.Context, collectionID int64, vchannel string, releaseSegmentIDs []int64) (bool, error) {
ret := _m.Called(ctx, collectionID, vchannel, releaseSegmentIDs)
if len(ret) == 0 {
panic("no return value specified for PrepareReleaseManualFlush")
}
var r0 bool
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, int64, string, []int64) (bool, error)); ok {
return rf(ctx, collectionID, vchannel, releaseSegmentIDs)
}
if rf, ok := ret.Get(0).(func(context.Context, int64, string, []int64) bool); ok {
r0 = rf(ctx, collectionID, vchannel, releaseSegmentIDs)
} else {
r0 = ret.Get(0).(bool)
}
if rf, ok := ret.Get(1).(func(context.Context, int64, string, []int64) error); ok {
r1 = rf(ctx, collectionID, vchannel, releaseSegmentIDs)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockHandlerClient_PrepareReleaseManualFlush_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'PrepareReleaseManualFlush'
type MockHandlerClient_PrepareReleaseManualFlush_Call struct {
*mock.Call
}
// PrepareReleaseManualFlush is a helper method to define mock.On call
// - ctx context.Context
// - collectionID int64
// - vchannel string
// - releaseSegmentIDs []int64
func (_e *MockHandlerClient_Expecter) PrepareReleaseManualFlush(ctx interface{}, collectionID interface{}, vchannel interface{}, releaseSegmentIDs interface{}) *MockHandlerClient_PrepareReleaseManualFlush_Call {
return &MockHandlerClient_PrepareReleaseManualFlush_Call{Call: _e.mock.On("PrepareReleaseManualFlush", ctx, collectionID, vchannel, releaseSegmentIDs)}
}
func (_c *MockHandlerClient_PrepareReleaseManualFlush_Call) Run(run func(ctx context.Context, collectionID int64, vchannel string, releaseSegmentIDs []int64)) *MockHandlerClient_PrepareReleaseManualFlush_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(int64), args[2].(string), args[3].([]int64))
})
return _c
}
func (_c *MockHandlerClient_PrepareReleaseManualFlush_Call) Return(_a0 bool, _a1 error) *MockHandlerClient_PrepareReleaseManualFlush_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockHandlerClient_PrepareReleaseManualFlush_Call) RunAndReturn(run func(context.Context, int64, string, []int64) (bool, error)) *MockHandlerClient_PrepareReleaseManualFlush_Call {
_c.Call.Return(run)
return _c
}
// NewMockHandlerClient creates a new instance of MockHandlerClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewMockHandlerClient(t interface {
internal/proxy/task.go
+21
View File
@@ -205,6 +205,13 @@ type dmlTask interface {
type BaseInsertTask = msgstream.InsertMsg
func validateTextStorageV3Enabled(schema *schemapb.CollectionSchema) error {
if err := typeutil.ValidateTextRequiresStorageV3(schema, Params.CommonCfg.UseLoonFFI.GetAsBool()); err != nil {
return merr.WrapErrParameterInvalidMsg("%s", err.Error())
}
return nil
}
type createCollectionTask struct {
baseTask
Condition
@@ -450,6 +457,9 @@ func (t *createCollectionTask) PreExecute(ctx context.Context) error {
if err := typeutil.NormalizeAndValidateExternalCollectionSchema(t.schema); err != nil {
return err
}
if err := validateTextStorageV3Enabled(t.schema); err != nil {
return err
}
// External collections must be single-shard: the refresh mechanism assigns all
// segments to VChannelNames[0], so multiple shards would leave segments orphaned.
@@ -944,6 +954,11 @@ func validateAddFieldRequest(schema *schemapb.CollectionSchema, newFieldSchema *
return merr.WrapErrParameterInvalidMsg("%s", err.Error())
}
}
schemaWithNewField := proto.Clone(schema).(*schemapb.CollectionSchema)
schemaWithNewField.Fields = append(schemaWithNewField.GetFields(), proto.Clone(newFieldSchema).(*schemapb.FieldSchema))
if err := validateTextStorageV3Enabled(schemaWithNewField); err != nil {
return err
}
if funcutil.SliceContain([]string{common.RowIDFieldName, common.TimeStampFieldName, common.MetaFieldName, common.NamespaceFieldName, common.VirtualPKFieldName}, newFieldSchema.GetName()) {
return merr.WrapErrParameterInvalidMsg(fmt.Sprintf("not support to add system field, field name = %s", newFieldSchema.GetName()))
}
@@ -3086,6 +3101,9 @@ func (t *loadCollectionTask) Execute(ctx context.Context) (err error) {
if err != nil {
return err
}
if err := validateTextStorageV3Enabled(collSchema.CollectionSchema); err != nil {
return err
}
// prepare load field list
loadFields, err := collSchema.GetLoadFieldIDs(t.GetLoadFields(), t.GetSkipLoadDynamicField())
if err != nil {
@@ -3344,6 +3362,9 @@ func (t *loadPartitionsTask) Execute(ctx context.Context) error {
if err != nil {
return err
}
if err := validateTextStorageV3Enabled(collSchema.CollectionSchema); err != nil {
return err
}
// prepare load field list
loadFields, err := collSchema.GetLoadFieldIDs(t.GetLoadFields(), t.GetSkipLoadDynamicField())
if err != nil {
internal/proxy/task_delete.go
+3
View File
@@ -287,6 +287,9 @@ func (dr *deleteRunner) Init(ctx context.Context) error {
if err != nil {
return ErrWithLog(log, "Failed to get collection schema", err)
}
if err := validateTextStorageV3Enabled(dr.schema.CollectionSchema); err != nil {
return ErrWithLog(log, "TEXT field requires StorageV3", err)
}
colInfo, err := globalMetaCache.GetCollectionInfo(ctx, dr.req.GetDbName(), collName, dr.collectionID)
if err != nil {
internal/proxy/task_flush_streaming.go
+2 -21
View File
@@ -52,19 +52,6 @@ func (t *flushTask) Execute(ctx context.Context) error {
return merr.WrapErrAsInputErrorWhen(err, merr.ErrCollectionNotFound, merr.ErrDatabaseNotFound)
}
// Check if collection has TEXT fields — TEXT collections use deferred flush via WAL.
schemaInfo, err := globalMetaCache.GetCollectionSchema(ctx, t.DbName, collName)
if err != nil {
return err
}
hasTextFields := false
for _, field := range schemaInfo.GetFields() {
if field.GetDataType() == schemapb.DataType_Text {
hasTextFields = true
break
}
}
vchannels, err := t.chMgr.getVChannels(collID)
if err != nil {
return err
@@ -77,9 +64,7 @@ func (t *flushTask) Execute(ctx context.Context) error {
if err != nil {
return err
}
if !hasTextFields {
onFlushSegmentIDs = append(onFlushSegmentIDs, segmentIDs...)
}
onFlushSegmentIDs = append(onFlushSegmentIDs, segmentIDs...)
}
// Ask datacoord to get flushed segment infos.
@@ -108,11 +93,7 @@ func (t *flushTask) Execute(ctx context.Context) error {
coll2Segments[collName] = &schemapb.LongArray{Data: onFlushSegmentIDs}
flushColl2Segments[collName] = &schemapb.LongArray{Data: resp.GetFlushSegmentIDs()}
coll2SealTimes[collName] = timeOfSeal.Unix()
if hasTextFields {
coll2FlushTs[collName] = 0
} else {
coll2FlushTs[collName] = flushTs
}
coll2FlushTs[collName] = flushTs
channelCps = resp.GetChannelCps()
}
t.result = &milvuspb.FlushResponse{
internal/proxy/task_import.go
+3
View File
@@ -105,6 +105,9 @@ func (it *importTask) PreExecute(ctx context.Context) error {
if schema.CollectionSchema == nil || len(schema.GetFields()) == 0 {
return merr.WrapErrImportFailed("collection schema has no fields")
}
if err := validateTextStorageV3Enabled(schema.CollectionSchema); err != nil {
return err
}
it.schema = schema
channels, err := node.chMgr.getVChannels(collectionID)
internal/proxy/task_insert.go
+3
View File
@@ -157,6 +157,9 @@ func (it *insertTask) PreExecute(ctx context.Context) error {
}
it.schema = schema.CollectionSchema
it.schemaVersion = schema.Version
if err := validateTextStorageV3Enabled(it.schema); err != nil {
return err
}
if err := genFunctionFields(ctx, it.insertMsg, schema, false); err != nil {
return err
internal/proxy/task_insert_test.go
+40
View File
@@ -77,6 +77,46 @@ func TestRepackInsertDataForStreamingServicePreservesExplicitZeroSchemaVersion(t
assert.Equal(t, int32(0), header.GetSchemaVersion())
}
func TestInsertTaskPreExecuteTextRequiresStorageV3(t *testing.T) {
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "false")
t.Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
})
oldCache := globalMetaCache
t.Cleanup(func() {
globalMetaCache = oldCache
})
const (
dbName = "db"
collectionName = "text_collection"
)
schema := newSchemaInfo(newTextSchemaForStorageV3Test(collectionName))
cache := NewMockCache(t)
cache.EXPECT().GetCollectionID(mock.Anything, dbName, collectionName).Return(int64(100), nil)
cache.EXPECT().GetCollectionInfo(mock.Anything, dbName, collectionName, int64(100)).Return(&collectionInfo{}, nil)
cache.EXPECT().GetCollectionSchema(mock.Anything, dbName, collectionName).Return(schema, nil)
globalMetaCache = cache
task := &insertTask{
ctx: context.Background(),
insertMsg: &BaseInsertTask{
InsertRequest: &msgpb.InsertRequest{
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Insert},
DbName: dbName,
CollectionName: collectionName,
NumRows: 1,
},
},
}
err := task.PreExecute(context.Background())
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "TEXT field requires StorageV3")
}
func TestInsertTask_CheckAligned(t *testing.T) {
var err error
internal/proxy/task_query.go
+3
View File
@@ -694,6 +694,9 @@ func (t *queryTask) PreExecute(ctx context.Context) error {
return err
}
t.schema = schema
if err := validateTextStorageV3Enabled(t.schema.CollectionSchema); err != nil {
return err
}
err = common.CheckNamespace(t.schema.CollectionSchema, t.request.Namespace)
if err != nil {
return err
internal/proxy/task_search.go
+3
View File
@@ -178,6 +178,9 @@ func (t *searchTask) PreExecute(ctx context.Context) error {
log.Warn("get collection schema failed", zap.Error(err))
return err
}
if err := validateTextStorageV3Enabled(t.schema.CollectionSchema); err != nil {
return err
}
collectionInfo, err2 := globalMetaCache.GetCollectionInfo(ctx, t.request.GetDbName(), collectionName, t.CollectionID)
if err2 != nil {
internal/proxy/task_search_test.go
+38
View File
@@ -81,6 +81,44 @@ func TestSearchTaskFillResultSkipsTopksInsufficientForSearchAggregation(t *testi
require.False(t, task.resultSizeInsufficient)
}
func TestSearchTaskPreExecuteTextRequiresStorageV3(t *testing.T) {
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "false")
t.Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
})
oldCache := globalMetaCache
t.Cleanup(func() {
globalMetaCache = oldCache
})
const (
dbName = "db"
collectionName = "text_collection"
collectionID = int64(100)
)
schema := newSchemaInfo(newTextSchemaForStorageV3Test(collectionName))
cache := NewMockCache(t)
cache.EXPECT().GetCollectionID(mock.Anything, dbName, collectionName).Return(collectionID, nil)
cache.EXPECT().GetCollectionSchema(mock.Anything, dbName, collectionName).Return(schema, nil)
globalMetaCache = cache
task := &searchTask{
ctx: context.Background(),
SearchRequest: &internalpb.SearchRequest{},
request: &milvuspb.SearchRequest{
DbName: dbName,
CollectionName: collectionName,
},
}
require.NoError(t, task.OnEnqueue())
err := task.PreExecute(context.Background())
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "TEXT field requires StorageV3")
}
func TestSearchTask_PostExecute(t *testing.T) {
var err error
internal/proxy/task_test.go
+87
View File
@@ -240,6 +240,24 @@ func constructCollectionSchemaByDataType(collectionName string, fieldName2DataTy
}
}
func newTextSchemaForStorageV3Test(collectionName string) *schemapb.CollectionSchema {
return &schemapb.CollectionSchema{
Name: collectionName,
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: testInt64Field, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_Text},
{
FieldID: 102,
Name: testFloatVecField,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: strconv.Itoa(testVecDim)},
},
},
},
}
}
func constructCollectionSchemaWithTTLField(collectionName string, ttlField string) *schemapb.CollectionSchema {
pk := &schemapb.FieldSchema{
FieldID: 100,
@@ -3980,6 +3998,43 @@ func Test_loadCollectionTask_Execute(t *testing.T) {
})
}
func TestLoadCollectionTaskExecuteTextRequiresStorageV3(t *testing.T) {
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "false")
t.Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
})
oldCache := globalMetaCache
t.Cleanup(func() {
globalMetaCache = oldCache
})
const (
dbName = "db"
collectionName = "text_collection"
collectionID = int64(100)
)
schema := newSchemaInfo(newTextSchemaForStorageV3Test(collectionName))
cache := NewMockCache(t)
cache.EXPECT().GetCollectionID(mock.Anything, dbName, collectionName).Return(collectionID, nil)
cache.EXPECT().GetCollectionSchema(mock.Anything, dbName, collectionName).Return(schema, nil)
globalMetaCache = cache
task := &loadCollectionTask{
LoadCollectionRequest: &milvuspb.LoadCollectionRequest{
Base: commonpbutil.NewMsgBase(),
DbName: dbName,
CollectionName: collectionName,
},
ctx: context.Background(),
}
err := task.Execute(context.Background())
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "TEXT field requires StorageV3")
}
func Test_loadPartitionTask_Execute(t *testing.T) {
qc := NewMixCoordMock()
@@ -7003,6 +7058,38 @@ func TestValidateAddFieldRequest(t *testing.T) {
assert.NoError(t, err)
})
t.Run("text field requires storage v3", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "false")
t.Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
})
schema := baseSchema()
newField := &schemapb.FieldSchema{
Name: "text_field",
DataType: schemapb.DataType_Text,
Nullable: true,
}
err := validateAddFieldRequest(schema, newField)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "TEXT field requires StorageV3")
})
t.Run("text field allowed with storage v3", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "true")
t.Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
})
schema := baseSchema()
newField := &schemapb.FieldSchema{
Name: "text_field",
DataType: schemapb.DataType_Text,
Nullable: true,
}
err := validateAddFieldRequest(schema, newField)
assert.NoError(t, err)
})
t.Run("field count exceeds MaxFieldNum", func(t *testing.T) {
schema := baseSchema()
// Build exactly MaxFieldNum fields so that the check triggers.
internal/proxy/task_upsert.go
+3
View File
@@ -1290,6 +1290,9 @@ func (it *upsertTask) PreExecute(ctx context.Context) error {
}
it.schema = schema
it.schemaVersion = schema.Version
if err := validateTextStorageV3Enabled(schema.CollectionSchema); err != nil {
return err
}
// Validate any FieldPartialUpdateOp directives attached to FieldData.
// A non-REPLACE op implicitly promotes the request to partial_update=true
internal/querycoordv2/task/task_test.go
+2 -1
View File
@@ -2102,7 +2102,8 @@ func (suite *TaskSuite) TestExecutor_MoveSegmentTask() {
suite.broker,
suite.target,
suite.cluster,
suite.nodeMgr)
suite.nodeMgr,
)
// Verify shard leader ID was set for load action in move task
executor.executeSegmentAction(moveTask, 0)
internal/querynodev2/delegator/delegator.go
+27 -40
View File
@@ -37,6 +37,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/distributed/streaming"
"github.com/milvus-io/milvus/internal/flushcommon/syncmgr"
"github.com/milvus-io/milvus/internal/querynodev2/cluster"
"github.com/milvus-io/milvus/internal/querynodev2/delegator/deletebuffer"
"github.com/milvus-io/milvus/internal/querynodev2/segments"
@@ -86,7 +87,6 @@ type ShardDelegator interface {
// data
ProcessInsert(insertRecords map[int64]*InsertData)
ProcessDelete(deleteData []*DeleteData, ts uint64)
ProcessManualFlush(ctx context.Context, flushTs uint64) error
LoadGrowing(ctx context.Context, infos []*querypb.SegmentLoadInfo, version int64) error
LoadL0(ctx context.Context, infos []*querypb.SegmentLoadInfo, version int64) error
LoadSegments(ctx context.Context, req *querypb.LoadSegmentsRequest) error
@@ -183,11 +183,13 @@ type shardDelegator struct {
// for slow down the delegator consumption and reduce the timetick dispatch frequency.
latestRequiredMVCCTimeTick *atomic.Uint64
// growing segment flush support for TEXT collections
// checkpointTracker tracks offset -> MsgPosition mapping for Growing Segments
// checkpointTracker is still updated by legacy growing ingest hooks. It no
// longer drives QueryNode-side metadata commit; WAL flusher owns commit.
checkpointTracker *segments.CheckpointTracker
// growingFlushManager manages periodic flush of Growing Segments (for TEXT collections)
growingFlushManager *segments.GrowingFlushManager
// growingSourceRegistration is the process-local registry lease for this
// delegator's optional growing-source source.
growingSourceRegistration *syncmgr.GrowingSourceRegistration
growingSourceProvider *delegatorGrowingSourceProvider
}
// getLogger returns the zap logger with pre-defined shard attributes.
@@ -268,9 +270,6 @@ func (sd *shardDelegator) prepareSearchFunction(req *internalpb.SearchRequest) (
// Start sets delegator to working state.
func (sd *shardDelegator) Start() {
sd.lifetime.SetState(lifetime.Working)
if sd.growingFlushManager != nil {
sd.growingFlushManager.Start(context.Background())
}
}
// Collection returns delegator collection id.
@@ -1327,17 +1326,16 @@ func (sd *shardDelegator) Close() {
sd.tsCond.L.Unlock()
sd.lifetime.Wait()
if sd.growingSourceProvider != nil {
sd.growingSourceProvider.Deactivate()
}
// Stop background snapshot loop before refunding candidates
sd.distribution.Close()
// Refund all sealed segment candidates in distribution
sd.distribution.RefundAllCandidates()
// stop growing flush manager
if sd.growingFlushManager != nil {
sd.growingFlushManager.Stop()
}
// clean idf oracle
if idfOracle := sd.getIDFOracle(); idfOracle != nil {
idfOracle.Close()
@@ -1481,25 +1479,17 @@ func NewShardDelegator(ctx context.Context, collectionID UniqueID, replicaID Uni
sd.publishIDFOracle(idfOracle)
}
// initialize GrowingFlushManager for TEXT collections
// this enables incremental flush of Growing Segments to preserve TEXT data
if sd.hasTextFields() {
// Register growing-source segments as optional local flush sources. Metadata
// commit is still owned by WAL flusher / WriteBuffer.
if sd.useGrowingSourceFlush() {
sd.checkpointTracker = segments.NewCheckpointTracker()
if binlogSaver != nil {
sd.growingFlushManager = segments.NewGrowingFlushManager(
collectionID,
channel,
collection.Schema(),
manager.Collection,
manager.Segment,
binlogSaver,
chunkManager,
sd.checkpointTracker,
)
log.Info("initialized GrowingFlushManager for TEXT collection")
} else {
log.Warn("binlogSaver is nil, GrowingFlushManager not initialized for TEXT collection")
}
sd.growingSourceProvider = newDelegatorGrowingSourceProvider(manager.Segment, func(ctx context.Context, fenceTs uint64) error {
_, err := sd.waitTSafe(ctx, fenceTs)
return err
})
sd.growingSourceRegistration = syncmgr.DefaultGrowingSourceRegistry().Register(sd.vchannelName, sd.growingSourceProvider)
sd.growingSourceProvider.SetRegistration(sd.growingSourceRegistration)
log.Info("registered growing-source source support")
}
sd.tsCond = syncutil.NewContextCond(&sync.Mutex{})
@@ -1508,17 +1498,14 @@ func NewShardDelegator(ctx context.Context, collectionID UniqueID, replicaID Uni
return sd, nil
}
// hasTextFields returns true if the collection has any TEXT type fields.
func (sd *shardDelegator) hasTextFields() bool {
if sd.collection == nil {
// useGrowingSourceFlush returns true when the collection should expose growing segments as a flush source.
func (sd *shardDelegator) useGrowingSourceFlush() bool {
if sd == nil || sd.collection == nil {
return false
}
for _, field := range sd.collection.Schema().GetFields() {
if field.GetDataType() == schemapb.DataType_Text {
return true
}
}
return false
return typeutil.UseGrowingSourceFlush(sd.collection.Schema(),
paramtable.Get().CommonCfg.UseLoonFFI.GetAsBool(),
paramtable.Get().CommonCfg.EnableGrowingSourceFlush.GetAsBool())
}
func (sd *shardDelegator) RunAnalyzer(ctx context.Context, req *querypb.RunAnalyzerRequest) ([]*milvuspb.AnalyzerResult, error) {
internal/querynodev2/delegator/delegator_data.go
+18 -109
View File
@@ -148,9 +148,9 @@ func (sd *shardDelegator) ProcessInsert(insertRecords map[int64]*InsertData) {
}
growing.UpdatePkCandidate(insertData.PrimaryKeys)
// record batch info for checkpoint tracking (TEXT collections only)
// record batch info for checkpoint tracking (growing-source collections only)
if sd.checkpointTracker != nil {
endOffset := growing.RowNum()
endOffset := growing.InsertCount()
sd.checkpointTracker.RecordBatch(
growing.ID(),
endOffset,
@@ -248,73 +248,6 @@ func (sd *shardDelegator) ProcessDelete(deleteData []*DeleteData, ts uint64) {
Observe(float64(tr.ElapseSpan().Milliseconds()))
}
// ProcessManualFlush handles manual flush request for TEXT collections.
// It triggers immediate flush of all unflushed data in Growing Segments.
// This is called when user explicitly requests flush via collection.flush().
func (sd *shardDelegator) ProcessManualFlush(ctx context.Context, flushTs uint64) error {
// Only process for TEXT collections that have GrowingFlushManager
if sd.growingFlushManager == nil || sd.checkpointTracker == nil {
return nil
}
log := sd.getLogger(ctx).With(zap.Uint64("flushTs", flushTs))
log.Info("processing manual flush for TEXT collection")
// Get all tracked segment IDs
segmentIDs := sd.checkpointTracker.GetSegmentIDs()
if len(segmentIDs) == 0 {
log.Debug("no segments to flush")
return nil
}
// Submit to the dynamic CGO pool to reuse the same concurrency control
// as other CGO operations (pool size = CPU cores).
pool := segments.GetDynamicPool()
futures := make([]*conc.Future[any], 0, len(segmentIDs))
for _, segID := range segmentIDs {
segment := sd.segmentManager.GetGrowing(segID)
if segment == nil {
continue
}
id := segID
future := pool.Submit(func() (any, error) {
if err := sd.growingFlushManager.ForceSyncAndSeal(ctx, id); err != nil {
log.Warn("failed to ForceSyncAndSeal during manual flush",
zap.Int64("segmentID", id),
zap.Error(err))
return nil, err
}
log.Info("ForceSyncAndSeal completed during manual flush",
zap.Int64("segmentID", id))
return nil, nil
})
futures = append(futures, future)
}
// wait for all flush tasks to complete and collect errors
conc.AwaitAll(futures...)
var firstErr error
failCount := 0
for _, f := range futures {
if _, err := f.Await(); err != nil {
failCount++
if firstErr == nil {
firstErr = err
}
}
}
if failCount > 0 {
log.Warn("manual flush completed with errors",
zap.Int("segmentCount", len(segmentIDs)),
zap.Int("failCount", failCount),
zap.Error(firstErr))
return fmt.Errorf("manual flush failed for %d/%d segments: %w", failCount, len(segmentIDs), firstErr)
}
log.Info("manual flush completed", zap.Int("segmentCount", len(segmentIDs)))
return nil
}
type BatchApplyRet = struct {
DeleteDataIdx int
StartIdx int
@@ -480,14 +413,18 @@ func (sd *shardDelegator) LoadGrowing(ctx context.Context, infos []*querypb.Segm
segmentIDs = lo.Map(loaded, func(segment segments.Segment, _ int) int64 { return segment.ID() })
log.Info("load growing segments done", zap.Int64s("segmentIDs", segmentIDs))
// initialize checkpoint tracking for recovered growing segments (TEXT collections only)
// initialize checkpoint tracking for recovered growing segments (growing-source collections only)
if sd.checkpointTracker != nil {
for _, segment := range loaded {
// the segment was recovered from binlog, so the current row count is the flushed offset
flushedOffset := segment.RowNum()
manifest := segment.LoadInfo().GetManifestPath()
if manifest == "" && flushedOffset > 0 {
return fmt.Errorf("recovered growing segment %d has %d flushed rows but no manifest path, cannot safely resume flush", segment.ID(), flushedOffset)
log.Info("initialize checkpoint tracker for legacy recovered growing segment",
zap.Int64("segmentID", segment.ID()),
zap.Int64("flushedOffset", flushedOffset))
sd.checkpointTracker.InitSegment(segment.ID(), flushedOffset)
continue
}
sd.checkpointTracker.InitSegmentWithManifest(segment.ID(), flushedOffset, manifest)
log.Info("initialized checkpoint tracker for recovered growing segment",
@@ -1184,6 +1121,7 @@ func (sd *shardDelegator) ReleaseSegments(ctx context.Context, req *querypb.Rele
case querypb.DataScope_Historical:
sealed = lo.Map(req.GetSegmentIDs(), convertSealed)
}
signal := sd.distribution.RemoveDistributions(sealed, growing)
// wait cleared signal
<-signal
@@ -1206,50 +1144,16 @@ func (sd *shardDelegator) ReleaseSegments(ctx context.Context, req *querypb.Rele
// - Sealed segment candidates (BloomFilterSet) are refunded in RemoveDistributions
// - Growing segment candidates (LocalSegment) are managed by segmentManager.Release()
if len(growing) > 0 {
// For TEXT collections: best-effort flush before releasing growing segments.
// Flush failure does NOT block release because:
// 1. Data is still in WAL — channel checkpoint clamping ensures WAL won't
// be truncated beyond unflushed growing segment data.
// 2. On recovery, WAL replays from the clamped checkpoint, restoring data.
// 3. Blocking release on flush failure would prevent QueryCoord balance
// and QueryNode graceful shutdown.
if sd.growingFlushManager != nil && sd.checkpointTracker != nil {
for _, entry := range growing {
segID := entry.SegmentID
segment := sd.segmentManager.GetGrowing(segID)
if segment == nil {
continue
}
// growing-source data is not flushed directly from QueryNode during
// release. WAL flusher owns metadata commits; if a source-backed sync
// has not committed yet, WriteBuffer checkpoint pin keeps WAL replayable.
flushedOffset := sd.checkpointTracker.GetFlushedOffset(segID)
currentOffset := segment.RowNum()
if flushedOffset >= currentOffset {
continue
}
log := sd.getLogger(ctx).With(zap.Int64("segmentID", segID))
log.Info("best-effort flush before release",
zap.Int64("unflushedRows", currentOffset-flushedOffset))
if err := sd.growingFlushManager.ForceSync(ctx, segID); err != nil {
log.Warn("best-effort flush failed before release, data will be recovered from WAL",
zap.Error(err))
}
}
}
// clean up checkpoint tracking for released growing segments (TEXT collections only)
// clean up checkpoint tracking for released growing segments (growing-source collections only)
if sd.checkpointTracker != nil {
for _, entry := range growing {
sd.checkpointTracker.RemoveSegment(entry.SegmentID)
}
}
// clean up sealed segment tracking to prevent memory leak (TEXT collections only)
if sd.growingFlushManager != nil {
for _, entry := range growing {
sd.growingFlushManager.RemoveSealedSegment(entry.SegmentID)
}
}
}
var releaseErr error
@@ -1273,6 +1177,11 @@ func (sd *shardDelegator) ReleaseSegments(ctx context.Context, req *querypb.Rele
if releaseErr != nil {
return releaseErr
}
if len(growing) > 0 && sd.growingSourceProvider != nil {
for _, entry := range growing {
sd.growingSourceProvider.ClearReleasePrepared(entry.SegmentID)
}
}
return nil
}
internal/querynodev2/delegator/delegator_data_test.go
+373 -12
View File
@@ -38,6 +38,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/flushcommon/syncmgr"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/mocks/util/mock_segcore"
"github.com/milvus-io/milvus/internal/querynodev2/cluster"
@@ -190,6 +191,40 @@ func (s *DelegatorDataSuite) genNormalCollection() {
})
}
func (s *DelegatorDataSuite) genTextCollection() {
s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
Name: "TestTextCollection",
Fields: []*schemapb.FieldSchema{
{
Name: "id",
FieldID: 100,
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
},
{
Name: "text",
FieldID: 101,
DataType: schemapb.DataType_Text,
},
{
Name: "vector",
FieldID: 102,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
}, nil, &querypb.LoadMetaInfo{
LoadType: querypb.LoadType_LoadCollection,
PartitionIDs: []int64{1001},
SchemaVersion: tsoutil.ComposeTSByTime(time.Now(), 0),
})
}
func (s *DelegatorDataSuite) genCollectionWithFunction() {
s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
Name: "TestCollection",
@@ -230,6 +265,11 @@ func (s *DelegatorDataSuite) genCollectionWithFunction() {
}
func (s *DelegatorDataSuite) SetupTest() {
paramtable.Get().Save(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key, "false")
s.T().Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key)
})
s.workerManager = &cluster.MockManager{}
s.manager = segments.NewManager()
s.loader = &segments.MockLoader{}
@@ -249,6 +289,15 @@ func (s *DelegatorDataSuite) SetupTest() {
s.delegator = sd
}
func (s *DelegatorDataSuite) enableGrowingSourceFlush() {
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "true")
paramtable.Get().Save(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key, "true")
s.T().Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
paramtable.Get().Reset(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key)
})
}
func (s *DelegatorDataSuite) TestProcessInsert() {
s.Run("normal_insert", func() {
s.delegator.ProcessInsert(map[int64]*InsertData{
@@ -342,6 +391,7 @@ func (s *DelegatorDataSuite) TestProcessDelete() {
ms.EXPECT().Partition().Return(info.GetPartitionID())
ms.EXPECT().Indexes().Return(nil)
ms.EXPECT().RowNum().Return(info.GetNumOfRows())
ms.EXPECT().LoadInfo().Return(info)
ms.EXPECT().Delete(mock.Anything, mock.Anything, mock.Anything).Return(nil)
ms.EXPECT().MayPkExist(mock.Anything).RunAndReturn(func(lc *storage.LocationsCache) bool {
return lc.GetPk().EQ(storage.NewInt64PrimaryKey(10))
@@ -581,12 +631,51 @@ func (s *DelegatorDataSuite) TestLoadGrowingWithBM25() {
mockSegment.EXPECT().ID().Return(int64(111))
// Note: Type() is no longer called since pkOracle was removed
// Candidate is now stored directly in SegmentEntry
mockSegment.EXPECT().RowNum().Return(int64(0)).Maybe()
mockSegment.EXPECT().LoadInfo().Return(&querypb.SegmentLoadInfo{SegmentID: 111}).Maybe()
mockSegment.EXPECT().GetBM25Stats().Return(map[int64]*storage.BM25Stats{})
err := s.delegator.LoadGrowing(context.Background(), []*querypb.SegmentLoadInfo{{SegmentID: 1}}, 1)
s.NoError(err)
}
func (s *DelegatorDataSuite) TestLoadLegacyGrowingSourceWithoutManifest() {
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "true")
defer paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
s.manager = segments.NewManager()
s.genTextCollection()
delegator, err := NewShardDelegator(context.Background(), s.collectionID, s.replicaID, s.vchannelName, s.version, s.workerManager, s.manager, s.loader, 10000, nil, s.chunkManager, NewChannelQueryView(nil, nil, nil, initialTargetVersion), nil)
s.Require().NoError(err)
sd, ok := delegator.(*shardDelegator)
s.Require().True(ok)
s.Require().NotNil(sd.checkpointTracker)
s.delegator = sd
loadInfo := &querypb.SegmentLoadInfo{
SegmentID: 1001,
CollectionID: s.collectionID,
PartitionID: 500,
InsertChannel: s.vchannelName,
NumOfRows: 50,
}
s.loader.EXPECT().
Load(mock.Anything, s.collectionID, segments.SegmentTypeGrowing, int64(0), mock.Anything).
Call.Return(func(ctx context.Context, collectionID int64, segmentType segments.SegmentType, version int64, infos ...*querypb.SegmentLoadInfo) []segments.Segment {
ms := segments.NewMockSegment(s.T())
ms.EXPECT().ID().Return(loadInfo.GetSegmentID())
ms.EXPECT().Partition().Return(loadInfo.GetPartitionID())
ms.EXPECT().RowNum().Return(loadInfo.GetNumOfRows())
ms.EXPECT().LoadInfo().Return(loadInfo)
return []segments.Segment{ms}
}, nil)
err = s.delegator.LoadGrowing(context.Background(), []*querypb.SegmentLoadInfo{loadInfo}, 0)
s.Require().NoError(err)
s.Equal(int64(50), s.delegator.checkpointTracker.GetFlushedOffset(loadInfo.GetSegmentID()))
s.Empty(s.delegator.checkpointTracker.GetAcknowledgedManifest(loadInfo.GetSegmentID()))
}
func (s *DelegatorDataSuite) TestLoadSegmentsWithBm25() {
s.genCollectionWithFunction()
s.Run("normal_run", func() {
@@ -1483,6 +1572,7 @@ func (s *DelegatorDataSuite) TestReleaseSegment() {
ms.EXPECT().Collection().Return(info.GetCollectionID())
ms.EXPECT().Indexes().Return(nil)
ms.EXPECT().RowNum().Return(info.GetNumOfRows())
ms.EXPECT().LoadInfo().Return(info)
ms.EXPECT().Delete(mock.Anything, mock.Anything, mock.Anything).Return(nil)
ms.EXPECT().MayPkExist(mock.Anything).Call.Return(func(pk storage.PrimaryKey) bool {
return pk.EQ(storage.NewInt64PrimaryKey(10))
@@ -1619,6 +1709,289 @@ func (s *DelegatorDataSuite) TestReleaseSegment() {
s.NoError(err)
}
func (s *DelegatorDataSuite) TestReleaseGrowingSourceAfterPreparedHandoff() {
ctx := context.Background()
s.workerManager = &cluster.MockManager{}
s.manager = segments.NewManager()
s.loader = &segments.MockLoader{}
s.genTextCollection()
s.enableGrowingSourceFlush()
delegator, err := NewShardDelegator(
ctx,
s.collectionID,
s.replicaID,
s.vchannelName,
s.version,
s.workerManager,
s.manager,
s.loader,
10000,
nil,
s.chunkManager,
NewChannelQueryView(nil, nil, nil, initialTargetVersion),
nil)
s.Require().NoError(err)
sd := delegator.(*shardDelegator)
defer sd.Close()
const (
segmentID = int64(1001)
partitionID = int64(500)
targetOffset = int64(10)
)
segment := segments.NewMockSegment(s.T())
segment.EXPECT().ID().Return(segmentID).Maybe()
segment.EXPECT().Version().Return(int64(1)).Maybe()
segment.EXPECT().Shard().Return(s.channel).Maybe()
segment.EXPECT().Collection().Return(s.collectionID).Maybe()
segment.EXPECT().Partition().Return(partitionID).Maybe()
segment.EXPECT().Type().Return(segments.SegmentTypeGrowing).Maybe()
segment.EXPECT().Level().Return(datapb.SegmentLevel_Legacy).Maybe()
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
segment.EXPECT().InsertCount().Return(targetOffset).Once()
segment.EXPECT().Unpin().Maybe()
s.manager.Segment.Put(ctx, segments.SegmentTypeGrowing, segment)
sd.distribution.AddGrowing(SegmentEntry{
NodeID: paramtable.GetNodeID(),
SegmentID: segmentID,
PartitionID: partitionID,
Version: 1,
})
worker := &cluster.MockWorker{}
worker.EXPECT().ReleaseSegments(mock.Anything, mock.AnythingOfType("*querypb.ReleaseSegmentsRequest")).Return(nil).Once()
s.workerManager.EXPECT().GetWorker(mock.Anything, paramtable.GetNodeID()).Return(worker, nil).Once()
err = sd.growingSourceProvider.PrepareGrowingSourceReleaseHandoff(ctx, 10000, []syncmgr.GrowingSourceReleaseHandoffSegment{
{
SegmentID: segmentID,
TargetOffset: targetOffset,
},
})
s.Require().NoError(err)
err = sd.ReleaseSegments(ctx, &querypb.ReleaseSegmentsRequest{
Base: commonpbutil.NewMsgBase(),
NodeID: paramtable.GetNodeID(),
CollectionID: s.collectionID,
SegmentIDs: []int64{segmentID},
Scope: querypb.DataScope_Streaming,
Shard: s.vchannelName,
}, false)
s.Require().NoError(err)
}
func (s *DelegatorDataSuite) TestReleaseGrowingSourceAfterFencePreparedHandoff() {
ctx := context.Background()
s.workerManager = &cluster.MockManager{}
s.manager = segments.NewManager()
s.loader = &segments.MockLoader{}
s.genTextCollection()
s.enableGrowingSourceFlush()
delegator, err := NewShardDelegator(
ctx,
s.collectionID,
s.replicaID,
s.vchannelName,
s.version,
s.workerManager,
s.manager,
s.loader,
10000,
nil,
s.chunkManager,
NewChannelQueryView(nil, nil, nil, initialTargetVersion),
nil)
s.Require().NoError(err)
sd := delegator.(*shardDelegator)
defer sd.Close()
const segmentID = int64(1001)
sd.distribution.AddGrowing(SegmentEntry{
NodeID: paramtable.GetNodeID(),
SegmentID: segmentID,
PartitionID: 500,
Version: 1,
})
err = sd.growingSourceProvider.PrepareGrowingSourceReleaseHandoff(ctx, 10000, []syncmgr.GrowingSourceReleaseHandoffSegment{
{SegmentID: segmentID},
})
s.Require().NoError(err)
worker := &cluster.MockWorker{}
worker.EXPECT().ReleaseSegments(mock.Anything, mock.AnythingOfType("*querypb.ReleaseSegmentsRequest")).Return(nil).Once()
s.workerManager.EXPECT().GetWorker(mock.Anything, paramtable.GetNodeID()).Return(worker, nil).Once()
err = sd.ReleaseSegments(ctx, &querypb.ReleaseSegmentsRequest{
Base: commonpbutil.NewMsgBase(),
NodeID: paramtable.GetNodeID(),
CollectionID: s.collectionID,
SegmentIDs: []int64{segmentID},
Scope: querypb.DataScope_Streaming,
Shard: s.vchannelName,
Checkpoint: &msgpb.MsgPosition{Timestamp: 10000},
}, false)
s.Require().NoError(err)
}
func (s *DelegatorDataSuite) TestReleaseGrowingSourceAfterNoRetainPreparedHandoff() {
ctx := context.Background()
s.workerManager = &cluster.MockManager{}
s.manager = segments.NewManager()
s.loader = &segments.MockLoader{}
s.genTextCollection()
s.enableGrowingSourceFlush()
delegator, err := NewShardDelegator(
ctx,
s.collectionID,
s.replicaID,
s.vchannelName,
s.version,
s.workerManager,
s.manager,
s.loader,
10000,
nil,
s.chunkManager,
NewChannelQueryView(nil, nil, nil, initialTargetVersion),
nil)
s.Require().NoError(err)
sd := delegator.(*shardDelegator)
defer sd.Close()
const segmentID = int64(1001)
sd.distribution.AddGrowing(SegmentEntry{
NodeID: paramtable.GetNodeID(),
SegmentID: segmentID,
PartitionID: 500,
Version: 1,
})
err = sd.growingSourceProvider.PrepareGrowingSourceReleaseHandoff(ctx, 10000, []syncmgr.GrowingSourceReleaseHandoffSegment{
{SegmentID: segmentID},
})
s.Require().NoError(err)
worker := &cluster.MockWorker{}
worker.EXPECT().ReleaseSegments(mock.Anything, mock.AnythingOfType("*querypb.ReleaseSegmentsRequest")).Return(nil).Once()
s.workerManager.EXPECT().GetWorker(mock.Anything, paramtable.GetNodeID()).Return(worker, nil).Once()
err = sd.ReleaseSegments(ctx, &querypb.ReleaseSegmentsRequest{
Base: commonpbutil.NewMsgBase(),
NodeID: paramtable.GetNodeID(),
CollectionID: s.collectionID,
SegmentIDs: []int64{segmentID},
Scope: querypb.DataScope_Streaming,
Shard: s.vchannelName,
Checkpoint: &msgpb.MsgPosition{Timestamp: 10000},
}, false)
s.Require().NoError(err)
}
func (s *DelegatorDataSuite) TestReleaseGrowingSourceWithoutPreparedHandoff() {
ctx := context.Background()
s.workerManager = &cluster.MockManager{}
s.manager = segments.NewManager()
s.loader = &segments.MockLoader{}
s.genTextCollection()
s.enableGrowingSourceFlush()
delegator, err := NewShardDelegator(
ctx,
s.collectionID,
s.replicaID,
s.vchannelName,
s.version,
s.workerManager,
s.manager,
s.loader,
10000,
nil,
s.chunkManager,
NewChannelQueryView(nil, nil, nil, initialTargetVersion),
nil)
s.Require().NoError(err)
sd := delegator.(*shardDelegator)
defer sd.Close()
const segmentID = int64(1001)
sd.distribution.AddGrowing(SegmentEntry{
NodeID: paramtable.GetNodeID(),
SegmentID: segmentID,
PartitionID: 500,
Version: 1,
})
worker := &cluster.MockWorker{}
worker.EXPECT().ReleaseSegments(mock.Anything, mock.AnythingOfType("*querypb.ReleaseSegmentsRequest")).Return(nil).Once()
s.workerManager.EXPECT().GetWorker(mock.Anything, paramtable.GetNodeID()).Return(worker, nil).Once()
err = sd.ReleaseSegments(ctx, &querypb.ReleaseSegmentsRequest{
Base: commonpbutil.NewMsgBase(),
NodeID: paramtable.GetNodeID(),
CollectionID: s.collectionID,
SegmentIDs: []int64{segmentID},
Scope: querypb.DataScope_Streaming,
Shard: s.vchannelName,
Checkpoint: &msgpb.MsgPosition{Timestamp: 10000},
}, false)
s.Require().NoError(err)
}
func (s *DelegatorDataSuite) TestReleaseGrowingSourceDroppedChannelWithoutPreparedHandoff() {
ctx := context.Background()
s.workerManager = &cluster.MockManager{}
s.manager = segments.NewManager()
s.loader = &segments.MockLoader{}
s.genTextCollection()
s.enableGrowingSourceFlush()
delegator, err := NewShardDelegator(
ctx,
s.collectionID,
s.replicaID,
s.vchannelName,
s.version,
s.workerManager,
s.manager,
s.loader,
10000,
nil,
s.chunkManager,
NewChannelQueryView(nil, nil, nil, initialTargetVersion),
nil)
s.Require().NoError(err)
sd := delegator.(*shardDelegator)
defer sd.Close()
const segmentID = int64(1001)
sd.distribution.AddGrowing(SegmentEntry{
NodeID: paramtable.GetNodeID(),
SegmentID: segmentID,
PartitionID: 500,
Version: 1,
})
worker := &cluster.MockWorker{}
worker.EXPECT().ReleaseSegments(mock.Anything, mock.AnythingOfType("*querypb.ReleaseSegmentsRequest")).Return(nil).Once()
s.workerManager.EXPECT().GetWorker(mock.Anything, paramtable.GetNodeID()).Return(worker, nil).Once()
err = sd.ReleaseSegments(ctx, &querypb.ReleaseSegmentsRequest{
Base: commonpbutil.NewMsgBase(),
NodeID: paramtable.GetNodeID(),
CollectionID: s.collectionID,
SegmentIDs: []int64{segmentID},
Scope: querypb.DataScope_Streaming,
Shard: s.vchannelName,
Checkpoint: &msgpb.MsgPosition{Timestamp: typeutil.MaxTimestamp},
}, false)
s.Require().NoError(err)
}
func (s *DelegatorDataSuite) TestLoadPartitionStats() {
segStats := make(map[UniqueID]storage.SegmentStats)
centroid := []float32{1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0}
@@ -1933,18 +2306,6 @@ func (s *DelegatorDataSuite) TestDelegatorData_ExcludeSegments() {
s.True(s.delegator.VerifyExcludedSegments(1, 5))
}
func (s *DelegatorDataSuite) TestProcessManualFlush_NilManager() {
// When growingFlushManager is nil (non-TEXT collection), should return immediately without panic
s.Nil(s.delegator.growingFlushManager)
s.delegator.ProcessManualFlush(context.Background(), 1000)
}
func (s *DelegatorDataSuite) TestProcessManualFlush_NilTracker() {
// When checkpointTracker is nil, should return immediately without panic
s.delegator.checkpointTracker = nil
s.delegator.ProcessManualFlush(context.Background(), 1000)
}
func TestSegmentEffectiveTs(t *testing.T) {
// Import segment: commit_timestamp takes precedence over start_position.Timestamp
info := &querypb.SegmentLoadInfo{
internal/querynodev2/delegator/delegator_test.go
+6
View File
@@ -91,6 +91,11 @@ func (s *DelegatorSuite) TearDownSuite() {
}
func (s *DelegatorSuite) SetupTest() {
paramtable.Get().Save(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key, "false")
s.T().Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key)
})
s.collectionID = 1000
s.partitionIDs = []int64{500, 501}
s.replicaID = 65535
@@ -110,6 +115,7 @@ func (s *DelegatorSuite) SetupTest() {
ms.EXPECT().Collection().Return(info.GetCollectionID())
ms.EXPECT().Indexes().Return(nil)
ms.EXPECT().RowNum().Return(info.GetNumOfRows())
ms.EXPECT().LoadInfo().Return(info)
ms.EXPECT().Delete(mock.Anything, mock.Anything, mock.Anything).Return(nil)
return ms
})
internal/querynodev2/delegator/delta_forward_test.go
+12
View File
@@ -67,6 +67,11 @@ func (s *StreamingForwardSuite) SetupSuite() {
}
func (s *StreamingForwardSuite) SetupTest() {
paramtable.Get().Save(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key, "false")
s.T().Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key)
})
s.collectionID = 1000
s.partitionIDs = []int64{500, 501}
s.replicaID = 65535
@@ -86,6 +91,7 @@ func (s *StreamingForwardSuite) SetupTest() {
ms.EXPECT().Collection().Return(info.GetCollectionID())
ms.EXPECT().Indexes().Return(nil)
ms.EXPECT().RowNum().Return(info.GetNumOfRows())
ms.EXPECT().LoadInfo().Return(info)
ms.EXPECT().Delete(mock.Anything, mock.Anything, mock.Anything).Return(nil)
return ms
})
@@ -300,6 +306,11 @@ func (s *GrowingMergeL0Suite) SetupSuite() {
}
func (s *GrowingMergeL0Suite) SetupTest() {
paramtable.Get().Save(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key, "false")
s.T().Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key)
})
s.collectionID = 1000
s.partitionIDs = []int64{500, 501}
s.replicaID = 65535
@@ -319,6 +330,7 @@ func (s *GrowingMergeL0Suite) SetupTest() {
ms.EXPECT().Collection().Return(info.GetCollectionID())
ms.EXPECT().Indexes().Return(nil)
ms.EXPECT().RowNum().Return(info.GetNumOfRows())
ms.EXPECT().LoadInfo().Return(info)
ms.EXPECT().Delete(mock.Anything, mock.Anything, mock.Anything).Return(nil)
return ms
})
internal/querynodev2/delegator/growing_flush_source.go
+562
View File
@@ -0,0 +1,562 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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.
package delegator
import (
"context"
"sync"
"github.com/cockroachdb/errors"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus/internal/flushcommon/syncmgr"
"github.com/milvus-io/milvus/internal/querynodev2/segments"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
)
var errGrowingSourceProviderClosed = errors.New("growing source provider is closed")
type delegatorGrowingSourceProvider struct {
segmentManager segments.SegmentManager
waitFence func(context.Context, uint64) error
mu sync.Mutex
cond *sync.Cond
closing bool
deactivated bool
registration *syncmgr.GrowingSourceRegistration
active int
retained map[int64]*retainedGrowingFlushSource
releaseAllowed map[int64]uint64
releasePrepared map[int64]int64
handoffOnly bool
handoffAllowed map[int64]struct{}
}
func newDelegatorGrowingSourceProvider(segmentManager segments.SegmentManager, waitFence func(context.Context, uint64) error) *delegatorGrowingSourceProvider {
provider := &delegatorGrowingSourceProvider{
segmentManager: segmentManager,
waitFence: waitFence,
retained: make(map[int64]*retainedGrowingFlushSource),
releaseAllowed: make(map[int64]uint64),
releasePrepared: make(map[int64]int64),
handoffAllowed: make(map[int64]struct{}),
}
provider.cond = sync.NewCond(&provider.mu)
return provider
}
func (p *delegatorGrowingSourceProvider) SetRegistration(registration *syncmgr.GrowingSourceRegistration) {
p.mu.Lock()
defer p.mu.Unlock()
p.registration = registration
}
func (p *delegatorGrowingSourceProvider) GetGrowingFlushSource(segmentID int64, targetOffset int64, _ *msgpb.MsgPosition) (syncmgr.GrowingFlushSource, syncmgr.GrowingSourceState) {
if !p.acquireLease(segmentID) {
return nil, syncmgr.GrowingSourceUnavailable
}
segment := p.segmentManager.GetGrowing(segmentID)
retained := false
if segment == nil {
var ok bool
segment, ok = p.getRetained(segmentID)
if !ok {
p.releaseLease()
return nil, syncmgr.GrowingSourceUnavailable
}
retained = true
} else if p.isDeactivated() {
p.releaseLease()
return nil, syncmgr.GrowingSourceUnavailable
}
if err := segment.PinIfNotReleased(); err != nil {
p.releaseLease()
return nil, syncmgr.GrowingSourceUnavailable
}
source := &delegatorGrowingFlushSource{segmentID: segmentID, segment: segment, provider: p, targetOffset: targetOffset, retained: retained}
if p.currentOffset(segment) < targetOffset {
return source, syncmgr.GrowingSourcePending
}
return source, syncmgr.GrowingSourceUsable
}
func (p *delegatorGrowingSourceProvider) PrepareGrowingSourceReleaseHandoff(ctx context.Context, fenceTs uint64, segments []syncmgr.GrowingSourceReleaseHandoffSegment) error {
if p.isDeactivated() {
return p.prepareDeactivatedGrowingSourceReleaseHandoff(fenceTs, segments)
}
handoffSnapshot := p.enterHandoffOnly(segments)
if p.waitFence != nil && fenceTs > 0 {
if err := p.waitFence(ctx, fenceTs); err != nil {
p.rollbackHandoffOnly(handoffSnapshot)
return err
}
}
snapshot := p.snapshotRetained(segments)
allowedSegments := make([]syncmgr.GrowingSourceReleaseHandoffSegment, 0, len(segments))
preparedSegments := make([]syncmgr.GrowingSourceReleaseHandoffSegment, 0, len(segments))
for _, segment := range segments {
allowedSegments = append(allowedSegments, segment)
if segment.TargetOffset <= 0 {
continue
}
if err := p.registerRetained(segment.SegmentID, segment.TargetOffset); err != nil {
if errors.Is(err, merr.ErrSegmentNotFound) {
continue
}
p.rollbackRetained(snapshot)
p.rollbackHandoffOnly(handoffSnapshot)
return err
}
preparedSegments = append(preparedSegments, segment)
}
p.markReleaseAllowed(fenceTs, allowedSegments)
p.markReleasePrepared(fenceTs, preparedSegments)
return nil
}
func (p *delegatorGrowingSourceProvider) prepareDeactivatedGrowingSourceReleaseHandoff(fenceTs uint64, segments []syncmgr.GrowingSourceReleaseHandoffSegment) error {
p.mu.Lock()
defer p.mu.Unlock()
for _, segment := range segments {
retained, ok := p.retained[segment.SegmentID]
if !ok {
continue
}
if segment.TargetOffset > retained.targetOffset {
continue
}
if current, ok := p.releaseAllowed[segment.SegmentID]; !ok || current < fenceTs {
p.releaseAllowed[segment.SegmentID] = fenceTs
}
if segment.TargetOffset <= 0 {
continue
}
if current, ok := p.releasePrepared[segment.SegmentID]; !ok || current < segment.TargetOffset {
p.releasePrepared[segment.SegmentID] = segment.TargetOffset
}
}
return nil
}
func (p *delegatorGrowingSourceProvider) registerRetained(segmentID int64, targetOffset int64) error {
p.mu.Lock()
if p.closing {
p.mu.Unlock()
return errGrowingSourceProviderClosed
}
if retained, ok := p.retained[segmentID]; ok {
if retained.targetOffset < targetOffset {
retained.targetOffset = targetOffset
}
p.mu.Unlock()
return nil
}
p.mu.Unlock()
segment := p.segmentManager.GetGrowing(segmentID)
if segment == nil {
return merr.WrapErrSegmentNotFound(segmentID)
}
if err := segment.PinIfNotReleased(); err != nil {
return err
}
currentOffset := p.currentOffset(segment)
if currentOffset < targetOffset {
segment.Unpin()
return errors.Errorf("growing-source segment %d is behind target offset, current=%d target=%d", segmentID, currentOffset, targetOffset)
}
p.mu.Lock()
defer p.mu.Unlock()
if p.closing {
segment.Unpin()
return errGrowingSourceProviderClosed
}
if retained, ok := p.retained[segmentID]; ok {
if retained.targetOffset < targetOffset {
retained.targetOffset = targetOffset
}
segment.Unpin()
return nil
}
p.retained[segmentID] = &retainedGrowingFlushSource{
segment: segment,
targetOffset: targetOffset,
}
return nil
}
type retainedSnapshot struct {
existed bool
source *retainedGrowingFlushSource
targetOffset int64
committedOffset int64
detached bool
}
func (p *delegatorGrowingSourceProvider) snapshotRetained(segments []syncmgr.GrowingSourceReleaseHandoffSegment) map[int64]retainedSnapshot {
p.mu.Lock()
defer p.mu.Unlock()
snapshot := make(map[int64]retainedSnapshot, len(segments))
for _, segment := range segments {
if _, ok := snapshot[segment.SegmentID]; ok {
continue
}
retained, existed := p.retained[segment.SegmentID]
entry := retainedSnapshot{existed: existed, source: retained}
if existed {
entry.targetOffset = retained.targetOffset
entry.committedOffset = retained.committedOffset
entry.detached = retained.detached
}
snapshot[segment.SegmentID] = entry
}
return snapshot
}
func (p *delegatorGrowingSourceProvider) rollbackRetained(snapshot map[int64]retainedSnapshot) {
var toUnpin []segments.Segment
p.mu.Lock()
for segmentID, entry := range snapshot {
current, exists := p.retained[segmentID]
if entry.existed {
p.retained[segmentID] = entry.source
entry.source.targetOffset = entry.targetOffset
entry.source.committedOffset = entry.committedOffset
entry.source.detached = entry.detached
if exists && current != entry.source {
toUnpin = append(toUnpin, current.segment)
}
continue
}
if exists {
delete(p.retained, segmentID)
toUnpin = append(toUnpin, current.segment)
}
}
p.mu.Unlock()
for _, segment := range toUnpin {
segment.Unpin()
}
}
type handoffSnapshot struct {
enabled bool
allowed map[int64]struct{}
}
func (p *delegatorGrowingSourceProvider) enterHandoffOnly(segments []syncmgr.GrowingSourceReleaseHandoffSegment) handoffSnapshot {
p.mu.Lock()
defer p.mu.Unlock()
snapshot := handoffSnapshot{
enabled: p.handoffOnly,
allowed: make(map[int64]struct{}, len(p.handoffAllowed)),
}
for segmentID := range p.handoffAllowed {
snapshot.allowed[segmentID] = struct{}{}
}
p.handoffOnly = true
for _, segment := range segments {
p.handoffAllowed[segment.SegmentID] = struct{}{}
}
return snapshot
}
func (p *delegatorGrowingSourceProvider) rollbackHandoffOnly(snapshot handoffSnapshot) {
p.mu.Lock()
defer p.mu.Unlock()
p.handoffOnly = snapshot.enabled
p.handoffAllowed = snapshot.allowed
}
func (p *delegatorGrowingSourceProvider) markReleaseAllowed(fenceTs uint64, segments []syncmgr.GrowingSourceReleaseHandoffSegment) {
p.mu.Lock()
defer p.mu.Unlock()
for _, segment := range segments {
if current, ok := p.releaseAllowed[segment.SegmentID]; !ok || current < fenceTs {
p.releaseAllowed[segment.SegmentID] = fenceTs
}
}
}
func (p *delegatorGrowingSourceProvider) markReleasePrepared(fenceTs uint64, segments []syncmgr.GrowingSourceReleaseHandoffSegment) {
p.mu.Lock()
defer p.mu.Unlock()
for _, segment := range segments {
if current, ok := p.releaseAllowed[segment.SegmentID]; !ok || current < fenceTs {
p.releaseAllowed[segment.SegmentID] = fenceTs
}
if current, ok := p.releasePrepared[segment.SegmentID]; !ok || current < segment.TargetOffset {
p.releasePrepared[segment.SegmentID] = segment.TargetOffset
}
}
}
func (p *delegatorGrowingSourceProvider) IsReleaseAllowed(segmentID int64, checkpointTs uint64) bool {
p.mu.Lock()
defer p.mu.Unlock()
return p.isReleaseAllowedLocked(segmentID, checkpointTs)
}
func (p *delegatorGrowingSourceProvider) IsReleasePrepared(segmentID int64, checkpointTs uint64) bool {
p.mu.Lock()
defer p.mu.Unlock()
if _, ok := p.releasePrepared[segmentID]; ok {
return p.isReleaseAllowedLocked(segmentID, checkpointTs)
}
return false
}
func (p *delegatorGrowingSourceProvider) isReleaseAllowedLocked(segmentID int64, checkpointTs uint64) bool {
fenceTs, ok := p.releaseAllowed[segmentID]
if !ok {
return false
}
return fenceTs == 0 || checkpointTs == 0 || checkpointTs <= fenceTs
}
func (p *delegatorGrowingSourceProvider) ClearReleasePrepared(segmentID int64) {
p.mu.Lock()
defer p.mu.Unlock()
delete(p.releasePrepared, segmentID)
delete(p.releaseAllowed, segmentID)
delete(p.handoffAllowed, segmentID)
}
func (p *delegatorGrowingSourceProvider) ReleasePreparedSegments() []int64 {
p.mu.Lock()
defer p.mu.Unlock()
segments := make([]int64, 0, len(p.releasePrepared))
for segmentID := range p.releasePrepared {
segments = append(segments, segmentID)
}
return segments
}
func (p *delegatorGrowingSourceProvider) MarkReleaseDetached(segmentID int64) {
p.mu.Lock()
retained, ok := p.retained[segmentID]
if !ok {
p.mu.Unlock()
return
}
retained.detached = true
registration, released := p.tryReleaseRetainedLocked(segmentID, retained)
p.mu.Unlock()
p.releaseRetained(registration, released)
}
func (p *delegatorGrowingSourceProvider) getRetained(segmentID int64) (segments.Segment, bool) {
p.mu.Lock()
defer p.mu.Unlock()
retained, ok := p.retained[segmentID]
if !ok {
return nil, false
}
return retained.segment, true
}
func (p *delegatorGrowingSourceProvider) releaseRetainedIfComplete(segmentID int64, targetOffset int64) {
p.mu.Lock()
retained, ok := p.retained[segmentID]
if !ok {
p.mu.Unlock()
return
}
if retained.committedOffset < targetOffset {
retained.committedOffset = targetOffset
}
registration, released := p.tryReleaseRetainedLocked(segmentID, retained)
p.mu.Unlock()
p.releaseRetained(registration, released)
}
func (p *delegatorGrowingSourceProvider) tryReleaseRetainedLocked(segmentID int64, retained *retainedGrowingFlushSource) (*syncmgr.GrowingSourceRegistration, *retainedGrowingFlushSource) {
if retained == nil ||
!retained.detached ||
retained.committedOffset < retained.targetOffset {
return nil, nil
}
delete(p.retained, segmentID)
return p.unregisterIfInactiveLocked(), retained
}
func (p *delegatorGrowingSourceProvider) releaseRetained(registration *syncmgr.GrowingSourceRegistration, retained *retainedGrowingFlushSource) {
if retained == nil {
syncmgr.DefaultGrowingSourceRegistry().Unregister(registration)
return
}
// Drop the retained pin first so Release can drain, then route through the
// segment manager's managed release path. The segment was already removed
// from the active maps by Detach, so release() will reconcile the
// on-releasing set, the segment gauge and the release callback. Calling
// segment.Release() directly here would leak the on-releasing set entry and
// keep Exist() reporting the segment forever.
retained.segment.Unpin()
p.segmentManager.ReleaseDetached(context.Background(), retained.segment)
syncmgr.DefaultGrowingSourceRegistry().Unregister(registration)
}
func (p *delegatorGrowingSourceProvider) acquireLease(segmentID int64) bool {
p.mu.Lock()
defer p.mu.Unlock()
if p.closing {
return false
}
if p.handoffOnly {
_, allowed := p.handoffAllowed[segmentID]
_, retained := p.retained[segmentID]
if !allowed && !retained {
return false
}
}
p.active++
return true
}
func (p *delegatorGrowingSourceProvider) releaseLease() {
p.mu.Lock()
defer p.mu.Unlock()
if p.active > 0 {
p.active--
}
if p.closing && p.active == 0 {
p.cond.Broadcast()
}
}
func (p *delegatorGrowingSourceProvider) isDeactivated() bool {
p.mu.Lock()
defer p.mu.Unlock()
return p.deactivated
}
func (p *delegatorGrowingSourceProvider) Deactivate() {
p.mu.Lock()
p.deactivated = true
registration := p.unregisterIfInactiveLocked()
p.mu.Unlock()
syncmgr.DefaultGrowingSourceRegistry().Unregister(registration)
}
func (p *delegatorGrowingSourceProvider) Close() {
p.mu.Lock()
p.closing = true
for p.active > 0 {
p.cond.Wait()
}
retained := p.retained
p.retained = make(map[int64]*retainedGrowingFlushSource)
p.releaseAllowed = make(map[int64]uint64)
p.releasePrepared = make(map[int64]int64)
p.handoffOnly = false
p.handoffAllowed = make(map[int64]struct{})
registration := p.registration
p.registration = nil
p.mu.Unlock()
for _, source := range retained {
source.segment.Unpin()
}
syncmgr.DefaultGrowingSourceRegistry().Unregister(registration)
}
func (p *delegatorGrowingSourceProvider) unregisterIfInactiveLocked() *syncmgr.GrowingSourceRegistration {
if !p.deactivated || len(p.retained) > 0 {
return nil
}
registration := p.registration
p.registration = nil
p.releaseAllowed = make(map[int64]uint64)
p.releasePrepared = make(map[int64]int64)
p.handoffOnly = false
p.handoffAllowed = make(map[int64]struct{})
return registration
}
func (p *delegatorGrowingSourceProvider) currentOffset(segment segments.Segment) int64 {
if segment == nil {
return 0
}
return segment.InsertCount()
}
type retainedGrowingFlushSource struct {
segment segments.Segment
targetOffset int64
committedOffset int64
detached bool
}
type delegatorGrowingFlushSource struct {
segmentID int64
segment segments.Segment
provider *delegatorGrowingSourceProvider
targetOffset int64
retained bool
once sync.Once
}
func (s *delegatorGrowingFlushSource) CurrentOffset() int64 {
if s.provider != nil {
return s.provider.currentOffset(s.segment)
}
if s.segment == nil {
return 0
}
return s.segment.InsertCount()
}
func (s *delegatorGrowingFlushSource) FlushGrowingData(ctx context.Context, startOffset, endOffset int64, config *syncmgr.GrowingFlushConfig) (*syncmgr.GrowingFlushResult, error) {
result, err := s.segment.FlushData(ctx, startOffset, endOffset, &segments.FlushConfig{
SegmentBasePath: config.SegmentBasePath,
PartitionBasePath: config.PartitionBasePath,
CollectionID: config.CollectionID,
PartitionID: config.PartitionID,
TextFieldIDs: config.TextFieldIDs,
TextLobPaths: config.TextLobPaths,
ReadVersion: config.ReadVersion,
})
if err != nil || result == nil {
return nil, err
}
return &syncmgr.GrowingFlushResult{
ManifestPath: result.ManifestPath,
NumRows: result.NumRows,
}, nil
}
func (s *delegatorGrowingFlushSource) Release() {
s.once.Do(func() {
if s.segment != nil {
s.segment.Unpin()
}
if s.provider != nil {
s.provider.releaseLease()
}
})
}
func (s *delegatorGrowingFlushSource) CommitGrowingFlush(targetOffset int64) {
if s.retained && s.provider != nil {
s.provider.releaseRetainedIfComplete(s.segmentID, targetOffset)
}
}
internal/querynodev2/delegator/growing_flush_source_test.go
+321
View File
@@ -0,0 +1,321 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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.
package delegator
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus/internal/flushcommon/syncmgr"
"github.com/milvus-io/milvus/internal/querynodev2/segments"
)
func TestDelegatorGrowingSourceProviderCloseWaitsForSourceRelease(t *testing.T) {
segmentManager := segments.NewMockSegmentManager(t)
segment := segments.NewMockSegment(t)
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
segment.EXPECT().InsertCount().Return(int64(10)).Once()
segment.EXPECT().Unpin().Once()
source, state := provider.GetGrowingFlushSource(1001, 10, nil)
require.Equal(t, syncmgr.GrowingSourceUsable, state)
require.NotNil(t, source)
closeDone := make(chan struct{})
go func() {
provider.Close()
close(closeDone)
}()
select {
case <-closeDone:
t.Fatal("provider close should wait for active growing source lease")
case <-time.After(50 * time.Millisecond):
}
source.Release()
require.Eventually(t, func() bool {
select {
case <-closeDone:
return true
default:
return false
}
}, time.Second, 10*time.Millisecond)
source, state = provider.GetGrowingFlushSource(1001, 10, nil)
require.Equal(t, syncmgr.GrowingSourceUnavailable, state)
require.Nil(t, source)
}
func TestDelegatorGrowingSourceProviderRetainedSource(t *testing.T) {
segmentManager := segments.NewMockSegmentManager(t)
segment := segments.NewMockSegment(t)
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
segment.EXPECT().InsertCount().Return(int64(10)).Once()
err := provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 0, []syncmgr.GrowingSourceReleaseHandoffSegment{
{SegmentID: 1001, TargetOffset: 10},
})
require.NoError(t, err)
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(nil).Twice()
segment.EXPECT().PinIfNotReleased().Return(nil).Twice()
segment.EXPECT().InsertCount().Return(int64(10)).Twice()
segment.EXPECT().Unpin().Times(3)
segmentManager.EXPECT().ReleaseDetached(context.Background(), segment).Once()
source, state := provider.GetGrowingFlushSource(1001, 10, nil)
require.Equal(t, syncmgr.GrowingSourceUsable, state)
require.NotNil(t, source)
source.Release()
source, state = provider.GetGrowingFlushSource(1001, 10, nil)
require.Equal(t, syncmgr.GrowingSourceUsable, state)
require.NotNil(t, source)
source.Release()
source.(syncmgr.GrowingFlushSourceCommitter).CommitGrowingFlush(10)
segmentManager.AssertNotCalled(t, "ReleaseDetached", context.Background(), segment)
provider.MarkReleaseDetached(1001)
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(nil).Once()
source, state = provider.GetGrowingFlushSource(1001, 10, nil)
require.Equal(t, syncmgr.GrowingSourceUnavailable, state)
require.Nil(t, source)
}
func TestDelegatorGrowingSourceProviderUsesInsertCountAsOffset(t *testing.T) {
segmentManager := segments.NewMockSegmentManager(t)
segment := segments.NewMockSegment(t)
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
segment.EXPECT().InsertCount().Return(int64(12)).Twice()
segment.EXPECT().Unpin().Once()
source, state := provider.GetGrowingFlushSource(1001, 12, nil)
require.Equal(t, syncmgr.GrowingSourceUsable, state)
require.NotNil(t, source)
require.EqualValues(t, 12, source.CurrentOffset())
source.Release()
}
func TestDelegatorGrowingSourceProviderPrepareWaitsFence(t *testing.T) {
segmentManager := segments.NewMockSegmentManager(t)
segment := segments.NewMockSegment(t)
waitedFence := uint64(0)
provider := newDelegatorGrowingSourceProvider(segmentManager, func(ctx context.Context, fenceTs uint64) error {
waitedFence = fenceTs
return nil
})
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
segment.EXPECT().InsertCount().Return(int64(10)).Once()
err := provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 200, []syncmgr.GrowingSourceReleaseHandoffSegment{
{SegmentID: 1001, TargetOffset: 10},
{SegmentID: 1002},
})
require.NoError(t, err)
require.EqualValues(t, 200, waitedFence)
require.True(t, provider.IsReleasePrepared(1001, 0))
require.True(t, provider.IsReleaseAllowed(1001, 0))
require.False(t, provider.IsReleasePrepared(1002, 200))
require.True(t, provider.IsReleaseAllowed(1002, 200))
require.False(t, provider.IsReleaseAllowed(1002, 201))
require.False(t, provider.IsReleasePrepared(1003, 200))
require.False(t, provider.IsReleaseAllowed(1003, 200))
segment.EXPECT().Unpin().Once()
segmentManager.EXPECT().ReleaseDetached(context.Background(), segment).Once()
provider.releaseRetainedIfComplete(1001, 10)
segmentManager.AssertNotCalled(t, "ReleaseDetached", context.Background(), segment)
provider.MarkReleaseDetached(1001)
}
func TestDelegatorGrowingSourceProviderHandoffOnlyRejectsNewSegmentsBeforeFence(t *testing.T) {
segmentManager := segments.NewMockSegmentManager(t)
segment := segments.NewMockSegment(t)
waitFenceEntered := make(chan struct{})
releaseFence := make(chan struct{})
provider := newDelegatorGrowingSourceProvider(segmentManager, func(ctx context.Context, fenceTs uint64) error {
close(waitFenceEntered)
<-releaseFence
return nil
})
prepareDone := make(chan error, 1)
go func() {
prepareDone <- provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 200, []syncmgr.GrowingSourceReleaseHandoffSegment{
{SegmentID: 1001, TargetOffset: 10},
})
}()
select {
case <-waitFenceEntered:
case <-time.After(time.Second):
t.Fatal("prepare did not enter fence wait")
}
source, state := provider.GetGrowingFlushSource(1002, 1, nil)
require.Equal(t, syncmgr.GrowingSourceUnavailable, state)
require.Nil(t, source)
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
segment.EXPECT().InsertCount().Return(int64(10)).Once()
close(releaseFence)
require.NoError(t, <-prepareDone)
segment.EXPECT().Unpin().Once()
provider.Close()
}
func TestDelegatorGrowingSourceProviderDeactivatedOnlyServesRetainedSources(t *testing.T) {
segmentManager := segments.NewMockSegmentManager(t)
segment := segments.NewMockSegment(t)
waitCount := 0
provider := newDelegatorGrowingSourceProvider(segmentManager, func(ctx context.Context, fenceTs uint64) error {
waitCount++
return nil
})
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
segment.EXPECT().InsertCount().Return(int64(10)).Once()
err := provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 100, []syncmgr.GrowingSourceReleaseHandoffSegment{
{SegmentID: 1001, TargetOffset: 10},
})
require.NoError(t, err)
require.Equal(t, 1, waitCount)
provider.Deactivate()
err = provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 200, []syncmgr.GrowingSourceReleaseHandoffSegment{
{SegmentID: 1002, TargetOffset: 5},
})
require.NoError(t, err)
require.Equal(t, 1, waitCount)
require.False(t, provider.IsReleaseAllowed(1002, 200))
require.False(t, provider.IsReleasePrepared(1002, 200))
require.True(t, provider.IsReleaseAllowed(1001, 100))
require.True(t, provider.IsReleasePrepared(1001, 100))
err = provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 200, []syncmgr.GrowingSourceReleaseHandoffSegment{
{SegmentID: 1001, TargetOffset: 11},
})
require.NoError(t, err)
require.False(t, provider.IsReleasePrepared(1001, 200))
segment.EXPECT().Unpin().Once()
segmentManager.EXPECT().ReleaseDetached(context.Background(), segment).Once()
provider.releaseRetainedIfComplete(1001, 10)
segmentManager.AssertNotCalled(t, "ReleaseDetached", context.Background(), segment)
provider.MarkReleaseDetached(1001)
}
func TestDelegatorGrowingSourceProviderReleasesWhenDetachedBeforeCommit(t *testing.T) {
segmentManager := segments.NewMockSegmentManager(t)
segment := segments.NewMockSegment(t)
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
segment.EXPECT().InsertCount().Return(int64(10)).Once()
err := provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 0, []syncmgr.GrowingSourceReleaseHandoffSegment{
{SegmentID: 1001, TargetOffset: 10},
})
require.NoError(t, err)
provider.MarkReleaseDetached(1001)
segmentManager.AssertNotCalled(t, "ReleaseDetached", context.Background(), segment)
segment.EXPECT().Unpin().Once()
segmentManager.EXPECT().ReleaseDetached(context.Background(), segment).Once()
provider.releaseRetainedIfComplete(1001, 10)
}
func TestDelegatorGrowingSourceProviderPrepareRollbackOnFailure(t *testing.T) {
segmentManager := segments.NewMockSegmentManager(t)
segment := segments.NewMockSegment(t)
behindSegment := segments.NewMockSegment(t)
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
segment.EXPECT().InsertCount().Return(int64(10)).Once()
segmentManager.EXPECT().GetGrowing(int64(1002)).Return(behindSegment).Once()
behindSegment.EXPECT().PinIfNotReleased().Return(nil).Once()
behindSegment.EXPECT().InsertCount().Return(int64(5)).Once()
behindSegment.EXPECT().Unpin().Once()
segment.EXPECT().Unpin().Once()
err := provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 200, []syncmgr.GrowingSourceReleaseHandoffSegment{
{SegmentID: 1001, TargetOffset: 10},
{SegmentID: 1002, TargetOffset: 10},
})
require.Error(t, err)
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(nil).Once()
source, state := provider.GetGrowingFlushSource(1001, 10, nil)
require.Equal(t, syncmgr.GrowingSourceUnavailable, state)
require.Nil(t, source)
}
func TestDelegatorGrowingSourceProviderRegisterRetainedSegmentNotFound(t *testing.T) {
segmentManager := segments.NewMockSegmentManager(t)
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
registry := syncmgr.NewGrowingSourceRegistry()
registration := registry.Register("ch", provider)
defer registry.Unregister(registration)
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(nil).Once()
err := registry.PrepareGrowingSourceReleaseHandoff(context.Background(), "ch", 0, []syncmgr.GrowingSourceReleaseHandoffSegment{
{SegmentID: 1001, TargetOffset: 10},
})
require.Error(t, err)
}
func TestDelegatorGrowingSourceProviderRegisterRetainedBehindTarget(t *testing.T) {
segmentManager := segments.NewMockSegmentManager(t)
segment := segments.NewMockSegment(t)
provider := newDelegatorGrowingSourceProvider(segmentManager, nil)
segmentManager.EXPECT().GetGrowing(int64(1001)).Return(segment).Once()
segment.EXPECT().PinIfNotReleased().Return(nil).Once()
segment.EXPECT().InsertCount().Return(int64(5)).Once()
segment.EXPECT().Unpin().Once()
err := provider.PrepareGrowingSourceReleaseHandoff(context.Background(), 0, []syncmgr.GrowingSourceReleaseHandoffSegment{
{SegmentID: 1001, TargetOffset: 10},
})
require.Error(t, err)
}
internal/querynodev2/delegator/mock_delegator.go
-47
View File
@@ -860,53 +860,6 @@ func (_c *MockShardDelegator_ProcessInsert_Call) RunAndReturn(run func(map[int64
return _c
}
// ProcessManualFlush provides a mock function with given fields: ctx, flushTs
func (_m *MockShardDelegator) ProcessManualFlush(ctx context.Context, flushTs uint64) error {
ret := _m.Called(ctx, flushTs)
if len(ret) == 0 {
panic("no return value specified for ProcessManualFlush")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, uint64) error); ok {
r0 = rf(ctx, flushTs)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockShardDelegator_ProcessManualFlush_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ProcessManualFlush'
type MockShardDelegator_ProcessManualFlush_Call struct {
*mock.Call
}
// ProcessManualFlush is a helper method to define mock.On call
// - ctx context.Context
// - flushTs uint64
func (_e *MockShardDelegator_Expecter) ProcessManualFlush(ctx interface{}, flushTs interface{}) *MockShardDelegator_ProcessManualFlush_Call {
return &MockShardDelegator_ProcessManualFlush_Call{Call: _e.mock.On("ProcessManualFlush", ctx, flushTs)}
}
func (_c *MockShardDelegator_ProcessManualFlush_Call) Run(run func(ctx context.Context, flushTs uint64)) *MockShardDelegator_ProcessManualFlush_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(uint64))
})
return _c
}
func (_c *MockShardDelegator_ProcessManualFlush_Call) Return(_a0 error) *MockShardDelegator_ProcessManualFlush_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockShardDelegator_ProcessManualFlush_Call) RunAndReturn(run func(context.Context, uint64) error) *MockShardDelegator_ProcessManualFlush_Call {
_c.Call.Return(run)
return _c
}
// Query provides a mock function with given fields: ctx, req
func (_m *MockShardDelegator) Query(ctx context.Context, req *querypb.QueryRequest) ([]*internalpb.RetrieveResults, error) {
ret := _m.Called(ctx, req)
internal/querynodev2/pipeline/filter_node.go
+3 -4
View File
@@ -150,15 +150,14 @@ func (fNode *filterNode) filtrate(c *Collection, msg msgstream.TsMsg) error {
}
return nil
case commonpb.MsgType_ManualFlush:
// Handle ManualFlush message for TEXT collections
// This triggers immediate flush of Growing Segment data
// ManualFlush is handled by StreamingNode WAL flusher (fence + persist).
// QueryNode only consumes the barrier so the pipeline advances.
manualFlushMsg := msg.(*adaptor.ManualFlushMessageBody)
header := manualFlushMsg.ManualFlushMessage.Header()
if header.GetCollectionId() != fNode.collectionID {
return merr.WrapErrCollectionNotFound(header.GetCollectionId())
}
// Trigger flush via delegator
return fNode.delegator.ProcessManualFlush(context.Background(), header.GetFlushTs())
return nil
default:
return merr.WrapErrParameterInvalid("msgType is Insert or Delete", "not")
}
internal/querynodev2/pipeline/message.go
+1 -2
View File
@@ -71,8 +71,7 @@ func (msg *insertNodeMsg) append(taskMsg msgstream.TsMsg) error {
msg.schemaVersion = taskMsg.BeginTs()
}
case commonpb.MsgType_ManualFlush:
// ManualFlush is already handled in filterNode.filtrate() via delegator.ProcessManualFlush().
// No-op here since ManualFlush doesn't carry insert/delete data.
// ManualFlush is consumed in filterNode.filtrate(); no insert/delete payload here.
default:
return merr.WrapErrParameterInvalid("msgType is Insert or Delete", "not")
}
internal/querynodev2/segments/checkpoint_tracker.go
+2 -2
View File
@@ -36,8 +36,8 @@ type BatchInfo struct {
}
// CheckpointTracker tracks the mapping between segment row offsets and WAL positions.
// it is used by GrowingFlushManager to determine the correct checkpoint when flushing
// a range of rows from a Growing Segment.
// It is legacy QueryNode-side state kept for compatibility while TEXT flush
// ownership moves to the WAL flusher / WriteBuffer path.
//
// design:
// - each Insert operation records a BatchInfo with (endOffset, position)
internal/querynodev2/segments/growing_flush_manager.go
-589
View File
@@ -1,589 +0,0 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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.
package segments
import (
"context"
"fmt"
"strconv"
"sync"
"time"
"github.com/cockroachdb/errors"
"go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/storagev2/packed"
milvuscommon "github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/log"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/util/metautil"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
)
// BinlogSaver is a minimal interface for saving binlog paths to DataCoord.
// This avoids depending on the full broker.Broker interface.
type BinlogSaver interface {
SaveBinlogPaths(ctx context.Context, req *datapb.SaveBinlogPathsRequest) error
}
// GrowingFlushManager manages the incremental flush of Growing Segments for TEXT collections.
// it periodically checks sync policies and triggers flush operations when conditions are met.
//
// key responsibilities:
// - track Growing Segments and their unflushed data via CheckpointTracker
// - apply sync policies to determine which segments to flush
// - extract data from Growing Segments and write to binlog storage
// - update metadata via DataCoord's SaveBinlogPaths API
// - update channel checkpoints
type GrowingFlushManager struct {
collectionID int64
channelName string
schema *schemapb.CollectionSchema
// dependencies
collectionManager CollectionManager
segmentManager SegmentManager
binlogSaver BinlogSaver
chunkManager storage.ChunkManager
checkpointTracker *CheckpointTracker
// sync policies
policies []GrowingSyncPolicy
// configuration
flushInterval time.Duration
// sync mutex to prevent concurrent sync on the same segment
syncMu sync.Map // map[int64]*sync.Mutex - segmentID -> mutex
// sealedSegments tracks segments that have already been sealed via ForceSyncAndSeal.
// This prevents duplicate SaveBinlogPaths(isFlush=true) requests when the user
// calls flush() multiple times, which can resurrect Dropped segments and cause
// duplicate data after compaction.
sealedSegments sync.Map // map[int64]struct{} - segmentID -> struct{}
// control
stopCh chan struct{}
stopOnce sync.Once
wg sync.WaitGroup
// logger
logger *log.MLogger
}
// GrowingFlushManagerOption is a function to configure GrowingFlushManager.
type GrowingFlushManagerOption func(*GrowingFlushManager)
// WithFlushInterval sets the flush check interval.
func WithFlushInterval(interval time.Duration) GrowingFlushManagerOption {
return func(m *GrowingFlushManager) {
m.flushInterval = interval
}
}
// WithSyncPolicies sets custom sync policies.
func WithSyncPolicies(policies ...GrowingSyncPolicy) GrowingFlushManagerOption {
return func(m *GrowingFlushManager) {
m.policies = policies
}
}
// NewGrowingFlushManager creates a new GrowingFlushManager instance.
func NewGrowingFlushManager(
collectionID int64,
channelName string,
schema *schemapb.CollectionSchema,
collectionManager CollectionManager,
segmentManager SegmentManager,
binlogSaver BinlogSaver,
chunkManager storage.ChunkManager,
checkpointTracker *CheckpointTracker,
opts ...GrowingFlushManagerOption,
) *GrowingFlushManager {
m := &GrowingFlushManager{
collectionID: collectionID,
channelName: channelName,
schema: schema,
collectionManager: collectionManager,
segmentManager: segmentManager,
binlogSaver: binlogSaver,
chunkManager: chunkManager,
checkpointTracker: checkpointTracker,
stopCh: make(chan struct{}),
logger: log.With(
zap.Int64("collectionID", collectionID),
zap.String("channel", channelName),
),
}
// apply options
for _, opt := range opts {
opt(m)
}
// set defaults
if m.flushInterval == 0 {
m.flushInterval = 100 * time.Millisecond
}
if len(m.policies) == 0 {
m.policies = m.defaultPolicies()
}
return m
}
// defaultPolicies returns the default sync policies.
func (m *GrowingFlushManager) defaultPolicies() []GrowingSyncPolicy {
staleDuration := paramtable.Get().DataNodeCfg.SyncPeriod.GetAsDuration(time.Second)
return []GrowingSyncPolicy{
GetGrowingFullBufferPolicy(),
GetGrowingStaleBufferPolicy(staleDuration),
}
}
// Start begins the background flush checking loop.
func (m *GrowingFlushManager) Start(ctx context.Context) {
m.wg.Add(1)
go m.backgroundLoop(ctx)
m.logger.Info("GrowingFlushManager started")
}
// Stop signals the manager to stop and waits for completion.
// Safe to call multiple times.
func (m *GrowingFlushManager) Stop() {
m.stopOnce.Do(func() {
close(m.stopCh)
})
m.wg.Wait()
m.logger.Info("GrowingFlushManager stopped")
}
// backgroundLoop periodically checks sync policies and triggers flushes.
func (m *GrowingFlushManager) backgroundLoop(ctx context.Context) {
defer m.wg.Done()
ticker := time.NewTicker(m.flushInterval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-m.stopCh:
return
case <-ticker.C:
m.checkAndTriggerSync(ctx)
}
}
}
// checkAndTriggerSync checks all policies and triggers sync for matching segments.
func (m *GrowingFlushManager) checkAndTriggerSync(ctx context.Context) {
// get all growing segments
segments := m.getGrowingSegments()
if len(segments) == 0 {
return
}
// wrap segments as GrowingSegmentBuffer for policy evaluation
buffers := m.wrapAsBuffers(segments)
currentTs := tsoutil.ComposeTSByTime(time.Now(), 0)
// track which segments to sync (dedup across policies)
syncSet := make(map[int64]string) // segmentID -> reason
// evaluate each policy
for _, policy := range m.policies {
segIDs := policy.SelectSegments(buffers, currentTs)
for _, segID := range segIDs {
if _, exists := syncSet[segID]; !exists {
syncSet[segID] = policy.Reason()
}
}
}
// trigger sync for each selected segment
for segID, reason := range syncSet {
segment := m.findSegment(segments, segID)
if segment == nil {
continue
}
m.triggerSync(ctx, segment, reason)
}
}
// getGrowingSegments returns all growing segments from the segment manager.
func (m *GrowingFlushManager) getGrowingSegments() []Segment {
// get segment IDs being tracked
trackedIDs := m.checkpointTracker.GetSegmentIDs()
if len(trackedIDs) == 0 {
return nil
}
segments := make([]Segment, 0, len(trackedIDs))
for _, segID := range trackedIDs {
seg := m.segmentManager.GetGrowing(segID)
if seg != nil {
segments = append(segments, seg)
}
}
return segments
}
// wrapAsBuffers wraps segments as GrowingSegmentBuffer for policy evaluation.
func (m *GrowingFlushManager) wrapAsBuffers(segments []Segment) []*GrowingSegmentBuffer {
buffers := make([]*GrowingSegmentBuffer, 0, len(segments))
for _, seg := range segments {
buf := NewGrowingSegmentBuffer(seg, m.checkpointTracker)
if buf.HasUnflushedData() {
buffers = append(buffers, buf)
}
}
return buffers
}
// findSegment finds a segment by ID from the list.
func (m *GrowingFlushManager) findSegment(segments []Segment, segID int64) Segment {
for _, seg := range segments {
if seg.ID() == segID {
return seg
}
}
return nil
}
// triggerSync initiates a sync operation for the given segment.
func (m *GrowingFlushManager) triggerSync(ctx context.Context, segment Segment, reason string) {
segID := segment.ID()
log := m.logger.With(
zap.Int64("segmentID", segID),
zap.String("reason", reason),
)
// get or create mutex for this segment to prevent concurrent sync
muInterface, _ := m.syncMu.LoadOrStore(segID, &sync.Mutex{})
mu := muInterface.(*sync.Mutex)
// try to acquire lock without blocking
if !mu.TryLock() {
log.Debug("sync already in progress, skipping")
return
}
log.Info("triggering growing segment sync")
// async sync to avoid blocking the check loop
m.wg.Add(1)
go func() {
defer m.wg.Done()
defer mu.Unlock()
if err := m.doSync(ctx, segment, false); err != nil {
log.Error("growing segment sync failed", zap.Error(err))
}
}()
}
// doSync performs the actual sync operation for a segment.
// when isSeal is true, this is the final flush for the segment (user flush/seal),
// and SaveBinlogPaths will set Flushed=true to trigger DataCoord handoff.
func (m *GrowingFlushManager) doSync(ctx context.Context, segment Segment, isSeal bool) error {
segID := segment.ID()
log := m.logger.With(zap.Int64("segmentID", segID))
// 1. get unflushed data range
// Use GetMaxEndOffset (from CheckpointTracker) instead of segment.RowNum() to
// determine how far to flush. This ensures we only flush data that has a
flushedOffset := m.checkpointTracker.GetFlushedOffset(segID)
currentOffset := m.checkpointTracker.GetMaxEndOffset(segID)
if flushedOffset >= currentOffset {
if isSeal {
lastManifest := m.checkpointTracker.GetAcknowledgedManifest(segID)
if lastManifest == "" {
log.Info("empty segment with no manifest, skip seal",
zap.Int64("flushedOffset", flushedOffset))
return nil
}
log.Info("no unflushed data, but sealing segment",
zap.String("lastManifest", lastManifest),
zap.Int64("flushedRows", flushedOffset))
return m.saveManifestAndCheckpoint(ctx, segment, lastManifest, nil, flushedOffset, true)
}
log.Debug("no unflushed data")
return nil
}
unflushedRows := currentOffset - flushedOffset
log.Info("starting sync",
zap.Int64("flushedOffset", flushedOffset),
zap.Int64("currentOffset", currentOffset),
zap.Int64("unflushedRows", unflushedRows),
)
// 2. get checkpoint position for this sync
checkpoint := m.checkpointTracker.GetCheckpoint(segID, currentOffset)
// 3. flush data directly via C++ milvus-storage (unified interface)
// Use the last acknowledged manifest version as read_version.
flushConfig, err := m.buildFlushConfig(segment)
if err != nil {
return err
}
flushConfig.ReadVersion = m.checkpointTracker.GetAcknowledgedVersion(segID)
flushResult, err := segment.FlushData(ctx, flushedOffset, currentOffset, flushConfig)
if err != nil {
return errors.Wrap(err, "failed to flush data via C++ milvus-storage")
}
if flushResult == nil || flushResult.ManifestPath == "" {
if isSeal {
lastManifest := m.checkpointTracker.GetAcknowledgedManifest(segID)
if lastManifest == "" {
log.Warn("FlushData returned no manifest and no prior manifest exists, skip seal")
return nil
}
// Use flushedOffset (not currentOffset) — the last manifest only contains
// data up to flushedOffset. Using currentOffset would over-report rows,
// causing sort compaction to fail with "unexpected row count".
log.Warn("no data flushed but seal requested, using last manifest",
zap.String("lastManifest", lastManifest),
zap.Int64("flushedOffset", flushedOffset),
zap.Int64("currentOffset", currentOffset))
return m.saveManifestAndCheckpoint(ctx, segment, lastManifest, checkpoint, flushedOffset, true)
}
log.Debug("no data flushed or no manifest path returned")
return nil
}
// 4. save manifest path and checkpoint to DataCoord
// In Storage V3 FFI mode, only manifest path is needed (all file info is in manifest)
if err := m.saveManifestAndCheckpoint(ctx, segment, flushResult.ManifestPath, checkpoint, currentOffset, isSeal); err != nil {
return err
}
// 5. update flushed offset AND acknowledged version
// Both updates happen ONLY after SaveBinlogPaths succeeds.
m.checkpointTracker.UpdateFlushedOffset(segID, currentOffset)
m.checkpointTracker.UpdateAcknowledgedManifest(segID, flushResult.ManifestPath)
log.Info("sync completed",
zap.Int64("flushedRows", flushResult.NumRows),
zap.String("manifest", flushResult.ManifestPath),
)
return nil
}
// buildFlushConfig builds FlushConfig for C++ FFI call.
// Go layer only constructs paths and config, all logic is in C++ side.
func (m *GrowingFlushManager) buildFlushConfig(segment Segment) (*FlushConfig, error) {
// Build paths for C++ side
// Segment path: {root}/insert_log/{coll}/{part}/{seg}
segmentBasePath := m.chunkManager.RootPath() + "/insert_log/" +
metautil.JoinIDPath(m.collectionID, segment.Partition(), segment.ID())
// Partition path: {root}/insert_log/{coll}/{part}
partitionBasePath := m.chunkManager.RootPath() + "/insert_log/" +
metautil.JoinIDPath(m.collectionID, segment.Partition())
// Collect TEXT field IDs and LOB paths from schema
var textFieldIDs []int64
var textLobPaths []string
if m.schema != nil {
for _, field := range m.schema.GetFields() {
if field.GetDataType() == schemapb.DataType_Text {
fieldID := field.GetFieldID()
textFieldIDs = append(textFieldIDs, fieldID)
// LOB path: {partition_path}/lobs/{field_id}
lobPath := fmt.Sprintf("%s/lobs/%d", partitionBasePath, fieldID)
textLobPaths = append(textLobPaths, lobPath)
}
}
}
globalWriterFormat := paramtable.Get().DataNodeCfg.StorageFormat.GetValue()
writerFormat, err := m.resolveFlushWriterFormat(segment, globalWriterFormat)
if err != nil {
return nil, err
}
return &FlushConfig{
SegmentBasePath: segmentBasePath,
PartitionBasePath: partitionBasePath,
CollectionID: m.collectionID,
PartitionID: segment.Partition(),
TextFieldIDs: textFieldIDs,
TextLobPaths: textLobPaths,
WriterFormat: writerFormat,
}, nil
}
func (m *GrowingFlushManager) resolveFlushWriterFormat(segment Segment, globalWriterFormat string) (string, error) {
if m.checkpointTracker == nil {
return globalWriterFormat, nil
}
manifest := m.checkpointTracker.GetAcknowledgedManifest(segment.ID())
if manifest == "" {
return globalWriterFormat, nil
}
format, err := packed.ResolveManifestSingleWriterFormat(manifest, packed.CreateStorageConfig(), m.flushWriterColumns(), globalWriterFormat)
if err != nil {
return "", err
}
return format, nil
}
func (m *GrowingFlushManager) flushWriterColumns() []string {
seen := make(map[string]struct{})
columns := []string{
strconv.FormatInt(milvuscommon.RowIDField, 10),
strconv.FormatInt(milvuscommon.TimeStampField, 10),
}
seen[columns[0]] = struct{}{}
seen[columns[1]] = struct{}{}
if m.schema == nil {
return columns
}
for _, field := range m.schema.GetFields() {
column := strconv.FormatInt(field.GetFieldID(), 10)
if _, ok := seen[column]; ok {
continue
}
seen[column] = struct{}{}
columns = append(columns, column)
}
return columns
}
// saveManifestAndCheckpoint saves manifest path and checkpoint to DataCoord.
// In Storage V3 FFI mode, only manifest path is needed - all file information
// (binlog paths, column groups, etc.) is contained in the manifest file.
// C++ side has already written all files and created the manifest.
func (m *GrowingFlushManager) saveManifestAndCheckpoint(
ctx context.Context,
segment Segment,
manifestPath string,
checkpoint *msgpb.MsgPosition,
numRows int64,
isSeal bool,
) error {
if m.binlogSaver == nil {
return errors.New("binlogSaver is nil, cannot save manifest and checkpoint")
}
segID := segment.ID()
// In Storage V3 FFI mode, SaveBinlogPaths is only used to:
// 1. Save manifest path (UpdateManifest)
// 2. Save checkpoint (UpdateCheckPoint)
// 3. Save start position (if needed)
// All file paths are in manifest, no need to pass Field2BinlogPaths
req := &datapb.SaveBinlogPathsRequest{
Base: &commonpb.MsgBase{},
SegmentID: segID,
CollectionID: m.collectionID,
PartitionID: segment.Partition(),
Flushed: isSeal, // true when sealing, false for incremental flush
WithFullBinlogs: false, // incremental mode
ManifestPath: manifestPath, // Storage V3 manifest path (from C++ side)
StorageVersion: storage.StorageV3,
}
if checkpoint != nil {
req.CheckPoints = []*datapb.CheckPoint{{
SegmentID: segID,
Position: checkpoint,
NumOfRows: numRows,
}}
}
return m.binlogSaver.SaveBinlogPaths(ctx, req)
}
// ForceSync forces an immediate sync of all unflushed data for a segment.
// this is typically called when a segment is being sealed or dropped.
func (m *GrowingFlushManager) ForceSync(ctx context.Context, segID int64) error {
segment := m.segmentManager.GetGrowing(segID)
if segment == nil {
return nil // segment not found, nothing to sync
}
// acquire per-segment mutex to prevent concurrent doSync with background triggerSync
muInterface, _ := m.syncMu.LoadOrStore(segID, &sync.Mutex{})
mu := muInterface.(*sync.Mutex)
mu.Lock()
defer mu.Unlock()
return m.doSync(ctx, segment, false)
}
// ForceSyncAndSeal forces an immediate sync and marks the segment as flushed/sealed.
// This is called when the segment is being sealed (user flush command).
// After this call, DataCoord will mark the segment as Flushed and trigger handoff.
//
// This method is idempotent: once a segment has been successfully sealed, subsequent
// calls are no-ops. This prevents duplicate SaveBinlogPaths(isFlush=true) requests
// when flush() is called multiple times, which would otherwise resurrect compacted
// segments via the Dropped→Flushed state transition race and cause duplicate data.
func (m *GrowingFlushManager) ForceSyncAndSeal(ctx context.Context, segID int64) error {
// Idempotency check: skip segments that have already been sealed.
if _, sealed := m.sealedSegments.Load(segID); sealed {
m.logger.Debug("segment already sealed, skipping duplicate ForceSyncAndSeal",
zap.Int64("segmentID", segID))
return nil
}
segment := m.segmentManager.GetGrowing(segID)
if segment == nil {
return nil // segment not found, nothing to sync
}
// acquire per-segment mutex to prevent concurrent doSync with background triggerSync
muInterface, _ := m.syncMu.LoadOrStore(segID, &sync.Mutex{})
mu := muInterface.(*sync.Mutex)
mu.Lock()
defer mu.Unlock()
// Double-check after acquiring lock, another goroutine may have sealed it.
if _, sealed := m.sealedSegments.Load(segID); sealed {
return nil
}
if err := m.doSync(ctx, segment, true); err != nil {
return err
}
m.sealedSegments.Store(segID, struct{}{})
return nil
}
// RemoveSealedSegment clears the sealed state for a segment.
// This should be called when the segment is released from QueryNode.
func (m *GrowingFlushManager) RemoveSealedSegment(segID int64) {
m.sealedSegments.Delete(segID)
}
// GetCheckpointTracker returns the checkpoint tracker for external access.
func (m *GrowingFlushManager) GetCheckpointTracker() *CheckpointTracker {
return m.checkpointTracker
}
internal/querynodev2/segments/growing_flush_manager_test.go
-931
View File
@@ -1,931 +0,0 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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.
package segments
import (
"context"
"sync"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/mocks/mock_storage"
"github.com/milvus-io/milvus/internal/storagev2/packed"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
func TestNewGrowingFlushManager(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_Text},
},
}
tracker := NewCheckpointTracker()
t.Run("default configuration", func(t *testing.T) {
m := NewGrowingFlushManager(
1, // collectionID
"test-channel", // channelName
schema, // schema
nil, // collectionManager
nil, // segmentManager
nil, // broker
nil, // chunkManager
tracker, // checkpointTracker
)
assert.NotNil(t, m)
assert.Equal(t, int64(1), m.collectionID)
assert.Equal(t, "test-channel", m.channelName)
assert.Equal(t, schema, m.schema)
assert.Equal(t, tracker, m.checkpointTracker)
assert.Equal(t, 100*time.Millisecond, m.flushInterval)
assert.True(t, len(m.policies) > 0) // Default policies should be set
})
t.Run("custom configuration", func(t *testing.T) {
customPolicy := GrowingSyncPolicyFunc{
fn: func(buffers []*GrowingSegmentBuffer, ts uint64) []int64 { return nil },
reason: "custom",
}
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
nil,
nil,
nil,
tracker,
WithFlushInterval(500*time.Millisecond),
WithSyncPolicies(customPolicy),
)
assert.Equal(t, 500*time.Millisecond, m.flushInterval)
assert.Len(t, m.policies, 1)
assert.Equal(t, "custom", m.policies[0].Reason())
})
}
func TestGrowingFlushManager_StartStop(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
},
}
tracker := NewCheckpointTracker()
segMgr := NewMockSegmentManager(t)
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
segMgr,
nil,
nil,
tracker,
WithFlushInterval(10*time.Millisecond),
)
ctx := context.Background()
// Start the manager
m.Start(ctx)
// Let it run briefly
time.Sleep(50 * time.Millisecond)
// Stop should complete without blocking
done := make(chan struct{})
go func() {
m.Stop()
close(done)
}()
select {
case <-done:
// Success - stopped cleanly
case <-time.After(time.Second):
t.Fatal("Stop() did not complete in time")
}
}
func TestGrowingFlushManager_WrapAsBuffers(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
}
tracker := NewCheckpointTracker()
// Record batches so segments have unflushed data (needed by HasUnflushedData check)
pos := &msgpb.MsgPosition{Timestamp: 100}
tracker.RecordBatch(1001, 50, pos)
tracker.RecordBatch(1002, 50, pos)
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
nil,
nil,
nil,
tracker,
)
// Create mock segments with all necessary expectations
seg1 := NewMockSegment(t)
seg1.EXPECT().ID().Return(int64(1001)).Maybe()
seg1.EXPECT().RowNum().Return(int64(50)).Maybe()
seg2 := NewMockSegment(t)
seg2.EXPECT().ID().Return(int64(1002)).Maybe()
seg2.EXPECT().RowNum().Return(int64(50)).Maybe()
segments := []Segment{seg1, seg2}
buffers := m.wrapAsBuffers(segments)
assert.Len(t, buffers, 2)
assert.Equal(t, int64(1001), buffers[0].SegmentID())
assert.Equal(t, int64(1002), buffers[1].SegmentID())
}
func TestGrowingFlushManager_GetGrowingSegments(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
}
tracker := NewCheckpointTracker()
segMgr := NewMockSegmentManager(t)
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
segMgr,
nil,
nil,
tracker,
)
t.Run("no tracked segments", func(t *testing.T) {
segments := m.getGrowingSegments()
assert.Len(t, segments, 0)
})
t.Run("with tracked segments", func(t *testing.T) {
// Record batches to track segments
pos := &msgpb.MsgPosition{Timestamp: 100}
tracker.RecordBatch(1001, 50, pos)
tracker.RecordBatch(1002, 50, pos)
// Setup mock expectations
seg1 := NewMockSegment(t)
seg2 := NewMockSegment(t)
segMgr.EXPECT().GetGrowing(int64(1001)).Return(seg1).Once()
segMgr.EXPECT().GetGrowing(int64(1002)).Return(seg2).Once()
segments := m.getGrowingSegments()
assert.Len(t, segments, 2)
})
t.Run("segment not found in manager", func(t *testing.T) {
// Remove old tracker data and create fresh tracker
tracker = NewCheckpointTracker()
m.checkpointTracker = tracker
pos := &msgpb.MsgPosition{Timestamp: 100}
tracker.RecordBatch(9999, 50, pos)
// Create a fresh mock for this test
segMgr2 := NewMockSegmentManager(t)
m.segmentManager = segMgr2
// Segment doesn't exist
segMgr2.EXPECT().GetGrowing(int64(9999)).Return(nil).Once()
segments := m.getGrowingSegments()
assert.Len(t, segments, 0)
})
}
func TestGrowingFlushManager_FindSegment(t *testing.T) {
paramtable.Init()
m := &GrowingFlushManager{}
seg1 := NewMockSegment(t)
seg1.EXPECT().ID().Return(int64(1001))
seg2 := NewMockSegment(t)
seg2.EXPECT().ID().Return(int64(1002))
segments := []Segment{seg1, seg2}
t.Run("found", func(t *testing.T) {
found := m.findSegment(segments, 1001)
assert.Equal(t, seg1, found)
})
t.Run("not found", func(t *testing.T) {
found := m.findSegment(segments, 9999)
assert.Nil(t, found)
})
}
func TestGrowingFlushManager_CheckAndTriggerSync(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
}
tracker := NewCheckpointTracker()
segMgr := NewMockSegmentManager(t)
// Create a custom policy that always selects segment 1001
customPolicy := GrowingSyncPolicyFunc{
fn: func(buffers []*GrowingSegmentBuffer, ts uint64) []int64 {
for _, buf := range buffers {
if buf.SegmentID() == 1001 {
return []int64{1001}
}
}
return nil
},
reason: "test policy",
}
// Create mock chunk manager for buildFlushConfig
mockCM := mock_storage.NewMockChunkManager(t)
mockCM.EXPECT().RootPath().Return("/test/root").Maybe()
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
segMgr,
nil, // broker - needed for triggerSync but we won't call it in this test
mockCM, // chunkManager - needed for buildFlushConfig
tracker,
WithSyncPolicies(customPolicy),
)
// Record batches to track segments
pos := &msgpb.MsgPosition{Timestamp: 100}
tracker.RecordBatch(1001, 50, pos)
// Setup mock segment
seg := NewMockSegment(t)
seg.EXPECT().ID().Return(int64(1001)).Maybe()
seg.EXPECT().MemSize().Return(int64(1024)).Maybe()
seg.EXPECT().RowNum().Return(int64(100)).Maybe()
seg.EXPECT().Partition().Return(int64(10)).Maybe()
// FlushData will be called in doSync goroutine
seg.EXPECT().FlushData(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, nil).Maybe()
segMgr.EXPECT().GetGrowing(int64(1001)).Return(seg).Maybe()
// checkAndTriggerSync will trigger sync in a goroutine.
// The actual sync will fail gracefully (no FlushResult returned)
ctx := context.Background()
m.checkAndTriggerSync(ctx)
// Wait briefly for the async goroutine to run
time.Sleep(50 * time.Millisecond)
}
func TestGrowingFlushManager_DefaultPolicies(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
}
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
nil,
nil,
nil,
NewCheckpointTracker(),
)
policies := m.defaultPolicies()
// Should have at least FullBufferPolicy and StaleBufferPolicy
assert.True(t, len(policies) >= 2)
reasons := make([]string, len(policies))
for i, p := range policies {
reasons[i] = p.Reason()
}
assert.Contains(t, reasons, "buffer full")
assert.Contains(t, reasons, "buffer stale")
}
// Test segment mutex for concurrent sync prevention
func TestGrowingFlushManager_SegmentMutex(t *testing.T) {
paramtable.Init()
m := &GrowingFlushManager{}
segID := int64(1001)
// Get or create mutex
mutexI, _ := m.syncMu.LoadOrStore(segID, &sync.Mutex{})
mu := mutexI.(*sync.Mutex)
// First lock should succeed
locked := mu.TryLock()
assert.True(t, locked)
// Second lock should fail (mutex is held)
locked2 := mu.TryLock()
assert.False(t, locked2)
// After unlock, lock should succeed again
mu.Unlock()
locked3 := mu.TryLock()
assert.True(t, locked3)
mu.Unlock()
}
// Mock helpers for testing
type testSyncPolicy struct {
mock.Mock
}
func (p *testSyncPolicy) SelectSegments(buffers []*GrowingSegmentBuffer, ts uint64) []int64 {
args := p.Called(buffers, ts)
return args.Get(0).([]int64)
}
func (p *testSyncPolicy) Reason() string {
args := p.Called()
return args.String(0)
}
func TestGrowingFlushManager_BuildFlushConfig(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_VarChar},
},
}
tracker := NewCheckpointTracker()
// Create mock chunk manager
mockCM := mock_storage.NewMockChunkManager(t)
mockCM.EXPECT().RootPath().Return("/test/root").Maybe()
m := NewGrowingFlushManager(
1, // collectionID
"test-channel", // channelName
schema, // schema
nil, // collectionManager
nil, // segmentManager
nil, // broker
mockCM, // chunkManager
tracker, // checkpointTracker
)
// Create mock segment
seg := NewMockSegment(t)
seg.EXPECT().ID().Return(int64(1001)).Maybe()
seg.EXPECT().Partition().Return(int64(10)).Maybe()
config, err := m.buildFlushConfig(seg)
assert.NoError(t, err)
assert.NotNil(t, config)
assert.Contains(t, config.SegmentBasePath, "/test/root")
assert.Contains(t, config.SegmentBasePath, "insert_log")
assert.Contains(t, config.PartitionBasePath, "/test/root")
assert.Equal(t, int64(1), config.CollectionID)
assert.Equal(t, int64(10), config.PartitionID)
assert.Equal(t, "parquet", config.WriterFormat)
}
func TestGrowingFlushManager_BuildFlushConfigKeepsGlobalFormatAndUsesManifestForSinglePolicy(t *testing.T) {
paramtable.Init()
params := paramtable.Get()
rootPath := t.TempDir()
assert.NoError(t, params.Save(params.CommonCfg.StorageType.Key, "local"))
assert.NoError(t, params.Save(params.LocalStorageCfg.Path.Key, rootPath))
assert.NoError(t, params.Save(params.DataNodeCfg.StorageFormat.Key, "parquet"))
t.Cleanup(func() {
_ = params.Reset(params.CommonCfg.StorageType.Key)
_ = params.Reset(params.LocalStorageCfg.Path.Key)
_ = params.Reset(params.DataNodeCfg.StorageFormat.Key)
})
manifestPath, err := packed.CreateManifestForSegment(
"files/querynode_format/segment",
[]string{"0", "1", "100", "101"},
"parquet",
[]packed.Fragment{{FilePath: "files/querynode_format/segment/_data/0.parquet", StartRow: 0, EndRow: 1, RowCount: 1}},
packed.CreateStorageConfig(),
)
assert.NoError(t, err)
assert.NoError(t, params.Save(params.DataNodeCfg.StorageFormat.Key, "vortex"))
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_VarChar},
},
}
tracker := NewCheckpointTracker()
tracker.InitSegmentWithManifest(1001, 1, manifestPath)
mockCM := mock_storage.NewMockChunkManager(t)
mockCM.EXPECT().RootPath().Return(rootPath).Maybe()
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
nil,
nil,
mockCM,
tracker,
)
seg := NewMockSegment(t)
seg.EXPECT().ID().Return(int64(1001)).Maybe()
seg.EXPECT().Partition().Return(int64(10)).Maybe()
config, err := m.buildFlushConfig(seg)
assert.NoError(t, err)
assert.NotNil(t, config)
assert.Equal(t, "parquet", config.WriterFormat)
}
func TestGrowingFlushManager_DoSync_NoUnflushedData(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
}
tracker := NewCheckpointTracker()
segMgr := NewMockSegmentManager(t)
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
segMgr,
nil,
nil,
tracker,
)
// Create mock segment with no unflushed data
seg := NewMockSegment(t)
seg.EXPECT().ID().Return(int64(1001)).Maybe()
seg.EXPECT().RowNum().Return(int64(100)).Maybe()
// Set flushed offset to equal row count (no unflushed data)
tracker.UpdateFlushedOffset(1001, 100)
// doSync should return nil when no data to flush
err := m.doSync(context.Background(), seg, false)
assert.NoError(t, err)
}
func TestGrowingFlushManager_DoSync_WithUnflushedData(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
}
tracker := NewCheckpointTracker()
mockCM := mock_storage.NewMockChunkManager(t)
mockCM.EXPECT().RootPath().Return("/test/root").Maybe()
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
nil,
nil,
mockCM,
tracker,
)
// Create mock segment with unflushed data
seg := NewMockSegment(t)
seg.EXPECT().ID().Return(int64(1001)).Maybe()
seg.EXPECT().Partition().Return(int64(10)).Maybe()
seg.EXPECT().RowNum().Return(int64(100)).Maybe()
// Set flushed offset less than row count (has unflushed data)
tracker.UpdateFlushedOffset(1001, 50)
// Record batch to get checkpoint
pos := &msgpb.MsgPosition{Timestamp: 200}
tracker.RecordBatch(1001, 100, pos)
// FlushData will be called
flushResult := &FlushResult{
ManifestPath: "/test/manifest.json",
NumRows: 50,
}
seg.EXPECT().FlushData(mock.Anything, int64(50), int64(100), mock.Anything).Return(flushResult, nil).Once()
// Without broker, saveManifestAndCheckpoint will fail
// but the test should still exercise the main logic
err := m.doSync(context.Background(), seg, false)
// Expected to fail because broker is nil
assert.Error(t, err)
}
func TestGrowingFlushManager_ForceSync(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
}
tracker := NewCheckpointTracker()
segMgr := NewMockSegmentManager(t)
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
segMgr,
nil,
nil,
tracker,
)
t.Run("segment not found", func(t *testing.T) {
segMgr.EXPECT().GetGrowing(int64(9999)).Return(nil).Once()
err := m.ForceSync(context.Background(), 9999)
assert.NoError(t, err) // should return nil when segment not found
})
t.Run("segment exists but no unflushed data", func(t *testing.T) {
seg := NewMockSegment(t)
seg.EXPECT().ID().Return(int64(1001)).Maybe()
seg.EXPECT().RowNum().Return(int64(100)).Maybe()
segMgr.EXPECT().GetGrowing(int64(1001)).Return(seg).Once()
// Set flushed offset equal to row count
tracker.UpdateFlushedOffset(1001, 100)
err := m.ForceSync(context.Background(), 1001)
assert.NoError(t, err)
})
}
func TestGrowingFlushManager_TriggerSync_AlreadyInProgress(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
}
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
nil,
nil,
nil,
NewCheckpointTracker(),
)
// Create mock segment
seg := NewMockSegment(t)
seg.EXPECT().ID().Return(int64(1001)).Maybe()
// Acquire the mutex first to simulate sync in progress
muInterface, _ := m.syncMu.LoadOrStore(int64(1001), &sync.Mutex{})
mu := muInterface.(*sync.Mutex)
mu.Lock()
defer mu.Unlock()
// triggerSync should skip because sync is already in progress
// (this is tested by not expecting any FlushData calls)
m.triggerSync(context.Background(), seg, "test")
// Give time for goroutine to potentially start
time.Sleep(10 * time.Millisecond)
// If FlushData was not called, the test passes
// The mock would fail if FlushData was called without expectation
}
func TestGrowingFlushManager_GetCheckpointTracker(t *testing.T) {
paramtable.Init()
tracker := NewCheckpointTracker()
m := NewGrowingFlushManager(
1,
"test-channel",
nil,
nil,
nil,
nil,
nil,
tracker,
)
assert.Equal(t, tracker, m.GetCheckpointTracker())
}
// MockBinlogSaver is a testify mock implementing the BinlogSaver interface.
type MockBinlogSaver struct {
mock.Mock
}
func (m *MockBinlogSaver) SaveBinlogPaths(ctx context.Context, req *datapb.SaveBinlogPathsRequest) error {
args := m.Called(ctx, req)
return args.Error(0)
}
func TestGrowingFlushManager_ForceSyncAndSeal(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
},
}
t.Run("segment not found returns nil", func(t *testing.T) {
tracker := NewCheckpointTracker()
segMgr := NewMockSegmentManager(t)
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
segMgr,
nil,
nil,
tracker,
)
segMgr.EXPECT().GetGrowing(int64(9999)).Return(nil).Once()
err := m.ForceSyncAndSeal(context.Background(), 9999)
assert.NoError(t, err)
})
t.Run("calls doSync with isSeal true", func(t *testing.T) {
tracker := NewCheckpointTracker()
segMgr := NewMockSegmentManager(t)
mockCM := mock_storage.NewMockChunkManager(t)
mockCM.EXPECT().RootPath().Return("/test/root").Maybe()
binlogSaver := &MockBinlogSaver{}
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
segMgr,
binlogSaver,
mockCM,
tracker,
)
seg := NewMockSegment(t)
seg.EXPECT().ID().Return(int64(1001)).Maybe()
seg.EXPECT().Partition().Return(int64(10)).Maybe()
seg.EXPECT().RowNum().Return(int64(100)).Maybe()
segMgr.EXPECT().GetGrowing(int64(1001)).Return(seg).Once()
// Set up tracker with unflushed data: flushed offset = 0, record batch up to 100
pos := &msgpb.MsgPosition{ChannelName: "test-channel", MsgID: []byte{1}, Timestamp: 100}
tracker.RecordBatch(1001, 100, pos)
// Mock FlushData to return a valid result
flushResult := &FlushResult{
ManifestPath: "/test/manifest.json",
NumRows: 100,
}
seg.EXPECT().FlushData(mock.Anything, int64(0), int64(100), mock.Anything).Return(flushResult, nil).Once()
// Capture the SaveBinlogPaths request and assert Flushed=true
binlogSaver.On("SaveBinlogPaths", mock.Anything, mock.MatchedBy(func(req *datapb.SaveBinlogPathsRequest) bool {
return req.Flushed == true
})).Return(nil).Once()
err := m.ForceSyncAndSeal(context.Background(), 1001)
assert.NoError(t, err)
binlogSaver.AssertExpectations(t)
})
}
func TestGrowingFlushManager_DoSync_isSeal(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
},
}
t.Run("isSeal true produces Flushed true", func(t *testing.T) {
tracker := NewCheckpointTracker()
mockCM := mock_storage.NewMockChunkManager(t)
mockCM.EXPECT().RootPath().Return("/test/root").Maybe()
binlogSaver := &MockBinlogSaver{}
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
nil,
binlogSaver,
mockCM,
tracker,
)
seg := NewMockSegment(t)
seg.EXPECT().ID().Return(int64(2001)).Maybe()
seg.EXPECT().Partition().Return(int64(10)).Maybe()
seg.EXPECT().RowNum().Return(int64(50)).Maybe()
// Set up tracker: flushed offset = 0, record batch up to 50
pos := &msgpb.MsgPosition{ChannelName: "test-channel", MsgID: []byte{1}, Timestamp: 200}
tracker.RecordBatch(2001, 50, pos)
flushResult := &FlushResult{
ManifestPath: "/test/manifest_seal.json",
NumRows: 50,
}
seg.EXPECT().FlushData(mock.Anything, int64(0), int64(50), mock.Anything).Return(flushResult, nil).Once()
binlogSaver.On("SaveBinlogPaths", mock.Anything, mock.MatchedBy(func(req *datapb.SaveBinlogPathsRequest) bool {
return req.Flushed == true
})).Return(nil).Once()
err := m.doSync(context.Background(), seg, true)
assert.NoError(t, err)
binlogSaver.AssertExpectations(t)
})
t.Run("isSeal false produces Flushed false", func(t *testing.T) {
tracker := NewCheckpointTracker()
mockCM := mock_storage.NewMockChunkManager(t)
mockCM.EXPECT().RootPath().Return("/test/root").Maybe()
binlogSaver := &MockBinlogSaver{}
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
nil,
binlogSaver,
mockCM,
tracker,
)
seg := NewMockSegment(t)
seg.EXPECT().ID().Return(int64(3001)).Maybe()
seg.EXPECT().Partition().Return(int64(10)).Maybe()
seg.EXPECT().RowNum().Return(int64(80)).Maybe()
// Set up tracker: flushed offset = 0, record batch up to 80
pos := &msgpb.MsgPosition{ChannelName: "test-channel", MsgID: []byte{1}, Timestamp: 300}
tracker.RecordBatch(3001, 80, pos)
flushResult := &FlushResult{
ManifestPath: "/test/manifest_incr.json",
NumRows: 80,
}
seg.EXPECT().FlushData(mock.Anything, int64(0), int64(80), mock.Anything).Return(flushResult, nil).Once()
binlogSaver.On("SaveBinlogPaths", mock.Anything, mock.MatchedBy(func(req *datapb.SaveBinlogPathsRequest) bool {
return req.Flushed == false
})).Return(nil).Once()
err := m.doSync(context.Background(), seg, false)
assert.NoError(t, err)
binlogSaver.AssertExpectations(t)
})
}
func TestGrowingFlushManager_DoSync_ChannelEmpty(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
},
}
tracker := NewCheckpointTracker()
mockCM := mock_storage.NewMockChunkManager(t)
mockCM.EXPECT().RootPath().Return("/test/root").Maybe()
binlogSaver := &MockBinlogSaver{}
m := NewGrowingFlushManager(
1,
"test-channel",
schema,
nil,
nil,
binlogSaver,
mockCM,
tracker,
)
seg := NewMockSegment(t)
seg.EXPECT().ID().Return(int64(4001)).Maybe()
seg.EXPECT().Partition().Return(int64(10)).Maybe()
seg.EXPECT().RowNum().Return(int64(60)).Maybe()
// Set up tracker: flushed offset = 0, record batch up to 60
pos := &msgpb.MsgPosition{ChannelName: "test-channel", MsgID: []byte{1}, Timestamp: 400}
tracker.RecordBatch(4001, 60, pos)
flushResult := &FlushResult{
ManifestPath: "/test/manifest_channel.json",
NumRows: 60,
}
seg.EXPECT().FlushData(mock.Anything, int64(0), int64(60), mock.Anything).Return(flushResult, nil).Once()
// Assert that Channel field is empty string in SaveBinlogPaths request
binlogSaver.On("SaveBinlogPaths", mock.Anything, mock.MatchedBy(func(req *datapb.SaveBinlogPathsRequest) bool {
return req.Channel == ""
})).Return(nil).Once()
err := m.doSync(context.Background(), seg, false)
assert.NoError(t, err)
binlogSaver.AssertExpectations(t)
}
internal/querynodev2/segments/growing_segment_buffer.go
-307
View File
@@ -1,307 +0,0 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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.
package segments
import (
"container/heap"
"math/rand"
"sort"
"time"
"github.com/samber/lo"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
// GrowingSegmentBuffer wraps a Growing Segment to provide the interface needed
// for sync policy selection. this is an adapter that allows reusing the sync
// policy algorithms from flushcommon/writebuffer with Growing Segments.
//
// note: we can't directly use writebuffer.SyncPolicy because it expects
// *writebuffer.segmentBuffer which is a private type. instead, we define
// our own GrowingSyncPolicy interface with the same algorithm logic.
type GrowingSegmentBuffer struct {
segment Segment
checkpointTracker *CheckpointTracker
}
// NewGrowingSegmentBuffer creates a new GrowingSegmentBuffer wrapping the given segment.
func NewGrowingSegmentBuffer(segment Segment, tracker *CheckpointTracker) *GrowingSegmentBuffer {
return &GrowingSegmentBuffer{
segment: segment,
checkpointTracker: tracker,
}
}
// SegmentID returns the segment ID.
func (b *GrowingSegmentBuffer) SegmentID() int64 {
return b.segment.ID()
}
// Segment returns the underlying segment.
func (b *GrowingSegmentBuffer) Segment() Segment {
return b.segment
}
// IsFull returns true if the segment's memory usage exceeds the threshold.
// the threshold is configured via paramtable.
func (b *GrowingSegmentBuffer) IsFull() bool {
threshold := paramtable.Get().DataNodeCfg.FlushInsertBufferSize.GetAsInt64()
if threshold <= 0 {
threshold = 16 * 1024 * 1024 // 16MB default
}
return b.segment.MemSize() >= threshold
}
// MinTimestamp returns the minimum timestamp of unflushed data.
func (b *GrowingSegmentBuffer) MinTimestamp() typeutil.Timestamp {
if b.checkpointTracker == nil {
return 0
}
return b.checkpointTracker.GetMinTimestamp(b.segment.ID())
}
// MemorySize returns the current memory usage of the segment.
func (b *GrowingSegmentBuffer) MemorySize() int64 {
return b.segment.MemSize()
}
// RowCount returns the current row count of the segment.
func (b *GrowingSegmentBuffer) RowCount() int64 {
return b.segment.RowNum()
}
// FlushedOffset returns the last flushed row offset.
func (b *GrowingSegmentBuffer) FlushedOffset() int64 {
if b.checkpointTracker == nil {
return 0
}
return b.checkpointTracker.GetFlushedOffset(b.segment.ID())
}
// UnflushedRowCount returns the number of rows that haven't been flushed yet.
func (b *GrowingSegmentBuffer) UnflushedRowCount() int64 {
return b.RowCount() - b.FlushedOffset()
}
// HasUnflushedData returns true if there are unflushed rows.
func (b *GrowingSegmentBuffer) HasUnflushedData() bool {
return b.UnflushedRowCount() > 0
}
// ============================================================================
// GrowingSyncPolicy - sync policy interface for Growing Segments
// ============================================================================
// GrowingSyncPolicy defines the interface for selecting which Growing Segments
// should be synced. this mirrors writebuffer.SyncPolicy but works with
// GrowingSegmentBuffer instead of the private segmentBuffer type.
type GrowingSyncPolicy interface {
// SelectSegments returns the IDs of segments that should be synced.
SelectSegments(buffers []*GrowingSegmentBuffer, ts typeutil.Timestamp) []int64
// Reason returns a human-readable reason for this policy.
Reason() string
}
// GrowingSyncPolicyFunc is a function type that implements GrowingSyncPolicy.
type GrowingSyncPolicyFunc struct {
fn func(buffers []*GrowingSegmentBuffer, ts typeutil.Timestamp) []int64
reason string
}
// SelectSegments implements GrowingSyncPolicy.
func (f GrowingSyncPolicyFunc) SelectSegments(buffers []*GrowingSegmentBuffer, ts typeutil.Timestamp) []int64 {
return f.fn(buffers, ts)
}
// Reason implements GrowingSyncPolicy.
func (f GrowingSyncPolicyFunc) Reason() string {
return f.reason
}
// ============================================================================
// sync policy factory functions
// these mirror the policies in flushcommon/writebuffer/sync_policy.go
// ============================================================================
// GetGrowingFullBufferPolicy returns a policy that selects segments whose buffer is full.
// this mirrors writebuffer.GetFullBufferPolicy.
func GetGrowingFullBufferPolicy() GrowingSyncPolicy {
return GrowingSyncPolicyFunc{
fn: func(buffers []*GrowingSegmentBuffer, _ typeutil.Timestamp) []int64 {
return lo.FilterMap(buffers, func(buf *GrowingSegmentBuffer, _ int) (int64, bool) {
return buf.SegmentID(), buf.IsFull()
})
},
reason: "buffer full",
}
}
// GetGrowingStaleBufferPolicy returns a policy that selects segments with stale data.
// this mirrors writebuffer.GetSyncStaleBufferPolicy.
func GetGrowingStaleBufferPolicy(staleDuration time.Duration) GrowingSyncPolicy {
return GrowingSyncPolicyFunc{
fn: func(buffers []*GrowingSegmentBuffer, ts typeutil.Timestamp) []int64 {
current := tsoutil.PhysicalTime(ts)
return lo.FilterMap(buffers, func(buf *GrowingSegmentBuffer, _ int) (int64, bool) {
minTs := buf.MinTimestamp()
if minTs == 0 {
return 0, false // no unflushed data
}
start := tsoutil.PhysicalTime(minTs)
// add jitter to avoid thundering herd
jitter := time.Duration(rand.Float64() * 0.1 * float64(staleDuration))
return buf.SegmentID(), current.Sub(start) > staleDuration+jitter
})
},
reason: "buffer stale",
}
}
// GetGrowingOldestBufferPolicy returns a policy that selects the N oldest buffers.
// this mirrors writebuffer.GetOldestBufferPolicy.
func GetGrowingOldestBufferPolicy(num int) GrowingSyncPolicy {
return GrowingSyncPolicyFunc{
fn: func(buffers []*GrowingSegmentBuffer, _ typeutil.Timestamp) []int64 {
if len(buffers) <= num {
// return all segments with unflushed data
return lo.FilterMap(buffers, func(buf *GrowingSegmentBuffer, _ int) (int64, bool) {
return buf.SegmentID(), buf.HasUnflushedData()
})
}
// use a max-heap to find the N oldest buffers
h := &growingSegmentHeap{}
heap.Init(h)
for _, buf := range buffers {
if !buf.HasUnflushedData() {
continue
}
heap.Push(h, buf)
if h.Len() > num {
heap.Pop(h)
}
}
return lo.Map(*h, func(buf *GrowingSegmentBuffer, _ int) int64 {
return buf.SegmentID()
})
},
reason: "oldest buffers",
}
}
// GetGrowingMemoryPressurePolicy returns a policy that syncs when total memory exceeds threshold.
func GetGrowingMemoryPressurePolicy(memoryThreshold int64) GrowingSyncPolicy {
return GrowingSyncPolicyFunc{
fn: func(buffers []*GrowingSegmentBuffer, _ typeutil.Timestamp) []int64 {
// calculate total unflushed memory
var totalMemory int64
for _, buf := range buffers {
if buf.HasUnflushedData() {
totalMemory += buf.MemorySize()
}
}
if totalMemory < memoryThreshold {
return nil
}
// memory pressure: sync the largest segments first
// use a max-heap by memory size
type memSegment struct {
id int64
memSize int64
}
segments := make([]memSegment, 0)
for _, buf := range buffers {
if buf.HasUnflushedData() {
segments = append(segments, memSegment{
id: buf.SegmentID(),
memSize: buf.MemorySize(),
})
}
}
// sort by memory size descending
sort.Slice(segments, func(i, j int) bool {
return segments[i].memSize > segments[j].memSize
})
// select segments until we've freed enough memory
var freedMemory int64
targetFree := totalMemory - memoryThreshold/2 // try to get to 50% of threshold
result := make([]int64, 0)
for _, seg := range segments {
if freedMemory >= targetFree {
break
}
result = append(result, seg.id)
freedMemory += seg.memSize
}
return result
},
reason: "memory pressure",
}
}
// GetGrowingRowCountPolicy returns a policy that syncs segments with more than N unflushed rows.
func GetGrowingRowCountPolicy(rowThreshold int64) GrowingSyncPolicy {
return GrowingSyncPolicyFunc{
fn: func(buffers []*GrowingSegmentBuffer, _ typeutil.Timestamp) []int64 {
return lo.FilterMap(buffers, func(buf *GrowingSegmentBuffer, _ int) (int64, bool) {
return buf.SegmentID(), buf.UnflushedRowCount() >= rowThreshold
})
},
reason: "row count threshold",
}
}
// ============================================================================
// heap implementation for oldest buffer policy
// ============================================================================
// growingSegmentHeap implements heap.Interface for selecting oldest segments.
// it's a max-heap based on MinTimestamp (larger timestamp = newer = pop first).
type growingSegmentHeap []*GrowingSegmentBuffer
func (h growingSegmentHeap) Len() int { return len(h) }
func (h growingSegmentHeap) Less(i, j int) bool {
// max-heap: larger timestamp comes first (gets popped first)
return h[i].MinTimestamp() > h[j].MinTimestamp()
}
func (h growingSegmentHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *growingSegmentHeap) Push(x any) {
*h = append(*h, x.(*GrowingSegmentBuffer))
}
func (h *growingSegmentHeap) Pop() any {
old := *h
n := len(old)
x := old[n-1]
*h = old[0 : n-1]
return x
}
internal/querynodev2/segments/growing_segment_buffer_test.go
-327
View File
@@ -1,327 +0,0 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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.
package segments
import (
"container/heap"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
)
func init() {
paramtable.Init()
}
func TestGrowingSegmentBuffer_Basic(t *testing.T) {
// Create mock segment
seg := NewMockSegment(t)
seg.EXPECT().ID().Return(int64(1001))
seg.EXPECT().MemSize().Return(int64(1024 * 1024)) // 1MB
seg.EXPECT().RowNum().Return(int64(100))
// Create buffer without checkpoint tracker
buf := NewGrowingSegmentBuffer(seg, nil)
assert.Equal(t, int64(1001), buf.SegmentID())
assert.Equal(t, seg, buf.Segment())
assert.Equal(t, int64(1024*1024), buf.MemorySize())
assert.Equal(t, int64(100), buf.RowCount())
assert.Equal(t, uint64(0), buf.MinTimestamp()) // nil tracker returns 0
assert.Equal(t, int64(0), buf.FlushedOffset()) // nil tracker returns 0
assert.Equal(t, int64(100), buf.UnflushedRowCount())
assert.True(t, buf.HasUnflushedData())
}
func TestGrowingSegmentBuffer_WithTracker(t *testing.T) {
tracker := NewCheckpointTracker()
segID := int64(1001)
// Record some batches
pos1 := &msgpb.MsgPosition{ChannelName: "ch1", MsgID: []byte{1}, Timestamp: 100}
pos2 := &msgpb.MsgPosition{ChannelName: "ch1", MsgID: []byte{2}, Timestamp: 200}
tracker.RecordBatch(segID, 50, pos1)
tracker.RecordBatch(segID, 100, pos2)
// Create mock segment
seg := NewMockSegment(t)
seg.EXPECT().ID().Return(segID).Maybe()
seg.EXPECT().MemSize().Return(int64(1024 * 1024)).Maybe()
seg.EXPECT().RowNum().Return(int64(100)).Maybe()
buf := NewGrowingSegmentBuffer(seg, tracker)
// Initially no flushed data
assert.Equal(t, uint64(100), buf.MinTimestamp()) // timestamp from first batch
assert.Equal(t, int64(0), buf.FlushedOffset())
assert.Equal(t, int64(100), buf.UnflushedRowCount())
assert.True(t, buf.HasUnflushedData())
// After flushing to offset 50
tracker.UpdateFlushedOffset(segID, 50)
assert.Equal(t, uint64(200), buf.MinTimestamp()) // timestamp from second batch
assert.Equal(t, int64(50), buf.FlushedOffset())
assert.Equal(t, int64(50), buf.UnflushedRowCount())
assert.True(t, buf.HasUnflushedData())
// After flushing all
tracker.UpdateFlushedOffset(segID, 100)
assert.Equal(t, uint64(0), buf.MinTimestamp()) // no more unflushed data
assert.Equal(t, int64(100), buf.FlushedOffset())
assert.Equal(t, int64(0), buf.UnflushedRowCount())
assert.False(t, buf.HasUnflushedData())
}
func TestGrowingSegmentBuffer_IsFull(t *testing.T) {
// Get the configured threshold - uses FlushInsertBufferSize (same as IsFull method)
threshold := paramtable.Get().DataNodeCfg.FlushInsertBufferSize.GetAsInt64()
if threshold <= 0 {
threshold = 16 * 1024 * 1024 // 16MB default
}
t.Run("not full", func(t *testing.T) {
seg := NewMockSegment(t)
seg.EXPECT().MemSize().Return(threshold - 1)
buf := NewGrowingSegmentBuffer(seg, nil)
assert.False(t, buf.IsFull())
})
t.Run("full", func(t *testing.T) {
seg := NewMockSegment(t)
seg.EXPECT().MemSize().Return(threshold)
buf := NewGrowingSegmentBuffer(seg, nil)
assert.True(t, buf.IsFull())
})
t.Run("over full", func(t *testing.T) {
seg := NewMockSegment(t)
seg.EXPECT().MemSize().Return(threshold + 1)
buf := NewGrowingSegmentBuffer(seg, nil)
assert.True(t, buf.IsFull())
})
}
// ============================================================================
// Sync Policy Tests
// ============================================================================
func createMockBuffer(id int64, memSize int64, rowCount int64, tracker *CheckpointTracker) *GrowingSegmentBuffer {
seg := &MockSegment{}
seg.On("ID").Return(id)
seg.On("MemSize").Return(memSize)
seg.On("RowNum").Return(rowCount)
return NewGrowingSegmentBuffer(seg, tracker)
}
func TestGetGrowingFullBufferPolicy(t *testing.T) {
threshold := paramtable.Get().DataNodeCfg.FlushInsertBufferSize.GetAsInt64()
if threshold <= 0 {
threshold = 16 * 1024 * 1024
}
policy := GetGrowingFullBufferPolicy()
assert.Equal(t, "buffer full", policy.Reason())
buffers := []*GrowingSegmentBuffer{
createMockBuffer(1, threshold-1, 100, nil), // not full
createMockBuffer(2, threshold, 100, nil), // full
createMockBuffer(3, threshold+1, 100, nil), // full
}
selected := policy.SelectSegments(buffers, 0)
assert.Len(t, selected, 2)
assert.Contains(t, selected, int64(2))
assert.Contains(t, selected, int64(3))
}
func TestGetGrowingRowCountPolicy(t *testing.T) {
tracker := NewCheckpointTracker()
// Create buffers with different unflushed row counts
buffers := []*GrowingSegmentBuffer{
createMockBuffer(1, 1024, 50, tracker), // 50 unflushed
createMockBuffer(2, 1024, 100, tracker), // 100 unflushed
createMockBuffer(3, 1024, 150, tracker), // 150 unflushed
}
policy := GetGrowingRowCountPolicy(100)
assert.Equal(t, "row count threshold", policy.Reason())
selected := policy.SelectSegments(buffers, 0)
assert.Len(t, selected, 2)
assert.Contains(t, selected, int64(2)) // 100 >= 100
assert.Contains(t, selected, int64(3)) // 150 >= 100
}
func TestGetGrowingStaleBufferPolicy(t *testing.T) {
tracker := NewCheckpointTracker()
staleDuration := 10 * time.Second
// Record batches with different timestamps
now := time.Now()
oldTs := tsoutil.ComposeTSByTime(now.Add(-20*time.Second), 0) // 20s ago - stale
recentTs := tsoutil.ComposeTSByTime(now.Add(-5*time.Second), 0) // 5s ago - not stale
currentTs := tsoutil.ComposeTSByTime(now, 0)
tracker.RecordBatch(1, 100, &msgpb.MsgPosition{Timestamp: oldTs})
tracker.RecordBatch(2, 100, &msgpb.MsgPosition{Timestamp: recentTs})
buffers := []*GrowingSegmentBuffer{
createMockBuffer(1, 1024, 100, tracker), // stale
createMockBuffer(2, 1024, 100, tracker), // not stale
createMockBuffer(3, 1024, 100, tracker), // no data (minTs = 0)
}
policy := GetGrowingStaleBufferPolicy(staleDuration)
assert.Equal(t, "buffer stale", policy.Reason())
// Use current time as the timestamp for comparison
selected := policy.SelectSegments(buffers, currentTs)
// Only segment 1 should be stale (20s > 10s + jitter)
// Segment 3 has no unflushed data (minTs = 0)
assert.Contains(t, selected, int64(1))
assert.NotContains(t, selected, int64(3)) // no unflushed data
}
func TestGetGrowingOldestBufferPolicy(t *testing.T) {
tracker := NewCheckpointTracker()
// Record batches with different timestamps
tracker.RecordBatch(1, 100, &msgpb.MsgPosition{Timestamp: 100})
tracker.RecordBatch(2, 100, &msgpb.MsgPosition{Timestamp: 200})
tracker.RecordBatch(3, 100, &msgpb.MsgPosition{Timestamp: 300})
tracker.RecordBatch(4, 100, &msgpb.MsgPosition{Timestamp: 400})
buffers := []*GrowingSegmentBuffer{
createMockBuffer(1, 1024, 100, tracker), // oldest (ts=100)
createMockBuffer(2, 1024, 100, tracker), // second oldest (ts=200)
createMockBuffer(3, 1024, 100, tracker), // third oldest (ts=300)
createMockBuffer(4, 1024, 100, tracker), // newest (ts=400)
}
t.Run("select 2 oldest", func(t *testing.T) {
policy := GetGrowingOldestBufferPolicy(2)
assert.Equal(t, "oldest buffers", policy.Reason())
selected := policy.SelectSegments(buffers, 0)
assert.Len(t, selected, 2)
assert.Contains(t, selected, int64(1))
assert.Contains(t, selected, int64(2))
})
t.Run("select more than available", func(t *testing.T) {
policy := GetGrowingOldestBufferPolicy(10)
selected := policy.SelectSegments(buffers, 0)
assert.Len(t, selected, 4) // returns all with unflushed data
})
t.Run("no unflushed data", func(t *testing.T) {
// Flush all data
tracker.UpdateFlushedOffset(1, 100)
tracker.UpdateFlushedOffset(2, 100)
tracker.UpdateFlushedOffset(3, 100)
tracker.UpdateFlushedOffset(4, 100)
policy := GetGrowingOldestBufferPolicy(2)
selected := policy.SelectSegments(buffers, 0)
assert.Len(t, selected, 0)
})
}
func TestGetGrowingMemoryPressurePolicy(t *testing.T) {
tracker := NewCheckpointTracker()
// Create buffers with different memory sizes
buffers := []*GrowingSegmentBuffer{
createMockBuffer(1, 10*1024*1024, 100, tracker), // 10MB
createMockBuffer(2, 20*1024*1024, 100, tracker), // 20MB
createMockBuffer(3, 30*1024*1024, 100, tracker), // 30MB
createMockBuffer(4, 40*1024*1024, 100, tracker), // 40MB
}
t.Run("below threshold", func(t *testing.T) {
// Total: 100MB, threshold: 200MB -> no pressure
policy := GetGrowingMemoryPressurePolicy(200 * 1024 * 1024)
assert.Equal(t, "memory pressure", policy.Reason())
selected := policy.SelectSegments(buffers, 0)
assert.Len(t, selected, 0)
})
t.Run("above threshold", func(t *testing.T) {
// Total: 100MB, threshold: 50MB -> pressure
// Target free: 100MB - 25MB = 75MB
policy := GetGrowingMemoryPressurePolicy(50 * 1024 * 1024)
selected := policy.SelectSegments(buffers, 0)
assert.True(t, len(selected) > 0)
// Should select largest segments first (4, 3, ...)
assert.Contains(t, selected, int64(4)) // 40MB - largest
})
}
func TestGrowingSegmentHeap(t *testing.T) {
tracker := NewCheckpointTracker()
// Record batches with different timestamps
tracker.RecordBatch(1, 100, &msgpb.MsgPosition{Timestamp: 300}) // newer
tracker.RecordBatch(2, 100, &msgpb.MsgPosition{Timestamp: 100}) // oldest
tracker.RecordBatch(3, 100, &msgpb.MsgPosition{Timestamp: 200}) // middle
buf1 := createMockBuffer(1, 1024, 100, tracker)
buf2 := createMockBuffer(2, 1024, 100, tracker)
buf3 := createMockBuffer(3, 1024, 100, tracker)
h := &growingSegmentHeap{}
heap.Init(h)
// Push all buffers using heap.Push to maintain heap property
heap.Push(h, buf1)
heap.Push(h, buf2)
heap.Push(h, buf3)
assert.Equal(t, 3, h.Len())
// Pop should return newest first (max-heap by timestamp)
popped := heap.Pop(h).(*GrowingSegmentBuffer)
assert.Equal(t, int64(1), popped.SegmentID()) // ts=300, newest
popped = heap.Pop(h).(*GrowingSegmentBuffer)
assert.Equal(t, int64(3), popped.SegmentID()) // ts=200, middle
popped = heap.Pop(h).(*GrowingSegmentBuffer)
assert.Equal(t, int64(2), popped.SegmentID()) // ts=100, oldest
}
func TestGrowingSyncPolicyFunc(t *testing.T) {
// Test that GrowingSyncPolicyFunc correctly implements the interface
customPolicy := GrowingSyncPolicyFunc{
fn: func(buffers []*GrowingSegmentBuffer, ts uint64) []int64 {
return []int64{1, 2, 3}
},
reason: "custom reason",
}
var policy GrowingSyncPolicy = customPolicy
assert.Equal(t, "custom reason", policy.Reason())
assert.Equal(t, []int64{1, 2, 3}, policy.SelectSegments(nil, 0))
}
internal/querynodev2/segments/manager.go
+37
View File
@@ -113,6 +113,12 @@ type SegmentManager interface {
RemoveBy(ctx context.Context, filters ...SegmentFilter) (int, int)
Clear(ctx context.Context)
// ReleaseDetached completes the release of a segment previously taken out of
// the active maps via DetachStreaming. It runs the same bookkeeping as the
// normal release path (release callback, segment teardown, metric Dec and
// clearing the on-releasing set).
ReleaseDetached(ctx context.Context, segment Segment)
// Deprecated: quick fix critical issue: #30857
// TODO: All Segment assigned to querynode should be managed by SegmentManager, including loading or releasing to perform a transaction.
Exist(segmentID typeutil.UniqueID, typ SegmentType) bool
@@ -657,6 +663,37 @@ func (mgr *segmentManager) Remove(ctx context.Context, segmentID typeutil.Unique
return removeGrowing, removeSealed
}
// DetachStreaming removes the given growing segment from the active maps
// WITHOUT releasing it. The segment is staged into the on-releasing set so
// Exist() keeps reporting it while growing-source flush handoff keeps the
// segment alive. The owner MUST call ReleaseDetached once it is done,
// otherwise the on-releasing set entry, the segment gauge and the release
// callback are never reconciled.
func (mgr *segmentManager) DetachStreaming(ctx context.Context, segmentID typeutil.UniqueID) int {
removeGrowing := 0
if mgr.removeSegmentWithType(SegmentTypeGrowing, segmentID) != nil {
removeGrowing = 1
}
log.Ctx(ctx).Info("detached segment from active segment manager",
zap.Int64("segmentID", segmentID),
zap.Int("growingCount", removeGrowing))
return removeGrowing
}
// ReleaseDetached completes the release of a segment previously taken out of the
// active maps via DetachStreaming. DetachStreaming deliberately skips release()
// so an out-of-band owner (growing-source flush handoff) can keep the segment
// alive; that owner MUST call ReleaseDetached when it is done. Otherwise the
// on-releasing set entry, the segment gauge (QueryNodeNumSegments) and the
// release callback are never reconciled, and Exist() would keep returning true
// for the segmentID.
func (mgr *segmentManager) ReleaseDetached(ctx context.Context, segment Segment) {
if segment == nil {
return
}
mgr.release(ctx, segment)
}
func (mgr *segmentManager) removeSegmentWithType(typ SegmentType, segmentID typeutil.UniqueID) Segment {
segment, ok := mgr.globalSegments.RemoveWithType(segmentID, typ)
if !ok {
internal/querynodev2/segments/mock_segment_manager.go
+92
View File
@@ -811,6 +811,98 @@ func (_c *MockSegmentManager_Remove_Call) RunAndReturn(run func(context.Context,
return _c
}
// Detach provides a mock function with given fields: ctx, segmentID, scope
func (_m *MockSegmentManager) Detach(ctx context.Context, segmentID int64, scope querypb.DataScope) (int, int) {
ret := _m.Called(ctx, segmentID, scope)
if len(ret) == 0 {
panic("no return value specified for Detach")
}
var r0 int
var r1 int
if rf, ok := ret.Get(0).(func(context.Context, int64, querypb.DataScope) (int, int)); ok {
return rf(ctx, segmentID, scope)
}
if rf, ok := ret.Get(0).(func(context.Context, int64, querypb.DataScope) int); ok {
r0 = rf(ctx, segmentID, scope)
} else {
r0 = ret.Get(0).(int)
}
if rf, ok := ret.Get(1).(func(context.Context, int64, querypb.DataScope) int); ok {
r1 = rf(ctx, segmentID, scope)
} else {
r1 = ret.Get(1).(int)
}
return r0, r1
}
// MockSegmentManager_Detach_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Detach'
type MockSegmentManager_Detach_Call struct {
*mock.Call
}
// Detach is a helper method to define mock.On call
// - ctx context.Context
// - segmentID int64
// - scope querypb.DataScope
func (_e *MockSegmentManager_Expecter) Detach(ctx interface{}, segmentID interface{}, scope interface{}) *MockSegmentManager_Detach_Call {
return &MockSegmentManager_Detach_Call{Call: _e.mock.On("Detach", ctx, segmentID, scope)}
}
func (_c *MockSegmentManager_Detach_Call) Run(run func(ctx context.Context, segmentID int64, scope querypb.DataScope)) *MockSegmentManager_Detach_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(int64), args[2].(querypb.DataScope))
})
return _c
}
func (_c *MockSegmentManager_Detach_Call) Return(_a0 int, _a1 int) *MockSegmentManager_Detach_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockSegmentManager_Detach_Call) RunAndReturn(run func(context.Context, int64, querypb.DataScope) (int, int)) *MockSegmentManager_Detach_Call {
_c.Call.Return(run)
return _c
}
// ReleaseDetached provides a mock function with given fields: ctx, segment
func (_m *MockSegmentManager) ReleaseDetached(ctx context.Context, segment Segment) {
_m.Called(ctx, segment)
}
// MockSegmentManager_ReleaseDetached_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ReleaseDetached'
type MockSegmentManager_ReleaseDetached_Call struct {
*mock.Call
}
// ReleaseDetached is a helper method to define mock.On call
// - ctx context.Context
// - segment Segment
func (_e *MockSegmentManager_Expecter) ReleaseDetached(ctx interface{}, segment interface{}) *MockSegmentManager_ReleaseDetached_Call {
return &MockSegmentManager_ReleaseDetached_Call{Call: _e.mock.On("ReleaseDetached", ctx, segment)}
}
func (_c *MockSegmentManager_ReleaseDetached_Call) Run(run func(ctx context.Context, segment Segment)) *MockSegmentManager_ReleaseDetached_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(Segment))
})
return _c
}
func (_c *MockSegmentManager_ReleaseDetached_Call) Return() *MockSegmentManager_ReleaseDetached_Call {
_c.Call.Return()
return _c
}
func (_c *MockSegmentManager_ReleaseDetached_Call) RunAndReturn(run func(context.Context, Segment)) *MockSegmentManager_ReleaseDetached_Call {
_c.Run(run)
return _c
}
// RemoveBy provides a mock function with given fields: ctx, filters
func (_m *MockSegmentManager) RemoveBy(ctx context.Context, filters ...SegmentFilter) (int, int) {
_va := make([]interface{}, len(filters))
internal/querynodev2/segments/segment_interface.go
+6
View File
@@ -30,6 +30,12 @@ import (
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
// BinlogSaver is a minimal interface for saving binlog paths to DataCoord.
// This avoids depending on the full broker.Broker interface.
type BinlogSaver interface {
SaveBinlogPaths(ctx context.Context, req *datapb.SaveBinlogPathsRequest) error
}
// ResourceUsage is used to estimate the resource usage of a sealed segment.
type ResourceUsage struct {
MemorySize uint64
internal/querynodev2/segments/segment_loader.go
+8 -5
View File
@@ -443,16 +443,19 @@ func (loader *segmentLoader) prepare(ctx context.Context, segmentType SegmentTyp
// filter out loaded & loading segments
infos := make([]*querypb.SegmentLoadInfo, 0, len(segments))
for _, segment := range segments {
// Not loaded & loading & releasing.
if !loader.manager.Segment.Exist(segment.GetSegmentID(), segmentType) &&
!loader.loadingSegments.Contain(segment.GetSegmentID()) {
// Only active loaded segments should be skipped here. SegmentManager.Exist()
// also reports detached/on-releasing segments, which are no longer active
// and must be allowed to load again.
isLoaded := loader.manager.Segment.GetWithType(segment.GetSegmentID(), segmentType) != nil
isLoading := loader.loadingSegments.Contain(segment.GetSegmentID())
if !isLoaded && !isLoading {
infos = append(infos, segment)
loader.loadingSegments.Insert(segment.GetSegmentID(), newLoadResult())
} else {
log.Info("skip loaded/loading segment",
zap.Int64("segmentID", segment.GetSegmentID()),
zap.Bool("isLoaded", len(loader.manager.Segment.GetBy(WithType(segmentType), WithID(segment.GetSegmentID()))) > 0),
zap.Bool("isLoading", loader.loadingSegments.Contain(segment.GetSegmentID())),
zap.Bool("isLoaded", isLoaded),
zap.Bool("isLoading", isLoading),
)
}
}
internal/querynodev2/segments/segment_loader_test.go
+55 -29
View File
@@ -46,6 +46,7 @@ import (
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/metric"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
type SegmentLoaderSuite struct {
@@ -1054,26 +1055,26 @@ func (suite *SegmentLoaderDetailSuite) TestWaitSegmentLoadDone() {
suite.Run("wait_success", func() {
idx := 0
var infos []*querypb.SegmentLoadInfo
suite.segmentManager.EXPECT().Exist(mock.Anything, mock.Anything).Return(false)
suite.segmentManager.EXPECT().GetWithType(suite.segmentID, SegmentTypeSealed).RunAndReturn(func(segmentID int64, segmentType commonpb.SegmentState) Segment {
defer func() { idx++ }()
if idx == 0 {
go func() {
<-time.After(time.Second)
suite.loader.notifyLoadFinish(infos...)
}()
}
return nil
})
suite.segmentManager.EXPECT().UpdateBy(mock.Anything, mock.Anything, mock.Anything).Return(0)
infos = suite.loader.prepare(context.Background(), SegmentTypeSealed, &querypb.SegmentLoadInfo{
loadInfo := &querypb.SegmentLoadInfo{
SegmentID: suite.segmentID,
PartitionID: suite.partitionID,
CollectionID: suite.collectionID,
NumOfRows: 100,
InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", suite.collectionID),
}
suite.segmentManager.EXPECT().GetWithType(suite.segmentID, SegmentTypeSealed).RunAndReturn(func(segmentID int64, segmentType commonpb.SegmentState) Segment {
defer func() { idx++ }()
if idx == 0 {
go func() {
<-time.After(time.Second)
suite.loader.notifyLoadFinish(loadInfo)
}()
}
return nil
})
suite.segmentManager.EXPECT().UpdateBy(mock.Anything, mock.Anything, mock.Anything).Return(0)
infos := suite.loader.prepare(context.Background(), SegmentTypeSealed, loadInfo)
suite.Len(infos, 1)
err := suite.loader.waitSegmentLoadDone(context.Background(), SegmentTypeSealed, []int64{suite.segmentID}, 0)
suite.NoError(err)
@@ -1083,26 +1084,26 @@ func (suite *SegmentLoaderDetailSuite) TestWaitSegmentLoadDone() {
suite.SetupTest()
var idx int
var infos []*querypb.SegmentLoadInfo
suite.segmentManager.EXPECT().Exist(mock.Anything, mock.Anything).Return(false)
suite.segmentManager.EXPECT().GetWithType(suite.segmentID, SegmentTypeSealed).RunAndReturn(func(segmentID int64, segmentType commonpb.SegmentState) Segment {
defer func() { idx++ }()
if idx == 0 {
go func() {
<-time.After(time.Second)
suite.loader.unregister(infos...)
}()
}
return nil
})
infos = suite.loader.prepare(context.Background(), SegmentTypeSealed, &querypb.SegmentLoadInfo{
loadInfo := &querypb.SegmentLoadInfo{
SegmentID: suite.segmentID,
PartitionID: suite.partitionID,
CollectionID: suite.collectionID,
NumOfRows: 100,
InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", suite.collectionID),
}
suite.segmentManager.EXPECT().GetWithType(suite.segmentID, SegmentTypeSealed).RunAndReturn(func(segmentID int64, segmentType commonpb.SegmentState) Segment {
defer func() { idx++ }()
if idx == 0 {
go func() {
<-time.After(time.Second)
suite.loader.unregister(loadInfo)
}()
}
return nil
})
infos := suite.loader.prepare(context.Background(), SegmentTypeSealed, loadInfo)
suite.Len(infos, 1)
err := suite.loader.waitSegmentLoadDone(context.Background(), SegmentTypeSealed, []int64{suite.segmentID}, 0)
suite.Error(err)
@@ -1111,7 +1112,6 @@ func (suite *SegmentLoaderDetailSuite) TestWaitSegmentLoadDone() {
suite.Run("wait_timeout", func() {
suite.SetupTest()
suite.segmentManager.EXPECT().Exist(mock.Anything, mock.Anything).Return(false)
suite.segmentManager.EXPECT().GetWithType(suite.segmentID, SegmentTypeSealed).RunAndReturn(func(segmentID int64, segmentType commonpb.SegmentState) Segment {
return nil
})
@@ -1132,6 +1132,32 @@ func (suite *SegmentLoaderDetailSuite) TestWaitSegmentLoadDone() {
})
}
func TestSegmentLoaderPrepareLoadsWhenSegmentIsNotActive(t *testing.T) {
segmentID := rand.Int63()
loadInfo := &querypb.SegmentLoadInfo{
SegmentID: segmentID,
PartitionID: rand.Int63(),
CollectionID: rand.Int63(),
NumOfRows: 100,
InsertChannel: "by-dev-rootcoord-dml_0_1v0",
}
segmentManager := NewMockSegmentManager(t)
loader := &segmentLoader{
manager: &Manager{
Segment: segmentManager,
},
loadingSegments: typeutil.NewConcurrentMap[int64, *loadResult](),
}
segmentManager.EXPECT().GetWithType(segmentID, SegmentTypeGrowing).Return(nil).Once()
infos := loader.prepare(context.Background(), SegmentTypeGrowing, loadInfo)
assert.Len(t, infos, 1)
assert.Equal(t, segmentID, infos[0].GetSegmentID())
assert.True(t, loader.loadingSegments.Contain(segmentID))
}
func TestConfigureUseTakeForOutput(t *testing.T) {
paramtable.Init()
internalKey := paramtable.Get().QueryNodeCfg.InternalCollectionUseTakeForOutput.Key
internal/querynodev2/server.go
+2 -2
View File
@@ -147,7 +147,7 @@ type QueryNode struct {
metricsRequest *metricsinfo.MetricsRequest
// binlogSaver for TEXT collection growing segment flush
// binlogSaver for growing-source segment flush
binlogSaver segments.BinlogSaver
}
@@ -569,7 +569,7 @@ func (node *QueryNode) SetEtcdClient(client *clientv3.Client) {
node.etcdCli = client
}
// SetBinlogSaver sets the BinlogSaver for TEXT collection growing segment flush.
// SetBinlogSaver sets the BinlogSaver for growing-source segment flush.
func (node *QueryNode) SetBinlogSaver(saver segments.BinlogSaver) {
node.binlogSaver = saver
}
internal/querynodev2/services.go
+93 -4
View File
@@ -23,6 +23,7 @@ import (
"sync"
"time"
"github.com/cockroachdb/errors"
"github.com/samber/lo"
"go.opentelemetry.io/otel/trace"
"go.uber.org/zap"
@@ -32,14 +33,19 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus/internal/distributed/streaming"
"github.com/milvus-io/milvus/internal/flushcommon/syncmgr"
"github.com/milvus-io/milvus/internal/querynodev2/delegator"
"github.com/milvus-io/milvus/internal/querynodev2/segments"
"github.com/milvus-io/milvus/internal/querynodev2/tasks"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/storagev2"
"github.com/milvus-io/milvus/internal/streamingnode/client/handler"
"github.com/milvus-io/milvus/internal/streamingnode/client/handler/registry"
"github.com/milvus-io/milvus/internal/util/analyzer"
"github.com/milvus-io/milvus/internal/util/fileresource"
"github.com/milvus-io/milvus/internal/util/searchutil/scheduler"
streamingstatus "github.com/milvus-io/milvus/internal/util/streamingutil/status"
"github.com/milvus-io/milvus/internal/util/streamrpc"
"github.com/milvus-io/milvus/internal/util/textmatch"
"github.com/milvus-io/milvus/pkg/v3/common"
@@ -61,6 +67,10 @@ import (
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
type segmentDetacher interface {
DetachStreaming(ctx context.Context, segmentID typeutil.UniqueID) int
}
// GetComponentStates returns information about whether the node is healthy
func (node *QueryNode) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest) (*milvuspb.ComponentStates, error) {
stats := &milvuspb.ComponentStates{
@@ -360,13 +370,52 @@ func (node *QueryNode) UnsubDmChannel(ctx context.Context, req *querypb.UnsubDmC
node.unsubscribingChannels.Insert(req.GetChannelName())
defer node.unsubscribingChannels.Remove(req.GetChannelName())
delegator, ok := node.delegators.GetAndRemove(req.GetChannelName())
_, ok := node.delegators.Get(req.GetChannelName())
if ok {
node.pipelineManager.Remove(req.GetChannelName())
growingSegmentIDs := node.localGrowingSegmentIDs(req.GetChannelName(), nil)
prepared, err := node.prepareReleaseManualFlush(ctx, req.GetCollectionID(), req.GetChannelName(), growingSegmentIDs)
prepareSkipped := false
if err != nil {
if isReleaseManualFlushPrepareUnavailable(err) {
log.Warn("release manual flush prepare unavailable before unsubscribing channel, continue unsubscribe",
zap.Int64s("segmentIDs", growingSegmentIDs),
zap.Error(err))
prepared = false
prepareSkipped = true
} else {
log.Warn("failed to prepare release manual flush before unsubscribing channel",
zap.Int64s("segmentIDs", growingSegmentIDs),
zap.Error(err))
return merr.Status(err), nil
}
}
if prepareSkipped {
log.Info("release manual flush prepare skipped before unsubscribing channel",
zap.Int64s("segmentIDs", growingSegmentIDs),
zap.Bool("prepared", prepared))
} else {
log.Info("release manual flush prepare result before unsubscribing channel",
zap.Int64s("segmentIDs", growingSegmentIDs),
zap.Bool("prepared", prepared))
}
delegator, ok := node.delegators.GetAndRemove(req.GetChannelName())
if !ok {
log.Info("channel already unsubscribed")
return merr.Success(), nil
}
node.pipelineManager.Remove(req.GetChannelName())
preparedGrowingSourceSegments := syncmgr.DefaultGrowingSourceRegistry().ReleasePreparedSegments(req.GetChannelName())
// close the delegator first to block all coming query/search requests
delegator.Close()
if detacher, ok := node.manager.Segment.(segmentDetacher); ok {
for _, segmentID := range preparedGrowingSourceSegments {
detacher.DetachStreaming(ctx, segmentID)
syncmgr.DefaultGrowingSourceRegistry().MarkReleaseDetached(req.GetChannelName(), segmentID)
syncmgr.DefaultGrowingSourceRegistry().ClearReleasePrepared(req.GetChannelName(), segmentID)
}
}
node.manager.Segment.RemoveBy(ctx, segments.WithChannel(req.GetChannelName()), segments.WithType(segments.SegmentTypeGrowing))
node.manager.Collection.Unref(req.GetCollectionID(), 1)
}
@@ -598,7 +647,7 @@ func (node *QueryNode) ReleaseSegments(ctx context.Context, req *querypb.Release
defer node.lifetime.Done()
if req.GetNeedTransfer() {
delegator, ok := node.delegators.Get(req.GetShard())
shardDelegator, ok := node.delegators.Get(req.GetShard())
if !ok {
msg := "failed to release segment, delegator not found"
log.Warn(msg)
@@ -607,7 +656,7 @@ func (node *QueryNode) ReleaseSegments(ctx context.Context, req *querypb.Release
}
req.NeedTransfer = false
err := delegator.ReleaseSegments(ctx, req, false)
err := shardDelegator.ReleaseSegments(ctx, req, false)
if err != nil {
log.Warn("delegator failed to release segment", zap.Error(err))
return merr.Status(err), nil
@@ -627,6 +676,46 @@ func (node *QueryNode) ReleaseSegments(ctx context.Context, req *querypb.Release
return merr.Success(), nil
}
func (node *QueryNode) localGrowingSegmentIDs(channel string, segmentIDs []int64) []int64 {
filters := []segments.SegmentFilter{
segments.WithChannel(channel),
segments.WithType(segments.SegmentTypeGrowing),
}
if len(segmentIDs) > 0 {
filters = append(filters, segments.WithIDs(segmentIDs...))
}
return lo.Map(node.manager.Segment.GetBy(filters...), func(segment segments.Segment, _ int) int64 {
return segment.ID()
})
}
func (node *QueryNode) prepareReleaseManualFlush(ctx context.Context, collectionID int64, channel string, segmentIDs []int64) (bool, error) {
segmentIDs = lo.Uniq(lo.Filter(segmentIDs, func(segmentID int64, _ int) bool {
return segmentID > 0 && !syncmgr.DefaultGrowingSourceRegistry().IsReleasePrepared(channel, segmentID, 0)
}))
if len(segmentIDs) == 0 {
return false, nil
}
wal := streaming.WAL()
if wal == nil {
return false, merr.WrapErrServiceUnavailable("streaming WAL is not initialized")
}
return wal.PrepareReleaseManualFlush(ctx, collectionID, channel, segmentIDs)
}
func isReleaseManualFlushPrepareUnavailable(err error) bool {
if err == nil {
return false
}
if errors.Is(err, merr.ErrServiceUnavailable) ||
errors.Is(err, handler.ErrClientClosed) ||
errors.Is(err, registry.ErrNoStreamingNodeDeployed) ||
errors.Is(err, registry.ErrNoReleaseManualFlushPreparer) {
return true
}
return streamingstatus.AsStreamingError(err).IsOnShutdown()
}
// GetSegmentInfo returns segment information of the collection on the queryNode, and the information includes memSize, numRow, indexName, indexID ...
func (node *QueryNode) GetSegmentInfo(ctx context.Context, in *querypb.GetSegmentInfoRequest) (*querypb.GetSegmentInfoResponse, error) {
if err := node.lifetime.Add(merr.IsHealthy); err != nil {
internal/querynodev2/services_test.go
+59
View File
@@ -28,6 +28,7 @@ import (
"github.com/cockroachdb/errors"
"github.com/samber/lo"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
clientv3 "go.etcd.io/etcd/client/v3"
@@ -45,7 +46,10 @@ import (
"github.com/milvus-io/milvus/internal/querynodev2/delegator"
"github.com/milvus-io/milvus/internal/querynodev2/segments"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/streamingnode/client/handler"
"github.com/milvus-io/milvus/internal/streamingnode/client/handler/registry"
"github.com/milvus-io/milvus/internal/util/dependency"
streamingstatus "github.com/milvus-io/milvus/internal/util/streamingutil/status"
"github.com/milvus-io/milvus/internal/util/streamingutil/util"
"github.com/milvus-io/milvus/internal/util/streamrpc"
"github.com/milvus-io/milvus/pkg/v3/common"
@@ -562,6 +566,61 @@ func (suite *ServiceSuite) TestUnsubDmChannels_Failed() {
suite.Equal(commonpb.ErrorCode_NotReadyServe, status.GetErrorCode())
}
func TestIsReleaseManualFlushPrepareUnavailable(t *testing.T) {
tests := []struct {
name string
err error
want bool
}{
{
name: "nil",
err: nil,
want: false,
},
{
name: "service unavailable",
err: merr.WrapErrServiceUnavailable("streaming WAL is not initialized"),
want: true,
},
{
name: "handler client closed",
err: handler.ErrClientClosed,
want: true,
},
{
name: "no streaming node deployed",
err: registry.ErrNoStreamingNodeDeployed,
want: true,
},
{
name: "no release manual flush preparer",
err: registry.ErrNoReleaseManualFlushPreparer,
want: true,
},
{
name: "streaming on shutdown",
err: streamingstatus.NewOnShutdownError("wal is on shutdown"),
want: true,
},
{
name: "generic prepare failure",
err: errors.New("prepare failed"),
want: false,
},
{
name: "streaming inner failure",
err: streamingstatus.NewInner("prepare failed"),
want: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
assert.Equal(t, tt.want, isReleaseManualFlushPrepareUnavailable(tt.err))
})
}
}
func (suite *ServiceSuite) genSegmentLoadInfos(schema *schemapb.CollectionSchema,
indexInfos []*indexpb.IndexInfo,
) []*querypb.SegmentLoadInfo {
internal/rootcoord/create_collection_task.go
+3
View File
@@ -198,6 +198,9 @@ func (t *createCollectionTask) validateSchema(ctx context.Context, schema *schem
if err := typeutil.NormalizeAndValidateExternalCollectionSchema(schema); err != nil {
return err
}
if err := typeutil.ValidateTextRequiresStorageV3(schema, Params.CommonCfg.UseLoonFFI.GetAsBool()); err != nil {
return merr.WrapErrParameterInvalidMsg("%s", err.Error())
}
// For external collections, validate the source URL scheme allowlist and
// the JSON spec structure (extfs allowlist + format whitelist) at the
internal/rootcoord/ddl_callbacks_alter_collection_add_field.go
+3
View File
@@ -79,6 +79,9 @@ func (c *Core) broadcastAlterCollectionForAddField(ctx context.Context, req *mil
if err := typeutil.ValidateExternalCollectionResolvedSchema(schema); err != nil {
return err
}
if err := typeutil.ValidateTextRequiresStorageV3(schema, Params.CommonCfg.UseLoonFFI.GetAsBool()); err != nil {
return merr.WrapErrParameterInvalidMsg("%s", err.Error())
}
cacheExpirations, err := c.getCacheExpireForCollection(ctx, req.GetDbName(), req.GetCollectionName())
if err != nil {
internal/rootcoord/ddl_callbacks_alter_collection_schema.go
+3
View File
@@ -185,6 +185,9 @@ func (c *Core) broadcastAlterCollectionSchemaAdd(ctx context.Context, broadcaste
if err := typeutil.ValidateExternalCollectionResolvedSchema(schema); err != nil {
return err
}
if err := typeutil.ValidateTextRequiresStorageV3(schema, Params.CommonCfg.UseLoonFFI.GetAsBool()); err != nil {
return merr.WrapErrParameterInvalidMsg("%s", err.Error())
}
// Broadcast.
cacheExpirations, err := c.getCacheExpireForCollection(ctx, req.GetDbName(), req.GetCollectionName())
internal/storage/record_writer.go
+1
View File
@@ -539,6 +539,7 @@ func NewPackedTextBatchWriter(
config := &packed.SegmentWriterConfig{
SegmentPath: basePath,
TextColumns: textColumnConfigs,
ColumnGroups: columnGroups,
WriterFormat: writerFormat,
SchemaBasedPattern: schemaBasedPattern,
SchemaBasedFormats: schemaBasedFormats,
internal/storagev2/packed/ffi_common.go
+1 -1
View File
@@ -61,7 +61,7 @@ var (
PropertyWriterPolicy = C.GoString(C.loon_properties_writer_policy)
PropertyWriterFormat = "writer.format"
PropertyWriterSchemaBasedPattern = C.GoString(C.loon_properties_writer_schema_base_patterns)
PropertyWriterSchemaBasedFormats = C.GoString(C.loon_properties_writer_schema_base_formats)
PropertyWriterSchemaBasedFormats = "writer.split.schema_based.formats"
// CMEK (Customer Managed Encryption Keys) writer properties
PropertyWriterEncEnable = C.GoString(C.loon_properties_writer_enc_enable) // Enable encryption for written data
internal/storagev2/packed/segment_writer_ffi.go
+41 -12
View File
@@ -31,7 +31,9 @@ import (
"github.com/apache/arrow/go/v17/arrow"
"github.com/apache/arrow/go/v17/arrow/cdata"
"github.com/samber/lo"
"github.com/milvus-io/milvus/internal/storagecommon"
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
)
@@ -51,6 +53,7 @@ type TextColumnConfig struct {
type SegmentWriterConfig struct {
SegmentPath string
TextColumns []TextColumnConfig
ColumnGroups []storagecommon.ColumnGroup
WriterFormat string
SchemaBasedPattern string
SchemaBasedFormats []string
@@ -80,18 +83,7 @@ func NewFFISegmentWriter(
return nil, fmt.Errorf("storageConfig must not be nil")
}
// create properties
extra := map[string]string{}
if config != nil && config.WriterFormat != "" {
extra[PropertyWriterFormat] = config.WriterFormat
}
if config != nil && config.SchemaBasedPattern != "" {
extra[PropertyWriterPolicy] = "schema_based"
extra[PropertyWriterSchemaBasedPattern] = config.SchemaBasedPattern
}
if config != nil && len(config.SchemaBasedFormats) > 0 {
extra[PropertyWriterSchemaBasedFormats] = strings.Join(config.SchemaBasedFormats, ",")
}
extra := segmentWriterProperties(schema, config)
cProperties, err := MakePropertiesFromStorageConfig(storageConfig, extra)
if err != nil {
return nil, err
@@ -116,6 +108,43 @@ func NewFFISegmentWriter(
}, nil
}
func segmentWriterProperties(schema *arrow.Schema, config *SegmentWriterConfig) map[string]string {
extra := map[string]string{
PropertyWriterFormat: "parquet",
}
if config == nil {
extra[PropertyWriterPolicy] = "single"
return extra
}
if config.WriterFormat != "" {
extra[PropertyWriterFormat] = config.WriterFormat
}
if len(config.SchemaBasedFormats) > 0 {
extra[PropertyWriterSchemaBasedFormats] = strings.Join(config.SchemaBasedFormats, ",")
}
if config.SchemaBasedPattern != "" {
extra[PropertyWriterPolicy] = "schema_based"
extra[PropertyWriterSchemaBasedPattern] = config.SchemaBasedPattern
return extra
}
columnGroups := config.ColumnGroups
if len(columnGroups) == 0 {
extra[PropertyWriterPolicy] = "single"
return extra
}
pattern := strings.Join(lo.Map(columnGroups, func(columnGroup storagecommon.ColumnGroup, _ int) string {
return strings.Join(lo.Map(columnGroup.Columns, func(index int, _ int) string {
return schema.Field(index).Name
}), "|")
}), ",")
extra[PropertyWriterPolicy] = "schema_based"
extra[PropertyWriterSchemaBasedPattern] = pattern
return extra
}
// Write writes a record batch to the segment writer.
func (w *FFISegmentWriter) Write(record arrow.Record) error {
var caa cdata.CArrowArray
internal/streamingnode/client/handler/handler_client.go
+5
View File
@@ -85,6 +85,11 @@ type HandlerClient interface {
// Returns an empty slice if no force promote has occurred.
GetSalvageCheckpoint(ctx context.Context, channelName string) ([]*wal.ReplicateCheckpoint, error)
// PrepareReleaseManualFlush prepares process-local release handoff.
// Returns false when the current process is not the local flush owner or
// the channel does not need growing-source retention.
PrepareReleaseManualFlush(ctx context.Context, collectionID int64, vchannel string, releaseSegmentIDs []int64) (bool, error)
// GetWALMetricsIfLocal gets the metrics of the local wal.
// It will only return the metrics of the local wal but not the remote wal.
GetWALMetricsIfLocal(ctx context.Context) (*types.StreamingNodeMetrics, error)
internal/streamingnode/client/handler/handler_client_impl.go
+48
View File
@@ -150,6 +150,54 @@ func (hc *handlerClientImpl) GetSalvageCheckpoint(ctx context.Context, pchannel
return cps.([]*wal.ReplicateCheckpoint), nil
}
// PrepareReleaseManualFlush appends a normal ManualFlush and prepares local growing-source retention.
func (hc *handlerClientImpl) PrepareReleaseManualFlush(ctx context.Context, collectionID int64, vchannel string, releaseSegmentIDs []int64) (bool, error) {
if !hc.lifetime.Add(typeutil.LifetimeStateWorking) {
return false, ErrClientClosed
}
defer hc.lifetime.Done()
pchannel := funcutil.ToPhysicalChannel(vchannel)
logger := log.With(
zap.String("pchannel", pchannel),
zap.String("vchannel", vchannel),
zap.Int64("collectionID", collectionID),
zap.Int64s("releaseSegmentIDs", releaseSegmentIDs),
zap.String("handler", "prepare release manual flush"),
)
result, err := hc.createHandlerAfterStreamingNodeReady(ctx, logger, pchannel, func(ctx context.Context, assign *types.PChannelInfoAssigned) (any, error) {
if assign.Channel.AccessMode != types.AccessModeRW {
logger.Info("skip release manual flush prepare because channel is not RW",
zap.String("accessMode", assign.Channel.AccessMode.String()))
return false, nil
}
_, err := registry.GetLocalAvailableWAL(assign.Channel)
if err != nil {
if errors.Is(err, registry.ErrNoStreamingNodeDeployed) ||
status.AsStreamingError(err).IsWrongStreamingNode() {
logger.Info("skip release manual flush prepare because channel is not owned by local streaming node",
zap.Error(err))
return false, nil
}
return false, err
}
preparer, err := registry.GetLocalReleaseManualFlushPreparer()
if err != nil {
return false, err
}
return preparer.PrepareReleaseManualFlush(ctx, assign.Channel, collectionID, vchannel, releaseSegmentIDs)
})
if err != nil {
return false, err
}
if result == nil {
return false, nil
}
return result.(bool), nil
}
// GetWALMetricsIfLocal gets the metrics of the local wal.
func (hc *handlerClientImpl) GetWALMetricsIfLocal(ctx context.Context) (*types.StreamingNodeMetrics, error) {
if !hc.lifetime.Add(typeutil.LifetimeStateWorking) {
internal/streamingnode/client/handler/handler_client_test.go
+34
View File
@@ -194,6 +194,40 @@ func TestHandlerClient_GetSalvageCheckpoint(t *testing.T) {
assert.Nil(t, cps)
}
func TestHandlerClient_PrepareReleaseManualFlush(t *testing.T) {
assignment := &types.PChannelInfoAssigned{
Channel: types.PChannelInfo{Name: "pchannel", Term: 1, AccessMode: types.AccessModeRO},
Node: types.StreamingNodeInfo{ServerID: 1, Address: "localhost"},
}
vchannel := "pchannel_100v0"
releaseSegmentIDs := []int64{1001}
service := mock_lazygrpc.NewMockService[streamingpb.StreamingNodeHandlerServiceClient](t)
service.EXPECT().Close().Return()
rb := mock_resolver.NewMockBuilder(t)
rb.EXPECT().Close().Run(func() {})
w := mock_assignment.NewMockWatcher(t)
w.EXPECT().Get(mock.Anything, "pchannel").Return(assignment)
w.EXPECT().Close().Run(func() {})
handler := &handlerClientImpl{
lifetime: typeutil.NewLifetime(),
service: service,
rb: rb,
watcher: w,
}
prepared, err := handler.PrepareReleaseManualFlush(context.Background(), 100, vchannel, releaseSegmentIDs)
assert.NoError(t, err)
assert.False(t, prepared)
handler.Close()
prepared, err = handler.PrepareReleaseManualFlush(context.Background(), 100, vchannel, releaseSegmentIDs)
assert.ErrorIs(t, err, ErrClientClosed)
assert.False(t, prepared)
}
func TestDial(t *testing.T) {
paramtable.Init()
internal/streamingnode/client/handler/registry/release_manual_flush_preparer.go
+43
View File
@@ -0,0 +1,43 @@
package registry
import (
"context"
"github.com/cockroachdb/errors"
"github.com/milvus-io/milvus/pkg/v3/log"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/syncutil"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
var (
releaseManualFlushPreparerRegistry = syncutil.NewFuture[ReleaseManualFlushPreparer]()
ErrNoReleaseManualFlushPreparer = errors.New("no release manual flush preparer")
)
// ReleaseManualFlushPreparer prepares process-local release handoff.
type ReleaseManualFlushPreparer interface {
PrepareReleaseManualFlush(ctx context.Context, pchannel types.PChannelInfo, collectionID int64, vchannel string, releaseSegmentIDs []int64) (prepared bool, err error)
}
// RegisterLocalReleaseManualFlushPreparer registers the process-local release handoff preparer.
func RegisterLocalReleaseManualFlushPreparer(preparer ReleaseManualFlushPreparer) {
if !paramtable.IsLocalComponentEnabled(typeutil.StreamingNodeRole) {
panic("unreachable: streaming node is not enabled but release manual flush preparer setup")
}
releaseManualFlushPreparerRegistry.Set(preparer)
log.Ctx(context.Background()).Info("register local release manual flush preparer done")
}
// GetLocalReleaseManualFlushPreparer returns the process-local release handoff preparer.
func GetLocalReleaseManualFlushPreparer() (ReleaseManualFlushPreparer, error) {
if !paramtable.IsLocalComponentEnabled(typeutil.StreamingNodeRole) {
return nil, ErrNoStreamingNodeDeployed
}
if !releaseManualFlushPreparerRegistry.Ready() {
return nil, ErrNoReleaseManualFlushPreparer
}
return releaseManualFlushPreparerRegistry.Get(), nil
}
internal/streamingnode/client/handler/registry/test_utility.go
+1
View File
@@ -7,4 +7,5 @@ import "github.com/milvus-io/milvus/pkg/v3/util/syncutil"
func ResetRegisterLocalWALManager() {
registry = syncutil.NewFuture[WALManager]()
releaseManualFlushPreparerRegistry = syncutil.NewFuture[ReleaseManualFlushPreparer]()
}
internal/streamingnode/server/flusher/flusherimpl/msg_handler_impl.go
-8
View File
@@ -50,14 +50,6 @@ func (impl *msgHandlerImpl) HandleCreateSegment(ctx context.Context, createSegme
}
logger := log.With(log.FieldMessage(createSegmentMsg))
// For TEXT collections, skip creating local metacache entry.
// Insert data is flushed by QueryNode's GrowingFlushManager, not StreamNode.
// This avoids useless empty SyncTasks and manifest path conflicts.
if impl.wbMgr.HasTextFields(vchannel) {
logger.Info("skip CreateNewGrowingSegment for TEXT collection, managed by QueryNode")
return nil
}
if err := impl.wbMgr.CreateNewGrowingSegment(ctx, vchannel, h.PartitionId, h.SegmentId, h.SchemaVersion); err != nil {
logger.Warn("fail to create new growing segment")
return err
internal/streamingnode/server/flusher/flusherimpl/msg_handler_impl_test.go
+5 -8
View File
@@ -103,12 +103,12 @@ func TestFlushMsgHandler_HandleManualFlush(t *testing.T) {
assert.NoError(t, err)
}
func TestFlushMsgHandler_HandleCreateSegment_TextSkip(t *testing.T) {
func TestFlushMsgHandler_HandleCreateSegment(t *testing.T) {
vchannel := "ch-0"
// Build a CreateSegment message with L0 level so that createNewGrowingSegment
// returns nil immediately (L0 skips MixCoordClient allocation).
// This lets us test the TEXT skip logic in HandleCreateSegment without
// This lets us test the growing-source skip logic in HandleCreateSegment without
// requiring the full resource.Resource() server context.
msg, err := message.NewCreateSegmentMessageBuilderV2().
WithHeader(&message.CreateSegmentMessageHeader{
@@ -130,20 +130,17 @@ func TestFlushMsgHandler_HandleCreateSegment_TextSkip(t *testing.T) {
)
assert.NoError(t, err)
t.Run("TEXT collection skips CreateNewGrowingSegment", func(t *testing.T) {
t.Run("growing-source collection creates WriteBuffer shell", func(t *testing.T) {
wbMgr := writebuffer.NewMockBufferManager(t)
wbMgr.EXPECT().HasTextFields(vchannel).Return(true)
// CreateNewGrowingSegment should NOT be called — no expectation set.
// If called, testify will fail the test.
wbMgr.EXPECT().CreateNewGrowingSegment(mock.Anything, vchannel, int64(10), int64(1001), mock.Anything).Return(nil)
handler := newMsgHandler(wbMgr)
err := handler.HandleCreateSegment(context.Background(), im)
assert.NoError(t, err)
})
t.Run("non-TEXT collection calls CreateNewGrowingSegment", func(t *testing.T) {
t.Run("non-growing-source collection calls CreateNewGrowingSegment", func(t *testing.T) {
wbMgr := writebuffer.NewMockBufferManager(t)
wbMgr.EXPECT().HasTextFields(vchannel).Return(false)
wbMgr.EXPECT().CreateNewGrowingSegment(mock.Anything, vchannel, int64(10), int64(1001), mock.Anything).Return(nil)
handler := newMsgHandler(wbMgr)
internal/streamingnode/server/server.go
+2
View File
@@ -78,6 +78,8 @@ func (s *Server) initBasicComponent() {
// initService initializes the grpc service.
func (s *Server) initService() {
writeBufferManager := resource.Resource().WriteBufferManager()
registry.RegisterLocalReleaseManualFlushPreparer(service.NewReleaseManualFlushPreparer(s.walManager, writeBufferManager))
s.handlerService = service.NewHandlerService(s.walManager)
s.managerService = service.NewManagerService(s.walManager)
s.registerGRPCService(s.grpcServer)
internal/streamingnode/server/service/handler_test.go
+133
View File
@@ -0,0 +1,133 @@
package service
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"google.golang.org/protobuf/types/known/anypb"
"github.com/milvus-io/milvus/internal/flushcommon/metacache"
"github.com/milvus-io/milvus/internal/flushcommon/writebuffer"
"github.com/milvus-io/milvus/internal/mocks/streamingnode/server/mock_wal"
"github.com/milvus-io/milvus/internal/mocks/streamingnode/server/mock_walmanager"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
)
type releaseManualFlushCheckBufferManager struct {
*writebuffer.MockBufferManager
needManualFlush bool
err error
}
func (m *releaseManualFlushCheckBufferManager) CheckReleaseManualFlushNeed(ctx context.Context, channel string, segmentIDs []int64) (bool, error) {
return m.needManualFlush, m.err
}
func TestReleaseManualFlushPreparer(t *testing.T) {
ctx := context.Background()
releaseSegmentIDs := []int64{1001}
affectedSegmentIDs := []int64{1002}
extra, err := anypb.New(&message.ManualFlushExtraResponse{SegmentIds: affectedSegmentIDs})
assert.NoError(t, err)
wal := mock_wal.NewMockWAL(t)
wal.EXPECT().Append(mock.Anything, mock.MatchedBy(func(msg message.MutableMessage) bool {
flushMsg, err := message.AsMutableManualFlushMessageV2(msg)
if err != nil {
return false
}
return flushMsg.VChannel() == "vchannel" &&
flushMsg.Header().GetCollectionId() == 10
})).Return(&types.AppendResult{
TimeTick: 200,
Extra: extra,
}, nil)
manager := mock_walmanager.NewMockManager(t)
manager.EXPECT().GetAvailableWAL(mock.Anything).Return(wal, nil)
wbManager := writebuffer.NewMockBufferManager(t)
wbManager.EXPECT().UseGrowingSourceFlush("vchannel").Return(true)
wbManager.EXPECT().
GetGrowingFlushProgress(mock.Anything, "vchannel", releaseSegmentIDs, uint64(200)).
Return([]writebuffer.GrowingFlushSegmentProgress{
{
SegmentID: 1001,
TargetOffset: 10,
NeedReleaseHandoff: true,
SourceMode: metacache.FlushSourceGrowing,
},
}, nil)
preparer := NewReleaseManualFlushPreparer(manager, wbManager)
prepared, err := preparer.PrepareReleaseManualFlush(ctx, types.PChannelInfo{Name: "pchannel", Term: 1}, 10, "vchannel", releaseSegmentIDs)
assert.NoError(t, err)
assert.True(t, prepared)
}
func TestReleaseManualFlushPreparerNoGrowingProgress(t *testing.T) {
ctx := context.Background()
releaseSegmentIDs := []int64{1001}
extra, err := anypb.New(&message.ManualFlushExtraResponse{})
assert.NoError(t, err)
wal := mock_wal.NewMockWAL(t)
wal.EXPECT().Append(mock.Anything, mock.Anything).Return(&types.AppendResult{
TimeTick: 200,
Extra: extra,
}, nil)
manager := mock_walmanager.NewMockManager(t)
manager.EXPECT().GetAvailableWAL(mock.Anything).Return(wal, nil)
wbManager := writebuffer.NewMockBufferManager(t)
wbManager.EXPECT().UseGrowingSourceFlush("vchannel").Return(true)
wbManager.EXPECT().
GetGrowingFlushProgress(mock.Anything, "vchannel", releaseSegmentIDs, uint64(200)).
Return([]writebuffer.GrowingFlushSegmentProgress{
{
SegmentID: 1001,
NeedReleaseHandoff: false,
SourceMode: metacache.FlushSourceUnknown,
},
}, nil)
preparer := NewReleaseManualFlushPreparer(manager, wbManager)
prepared, err := preparer.PrepareReleaseManualFlush(ctx, types.PChannelInfo{Name: "pchannel", Term: 1}, 10, "vchannel", releaseSegmentIDs)
assert.NoError(t, err)
assert.False(t, prepared)
}
func TestReleaseManualFlushPreparerSkipNotNeeded(t *testing.T) {
ctx := context.Background()
wbManager := &releaseManualFlushCheckBufferManager{
MockBufferManager: writebuffer.NewMockBufferManager(t),
needManualFlush: false,
}
wbManager.EXPECT().UseGrowingSourceFlush("vchannel").Return(true)
preparer := NewReleaseManualFlushPreparer(mock_walmanager.NewMockManager(t), wbManager)
prepared, err := preparer.PrepareReleaseManualFlush(ctx, types.PChannelInfo{Name: "pchannel", Term: 1}, 10, "vchannel", []int64{1001})
assert.NoError(t, err)
assert.False(t, prepared)
}
func TestReleaseManualFlushPreparerSkipNonGrowingSource(t *testing.T) {
ctx := context.Background()
wbManager := writebuffer.NewMockBufferManager(t)
wbManager.EXPECT().UseGrowingSourceFlush("vchannel").Return(false)
preparer := NewReleaseManualFlushPreparer(mock_walmanager.NewMockManager(t), wbManager)
prepared, err := preparer.PrepareReleaseManualFlush(ctx, types.PChannelInfo{Name: "pchannel", Term: 1}, 10, "vchannel", []int64{1001})
assert.NoError(t, err)
assert.False(t, prepared)
}
internal/streamingnode/server/service/release_manual_flush_preparer.go
+107
View File
@@ -0,0 +1,107 @@
package service
import (
"context"
"go.uber.org/zap"
"github.com/milvus-io/milvus/internal/flushcommon/writebuffer"
"github.com/milvus-io/milvus/internal/streamingnode/server/walmanager"
"github.com/milvus-io/milvus/internal/util/streamingutil/status"
"github.com/milvus-io/milvus/pkg/v3/log"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
)
// NewReleaseManualFlushPreparer creates a process-local release manual flush preparer.
func NewReleaseManualFlushPreparer(walManager walmanager.Manager, writeBufferManager writebuffer.BufferManager) *releaseManualFlushPreparer {
return &releaseManualFlushPreparer{
walManager: walManager,
writeBufferManager: writeBufferManager,
}
}
type releaseManualFlushPreparer struct {
walManager walmanager.Manager
writeBufferManager writebuffer.BufferManager
}
// PrepareReleaseManualFlush appends a normal ManualFlush and retains the requested local growing segments.
func (p *releaseManualFlushPreparer) PrepareReleaseManualFlush(ctx context.Context, pchannel types.PChannelInfo, collectionID int64, vchannel string, releaseSegmentIDs []int64) (bool, error) {
if p.writeBufferManager == nil {
return false, status.NewInner("write buffer manager is not initialized")
}
if vchannel == "" {
return false, status.NewInvaildArgument("vchannel is empty")
}
if collectionID == 0 {
return false, status.NewInvaildArgument("collection id is empty")
}
if len(releaseSegmentIDs) == 0 {
return false, nil
}
if !p.writeBufferManager.UseGrowingSourceFlush(vchannel) {
log.Ctx(ctx).Info("skip release manual flush prepare because channel does not use growing-source flush",
zap.String("vchannel", vchannel),
zap.Int64("collectionID", collectionID),
zap.Int64s("releaseSegmentIDs", releaseSegmentIDs))
return false, nil
}
if checker, ok := p.writeBufferManager.(writebuffer.ReleaseManualFlushNeedChecker); ok {
needManualFlush, err := checker.CheckReleaseManualFlushNeed(ctx, vchannel, releaseSegmentIDs)
if err != nil {
return false, err
}
if !needManualFlush {
log.Ctx(ctx).Info("skip release manual flush prepare because target segments do not need release handoff",
zap.String("vchannel", vchannel),
zap.Int64("collectionID", collectionID),
zap.Int64s("releaseSegmentIDs", releaseSegmentIDs))
return false, nil
}
}
wal, err := p.walManager.GetAvailableWAL(pchannel)
if err != nil {
return false, err
}
flushMsg, err := message.NewManualFlushMessageBuilderV2().
WithVChannel(vchannel).
WithHeader(&message.ManualFlushMessageHeader{
CollectionId: collectionID,
}).
WithBody(&message.ManualFlushMessageBody{}).
BuildMutable()
if err != nil {
return false, err
}
appendResult, err := wal.Append(ctx, flushMsg)
if err != nil {
return false, err
}
var flushMsgResponse message.ManualFlushExtraResponse
if err := appendResult.GetExtra(&flushMsgResponse); err != nil {
return false, err
}
progress, err := p.writeBufferManager.GetGrowingFlushProgress(ctx, vchannel, releaseSegmentIDs, appendResult.TimeTick)
if err != nil {
return false, err
}
prepared := false
for _, segmentProgress := range progress {
if segmentProgress.NeedReleaseHandoff {
prepared = true
break
}
}
log.Ctx(ctx).Info("prepared release manual flush",
zap.String("vchannel", vchannel),
zap.Int64("collectionID", collectionID),
zap.Uint64("flushTs", appendResult.TimeTick),
zap.Int64s("releaseSegmentIDs", releaseSegmentIDs),
zap.Int64s("affectedSegmentIDs", flushMsgResponse.GetSegmentIds()),
zap.Bool("retained", prepared),
zap.Any("progress", progress))
return prepared, nil
}
internal/streamingnode/server/wal/interceptors/shard/shards/partition_manager.go
+1 -1
View File
@@ -212,6 +212,6 @@ func (m *partitionManager) assignSegment(req *AssignSegmentRequest) (*AssignSegm
// There is no segment can be allocated for the insert request.
// Ask a new pending segment to insert.
m.asyncAllocSegment(req.SchemaVersion)
m.asyncAllocSegment(req.SchemaVersion, req.UseGrowingSourceFlush)
return nil, ErrWaitForNewSegment
}
internal/streamingnode/server/wal/interceptors/shard/shards/segment_alloc_worker.go
+13 -11
View File
@@ -18,7 +18,7 @@ import (
)
// asyncAllocSegment allocates a new growing segment asynchronously.
func (m *partitionManager) asyncAllocSegment(schemaVersion int32) {
func (m *partitionManager) asyncAllocSegment(schemaVersion int32, useGrowingSourceFlush bool) {
if m.onAllocating != nil {
m.Logger().Debug("segment alloc worker is already on allocating")
// manager is already on allocating.
@@ -27,12 +27,13 @@ func (m *partitionManager) asyncAllocSegment(schemaVersion int32) {
// Create a notifier to notify the waiter when the allocation is done.
m.onAllocating = make(chan struct{})
w := &segmentAllocWorker{
ctx: m.ctx,
collectionID: m.collectionID,
partitionID: m.partitionID,
vchannel: m.vchannel,
wal: m.wal.Get(),
schemaVersion: schemaVersion,
ctx: m.ctx,
collectionID: m.collectionID,
partitionID: m.partitionID,
vchannel: m.vchannel,
wal: m.wal.Get(),
schemaVersion: schemaVersion,
useGrowingSourceFlush: useGrowingSourceFlush,
}
w.SetLogger(m.Logger())
// It should always done asynchronously.
@@ -50,10 +51,11 @@ type segmentAllocWorker struct {
wal wal.WAL
// The following fields are preserved across retries to ensure the same segment
// configuration is used when rebuilding the message after a failed append.
segmentID uint64 // allocated segment ID
storageVersion int64 // storage version determined at first attempt
limitation segmentLimitation // segment limitation determined at first attempt
schemaVersion int32
segmentID uint64 // allocated segment ID
storageVersion int64 // storage version determined at first attempt
limitation segmentLimitation // segment limitation determined at first attempt
schemaVersion int32
useGrowingSourceFlush bool
}
// do is the main loop of the segment allocation worker.
internal/streamingnode/server/wal/interceptors/shard/shards/segment_workers_test.go
+36
View File
@@ -15,6 +15,7 @@ import (
"github.com/milvus-io/milvus/internal/mocks/streamingnode/server/mock_wal"
"github.com/milvus-io/milvus/internal/mocks/streamingnode/server/wal/interceptors/shard/mock_utils"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/streamingnode/server/resource"
"github.com/milvus-io/milvus/internal/streamingnode/server/wal"
"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors/shard/policy"
@@ -295,6 +296,41 @@ func TestSegmentAllocWorker_InitSegmentConfig(t *testing.T) {
assert.Equal(t, firstLimitation, w.limitation)
}
func TestSegmentAllocWorkerStorageVersionFollowsUseLoonFFI(t *testing.T) {
paramtable.Init()
resource.InitForTest(t)
param := paramtable.Get()
defer param.Reset(param.CommonCfg.UseLoonFFI.Key)
for name, tc := range map[string]struct {
useLoonFFI string
useGrowingSourceFlush bool
expected int64
}{
"v2_without_growing_source": {useLoonFFI: "false", useGrowingSourceFlush: false, expected: storage.StorageV2},
"v2_with_growing_source": {useLoonFFI: "false", useGrowingSourceFlush: true, expected: storage.StorageV2},
"v3_without_growing_source": {useLoonFFI: "true", useGrowingSourceFlush: false, expected: storage.StorageV3},
"v3_with_growing_source": {useLoonFFI: "true", useGrowingSourceFlush: true, expected: storage.StorageV3},
} {
t.Run(name, func(t *testing.T) {
param.Save(param.CommonCfg.UseLoonFFI.Key, tc.useLoonFFI)
w := &segmentAllocWorker{
ctx: context.Background(),
collectionID: 1,
partitionID: 2,
vchannel: "v1",
wal: mock_wal.NewMockWAL(t),
useGrowingSourceFlush: tc.useGrowingSourceFlush,
}
w.SetLogger(log.With())
err := w.initSegmentConfig()
assert.NoError(t, err)
assert.Equal(t, tc.expected, w.storageVersion)
})
}
}
// TestSegmentFlushWorker_WaitForTxnManagerRecoverDone tests the txn manager wait behavior
func TestSegmentFlushWorker_WaitForTxnManagerRecoverDone(t *testing.T) {
paramtable.Init()
internal/streamingnode/server/wal/interceptors/shard/shards/shard_manager.go
+39 -12
View File
@@ -18,7 +18,9 @@ import (
"github.com/milvus-io/milvus/pkg/v3/log"
"github.com/milvus-io/milvus/pkg/v3/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/syncutil"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
var (
@@ -120,8 +122,8 @@ func newSegmentAllocManagersFromRecovery(pchannel types.PChannelInfo, recoverInf
// recover the segment infos from the streaming node segment assignment meta storage
partitionToSegmentManagers := make(map[PartitionUniqueKey]map[int64]*segmentAllocManager)
growingBelongs := make(map[int64]stats.SegmentBelongs)
seenSegments := make(map[int64]struct{}, len(recoverInfos.SegmentAssignments))
for _, rawMeta := range recoverInfos.SegmentAssignments {
m := newSegmentAllocManagerFromProto(pchannel, rawMeta)
coll, ok := collections[rawMeta.GetCollectionId()]
if !ok {
panic(fmt.Sprintf("segment assignment meta is dirty, collection not found, %d", rawMeta.GetCollectionId()))
@@ -129,24 +131,33 @@ func newSegmentAllocManagersFromRecovery(pchannel types.PChannelInfo, recoverInf
if _, ok := coll.PartitionIDs[rawMeta.GetPartitionId()]; !ok {
panic(fmt.Sprintf("segment assignment meta is dirty, partition not found, partition not found, %d", rawMeta.GetPartitionId()))
}
if _, ok := growingBelongs[rawMeta.GetSegmentId()]; ok {
if _, ok := seenSegments[rawMeta.GetSegmentId()]; ok {
panic(fmt.Sprintf("segment assignment meta is dirty, segment repeated, %d", rawMeta.GetSegmentId()))
}
growingBelongs[m.GetSegmentID()] = stats.SegmentBelongs{
PChannel: pchannel.Name,
VChannel: m.GetVChannel(),
CollectionID: rawMeta.GetCollectionId(),
PartitionID: rawMeta.GetPartitionId(),
SegmentID: m.GetSegmentID(),
}
seenSegments[rawMeta.GetSegmentId()] = struct{}{}
uniqueKey := PartitionUniqueKey{
CollectionID: rawMeta.GetCollectionId(),
PartitionID: rawMeta.GetPartitionId(),
}
if _, ok := partitionToSegmentManagers[uniqueKey]; !ok {
partitionToSegmentManagers[uniqueKey] = make(map[int64]*segmentAllocManager, 2)
switch rawMeta.GetState() {
case streamingpb.SegmentAssignmentState_SEGMENT_ASSIGNMENT_STATE_GROWING:
m := newSegmentAllocManagerFromProto(pchannel, rawMeta)
growingBelongs[m.GetSegmentID()] = stats.SegmentBelongs{
PChannel: pchannel.Name,
VChannel: m.GetVChannel(),
CollectionID: rawMeta.GetCollectionId(),
PartitionID: rawMeta.GetPartitionId(),
SegmentID: m.GetSegmentID(),
}
if _, ok := partitionToSegmentManagers[uniqueKey]; !ok {
partitionToSegmentManagers[uniqueKey] = make(map[int64]*segmentAllocManager, 2)
}
partitionToSegmentManagers[uniqueKey][rawMeta.GetSegmentId()] = m
case streamingpb.SegmentAssignmentState_SEGMENT_ASSIGNMENT_STATE_FLUSHED:
continue
default:
panic(fmt.Sprintf("segment assignment meta has unknown state, segment %d state %s", rawMeta.GetSegmentId(), rawMeta.GetState()))
}
partitionToSegmentManagers[uniqueKey][rawMeta.GetSegmentId()] = m
}
return partitionToSegmentManagers, growingBelongs
}
@@ -212,6 +223,22 @@ func (c *CollectionInfo) SchemaVersion() int32 {
return s.GetVersion()
}
func (c *CollectionInfo) UseGrowingSourceFlush() bool {
if c == nil || c.Schema == nil {
return false
}
return typeutil.UseGrowingSourceFlush(c.Schema.GetSchema(),
paramtable.Get().CommonCfg.UseLoonFFI.GetAsBool(),
paramtable.Get().CommonCfg.EnableGrowingSourceFlush.GetAsBool())
}
func (c *CollectionInfo) HasTextField() bool {
if c == nil || c.Schema == nil || c.Schema.GetSchema() == nil {
return false
}
return typeutil.HasTextField(c.Schema.GetSchema())
}
func (m *shardManagerImpl) Channel() types.PChannelInfo {
return m.pchannel
}
internal/streamingnode/server/wal/interceptors/shard/shards/shard_manager_segment.go
+8 -6
View File
@@ -13,12 +13,13 @@ import (
// AssignSegmentRequest is a request to allocate segment.
type AssignSegmentRequest struct {
CollectionID int64
PartitionID int64
ModifiedMetrics stats.ModifiedMetrics
TimeTick uint64
TxnSession TxnSession
SchemaVersion int32
CollectionID int64
PartitionID int64
ModifiedMetrics stats.ModifiedMetrics
TimeTick uint64
TxnSession TxnSession
SchemaVersion int32
UseGrowingSourceFlush bool
}
// AssignSegmentResult is a result of segment allocation.
@@ -142,6 +143,7 @@ func (m *shardManagerImpl) AssignSegment(req *AssignSegmentRequest) (*AssignSegm
// single place that resolves and stamps it before forwarding to partitionManager.
if info := m.collections[req.CollectionID]; info != nil {
req.SchemaVersion = info.SchemaVersion()
req.UseGrowingSourceFlush = info.UseGrowingSourceFlush()
}
result, err := pm.AssignSegment(req)
internal/streamingnode/server/wal/interceptors/shard/shards/shard_manager_test.go
+166
View File
@@ -1,6 +1,7 @@
package shards
import (
"context"
"testing"
"time"
@@ -12,6 +13,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/mocks/streamingnode/server/mock_wal"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/streamingnode/server/resource"
"github.com/milvus-io/milvus/internal/streamingnode/server/wal"
"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors/shard/policy"
@@ -292,6 +294,89 @@ func TestShardManager(t *testing.T) {
m.Close()
}
func TestShardManagerAssignSegmentTextUsesV3CreateSegmentWhenStorageV3Enabled(t *testing.T) {
paramtable.Init()
resource.InitForTest(t)
param := paramtable.Get()
param.Save(param.CommonCfg.UseLoonFFI.Key, "true")
defer param.Reset(param.CommonCfg.UseLoonFFI.Key)
channel := types.PChannelInfo{
Name: "test_text_channel",
Term: 1,
}
appendCh := make(chan *message.CreateSegmentMessageHeader, 1)
w := mock_wal.NewMockWAL(t)
w.EXPECT().Available().RunAndReturn(func() <-chan struct{} {
return make(chan struct{})
}).Maybe()
w.EXPECT().Append(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, msg message.MutableMessage) (*types.AppendResult, error) {
createMsg := message.MustAsMutableCreateSegmentMessageV2(msg)
appendCh <- createMsg.Header()
return &types.AppendResult{
MessageID: rmq.NewRmqID(1),
TimeTick: 1000,
}, nil
}).Once()
f := syncutil.NewFuture[wal.WAL]()
f.Set(w)
m := RecoverShardManager(&ShardManagerRecoverParam{
ChannelInfo: channel,
WAL: f,
InitialRecoverSnapshot: &recovery.RecoverySnapshot{
VChannels: map[string]*streamingpb.VChannelMeta{
"v_text": {
Vchannel: "v_text",
State: streamingpb.VChannelState_VCHANNEL_STATE_NORMAL,
CollectionInfo: &streamingpb.CollectionInfoOfVChannel{
CollectionId: 10,
Partitions: []*streamingpb.PartitionInfoOfVChannel{
{PartitionId: 20},
},
Schemas: []*streamingpb.CollectionSchemaOfVChannel{
{
Schema: &schemapb.CollectionSchema{
Name: "text_collection",
Version: 7,
Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_Int64},
{FieldID: 101, DataType: schemapb.DataType_Text},
},
},
},
},
},
},
},
SegmentAssignments: map[int64]*streamingpb.SegmentAssignmentMeta{},
Checkpoint: &recovery.WALCheckpoint{TimeTick: 100},
},
TxnManager: &mockedTxnManager{},
}).(*shardManagerImpl)
defer m.Close()
_, err := m.AssignSegment(&AssignSegmentRequest{
CollectionID: 10,
PartitionID: 20,
TimeTick: 200,
ModifiedMetrics: stats.ModifiedMetrics{
Rows: 1,
BinarySize: 1,
},
})
assert.ErrorIs(t, err, ErrWaitForNewSegment)
select {
case header := <-appendCh:
assert.Equal(t, storage.StorageV3, header.GetStorageVersion())
assert.Equal(t, int32(7), header.GetSchemaVersion())
case <-time.After(time.Second):
t.Fatal("timed out waiting for create segment message")
}
}
func TestShardManagerSchemaVersionCheck(t *testing.T) {
paramtable.Init()
resource.InitForTest(t)
@@ -369,6 +454,40 @@ func TestShardManagerSchemaVersionCheck(t *testing.T) {
assert.ErrorIs(t, err, ErrCollectionSchemaVersionNotMatch)
assert.Equal(t, int32(1), ver)
textSchema := &schemapb.CollectionSchema{
Name: "test_text_schema_collection",
Version: 4,
Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_Int64},
{FieldID: 101, DataType: schemapb.DataType_Text},
},
}
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "true")
defer paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
createMsgTextSchema := message.NewCreateCollectionMessageBuilderV1().
WithVChannel("v_text_schema").
WithHeader(&message.CreateCollectionMessageHeader{
CollectionId: 104,
PartitionIds: []int64{204},
}).
WithBody(&msgpb.CreateCollectionRequest{
CollectionSchema: textSchema,
}).
MustBuildMutable().
WithTimeTick(250).
WithLastConfirmedUseMessageID().
IntoImmutableMessage(rmq.NewRmqID(13))
m.CreateCollection(message.MustAsImmutableCreateCollectionMessageV1(createMsgTextSchema))
ver, err = m.CheckIfCollectionSchemaVersionMatch(&message.InsertMessageHeader{
CollectionId: 104,
SchemaVersion: proto.Int32(4),
})
assert.NoError(t, err)
assert.Equal(t, int32(4), ver)
assert.True(t, m.collections[104].HasTextField())
assert.True(t, m.collections[104].UseGrowingSourceFlush())
// Test 3: Create collection without schema (legacy), then check version
createMsgNoSchema := message.NewCreateCollectionMessageBuilderV1().
WithVChannel("v_no_schema").
@@ -705,3 +824,50 @@ func TestCollectionInfoSchemaVersion(t *testing.T) {
}
assert.Equal(t, int32(3), ci.SchemaVersion())
}
func TestCollectionInfoUseGrowingSourceFlush(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "false")
paramtable.Get().Save(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
defer paramtable.Get().Reset(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key)
assert.False(t, (*CollectionInfo)(nil).UseGrowingSourceFlush())
ci := &CollectionInfo{}
assert.False(t, ci.UseGrowingSourceFlush())
ci.Schema = &streamingpb.CollectionSchemaOfVChannel{}
assert.False(t, ci.UseGrowingSourceFlush())
ci.Schema = &streamingpb.CollectionSchemaOfVChannel{
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_Int64},
},
},
}
paramtable.Get().Save(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key, "true")
assert.False(t, ci.UseGrowingSourceFlush())
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "true")
assert.True(t, ci.UseGrowingSourceFlush())
}
func TestCollectionInfoUseGrowingSourceFlush_TextField(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "false")
paramtable.Get().Save(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
defer paramtable.Get().Reset(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key)
ci := &CollectionInfo{
Schema: &streamingpb.CollectionSchemaOfVChannel{
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_Int64},
{FieldID: 101, DataType: schemapb.DataType_Text},
},
},
},
}
assert.True(t, ci.HasTextField())
assert.False(t, ci.UseGrowingSourceFlush())
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "true")
assert.True(t, ci.UseGrowingSourceFlush())
}
internal/streamingnode/server/wal/recovery/recovery_storage_test.go
+44
View File
@@ -19,6 +19,7 @@ import (
"github.com/milvus-io/milvus/internal/streamingnode/server/resource"
"github.com/milvus-io/milvus/internal/streamingnode/server/wal/utility"
internaltypes "github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/pkg/v3/log"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/proto/messagespb"
"github.com/milvus-io/milvus/pkg/v3/proto/streamingpb"
@@ -180,6 +181,49 @@ func TestRecoveryStorage(t *testing.T) {
}
}
func TestRecoveryStorageManualFlushMarksSegmentsFlushed(t *testing.T) {
const segmentID = int64(1001)
r := &recoveryStorageImpl{
segments: map[int64]*segmentRecoveryInfo{
segmentID: {
meta: &streamingpb.SegmentAssignmentMeta{
CollectionId: 1,
PartitionId: 2,
SegmentId: segmentID,
Vchannel: "v1",
State: streamingpb.SegmentAssignmentState_SEGMENT_ASSIGNMENT_STATE_GROWING,
Stat: &streamingpb.SegmentAssignmentStat{
ModifiedRows: 10,
ModifiedBinarySize: 100,
CreateSegmentTimeTick: 100,
},
},
},
},
}
r.SetLogger(log.With())
msg := message.NewManualFlushMessageBuilderV2().
WithVChannel("v1").
WithHeader(&message.ManualFlushMessageHeader{
CollectionId: 1,
SegmentIds: []int64{segmentID},
}).
WithBody(&message.ManualFlushMessageBody{}).
MustBuildMutable().
WithTimeTick(200).
WithLastConfirmedUseMessageID().
IntoImmutableMessage(rmq.NewRmqID(2))
r.handleManualFlush(message.MustAsImmutableManualFlushMessageV2(msg))
segment := r.segments[segmentID]
snapshot, shouldBeRemoved := segment.ConsumeDirtyAndGetSnapshot()
require.NotNil(t, snapshot)
assert.True(t, shouldBeRemoved)
assert.Equal(t, streamingpb.SegmentAssignmentState_SEGMENT_ASSIGNMENT_STATE_FLUSHED, snapshot.GetState())
assert.EqualValues(t, 200, snapshot.GetCheckpointTimeTick())
}
type streamBuilder struct {
channel types.PChannelInfo
lastConfirmedMessageID int64
internal/streamingnode/server/wal/recovery/recovery_stream.go
+6 -4
View File
@@ -97,24 +97,26 @@ func (r *recoveryStorageImpl) getSnapshot() *RecoverySnapshot {
if !segment.IsGrowing() {
continue
}
// Defensive filtering: skip GROWING segments whose parent vchannel does not exist
// or is not active, or whose partition has been dropped. This can happen due to
// Defensive filtering: skip recoverable segment assignments whose parent vchannel
// does not exist or is not active, or whose partition has been dropped. This can happen due to
// non-atomic etcd persistence or Kafka offset compaction replaying CreateSegment
// for dropped collections/partitions.
if _, ok := vchannels[segment.meta.Vchannel]; !ok {
r.Logger().Warn("getSnapshot: skipping orphaned growing segment with non-active vchannel",
r.Logger().Warn("getSnapshot: skipping orphaned segment assignment with non-active vchannel",
zap.Int64("segmentID", segmentID),
zap.String("vchannel", segment.meta.Vchannel),
zap.Int64("collectionID", segment.meta.CollectionId),
zap.String("state", segment.meta.State.String()),
)
continue
}
if _, ok := activePartitions[segment.meta.PartitionId]; !ok {
r.Logger().Warn("getSnapshot: skipping orphaned growing segment with dropped partition",
r.Logger().Warn("getSnapshot: skipping orphaned segment assignment with dropped partition",
zap.Int64("segmentID", segmentID),
zap.String("vchannel", segment.meta.Vchannel),
zap.Int64("collectionID", segment.meta.CollectionId),
zap.Int64("partitionID", segment.meta.PartitionId),
zap.String("state", segment.meta.State.String()),
)
continue
}
internal/streamingnode/server/wal/recovery/segment_recovery_info.go
+3
View File
@@ -85,6 +85,9 @@ func (info *segmentRecoveryInfo) BinarySize() uint64 {
// ObserveInsert is called when an insert message is observed.
func (info *segmentRecoveryInfo) ObserveInsert(timetick uint64, assignment *messagespb.PartitionSegmentAssignment) {
if !info.IsGrowing() {
return
}
if timetick < info.meta.CheckpointTimeTick {
// the txn message will share the same time tick.
// so we only filter the time tick is less than the checkpoint time tick.
internal/streamingnode/server/wal/recovery/segment_recovery_info_test.go
-4
View File
@@ -77,10 +77,6 @@ func TestSegmentRecoveryInfo(t *testing.T) {
assert.Nil(t, snapshot)
assert.False(t, shouldBeRemoved)
// insert may came from same txn with same txn.
info.ObserveInsert(ts, assign)
assert.True(t, info.dirty)
ts += 1
info.ObserveFlush(ts)
snapshot, shouldBeRemoved = info.ConsumeDirtyAndGetSnapshot()

Some files were not shown because too many files have changed in this diff Show More