admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-06 17:18:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

enhance: add write rate limit for disk file writer (#43912)

Browse Source 
issue: #43040

---------

Signed-off-by: Shawn Wang <shawn.wang@zilliz.com>
This commit is contained in:
sparknack 2025-08-25 10:27:47 +08:00 committed by GitHub
parent d0e3a33c37
commit 4fae074d56
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
24 changed files with 611 additions and 87 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
configs/milvus.yaml
View File

@ -934,6 +934,16 @@ common:
# The default value is 0, which means the caller will perform write operations directly without using an additional writer thread pool.
# In this case, the maximum concurrency of disk write operations is determined by the caller's thread pool size.
diskWriteNumThreads: 0
diskWriteRateLimiter:
refillPeriodUs: 100000 # refill period in microseconds if disk rate limiter is enabled, default is 100000us (100ms)
avgKBps: 262144 # average kilobytes per second if disk rate limiter is enabled, default is 262144KB/s (256MB/s)
maxBurstKBps: 524288 # max burst kilobytes per second if disk rate limiter is enabled, default is 524288KB/s (512MB/s)
# amplification ratio for high priority tasks if disk rate limiter is enabled, value <= 0 means ratio limit is disabled.
# The ratio is the multiplication factor of the configured bandwidth.
# For example, if the rate limit is 100KB/s, and the high priority ratio is 2, then the high priority tasks will be limited to 200KB/s.
highPriorityRatio: -1
middlePriorityRatio: -1 # amplification ratio for middle priority tasks if disk rate limiter is enabled, value <= 0 means ratio limit is disabled
lowPriorityRatio: -1 # amplification ratio for low priority tasks if disk rate limiter is enabled, value <= 0 means ratio limit is disabled
security:
authorizationEnabled: false
# The superusers will ignore some system check processes,

16
internal/core/src/common/type_c.h
View File

@ -93,6 +93,22 @@ typedef struct CStorageConfig {
bool use_custom_part_upload;
} CStorageConfig;
typedef struct CDiskWriteRateLimiterConfig {
int64_t refill_period_us;
int64_t avg_bps;
int64_t max_burst_bps;
int32_t high_priority_ratio;
int32_t middle_priority_ratio;
int32_t low_priority_ratio;
} CDiskWriteRateLimiterConfig;
typedef struct CDiskWriteConfig {
const char* mode;
uint64_t buffer_size_kb;
int nr_threads;
CDiskWriteRateLimiterConfig rate_limiter_config;
} CDiskWriteConfig;
typedef struct CMmapConfig {
const char* cache_read_ahead_policy;
const char* mmap_path;

23
internal/core/src/index/BitmapIndex.cpp
View File

@ -30,6 +30,8 @@
#include "index/Meta.h"
#include "index/ScalarIndex.h"
#include "index/Utils.h"
#include "pb/common.pb.h"
#include "storage/ThreadPools.h"
#include "storage/Util.h"
#include "query/Utils.h"
@ -455,14 +457,16 @@ void
BitmapIndex<T>::MMapIndexData(const std::string& file_name,
const uint8_t* data_ptr,
size_t data_size,
size_t index_length) {
size_t index_length,
milvus::proto::common::LoadPriority priority) {
std::filesystem::create_directories(
std::filesystem::path(file_name).parent_path());
auto file_offset = 0;
std::map<T, std::pair<int32_t, int32_t>> bitmaps;
{
auto file_writer = storage::FileWriter(file_name);
auto file_writer = storage::FileWriter(
file_name, storage::io::GetPriorityFromLoadPriority(priority));
for (size_t i = 0; i < index_length; ++i) {
T key = ParseKey(&data_ptr);
@ -534,12 +538,17 @@ BitmapIndex<T>::LoadWithoutAssemble(const BinarySet& binary_set,
build_mode_ == BitmapIndexBuildMode::ROARING) {
auto mmap_filepath =
GetValueFromConfig<std::string>(config, MMAP_FILE_PATH);
auto priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
AssertInfo(mmap_filepath.has_value(),
"mmap filepath is empty when load index");
MMapIndexData(mmap_filepath.value(),
index_data_buffer->data.get(),
index_data_buffer->size,
index_length);
index_length,
priority);
} else {
DeserializeIndexData(index_data_buffer->data.get(), index_length);
}
@ -569,8 +578,12 @@ BitmapIndex<T>::Load(milvus::tracer::TraceContext ctx, const Config& config) {
GetValueFromConfig<std::vector<std::string>>(config, "index_files");
AssertInfo(index_files.has_value(),
"index file paths is empty when load bitmap index");
auto index_datas = file_manager_->LoadIndexToMemory(
index_files.value(), config[milvus::LOAD_PRIORITY]);
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
auto index_datas =
file_manager_->LoadIndexToMemory(index_files.value(), load_priority);
BinarySet binary_set;
AssembleIndexDatas(index_datas, binary_set);
// clear index_datas to free memory early

4
internal/core/src/index/BitmapIndex.h
View File

@ -23,6 +23,7 @@
#include "common/RegexQuery.h"
#include "index/ScalarIndex.h"
#include "pb/common.pb.h"
#include "storage/FileManager.h"
#include "storage/MemFileManagerImpl.h"
@ -252,7 +253,8 @@ class BitmapIndex : public ScalarIndex<T> {
MMapIndexData(const std::string& filepath,
const uint8_t* data,
size_t data_size,
size_t index_length);
size_t index_length,
milvus::proto::common::LoadPriority priority);
void
UnmapIndexData();

7
internal/core/src/index/HybridScalarIndex.cpp
View File

@ -364,9 +364,12 @@ HybridScalarIndex<T>::Load(milvus::tracer::TraceContext ctx,
auto index_type_file = GetRemoteIndexTypeFile(index_files.value());
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
auto index_datas = mem_file_manager_->LoadIndexToMemory(
std::vector<std::string>{index_type_file},
config[milvus::LOAD_PRIORITY]);
std::vector<std::string>{index_type_file}, load_priority);
BinarySet binary_set;
AssembleIndexDatas(index_datas, binary_set);
// clear index_datas to free memory early

15
internal/core/src/index/InvertedIndexTantivy.cpp
View File

@ -184,8 +184,11 @@ InvertedIndexTantivy<T>::Load(milvus::tracer::TraceContext ctx,
LoadIndexMetas(inverted_index_files, config);
RetainTantivyIndexFiles(inverted_index_files);
disk_file_manager_->CacheIndexToDisk(inverted_index_files,
config[milvus::LOAD_PRIORITY]);
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
disk_file_manager_->CacheIndexToDisk(inverted_index_files, load_priority);
auto prefix = disk_file_manager_->GetLocalIndexObjectPrefix();
path_ = prefix;
auto load_in_mmap =
@ -212,11 +215,15 @@ InvertedIndexTantivy<T>::LoadIndexMetas(
return boost::filesystem::path(file).filename().string() ==
INDEX_NULL_OFFSET_FILE_NAME;
});
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
if (null_offset_file_itr != index_files.end()) {
// null offset file is not sliced
auto index_datas = mem_file_manager_->LoadIndexToMemory(
{*null_offset_file_itr}, config[milvus::LOAD_PRIORITY]);
{*null_offset_file_itr}, load_priority);
auto null_offset_data =
std::move(index_datas.at(INDEX_NULL_OFFSET_FILE_NAME));
fill_null_offsets(null_offset_data->PayloadData(),
@ -233,7 +240,7 @@ InvertedIndexTantivy<T>::LoadIndexMetas(
if (null_offset_files.size() > 0) {
// null offset file is sliced
auto index_datas = mem_file_manager_->LoadIndexToMemory(
null_offset_files, config[milvus::LOAD_PRIORITY]);
null_offset_files, load_priority);
auto null_offsets_data = CompactIndexDatas(index_datas);
auto null_offsets_data_codecs =

8
internal/core/src/index/JsonInvertedIndex.cpp
View File

@ -107,11 +107,15 @@ JsonInvertedIndex<T>::LoadIndexMetas(
return boost::filesystem::path(file).filename().string() ==
INDEX_NON_EXIST_OFFSET_FILE_NAME;
});
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
if (non_exist_offset_file_itr != index_files.end()) {
// null offset file is not sliced
auto index_datas = this->mem_file_manager_->LoadIndexToMemory(
{*non_exist_offset_file_itr}, config[milvus::LOAD_PRIORITY]);
{*non_exist_offset_file_itr}, load_priority);
auto non_exist_offset_data =
std::move(index_datas.at(INDEX_NON_EXIST_OFFSET_FILE_NAME));
fill_non_exist_offset(non_exist_offset_data->PayloadData(),
@ -129,7 +133,7 @@ JsonInvertedIndex<T>::LoadIndexMetas(
if (non_exist_offset_files.size() > 0) {
// null offset file is sliced
auto index_datas = this->mem_file_manager_->LoadIndexToMemory(
non_exist_offset_files, config[milvus::LOAD_PRIORITY]);
non_exist_offset_files, load_priority);
auto non_exist_offset_data = CompactIndexDatas(index_datas);
auto non_exist_offset_data_codecs = std::move(

6
internal/core/src/index/JsonKeyStatsInvertedIndex.cpp
View File

@ -401,8 +401,12 @@ JsonKeyStatsInvertedIndex::Load(milvus::tracer::TraceContext ctx,
index_file = remote_prefix + "/" + index_file;
}
}
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
disk_file_manager_->CacheJsonKeyIndexToDisk(index_files.value(),
config[milvus::LOAD_PRIORITY]);
load_priority);
AssertInfo(
tantivy_index_exist(path_.c_str()), "index not exist: {}", path_);

11
internal/core/src/index/NgramInvertedIndex.cpp
View File

@ -129,6 +129,10 @@ NgramInvertedIndex::Load(milvus::tracer::TraceContext ctx,
AssertInfo(index_files.has_value(),
"index file paths is empty when load ngram index");
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
auto files_value = index_files.value();
auto it = std::find_if(
files_value.begin(), files_value.end(), [](const std::string& file) {
@ -140,8 +144,8 @@ NgramInvertedIndex::Load(milvus::tracer::TraceContext ctx,
std::vector<std::string> file;
file.push_back(*it);
files_value.erase(it);
auto index_datas = mem_file_manager_->LoadIndexToMemory(
file, config[milvus::LOAD_PRIORITY]);
auto index_datas =
mem_file_manager_->LoadIndexToMemory(file, load_priority);
BinarySet binary_set;
AssembleIndexDatas(index_datas, binary_set);
// clear index_datas to free memory early
@ -153,8 +157,7 @@ NgramInvertedIndex::Load(milvus::tracer::TraceContext ctx,
(size_t)index_valid_data->size);
}
disk_file_manager_->CacheNgramIndexToDisk(files_value,
config[milvus::LOAD_PRIORITY]);
disk_file_manager_->CacheNgramIndexToDisk(files_value, load_priority);
AssertInfo(
tantivy_index_exist(path_.c_str()), "index not exist: {}", path_);
auto load_in_mmap =

18
internal/core/src/index/ScalarIndexSort.cpp
View File

@ -30,6 +30,8 @@
#include "common/Types.h"
#include "index/Utils.h"
#include "index/ScalarIndexSort.h"
#include "pb/common.pb.h"
#include "storage/ThreadPools.h"
#include "storage/Util.h"
namespace milvus::index {
@ -213,7 +215,13 @@ ScalarIndexSort<T>::LoadWithoutAssemble(const BinarySet& index_binary,
auto aligned_size =
((index_data->size + ALIGNMENT - 1) / ALIGNMENT) * ALIGNMENT;
{
auto file_writer = storage::FileWriter(mmap_filepath_);
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
auto file_writer = storage::FileWriter(
mmap_filepath_,
storage::io::GetPriorityFromLoadPriority(load_priority));
file_writer.Write(index_data->data.get(), (size_t)index_data->size);
if (aligned_size > index_data->size) {
@ -294,8 +302,12 @@ ScalarIndexSort<T>::Load(milvus::tracer::TraceContext ctx,
GetValueFromConfig<std::vector<std::string>>(config, "index_files");
AssertInfo(index_files.has_value(),
"index file paths is empty when load disk ann index");
auto index_datas = file_manager_->LoadIndexToMemory(
index_files.value(), config[milvus::LOAD_PRIORITY]);
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
auto index_datas =
file_manager_->LoadIndexToMemory(index_files.value(), load_priority);
BinarySet binary_set;
AssembleIndexDatas(index_datas, binary_set);
// clear index_datas to free memory early

16
internal/core/src/index/StringIndexMarisa.cpp
View File

@ -37,6 +37,7 @@
#include "index/Utils.h"
#include "index/Index.h"
#include "marisa/base.h"
#include "storage/ThreadPools.h"
#include "storage/Util.h"
#include "storage/FileWriter.h"
@ -189,9 +190,14 @@ StringIndexMarisa::LoadWithoutAssemble(const BinarySet& set,
auto index = set.GetByName(MARISA_TRIE_INDEX);
auto len = index->size;
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
{
auto file_writer = storage::FileWriter(file_name);
auto file_writer = storage::FileWriter(
file_name, storage::io::GetPriorityFromLoadPriority(load_priority));
file_writer.Write(index->data.get(), len);
file_writer.Finish();
}
@ -225,8 +231,12 @@ StringIndexMarisa::Load(milvus::tracer::TraceContext ctx,
GetValueFromConfig<std::vector<std::string>>(config, "index_files");
AssertInfo(index_files.has_value(),
"index file paths is empty when load index");
auto index_datas = file_manager_->LoadIndexToMemory(
index_files.value(), config[milvus::LOAD_PRIORITY]);
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
auto index_datas =
file_manager_->LoadIndexToMemory(index_files.value(), load_priority);
BinarySet binary_set;
AssembleIndexDatas(index_datas, binary_set);
// clear index_datas to free memory early

11
internal/core/src/index/TextMatchIndex.cpp
View File

@ -144,12 +144,16 @@ TextMatchIndex::Load(const Config& config) {
return file.substr(file.find_last_of('/') + 1) ==
"index_null_offset";
});
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
if (it != files_value.end()) {
std::vector<std::string> file;
file.push_back(*it);
files_value.erase(it);
auto index_datas = mem_file_manager_->LoadIndexToMemory(
file, config[milvus::LOAD_PRIORITY]);
auto index_datas =
mem_file_manager_->LoadIndexToMemory(file, load_priority);
BinarySet binary_set;
AssembleIndexDatas(index_datas, binary_set);
// clear index_datas to free memory early
@ -160,8 +164,7 @@ TextMatchIndex::Load(const Config& config) {
index_valid_data->data.get(),
(size_t)index_valid_data->size);
}
disk_file_manager_->CacheTextLogToDisk(files_value,
config[milvus::LOAD_PRIORITY]);
disk_file_manager_->CacheTextLogToDisk(files_value, load_priority);
AssertInfo(
tantivy_index_exist(prefix.c_str()), "index not exist: {}", prefix);

7
internal/core/src/index/VectorDiskIndex.cpp
View File

@ -99,8 +99,11 @@ VectorDiskAnnIndex<T>::Load(milvus::tracer::TraceContext ctx,
"index file paths is empty when load disk ann index data");
// If index is loaded with stream, we don't need to cache index to disk
if (!index_.LoadIndexWithStream()) {
file_manager_->CacheIndexToDisk(index_files.value(),
config[milvus::LOAD_PRIORITY]);
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
file_manager_->CacheIndexToDisk(index_files.value(), load_priority);
}
read_file_span->End();
}

32
internal/core/src/index/VectorMemIndex.cpp
View File

@ -207,13 +207,18 @@ VectorMemIndex<T>::Load(milvus::tracer::TraceContext ctx,
milvus::tracer::GetTracer()->WithActiveSpan(read_file_span);
LOG_INFO("load with slice meta: {}", !slice_meta_filepath.empty());
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
if (!slice_meta_filepath
.empty()) { // load with the slice meta info, then we can load batch by batch
std::string index_file_prefix = slice_meta_filepath.substr(
0, slice_meta_filepath.find_last_of('/') + 1);
auto result = file_manager_->LoadIndexToMemory(
{slice_meta_filepath}, config[milvus::LOAD_PRIORITY]);
{slice_meta_filepath}, load_priority);
auto raw_slice_meta = std::move(result[INDEX_FILE_SLICE_META]);
Config meta_data = Config::parse(std::string(
reinterpret_cast<const char*>(raw_slice_meta->PayloadData()),
@ -230,8 +235,8 @@ VectorMemIndex<T>::Load(milvus::tracer::TraceContext ctx,
batch.push_back(index_file_prefix + file_name);
}
auto batch_data = file_manager_->LoadIndexToMemory(
batch, config[milvus::LOAD_PRIORITY]);
auto batch_data =
file_manager_->LoadIndexToMemory(batch, load_priority);
int64_t payload_size = 0;
index_data_codecs.insert({prefix, IndexDataCodec{}});
auto& index_data_codec = index_data_codecs.at(prefix);
@ -259,7 +264,7 @@ VectorMemIndex<T>::Load(milvus::tracer::TraceContext ctx,
auto result = file_manager_->LoadIndexToMemory(
std::vector<std::string>(pending_index_files.begin(),
pending_index_files.end()),
config[milvus::LOAD_PRIORITY]);
load_priority);
for (auto&& index_data : result) {
auto prefix = index_data.first;
index_data_codecs.insert({prefix, IndexDataCodec{}});
@ -589,7 +594,14 @@ void VectorMemIndex<T>::LoadFromFile(const Config& config) {
std::filesystem::create_directories(
std::filesystem::path(local_filepath.value()).parent_path());
auto file_writer = storage::FileWriter(local_filepath.value());
auto load_priority =
GetValueFromConfig<milvus::proto::common::LoadPriority>(
config, milvus::LOAD_PRIORITY)
.value_or(milvus::proto::common::LoadPriority::HIGH);
auto file_writer = storage::FileWriter(
local_filepath.value(),
storage::io::GetPriorityFromLoadPriority(load_priority));
auto index_files =
GetValueFromConfig<std::vector<std::string>>(config, "index_files");
@ -626,8 +638,8 @@ void VectorMemIndex<T>::LoadFromFile(const Config& config) {
std::vector<std::string> batch{};
batch.reserve(parallel_degree);
auto result = file_manager_->LoadIndexToMemory(
{slice_meta_filepath}, config[milvus::LOAD_PRIORITY]);
auto result = file_manager_->LoadIndexToMemory({slice_meta_filepath},
load_priority);
auto raw_slice_meta = std::move(result[INDEX_FILE_SLICE_META]);
Config meta_data = Config::parse(std::string(
reinterpret_cast<const char*>(raw_slice_meta->PayloadData()),
@ -639,8 +651,8 @@ void VectorMemIndex<T>::LoadFromFile(const Config& config) {
auto total_len = static_cast<size_t>(item[TOTAL_LEN]);
auto HandleBatch = [&](int index) {
auto start_load2_mem = std::chrono::system_clock::now();
auto batch_data = file_manager_->LoadIndexToMemory(
batch, config[milvus::LOAD_PRIORITY]);
auto batch_data =
file_manager_->LoadIndexToMemory(batch, load_priority);
load_duration_sum +=
(std::chrono::system_clock::now() - start_load2_mem);
for (int j = index - batch.size() + 1; j <= index; j++) {
@ -676,7 +688,7 @@ void VectorMemIndex<T>::LoadFromFile(const Config& config) {
auto result = file_manager_->LoadIndexToMemory(
std::vector<std::string>(pending_index_files.begin(),
pending_index_files.end()),
config[milvus::LOAD_PRIORITY]);
load_priority);
load_duration_sum +=
(std::chrono::system_clock::now() - start_load_files2_mem);
//2. write data into files

4
internal/core/src/storage/DiskFileManagerImpl.cpp
View File

@ -312,7 +312,9 @@ DiskFileManagerImpl::CacheIndexToDiskInternal(
uint64_t(DEFAULT_FIELD_MAX_MEMORY_LIMIT / FILE_SLICE_SIZE.load());
{
auto file_writer = storage::FileWriter(local_index_file_name);
auto file_writer = storage::FileWriter(
local_index_file_name,
storage::io::GetPriorityFromLoadPriority(priority));
auto appendIndexFiles = [&]() {
auto index_chunks_futures =
GetObjectData(rcm_.get(),

43
internal/core/src/storage/FileWriter.cpp
View File

@ -21,7 +21,10 @@
namespace milvus::storage {
FileWriter::FileWriter(std::string filename) : filename_(std::move(filename)) {
FileWriter::FileWriter(std::string filename, io::Priority priority)
: filename_(std::move(filename)),
priority_(priority),
rate_limiter_(io::WriteRateLimiter::GetInstance()) {
auto mode = GetMode();
use_direct_io_ = mode == WriteMode::DIRECT;
auto open_flags = O_CREAT | O_RDWR | O_TRUNC;
@ -115,12 +118,38 @@ void
FileWriter::PositionedWriteWithCheck(const void* data,
size_t nbyte,
size_t file_offset) {
if (!PositionedWrite(data, nbyte, file_offset)) {
Cleanup();
ThrowInfo(ErrorCode::FileWriteFailed,
"Failed to write to file: {}, error: {}",
filename_,
strerror(errno));
size_t bytes_to_write = nbyte;
int32_t empty_loops = 0;
int64_t total_wait_us = 0;
size_t alignment_bytes = use_direct_io_ ? ALIGNMENT_BYTES : 1;
while (bytes_to_write != 0) {
auto allowed_bytes =
rate_limiter_.Acquire(bytes_to_write, alignment_bytes, priority_);
if (allowed_bytes == 0) {
++empty_loops;
// if the empty loops is too large or the total wait time is too long, we should write the data directly
if (empty_loops > MAX_EMPTY_LOOPS || total_wait_us > MAX_WAIT_US) {
allowed_bytes = rate_limiter_.GetBytesPerPeriod();
empty_loops = 0;
total_wait_us = 0;
} else {
int64_t wait_us = (1 << (empty_loops / 10)) *
rate_limiter_.GetRateLimitPeriod();
std::this_thread::sleep_for(std::chrono::microseconds(wait_us));
total_wait_us += wait_us;
continue;
}
}
if (!PositionedWrite(data, allowed_bytes, file_offset)) {
Cleanup();
ThrowInfo(ErrorCode::FileWriteFailed,
"Failed to write to file: {}, error: {}",
filename_,
strerror(errno));
}
file_offset += allowed_bytes;
bytes_to_write -= allowed_bytes;
data = static_cast<const char*>(data) + allowed_bytes;
}
}

200
internal/core/src/storage/FileWriter.h
View File

@ -16,24 +16,200 @@
#pragma once
#include <algorithm>
#include <array>
#include <cassert>
#include <chrono>
#include <cstring>
#include <fcntl.h>
#include <folly/executors/CPUThreadPoolExecutor.h>
#include <mutex>
#include <stdexcept>
#include <string>
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <cstring>
#include <thread>
#include <folly/executors/CPUThreadPoolExecutor.h>
#include <folly/executors/SerialExecutor.h>
#include <unistd.h>
#include "common/EasyAssert.h"
#include "storage/PayloadWriter.h"
#include "log/Log.h"
#include "pb/common.pb.h"
#include "storage/ThreadPools.h"
namespace milvus::storage {
namespace io {
enum class Priority { HIGH = 0, MIDDLE = 1, LOW = 2, NR_PRIORITY = 3 };
inline Priority
GetPriorityFromLoadPriority(milvus::proto::common::LoadPriority priority) {
return priority == milvus::proto::common::LoadPriority::HIGH
? io::Priority::HIGH
: io::Priority::LOW;
}
class WriteRateLimiter {
public:
static WriteRateLimiter&
GetInstance() {
static WriteRateLimiter instance;
return instance;
}
void
Configure(int64_t refill_period_us,
int64_t avg_bps,
int64_t max_burst_bps,
int32_t high_priority_ratio,
int32_t middle_priority_ratio,
int32_t low_priority_ratio) {
if (refill_period_us <= 0 || avg_bps <= 0 || max_burst_bps <= 0 ||
avg_bps > max_burst_bps) {
ThrowInfo(ErrorCode::InvalidParameter,
"All parameters must be positive, but got: "
"refill_period_us: {}, "
"avg_bps: {}, max_burst_bps: {}",
refill_period_us,
avg_bps,
max_burst_bps);
}
std::unique_lock<std::mutex> lock(mutex_);
// avoid too small refill period, 1ms is used as the minimum refill period
refill_period_us_ = std::max<int64_t>(1000, refill_period_us);
refill_bytes_per_period_ = avg_bps * refill_period_us_ / 1000000;
if (refill_bytes_per_period_ <= 0) {
refill_bytes_per_period_ = 1;
}
expire_periods_ = max_burst_bps * refill_period_us_ / 1000000 /
refill_bytes_per_period_;
if (expire_periods_ <= 0) {
expire_periods_ = 1;
}
available_bytes_ = 0;
last_refill_time_ = std::chrono::steady_clock::now();
priority_ratio_ = {
high_priority_ratio, middle_priority_ratio, low_priority_ratio};
LOG_INFO(
"Disk rate limiter configured with refill_period_us: {}, "
"refill_bytes_per_period: {},avg_bps: {}, max_burst_bps: {}, "
"expire_periods: {}, high_priority_ratio: {}, "
"middle_priority_ratio: {}, low_priority_ratio: {}",
refill_period_us_,
refill_bytes_per_period_,
avg_bps,
max_burst_bps,
expire_periods_,
high_priority_ratio,
middle_priority_ratio,
low_priority_ratio);
}
size_t
Acquire(size_t bytes,
size_t alignment_bytes = 1,
Priority priority = Priority::MIDDLE) {
if (static_cast<int>(priority) >=
static_cast<int>(Priority::NR_PRIORITY)) {
ThrowInfo(ErrorCode::InvalidParameter,
"Invalid priority value: {}",
static_cast<int>(priority));
}
// if priority ratio is <= 0, no rate limit is applied, return the original bytes
if (priority_ratio_[static_cast<int>(priority)] <= 0) {
return bytes;
}
AssertInfo(alignment_bytes > 0 && bytes >= alignment_bytes &&
(bytes % alignment_bytes == 0),
"alignment_bytes must be positive and bytes must be "
"divisible by alignment_bytes");
std::unique_lock<std::mutex> lock(mutex_);
// recheck the amplification ratio after taking the lock
auto amplification_ratio = priority_ratio_[static_cast<int>(priority)];
if (amplification_ratio <= 0) {
return bytes;
}
// calculate the available bytes by delta periods
std::chrono::steady_clock::time_point now =
std::chrono::steady_clock::now();
// steady_clock is monotonic, so the time delta is always >= 0
auto delta_periods = static_cast<int>(
std::chrono::duration_cast<std::chrono::microseconds>(
now - last_refill_time_)
.count() /
refill_period_us_);
// early return if the time delta is less than the refill period and
// the available bytes is less than the alignment bytes
if (delta_periods == 0 && available_bytes_ < alignment_bytes) {
return 0;
}
if (delta_periods > expire_periods_) {
available_bytes_ += expire_periods_ * refill_bytes_per_period_;
} else {
available_bytes_ += delta_periods * refill_bytes_per_period_;
}
// keep the available bytes in the range of [0, refill_bytes_per_period_ * expire_periods_]
available_bytes_ = std::min(
available_bytes_,
static_cast<size_t>(refill_bytes_per_period_ * expire_periods_));
// calculate the allowed bytes with amplification ratio
auto ret = std::min(bytes, available_bytes_ * amplification_ratio);
// align the allowed bytes to the alignment bytes
ret = (ret / alignment_bytes) * alignment_bytes;
// update available_bytes_ by removing the amplification ratio, the updated value is always >= 0
available_bytes_ -= ret / amplification_ratio;
// update the last refill time only if delta_periods > 0
if (delta_periods > 0) {
last_refill_time_ = now;
}
return ret;
}
size_t
GetRateLimitPeriod() const {
return refill_period_us_;
}
size_t
GetBytesPerPeriod() const {
return refill_bytes_per_period_;
}
void
Reset() {
std::unique_lock<std::mutex> lock(mutex_);
available_bytes_ = refill_bytes_per_period_;
last_refill_time_ = std::chrono::steady_clock::now();
}
WriteRateLimiter(const WriteRateLimiter&) = delete;
WriteRateLimiter&
operator=(const WriteRateLimiter&) = delete;
~WriteRateLimiter() = default;
private:
WriteRateLimiter() = default;
// Set the default rate limit to a valid, reasonable value.
// These values should always be overridden by the yaml configuration, but
// if not, the default can still serve as a reasonable "no-limit" fallback.
int64_t refill_period_us_ = 100000; // 100ms
int64_t refill_bytes_per_period_ = 1024ll * 1024 * 1024; // 1GB
int32_t expire_periods_ = 10; // 10 periods
std::chrono::steady_clock::time_point last_refill_time_ =
std::chrono::steady_clock::now();
size_t available_bytes_ = 0;
std::array<int32_t, 3> priority_ratio_ = {-1, -1, -1};
std::mutex mutex_;
};
} // namespace io
/**
* FileWriter is a class that sequentially writes data to new files, designed specifically for saving temporary data downloaded from remote storage.
* It supports both buffered and direct I/O, and can use an additional thread pool to write data to files.
@ -42,7 +218,7 @@ namespace milvus::storage {
*
* The basic usage is:
*
* auto file_writer = FileWriter("path/to/file");
* auto file_writer = FileWriter("path/to/file", io::Priority::MIDDLE);
* file_writer.Write(data, size);
* ...
* file_writer.Write(data, size);
@ -57,8 +233,12 @@ class FileWriter {
static constexpr size_t MAX_BUFFER_SIZE = 64 * 1024 * 1024; // 64MB
static constexpr size_t MIN_BUFFER_SIZE = 4 * 1024; // 4KB
static constexpr size_t DEFAULT_BUFFER_SIZE = 64 * 1024; // 64KB
// for rate limiter
static constexpr int MAX_EMPTY_LOOPS = 20;
static constexpr int64_t MAX_WAIT_US = 5000000; // 5s
explicit FileWriter(std::string filename);
explicit FileWriter(std::string filename,
io::Priority priority = io::Priority::MIDDLE);
~FileWriter();
@ -120,6 +300,10 @@ class FileWriter {
mode_; // The write mode, which can be 'buffered' (default) or 'direct'.
static size_t
buffer_size_; // The buffer size used for direct I/O, which is only used when the write mode is 'direct'.
// for rate limiter
io::Priority priority_;
io::WriteRateLimiter& rate_limiter_;
};
class FileWriteWorkerPool {

18
internal/core/src/storage/storage_c.cpp
View File

@ -118,18 +118,16 @@ InitMmapManager(CMmapConfig c_mmap_config) {
}
CStatus
InitFileWriterConfig(const char* mode,
uint64_t buffer_size_kb,
int nr_threads) {
InitDiskFileWriterConfig(CDiskWriteConfig c_disk_write_config) {
try {
std::string mode_str(mode);
std::string mode_str(c_disk_write_config.mode);
if (mode_str == "direct") {
milvus::storage::FileWriter::SetMode(
milvus::storage::FileWriter::WriteMode::DIRECT);
// buffer size checking is done in FileWriter::SetBufferSize,
// and it will try to find a proper and valid buffer size
milvus::storage::FileWriter::SetBufferSize(
buffer_size_kb * 1024); // convert to bytes
c_disk_write_config.buffer_size_kb * 1024); // convert to bytes
} else if (mode_str == "buffered") {
milvus::storage::FileWriter::SetMode(
milvus::storage::FileWriter::WriteMode::BUFFERED);
@ -138,7 +136,15 @@ InitFileWriterConfig(const char* mode,
"Invalid mode");
}
milvus::storage::FileWriteWorkerPool::GetInstance().Configure(
nr_threads);
c_disk_write_config.nr_threads);
// configure rate limiter
milvus::storage::io::WriteRateLimiter::GetInstance().Configure(
c_disk_write_config.rate_limiter_config.refill_period_us,
c_disk_write_config.rate_limiter_config.avg_bps,
c_disk_write_config.rate_limiter_config.max_burst_bps,
c_disk_write_config.rate_limiter_config.high_priority_ratio,
c_disk_write_config.rate_limiter_config.middle_priority_ratio,
c_disk_write_config.rate_limiter_config.low_priority_ratio);
return milvus::SuccessCStatus();
} catch (std::exception& e) {
return milvus::FailureCStatus(&e);

2
internal/core/src/storage/storage_c.h
View File

@ -40,7 +40,7 @@ void
ResizeTheadPool(int64_t priority, float ratio);
CStatus
InitFileWriterConfig(const char* mode, uint64_t buffer_size_kb, int nr_threads);
InitDiskFileWriterConfig(CDiskWriteConfig c_disk_write_config);
#ifdef __cplusplus
};

125
internal/core/unittest/test_file_writer.cpp
View File

@ -13,7 +13,9 @@
#include <gmock/gmock.h>
#include <filesystem>
#include <fstream>
#include <random>
#include <chrono>
#include <cstdlib>
#include <cstring>
#include <thread>
#include "storage/FileWriter.h"
@ -32,6 +34,14 @@ class FileWriterTest : public testing::Test {
void
TearDown() override {
std::filesystem::remove_all(test_dir_);
// Reset rate limiter to disabled ratios to avoid test interference
auto& limiter = milvus::storage::io::WriteRateLimiter::GetInstance();
limiter.Configure(/*refill_period_us*/ 100000,
/*avg_bps*/ 8192 * 10,
/*max_burst_bps*/ 8192 * 40,
/*high*/ -1,
/*middle*/ -1,
/*low*/ -1);
}
std::filesystem::path test_dir_;
@ -250,11 +260,9 @@ TEST_F(FileWriterTest, MemoryAddressAlignedDataWriteWithDirectIO) {
std::string filename = (test_dir_ / "aligned_write.txt").string();
FileWriter writer(filename);
// Create 4KB aligned data using posix_memalign
void* aligned_data = nullptr;
if (posix_memalign(&aligned_data, 4096, kBufferSize) != 0) {
throw std::runtime_error("Failed to allocate aligned memory");
}
// Create 4KB aligned data using aligned_alloc
void* aligned_data = std::aligned_alloc(4096, kBufferSize);
ASSERT_NE(aligned_data, nullptr);
std::generate(static_cast<char*>(aligned_data),
static_cast<char*>(aligned_data) + kBufferSize,
std::rand);
@ -367,6 +375,98 @@ TEST_F(FileWriterTest, ExistingFileWithDirectIO) {
EXPECT_EQ(content, new_data);
}
// Test rate limiter basic behavior: alignment and refill period
TEST_F(FileWriterTest, RateLimiterAlignmentAndPeriods) {
using milvus::storage::io::Priority;
using milvus::storage::io::WriteRateLimiter;
// Configure: 100ms period, 8KB per period avg, 32KB burst, ratios enabled
auto& limiter = WriteRateLimiter::GetInstance();
limiter.Configure(/*refill_period_us*/ 100000,
/*avg_bps*/ 8192 * 10, // 8KB per 100ms
/*max_burst_bps*/ 8192 * 40, // 32KB burst
/*high*/ 1,
/*middle*/ 1,
/*low*/ 1);
// Wait one period to accumulate credits
std::this_thread::sleep_for(std::chrono::milliseconds(120));
// Request 8KB with 4KB alignment → expect a multiple of 4KB, <= 8KB
size_t allowed = limiter.Acquire(/*bytes*/ 8192,
/*alignment*/ 4096,
/*priority*/ Priority::MIDDLE);
EXPECT_GT(allowed, 0u);
EXPECT_LE(allowed, static_cast<size_t>(8192));
EXPECT_EQ(allowed % 4096, 0u);
}
// Test that buffered IO path writes correct data under throttling (no overlap)
TEST_F(FileWriterTest, FileWriterBufferedRateLimitedWriteCorrectness) {
using milvus::storage::io::Priority;
using milvus::storage::io::WriteRateLimiter;
FileWriter::SetMode(FileWriter::WriteMode::BUFFERED);
// Configure limiter to force multiple internal chunks
auto& limiter = WriteRateLimiter::GetInstance();
limiter.Configure(/*refill_period_us*/ 50000, // 50ms
/*avg_bps*/ 4096 * 20, // 4KB per 50ms
/*max_burst_bps*/ 4096 * 80, // 16KB burst
/*high*/ 1,
/*middle*/ 1,
/*low*/ 1);
// Prepare data larger than a few chunks
const size_t total_size = 12 * 4096;
std::vector<char> data(total_size);
std::generate(data.begin(), data.end(), std::rand);
std::string filename = (test_dir_ / "buffered_rate_limited.txt").string();
{
FileWriter writer(filename, Priority::MIDDLE);
writer.Write(data.data(), data.size());
writer.Finish();
}
// Verify file contents match exactly
std::ifstream file(filename, std::ios::binary);
std::vector<char> read_data((std::istreambuf_iterator<char>(file)),
std::istreambuf_iterator<char>());
EXPECT_EQ(read_data.size(), data.size());
EXPECT_EQ(read_data, data);
}
// Test that priority ratio impacts allowance (HIGH > MIDDLE)
TEST_F(FileWriterTest, RateLimiterPriorityRatioEffect) {
using milvus::storage::io::Priority;
using milvus::storage::io::WriteRateLimiter;
auto& limiter = WriteRateLimiter::GetInstance();
// 100ms period, 8KB per period, 32KB burst
limiter.Configure(/*refill_period_us*/ 100000,
/*avg_bps*/ 8192 * 10,
/*max_burst_bps*/ 8192 * 40,
/*high*/ 2,
/*middle*/ 1,
/*low*/ 1);
// Accumulate two periods of credits
std::this_thread::sleep_for(std::chrono::milliseconds(220));
// Request with same bytes and alignment; HIGH should allow more than MIDDLE
size_t req = 8 * 4096; // divisible by 4KB
size_t mid = limiter.Acquire(req, 4096, Priority::MIDDLE);
// Reset time/credits by waiting again for comparable conditions
std::this_thread::sleep_for(std::chrono::milliseconds(220));
size_t hig = limiter.Acquire(req, 4096, Priority::HIGH);
EXPECT_GT(hig, mid);
EXPECT_EQ(mid % 4096, 0u);
EXPECT_EQ(hig % 4096, 0u);
}
// Test config FileWriterConfig with very small buffer size
TEST_F(FileWriterTest, SmallBufferSizeWriteWithDirectIO) {
const size_t small_buffer_size = 64; // 64 bytes
@ -428,11 +528,9 @@ TEST_F(FileWriterTest, HalfAlignedDataWriteWithDirectIO) {
std::string filename = (test_dir_ / "half_aligned_buffer.txt").string();
FileWriter writer(filename);
char* aligned_buffer = nullptr;
int ret = posix_memalign(reinterpret_cast<void**>(&aligned_buffer),
kBufferSize,
aligned_buffer_size);
ASSERT_EQ(ret, 0);
char* aligned_buffer = static_cast<char*>(
std::aligned_alloc(kBufferSize, aligned_buffer_size));
ASSERT_NE(aligned_buffer, nullptr);
const size_t first_half_size = kBufferSize / 2;
const size_t rest_size = aligned_buffer_size - first_half_size;
@ -485,8 +583,9 @@ TEST_F(FileWriterTest, VeryLargeFileWriteWithDirectIO) {
const size_t large_size = 100 * 1024 * 1024; // 100MB
const size_t alignment = 4096; // 4KB alignment
char* aligned_data = nullptr;
ASSERT_EQ(posix_memalign((void**)&aligned_data, alignment, large_size), 0);
char* aligned_data =
static_cast<char*>(std::aligned_alloc(alignment, large_size));
ASSERT_NE(aligned_data, nullptr);
std::generate(aligned_data, aligned_data + large_size, std::rand);
writer.Write(aligned_data, large_size);

1
internal/core/unittest/test_stlsort_index.cpp
View File

@ -26,6 +26,7 @@ test_stlsort_for_range(
{
Config config;
config[milvus::index::ENABLE_MMAP] = enable_mmap;
config[milvus::LOAD_PRIORITY] = milvus::proto::common::LoadPriority::HIGH;
auto index = std::make_shared<index::ScalarIndexSort<int64_t>>();
index->Load(binary_set, config);

32
internal/querynodev2/server.go
View File

@ -200,16 +200,22 @@ func ResizeHighPriorityPool(evt *config.Event) {
}
}
func (node *QueryNode) ReconfigFileWriterParams(evt *config.Event) {
func (node *QueryNode) ReconfigDiskFileWriterParams(evt *config.Event) {
if evt.HasUpdated {
if err := initcore.InitFileWriterConfig(paramtable.Get()); err != nil {
if err := initcore.InitDiskFileWriterConfig(paramtable.Get()); err != nil {
log.Ctx(node.ctx).Warn("QueryNode failed to reconfigure file writer params", zap.Error(err))
return
}
log.Ctx(node.ctx).Info("QueryNode reconfig file writer params successfully",
zap.String("mode", paramtable.Get().CommonCfg.DiskWriteMode.GetValue()),
zap.Uint64("bufferSize", paramtable.Get().CommonCfg.DiskWriteBufferSizeKb.GetAsUint64()),
zap.Int("nrThreads", paramtable.Get().CommonCfg.DiskWriteNumThreads.GetAsInt()))
zap.Int("nrThreads", paramtable.Get().CommonCfg.DiskWriteNumThreads.GetAsInt()),
zap.Uint64("refillPeriodUs", paramtable.Get().CommonCfg.DiskWriteRateLimiterRefillPeriodUs.GetAsUint64()),
zap.Uint64("maxBurstKBps", paramtable.Get().CommonCfg.DiskWriteRateLimiterMaxBurstKBps.GetAsUint64()),
zap.Uint64("avgKBps", paramtable.Get().CommonCfg.DiskWriteRateLimiterAvgKBps.GetAsUint64()),
zap.Int("highPriorityRatio", paramtable.Get().CommonCfg.DiskWriteRateLimiterHighPriorityRatio.GetAsInt()),
zap.Int("middlePriorityRatio", paramtable.Get().CommonCfg.DiskWriteRateLimiterMiddlePriorityRatio.GetAsInt()),
zap.Int("lowPriorityRatio", paramtable.Get().CommonCfg.DiskWriteRateLimiterLowPriorityRatio.GetAsInt()))
}
}
@ -218,11 +224,23 @@ func (node *QueryNode) RegisterSegcoreConfigWatcher() {
pt.Watch(pt.CommonCfg.HighPriorityThreadCoreCoefficient.Key,
config.NewHandler("common.threadCoreCoefficient.highPriority", ResizeHighPriorityPool))
pt.Watch(pt.CommonCfg.DiskWriteMode.Key,
config.NewHandler("common.diskWriteMode", node.ReconfigFileWriterParams))
config.NewHandler("common.diskWriteMode", node.ReconfigDiskFileWriterParams))
pt.Watch(pt.CommonCfg.DiskWriteBufferSizeKb.Key,
config.NewHandler("common.diskWriteBufferSizeKb", node.ReconfigFileWriterParams))
config.NewHandler("common.diskWriteBufferSizeKb", node.ReconfigDiskFileWriterParams))
pt.Watch(pt.CommonCfg.DiskWriteNumThreads.Key,
config.NewHandler("common.diskWriteNumThreads", node.ReconfigFileWriterParams))
config.NewHandler("common.diskWriteNumThreads", node.ReconfigDiskFileWriterParams))
pt.Watch(pt.CommonCfg.DiskWriteRateLimiterRefillPeriodUs.Key,
config.NewHandler("common.diskWriteRateLimiter.refillPeriodUs", node.ReconfigDiskFileWriterParams))
pt.Watch(pt.CommonCfg.DiskWriteRateLimiterMaxBurstKBps.Key,
config.NewHandler("common.diskWriteRateLimiter.maxBurstKBps", node.ReconfigDiskFileWriterParams))
pt.Watch(pt.CommonCfg.DiskWriteRateLimiterAvgKBps.Key,
config.NewHandler("common.diskWriteRateLimiter.avgKBps", node.ReconfigDiskFileWriterParams))
pt.Watch(pt.CommonCfg.DiskWriteRateLimiterHighPriorityRatio.Key,
config.NewHandler("common.diskWriteRateLimiter.highPriorityRatio", node.ReconfigDiskFileWriterParams))
pt.Watch(pt.CommonCfg.DiskWriteRateLimiterMiddlePriorityRatio.Key,
config.NewHandler("common.diskWriteRateLimiter.middlePriorityRatio", node.ReconfigDiskFileWriterParams))
pt.Watch(pt.CommonCfg.DiskWriteRateLimiterLowPriorityRatio.Key,
config.NewHandler("common.diskWriteRateLimiter.lowPriorityRatio", node.ReconfigDiskFileWriterParams))
}
// InitSegcore set init params of segCore, such as chunkRows, SIMD type...
@ -302,7 +320,7 @@ func (node *QueryNode) InitSegcore() error {
return err
}
err = initcore.InitFileWriterConfig(paramtable.Get())
err = initcore.InitDiskFileWriterConfig(paramtable.Get())
if err != nil {
return err
}

26
internal/util/initcore/init_core.go
View File

@ -272,16 +272,36 @@ func InitMmapManager(params *paramtable.ComponentParam) error {
return HandleCStatus(&status, "InitMmapManager failed")
}
func InitFileWriterConfig(params *paramtable.ComponentParam) error {
func InitDiskFileWriterConfig(params *paramtable.ComponentParam) error {
mode := params.CommonCfg.DiskWriteMode.GetValue()
bufferSize := params.CommonCfg.DiskWriteBufferSizeKb.GetAsUint64()
numThreads := params.CommonCfg.DiskWriteNumThreads.GetAsInt()
refillPeriodUs := params.CommonCfg.DiskWriteRateLimiterRefillPeriodUs.GetAsInt64()
maxBurstKBps := params.CommonCfg.DiskWriteRateLimiterMaxBurstKBps.GetAsInt64()
avgKBps := params.CommonCfg.DiskWriteRateLimiterAvgKBps.GetAsInt64()
highPriorityRatio := params.CommonCfg.DiskWriteRateLimiterHighPriorityRatio.GetAsInt()
middlePriorityRatio := params.CommonCfg.DiskWriteRateLimiterMiddlePriorityRatio.GetAsInt()
lowPriorityRatio := params.CommonCfg.DiskWriteRateLimiterLowPriorityRatio.GetAsInt()
cMode := C.CString(mode)
cBufferSize := C.uint64_t(bufferSize)
cNumThreads := C.int(numThreads)
defer C.free(unsafe.Pointer(cMode))
status := C.InitFileWriterConfig(cMode, cBufferSize, cNumThreads)
return HandleCStatus(&status, "InitFileWriterConfig failed")
diskWriteRateLimiterConfig := C.CDiskWriteRateLimiterConfig{
refill_period_us: C.int64_t(refillPeriodUs),
avg_bps: C.int64_t(avgKBps * 1024),
max_burst_bps: C.int64_t(maxBurstKBps * 1024),
high_priority_ratio: C.int32_t(highPriorityRatio),
middle_priority_ratio: C.int32_t(middlePriorityRatio),
low_priority_ratio: C.int32_t(lowPriorityRatio),
}
diskWriteConfig := C.CDiskWriteConfig{
mode: cMode,
buffer_size_kb: cBufferSize,
nr_threads: cNumThreads,
rate_limiter_config: diskWriteRateLimiterConfig,
}
status := C.InitDiskFileWriterConfig(diskWriteConfig)
return HandleCStatus(&status, "InitDiskFileWriterConfig failed")
}
var coreParamCallbackInitOnce sync.Once

63
pkg/util/paramtable/component_param.go
View File

@ -246,6 +246,13 @@ type commonConfig struct {
DiskWriteBufferSizeKb ParamItem `refreshable:"true"`
DiskWriteNumThreads ParamItem `refreshable:"true"`
DiskWriteRateLimiterRefillPeriodUs ParamItem `refreshable:"true"`
DiskWriteRateLimiterAvgKBps ParamItem `refreshable:"true"`
DiskWriteRateLimiterMaxBurstKBps ParamItem `refreshable:"true"`
DiskWriteRateLimiterHighPriorityRatio ParamItem `refreshable:"true"`
DiskWriteRateLimiterMiddlePriorityRatio ParamItem `refreshable:"true"`
DiskWriteRateLimiterLowPriorityRatio ParamItem `refreshable:"true"`
AuthorizationEnabled ParamItem `refreshable:"false"`
SuperUsers ParamItem `refreshable:"true"`
DefaultRootPassword ParamItem `refreshable:"false"`
@ -686,6 +693,62 @@ In this case, the maximum concurrency of disk write operations is determined by
}
p.DiskWriteNumThreads.Init(base.mgr)
p.DiskWriteRateLimiterRefillPeriodUs = ParamItem{
Key: "common.diskWriteRateLimiter.refillPeriodUs",
Version: "2.6.0",
DefaultValue: "100000",
Doc: "refill period in microseconds if disk rate limiter is enabled, default is 100000us (100ms)",
Export: true,
}
p.DiskWriteRateLimiterRefillPeriodUs.Init(base.mgr)
p.DiskWriteRateLimiterAvgKBps = ParamItem{
Key: "common.diskWriteRateLimiter.avgKBps",
Version: "2.6.0",
DefaultValue: "262144",
Doc: "average kilobytes per second if disk rate limiter is enabled, default is 262144KB/s (256MB/s)",
Export: true,
}
p.DiskWriteRateLimiterAvgKBps.Init(base.mgr)
p.DiskWriteRateLimiterMaxBurstKBps = ParamItem{
Key: "common.diskWriteRateLimiter.maxBurstKBps",
Version: "2.6.0",
DefaultValue: "524288",
Doc: "max burst kilobytes per second if disk rate limiter is enabled, default is 524288KB/s (512MB/s)",
Export: true,
}
p.DiskWriteRateLimiterMaxBurstKBps.Init(base.mgr)
p.DiskWriteRateLimiterHighPriorityRatio = ParamItem{
Key: "common.diskWriteRateLimiter.highPriorityRatio",
Version: "2.6.0",
DefaultValue: "-1",
Doc: `amplification ratio for high priority tasks if disk rate limiter is enabled, value <= 0 means ratio limit is disabled.
The ratio is the multiplication factor of the configured bandwidth.
For example, if the rate limit is 100KB/s, and the high priority ratio is 2, then the high priority tasks will be limited to 200KB/s.`,
Export: true,
}
p.DiskWriteRateLimiterHighPriorityRatio.Init(base.mgr)
p.DiskWriteRateLimiterMiddlePriorityRatio = ParamItem{
Key: "common.diskWriteRateLimiter.middlePriorityRatio",
Version: "2.6.0",
DefaultValue: "-1",
Doc: "amplification ratio for middle priority tasks if disk rate limiter is enabled, value <= 0 means ratio limit is disabled",
Export: true,
}
p.DiskWriteRateLimiterMiddlePriorityRatio.Init(base.mgr)
p.DiskWriteRateLimiterLowPriorityRatio = ParamItem{
Key: "common.diskWriteRateLimiter.lowPriorityRatio",
Version: "2.6.0",
DefaultValue: "-1",
Doc: "amplification ratio for low priority tasks if disk rate limiter is enabled, value <= 0 means ratio limit is disabled",
Export: true,
}
p.DiskWriteRateLimiterLowPriorityRatio.Init(base.mgr)
p.BuildIndexThreadPoolRatio = ParamItem{
Key: "common.buildIndexThreadPoolRatio",
Version: "2.4.0",