mooncake-track
/
Mooncake
0
0
Fork 12
Code Issues Pull Requests 13 Actions Packages Projects Releases Wiki Activity

[Store] Implement a file interface based on io uring to optimize storage backend. (#1500) 

Co-authored-by: zhuxinjie-nz <240190801+zhuxinjie-nz@users.noreply.github.com>
This commit is contained in:
zhangzuo21 2026-02-25 14:18:22 +08:00 committed by GitHub
parent 53e750b2b5
commit 4ff4dc1e27
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 2404 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mooncake-store/benchmarks/CMakeLists.txt
View File

@ -6,6 +6,9 @@ target_link_libraries(allocator_bench PRIVATE cachelib_memory_allocator mooncake
add_executable(master_bench master_bench.cpp)
target_link_libraries(master_bench PRIVATE cachelib_memory_allocator mooncake_store)
# Add file interface benchmark executable
add_executable(file_interface_bench file_interface_bench.cpp)
target_link_libraries(file_interface_bench PRIVATE mooncake_store glog::glog)
# Add storage backend benchmark executable
# This benchmark tests all storage backends (OffsetAllocator, Bucket, FilePerKey)
# with realistic KV cache workloads for LLM inference scenarios

864
mooncake-store/benchmarks/file_interface_bench.cpp Normal file
View File

@ -0,0 +1,864 @@
#include <iostream>
#include <fstream>
#include <random>
#include <chrono>
#include <iomanip>
#include <memory>
#include <string>
#include <vector>
#include <utility>
#include <cstring>
#include <cerrno>
#include <cmath>
#include <limits>
#include <fcntl.h>
#include <unistd.h>
#include <glog/logging.h>
#include "file_interface.h"
using namespace mooncake;
// Aligned memory allocation for O_DIRECT
void* aligned_alloc_buffer(size_t size, size_t alignment = 4096) {
void* ptr = nullptr;
if (posix_memalign(&ptr, alignment, size) != 0) {
return nullptr;
}
return ptr;
}
void aligned_free_buffer(void* ptr) { free(ptr); }
struct BenchmarkConfig {
std::string file_path = "/tmp/file_bench_test.dat";
size_t data_size = 1024 * 1024 * 100; // 100 MB
bool use_uring = false;
unsigned uring_queue_depth = 32;
bool verify_data = true;
int iterations = 5; // Number of iterations for each test
bool use_direct_io =
false; // Use O_DIRECT with raw syscalls (bypasses StorageFile)
bool use_uring_direct_io =
false; // Use O_DIRECT with UringFile interface (with copy)
bool use_uring_direct_io_zero_copy =
false; // Use O_DIRECT with zero-copy aligned interface
bool use_registered_buffers =
false; // Register buffers with io_uring for optimal performance
size_t alignment = 4096; // Alignment for O_DIRECT (4KB)
size_t chunk_size =
1024 * 1024; // Chunk size for I/O operations (1MB default)
};
struct BenchmarkStats {
double min = std::numeric_limits<double>::max();
double max = 0;
double sum = 0;
double sum_sq = 0;
int count = 0;
void Add(double value) {
min = std::min(min, value);
max = std::max(max, value);
sum += value;
sum_sq += value * value;
count++;
}
double Mean() const { return count > 0 ? sum / count : 0; }
double StdDev() const {
if (count < 2) return 0;
double mean = Mean();
return std::sqrt((sum_sq / count) - (mean * mean));
}
void Print(const std::string& label) const {
std::cout << std::fixed << std::setprecision(2);
std::cout << label << ": " << Mean() << " MB/s (min=" << min
<< ", max=" << max << ", stddev=" << StdDev() << ")"
<< std::endl;
}
};
class FileInterfaceBenchmark {
public:
explicit FileInterfaceBenchmark(const BenchmarkConfig& config)
: config_(config) {}
void Run() {
std::cout << "=== File Interface Benchmark ===" << std::endl;
std::cout << "File path: " << config_.file_path << std::endl;
std::cout << "Data size: " << FormatSize(config_.data_size)
<< std::endl;
std::cout << "Use io_uring: " << (config_.use_uring ? "Yes" : "No")
<< std::endl;
if (config_.use_uring) {
std::cout << "io_uring queue depth: " << config_.uring_queue_depth
<< std::endl;
std::cout << "io_uring O_DIRECT: "
<< (config_.use_uring_direct_io ? "Yes (with copy)"
: "No")
<< std::endl;
std::cout << "io_uring O_DIRECT zero-copy: "
<< (config_.use_uring_direct_io_zero_copy ? "Yes" : "No")
<< std::endl;
std::cout << "io_uring registered buffers: "
<< (config_.use_registered_buffers ? "Yes" : "No")
<< std::endl;
}
std::cout << "Use O_DIRECT (raw syscalls): "
<< (config_.use_direct_io ? "Yes" : "No") << std::endl;
if (config_.use_direct_io) {
std::cout << "Alignment: " << config_.alignment << " bytes"
<< std::endl;
}
std::cout << "Chunk size: " << FormatSize(config_.chunk_size)
<< std::endl;
std::cout << "Iterations: " << config_.iterations << std::endl;
std::cout << std::endl;
// Ensure data size is aligned for O_DIRECT
size_t actual_size = config_.data_size;
if (config_.use_direct_io || config_.use_uring_direct_io_zero_copy) {
actual_size = ((config_.data_size + config_.alignment - 1) /
config_.alignment) *
config_.alignment;
if (actual_size != config_.data_size) {
std::cout << "Adjusted data size to " << FormatSize(actual_size)
<< " for O_DIRECT alignment" << std::endl;
}
}
// Generate test data
std::cout << "Generating test data..." << std::endl;
void* write_data_ptr = nullptr;
std::vector<char> write_data_vec;
if (config_.use_direct_io || config_.use_uring_direct_io_zero_copy) {
// Allocate aligned buffer for O_DIRECT (raw syscall or zero-copy
// mode)
write_data_ptr =
aligned_alloc_buffer(actual_size, config_.alignment);
if (!write_data_ptr) {
std::cerr << "Failed to allocate aligned buffer!" << std::endl;
return;
}
GenerateRandomDataAligned(write_data_ptr, actual_size);
} else {
// UringFile or PosixFile: use std::vector (UringFile will handle
// alignment internally)
write_data_vec = GenerateRandomData(actual_size);
write_data_ptr = write_data_vec.data();
}
std::cout << "Test data generated." << std::endl << std::endl;
BenchmarkStats write_stats, read_stats;
for (int iter = 0; iter < config_.iterations; ++iter) {
std::cout << "--- Iteration " << (iter + 1) << "/"
<< config_.iterations << " ---" << std::endl;
if (config_.use_uring_direct_io_zero_copy) {
// Zero-copy aligned I/O mode
auto write_result_aligned =
BenchmarkAlignedIO(write_data_ptr, actual_size, true);
if (!write_result_aligned) {
std::cerr << "Zero-copy aligned write failed!" << std::endl;
aligned_free_buffer(write_data_ptr);
return;
}
write_stats.Add(write_result_aligned->first);
// Sync and drop cache before read
SyncAndDropCache();
auto read_result_aligned =
BenchmarkAlignedIO(nullptr, actual_size, false);
if (!read_result_aligned) {
std::cerr << "Zero-copy aligned read failed!" << std::endl;
aligned_free_buffer(write_data_ptr);
return;
}
read_stats.Add(read_result_aligned->first);
// Verify data consistency (only first iteration)
if (config_.verify_data && iter == 0) {
bool verified =
VerifyData(write_data_ptr, read_result_aligned->second,
actual_size);
if (!verified) {
std::cerr << "Data verification FAILED!" << std::endl;
aligned_free_buffer(read_result_aligned->second);
aligned_free_buffer(write_data_ptr);
return;
}
std::cout << "Data verification PASSED" << std::endl;
}
// Free read buffer
aligned_free_buffer(read_result_aligned->second);
} else if (config_.use_direct_io) {
// Use direct I/O syscalls
auto write_result_direct =
BenchmarkDirectIO(write_data_ptr, actual_size, true);
if (!write_result_direct) {
std::cerr << "Direct I/O write benchmark failed!"
<< std::endl;
aligned_free_buffer(write_data_ptr);
return;
}
write_stats.Add(write_result_direct->first);
// Sync and drop cache before read
SyncAndDropCache();
auto read_result_direct =
BenchmarkDirectIO(nullptr, actual_size, false);
if (!read_result_direct) {
std::cerr << "Direct I/O read benchmark failed!"
<< std::endl;
aligned_free_buffer(write_data_ptr);
return;
}
read_stats.Add(read_result_direct->first);
// Verify data consistency (only first iteration)
if (config_.verify_data && iter == 0) {
bool verified =
VerifyData(write_data_ptr, read_result_direct->second,
actual_size);
if (!verified) {
std::cerr << "Data verification FAILED!" << std::endl;
aligned_free_buffer(read_result_direct->second);
aligned_free_buffer(write_data_ptr);
return;
}
std::cout << "Data verification PASSED" << std::endl;
}
// Free read buffer
aligned_free_buffer(read_result_direct->second);
} else {
// Use StorageFile interface
auto write_result = BenchmarkWrite(write_data_ptr, actual_size);
if (!write_result) {
std::cerr << "Write benchmark failed!" << std::endl;
return;
}
write_stats.Add(*write_result);
auto read_result = BenchmarkRead(actual_size);
if (!read_result) {
std::cerr << "Read benchmark failed!" << std::endl;
return;
}
read_stats.Add(read_result->first);
// Verify data consistency (only first iteration)
if (config_.verify_data && iter == 0) {
bool verified = VerifyData(
write_data_ptr, read_result->second, actual_size);
if (!verified) {
std::cerr << "Data verification FAILED!" << std::endl;
aligned_free_buffer(read_result->second);
return;
}
std::cout << "Data verification PASSED" << std::endl;
}
// Free read buffer (always allocated in BenchmarkRead now)
aligned_free_buffer(read_result->second);
}
// Cleanup for next iteration
unlink(config_.file_path.c_str());
std::cout << std::endl;
}
// Cleanup aligned buffer
if (config_.use_direct_io || config_.use_uring_direct_io_zero_copy) {
aligned_free_buffer(write_data_ptr);
}
// Print summary
std::cout << "=== Summary ===" << std::endl;
write_stats.Print("Write bandwidth");
read_stats.Print("Read bandwidth");
}
private:
std::vector<char> GenerateRandomData(size_t size) {
std::vector<char> data(size);
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_int_distribution<unsigned char> dist(0, 255);
for (size_t i = 0; i < size; ++i) {
data[i] = static_cast<char>(dist(gen));
}
return data;
}
void GenerateRandomDataAligned(void* ptr, size_t size) {
char* data = static_cast<char*>(ptr);
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_int_distribution<unsigned char> dist(0, 255);
for (size_t i = 0; i < size; ++i) {
data[i] = static_cast<char>(dist(gen));
}
}
std::unique_ptr<StorageFile> CreateFile(int fd) {
#ifdef USE_URING
if (config_.use_uring) {
return std::make_unique<UringFile>(config_.file_path, fd,
config_.uring_queue_depth,
config_.use_uring_direct_io);
}
#endif
return std::make_unique<PosixFile>(config_.file_path, fd);
}
bool SyncAndDropCache() {
#ifdef __linux__
// Flush dirty pages
::sync();
if (::geteuid() != 0) {
std::cerr << "[CACHE] Not root (euid=" << ::geteuid()
<< "). Cannot write /proc/sys/vm/drop_caches.\n"
<< " Run: sudo -E ./file_interface_bench ...\n";
return false;
}
std::ofstream drop("/proc/sys/vm/drop_caches");
if (!drop.is_open()) {
std::cerr << "[CACHE] Failed to open /proc/sys/vm/drop_caches: "
<< std::strerror(errno) << "\n";
return false;
}
drop << "3" << std::flush;
if (drop.fail()) {
std::cerr
<< "[CACHE] Failed to write to /proc/sys/vm/drop_caches\n";
return false;
}
drop.close();
std::cout << "Synced and dropped page cache" << std::endl;
return true;
#else
std::cerr << "[CACHE] drop_caches not supported on this platform\n";
return false;
#endif
}
std::optional<double> BenchmarkWrite(void* data, size_t size) {
std::cout << "=== Write Benchmark ===" << std::endl;
// Open file for writing
int flags = O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC;
if (config_.use_uring_direct_io) {
flags |= O_DIRECT;
std::cout << "Using UringFile with O_DIRECT" << std::endl;
}
int fd = open(config_.file_path.c_str(), flags, 0644);
if (fd < 0) {
std::cerr << "Failed to open file for writing: "
<< config_.file_path << " - " << std::strerror(errno)
<< std::endl;
return std::nullopt;
}
auto file = CreateFile(fd);
if (!file) {
std::cerr << "Failed to create file object" << std::endl;
close(fd);
return std::nullopt;
}
// Perform write operation
auto start = std::chrono::high_resolution_clock::now();
auto result = file->write(
std::span<const char>(static_cast<const char*>(data), size), size);
auto end = std::chrono::high_resolution_clock::now();
if (!result) {
std::cerr << "Write operation failed with error code: "
<< static_cast<int>(result.error()) << std::endl;
return std::nullopt;
}
size_t bytes_written = result.value();
auto duration_us =
std::chrono::duration_cast<std::chrono::microseconds>(end - start)
.count();
double duration_s = duration_us / 1000000.0;
double bandwidth_mbps =
(bytes_written / (1024.0 * 1024.0)) / duration_s;
std::cout << std::fixed << std::setprecision(2);
std::cout << "Bytes written: " << FormatSize(bytes_written)
<< std::endl;
std::cout << "Time: " << duration_us << " us (" << duration_s << " s)"
<< std::endl;
std::cout << "Bandwidth: " << bandwidth_mbps << " MB/s" << std::endl;
// Sync to disk and drop page cache
SyncAndDropCache();
return bandwidth_mbps;
}
// Zero-copy aligned I/O using UringFile's aligned interface
std::optional<std::pair<double, void*>> BenchmarkAlignedIO(void* write_data,
size_t size,
bool is_write) {
std::cout << "=== Zero-Copy Aligned I/O "
<< (is_write ? "Write" : "Read")
<< " Benchmark ===" << std::endl;
// Open file with O_DIRECT
int flags = is_write
? (O_WRONLY | O_CREAT | O_TRUNC | O_DIRECT | O_CLOEXEC)
: (O_RDONLY | O_DIRECT | O_CLOEXEC);
int fd = open(config_.file_path.c_str(), flags, 0644);
if (fd < 0) {
std::cerr << "Failed to open file with O_DIRECT: "
<< config_.file_path << " - " << std::strerror(errno)
<< std::endl;
return std::nullopt;
}
auto file = CreateFile(fd);
if (!file) {
std::cerr << "Failed to create file object" << std::endl;
close(fd);
return std::nullopt;
}
void* buffer = nullptr;
if (is_write) {
buffer = write_data;
} else {
buffer = aligned_alloc_buffer(size, config_.alignment);
if (!buffer) {
std::cerr << "Failed to allocate aligned buffer" << std::endl;
return std::nullopt;
}
}
// Register buffer if requested (NOT counted in I/O time)
auto* uring_file = dynamic_cast<UringFile*>(file.get());
if (config_.use_registered_buffers && uring_file) {
auto reg_start = std::chrono::high_resolution_clock::now();
bool registered = uring_file->register_buffer(buffer, size);
auto reg_end = std::chrono::high_resolution_clock::now();
auto reg_us = std::chrono::duration_cast<std::chrono::microseconds>(
reg_end - reg_start)
.count();
if (registered) {
std::cout << "Buffer registration: " << reg_us
<< " us (excluded from I/O time)" << std::endl;
} else {
std::cerr
<< "Warning: Buffer registration failed, continuing without"
<< std::endl;
}
}
// Perform zero-copy I/O operation
auto start = std::chrono::high_resolution_clock::now();
tl::expected<size_t, ErrorCode> result;
if (is_write) {
auto* uring_file = dynamic_cast<UringFile*>(file.get());
if (uring_file) {
result = uring_file->write_aligned(buffer, size, 0);
} else {
std::cerr << "Not a UringFile instance!" << std::endl;
if (!is_write) aligned_free_buffer(buffer);
return std::nullopt;
}
} else {
auto* uring_file = dynamic_cast<UringFile*>(file.get());
if (uring_file) {
result = uring_file->read_aligned(buffer, size, 0);
} else {
std::cerr << "Not a UringFile instance!" << std::endl;
aligned_free_buffer(buffer);
return std::nullopt;
}
}
auto end = std::chrono::high_resolution_clock::now();
// Unregister buffer if it was registered (NOT counted in I/O time)
if (config_.use_registered_buffers && uring_file &&
uring_file->is_buffer_registered()) {
auto unreg_start = std::chrono::high_resolution_clock::now();
uring_file->unregister_buffer();
auto unreg_end = std::chrono::high_resolution_clock::now();
auto unreg_us =
std::chrono::duration_cast<std::chrono::microseconds>(
unreg_end - unreg_start)
.count();
std::cout << "Buffer unregistration: " << unreg_us
<< " us (excluded from I/O time)" << std::endl;
}
if (!result) {
std::cerr << "Zero-copy I/O operation failed with error code: "
<< static_cast<int>(result.error()) << std::endl;
if (!is_write) aligned_free_buffer(buffer);
return std::nullopt;
}
size_t total_bytes = result.value();
auto duration_us =
std::chrono::duration_cast<std::chrono::microseconds>(end - start)
.count();
double duration_s = duration_us / 1000000.0;
double bandwidth_mbps = (total_bytes / (1024.0 * 1024.0)) / duration_s;
std::cout << std::fixed << std::setprecision(2);
std::cout << "Bytes " << (is_write ? "written" : "read") << ": "
<< FormatSize(total_bytes) << std::endl;
std::cout << "Time: " << duration_us << " us (" << duration_s << " s)"
<< std::endl;
std::cout << "Bandwidth: " << bandwidth_mbps << " MB/s" << std::endl;
std::cout << "Zero-copy: YES (no memory copy overhead)" << std::endl;
return std::make_pair(bandwidth_mbps, buffer);
}
// Direct I/O benchmark using raw syscalls (bypasses StorageFile interface)
std::optional<std::pair<double, void*>> BenchmarkDirectIO(void* write_data,
size_t size,
bool is_write) {
std::cout << "=== Direct I/O " << (is_write ? "Write" : "Read")
<< " Benchmark ===" << std::endl;
// Open file with O_DIRECT
int flags = is_write ? (O_WRONLY | O_CREAT | O_TRUNC | O_DIRECT)
: (O_RDONLY | O_DIRECT);
int fd = open(config_.file_path.c_str(), flags, 0644);
if (fd < 0) {
std::cerr << "Failed to open file with O_DIRECT: "
<< config_.file_path << " - " << std::strerror(errno)
<< std::endl;
return std::nullopt;
}
void* buffer = nullptr;
if (is_write) {
buffer = write_data;
} else {
buffer = aligned_alloc_buffer(size, config_.alignment);
if (!buffer) {
std::cerr << "Failed to allocate aligned buffer" << std::endl;
close(fd);
return std::nullopt;
}
}
// Perform I/O operation
auto start = std::chrono::high_resolution_clock::now();
size_t total_bytes = 0;
size_t remaining = size;
off_t offset = 0;
while (remaining > 0) {
size_t this_chunk = std::min(config_.chunk_size, remaining);
ssize_t ret;
if (is_write) {
ret = pwrite(fd, static_cast<char*>(buffer) + offset,
this_chunk, offset);
} else {
ret = pread(fd, static_cast<char*>(buffer) + offset, this_chunk,
offset);
}
if (ret < 0) {
std::cerr << "I/O error: " << std::strerror(errno) << std::endl;
if (!is_write) aligned_free_buffer(buffer);
close(fd);
return std::nullopt;
}
if (ret == 0) break; // EOF for read
total_bytes += ret;
offset += ret;
remaining -= ret;
}
auto end = std::chrono::high_resolution_clock::now();
close(fd);
auto duration_us =
std::chrono::duration_cast<std::chrono::microseconds>(end - start)
.count();
double duration_s = duration_us / 1000000.0;
double bandwidth_mbps = (total_bytes / (1024.0 * 1024.0)) / duration_s;
std::cout << std::fixed << std::setprecision(2);
std::cout << "Bytes " << (is_write ? "written" : "read") << ": "
<< FormatSize(total_bytes) << std::endl;
std::cout << "Time: " << duration_us << " us (" << duration_s << " s)"
<< std::endl;
std::cout << "Bandwidth: " << bandwidth_mbps << " MB/s" << std::endl;
return std::make_pair(bandwidth_mbps, buffer);
}
std::optional<std::pair<double, void*>> BenchmarkRead(size_t size) {
std::cout << "=== Read Benchmark ===" << std::endl;
// Open file for reading
int flags = O_RDONLY | O_CLOEXEC;
if (config_.use_uring_direct_io) {
flags |= O_DIRECT;
std::cout << "Using UringFile with O_DIRECT (with copy)"
<< std::endl;
}
int fd = open(config_.file_path.c_str(), flags);
if (fd < 0) {
std::cerr << "Failed to open file for reading: "
<< config_.file_path << " - " << std::strerror(errno)
<< std::endl;
return std::nullopt;
}
auto file = CreateFile(fd);
if (!file) {
std::cerr << "Failed to create file object" << std::endl;
close(fd);
return std::nullopt;
}
// Allocate a buffer that persists after function return
void* read_buffer = aligned_alloc_buffer(size, config_.alignment);
if (!read_buffer) {
std::cerr << "Failed to allocate read buffer" << std::endl;
return std::nullopt;
}
// Perform read operation into a temporary string
std::string read_data_str;
auto start = std::chrono::high_resolution_clock::now();
auto result = file->read(read_data_str, size);
auto end = std::chrono::high_resolution_clock::now();
if (!result) {
std::cerr << "Read operation failed with error code: "
<< static_cast<int>(result.error()) << std::endl;
aligned_free_buffer(read_buffer);
return std::nullopt;
}
size_t bytes_read = result.value();
// Copy data to persistent buffer
std::memcpy(read_buffer, read_data_str.data(), bytes_read);
auto duration_us =
std::chrono::duration_cast<std::chrono::microseconds>(end - start)
.count();
double duration_s = duration_us / 1000000.0;
double bandwidth_mbps = (bytes_read / (1024.0 * 1024.0)) / duration_s;
std::cout << std::fixed << std::setprecision(2);
std::cout << "Bytes read: " << FormatSize(bytes_read) << std::endl;
std::cout << "Time: " << duration_us << " us (" << duration_s << " s)"
<< std::endl;
std::cout << "Bandwidth: " << bandwidth_mbps << " MB/s" << std::endl;
return std::make_pair(bandwidth_mbps, read_buffer);
}
bool VerifyData(void* expected, void* actual, size_t size) {
bool match = std::memcmp(expected, actual, size) == 0;
if (!match) {
// Find first mismatch position
const char* exp_data = static_cast<const char*>(expected);
const char* act_data = static_cast<const char*>(actual);
for (size_t i = 0; i < size; ++i) {
if (exp_data[i] != act_data[i]) {
std::cerr << "First mismatch at offset " << i << ": "
<< "expected 0x" << std::hex
<< (int)(unsigned char)exp_data[i] << ", got 0x"
<< (int)(unsigned char)act_data[i] << std::dec
<< std::endl;
break;
}
}
}
return match;
}
static std::string FormatSize(size_t bytes) {
const char* units[] = {"B", "KB", "MB", "GB"};
int unit_index = 0;
double size = static_cast<double>(bytes);
while (size >= 1024.0 && unit_index < 3) {
size /= 1024.0;
unit_index++;
}
std::ostringstream oss;
oss << std::fixed << std::setprecision(2) << size << " "
<< units[unit_index];
return oss.str();
}
BenchmarkConfig config_;
};
void PrintUsage(const char* program_name) {
std::cout << "Usage: " << program_name << " [options]" << std::endl;
std::cout << "Options:" << std::endl;
std::cout << " --file <path> Path to test file (default: "
"/tmp/file_bench_test.dat)"
<< std::endl;
std::cout << " --size <bytes> Size of test data in bytes (default: "
"104857600 = 100MB)"
<< std::endl;
std::cout
<< " --size-mb <MB> Size of test data in MB (overrides --size)"
<< std::endl;
std::cout << " --chunk-size <KB> I/O chunk size in KB (default: 1024)"
<< std::endl;
#ifdef USE_URING
std::cout << " --use-uring Use io_uring for I/O operations"
<< std::endl;
std::cout << " --queue-depth <n> io_uring queue depth (default: 32)"
<< std::endl;
std::cout << " --uring-direct-io Enable O_DIRECT in UringFile "
"(with memory copy)"
<< std::endl;
std::cout << " --uring-direct-io-zerocopy Enable O_DIRECT with zero-copy "
"(requires aligned buffer)"
<< std::endl;
std::cout << " --use-registered-buffers Register buffers with io_uring "
"for optimal performance"
<< std::endl;
#endif
std::cout << " --direct-io Use O_DIRECT with raw syscalls "
"(bypasses StorageFile)"
<< std::endl;
std::cout << " --iterations <n> Number of iterations (default: 5)"
<< std::endl;
std::cout << " --no-verify Skip data verification" << std::endl;
std::cout << " --help Show this help message" << std::endl;
}
int main(int argc, char* argv[]) {
google::InitGoogleLogging(argv[0]);
BenchmarkConfig config;
// Parse command line arguments
for (int i = 1; i < argc; ++i) {
std::string arg = argv[i];
if (arg == "--help") {
PrintUsage(argv[0]);
return 0;
} else if (arg == "--file" && i + 1 < argc) {
config.file_path = argv[++i];
} else if (arg == "--size" && i + 1 < argc) {
config.data_size = std::stoull(argv[++i]);
} else if (arg == "--size-mb" && i + 1 < argc) {
config.data_size = std::stoull(argv[++i]) * 1024 * 1024;
} else if (arg == "--chunk-size" && i + 1 < argc) {
config.chunk_size =
std::stoull(argv[++i]) * 1024; // Convert KB to bytes
} else if (arg == "--use-uring") {
#ifdef USE_URING
config.use_uring = true;
#else
std::cerr << "Warning: io_uring support not compiled in"
<< std::endl;
#endif
} else if (arg == "--queue-depth" && i + 1 < argc) {
config.uring_queue_depth = std::stoul(argv[++i]);
} else if (arg == "--uring-direct-io") {
#ifdef USE_URING
config.use_uring_direct_io = true;
if (!config.use_uring) {
std::cerr << "Note: --uring-direct-io requires --use-uring, "
"enabling io_uring"
<< std::endl;
config.use_uring = true;
}
#else
std::cerr << "Warning: io_uring support not compiled in"
<< std::endl;
#endif
} else if (arg == "--uring-direct-io-zerocopy") {
#ifdef USE_URING
config.use_uring_direct_io_zero_copy = true;
if (!config.use_uring) {
std::cerr << "Note: --uring-direct-io-zerocopy requires "
"--use-uring, enabling io_uring"
<< std::endl;
config.use_uring = true;
}
#else
std::cerr << "Warning: io_uring support not compiled in"
<< std::endl;
#endif
} else if (arg == "--use-registered-buffers") {
#ifdef USE_URING
config.use_registered_buffers = true;
if (!config.use_uring) {
std::cerr << "Note: --use-registered-buffers requires "
"--use-uring, enabling io_uring"
<< std::endl;
config.use_uring = true;
}
if (!config.use_uring_direct_io_zero_copy) {
std::cerr << "Note: --use-registered-buffers works best with "
"--uring-direct-io-zerocopy"
<< std::endl;
}
#else
std::cerr << "Warning: io_uring support not compiled in"
<< std::endl;
#endif
} else if (arg == "--direct-io") {
config.use_direct_io = true;
} else if (arg == "--iterations" && i + 1 < argc) {
config.iterations = std::stoi(argv[++i]);
} else if (arg == "--no-verify") {
config.verify_data = false;
} else {
std::cerr << "Unknown argument: " << arg << std::endl;
PrintUsage(argv[0]);
return 1;
}
}
FileInterfaceBenchmark benchmark(config);
benchmark.Run();
return 0;
}

61
mooncake-store/include/aligned_client_buffer.hpp Normal file
View File

@ -0,0 +1,61 @@
#pragma once
#include "client_buffer.hpp"
namespace mooncake {
/**
* AlignedClientBufferAllocator extends ClientBufferAllocator to provide
* 4096-byte aligned memory allocation, which is required for O_DIRECT I/O
* operations with io_uring.
*
* This allocator ensures:
* - Base memory address is aligned to 4096 bytes (page size)
* - Can be registered with UringFile for zero-copy I/O
* - Maintains all features of the parent ClientBufferAllocator
*/
class AlignedClientBufferAllocator : public ClientBufferAllocator {
public:
// Alignment requirement for O_DIRECT I/O
static constexpr size_t kDirectIOAlignment = 4096;
/**
* Create an AlignedClientBufferAllocator with aligned memory
* @param size Total size of the buffer to allocate
* @param protocol Optional protocol string (unused, for compatibility)
* @param use_hugepage Whether to use huge pages for allocation
* @return Shared pointer to the allocator, or nullptr on failure
*/
static std::shared_ptr<AlignedClientBufferAllocator> create(
size_t size, const std::string& protocol = "",
bool use_hugepage = false);
/**
* Get the base pointer of the aligned buffer
* @return Base address of the allocated buffer
*/
[[nodiscard]] void* get_base_pointer() const { return getBase(); }
/**
* Get the total size of the aligned buffer
* @return Total size in bytes
*/
[[nodiscard]] size_t get_total_size() const { return size(); }
/**
* Destructor - properly frees the aligned memory
*/
~AlignedClientBufferAllocator();
private:
// Private constructor - use create() factory method
AlignedClientBufferAllocator(void* aligned_buffer, size_t size,
const std::string& protocol,
bool use_hugepage);
// Store whether we own the memory (for cleanup)
bool owns_memory_;
size_t allocated_size_; // Store the actual allocated size for cleanup
};
} // namespace mooncake

11
mooncake-store/include/client_buffer.hpp
View File

@ -51,19 +51,22 @@ class ClientBufferAllocator
[[nodiscard]] std::optional<BufferHandle> allocate(size_t size);
protected:
// Constructors accessible to derived classes
ClientBufferAllocator(void* addr, size_t size, const std::string& protocol);
void* buffer_;
bool use_hugepage_ = false;
private:
// Private constructors for different memory types
ClientBufferAllocator(size_t size, const std::string& protocol,
bool use_hugepage);
ClientBufferAllocator(void* addr, size_t size, const std::string& protocol);
std::shared_ptr<offset_allocator::OffsetAllocator> allocator_;
std::string protocol;
void* buffer_;
size_t buffer_size_;
bool is_external_memory_ = false;
bool use_hugepage_ = false;
};
/**

83
mooncake-store/include/file_interface.h
View File

@ -5,9 +5,13 @@
#include <sys/uio.h>
#include <cstdio>
#include "types.h"
#include "mutex.h"
#include <atomic>
#include <thread>
#include <sys/file.h>
#ifdef USE_URING
#include <liburing.h>
#endif
namespace mooncake {
class FileLockRAII {
@ -169,6 +173,85 @@ class PosixFile : public StorageFile {
off_t offset) override;
};
#ifdef USE_URING
class UringFile : public StorageFile {
public:
UringFile(const std::string &filename, int fd, unsigned queue_depth = 32,
bool use_direct_io = false);
~UringFile() override;
tl::expected<size_t, ErrorCode> write(const std::string &buffer,
size_t length) override;
tl::expected<size_t, ErrorCode> write(std::span<const char> data,
size_t length) override;
tl::expected<size_t, ErrorCode> read(std::string &buffer,
size_t length) override;
tl::expected<size_t, ErrorCode> vector_write(const iovec *iov, int iovcnt,
off_t offset) override;
tl::expected<size_t, ErrorCode> vector_read(const iovec *iov, int iovcnt,
off_t offset) override;
// Zero-copy interface for O_DIRECT: caller must provide aligned buffer
tl::expected<size_t, ErrorCode> read_aligned(void *buffer, size_t length,
off_t offset = 0);
tl::expected<size_t, ErrorCode> write_aligned(const void *buffer,
size_t length,
off_t offset = 0);
// Flush data to stable storage via
// io_uring_prep_fsync(IORING_FSYNC_DATASYNC). Must be called after write
// and before writing dependent metadata files.
tl::expected<void, ErrorCode> datasync();
// Buffer registration interface for high-performance I/O
// Register a single buffer with io_uring to avoid get_user_pages() overhead
// Returns true on success, false on failure
bool register_buffer(void *buffer, size_t length);
// Unregister previously registered buffer
void unregister_buffer();
// Check if a buffer is currently registered
bool is_buffer_registered() const { return buffer_registered_; }
private:
struct io_uring ring_;
bool ring_initialized_;
bool files_registered_;
bool buffer_registered_;
unsigned queue_depth_;
bool use_direct_io_;
static constexpr size_t ALIGNMENT_ =
4096; // O_DIRECT alignment requirement
// Registered buffer info
void *registered_buffer_;
size_t registered_buffer_size_;
struct iovec registered_iovec_;
/// Submit all pending SQEs and wait for exactly @p n completions.
/// Returns the total bytes transferred, or an error.
tl::expected<size_t, ErrorCode> submit_and_wait_n(int n);
/// Calculate optimal chunk size for parallel I/O based on:
/// - total_len: remaining bytes to transfer
/// - available_depth: number of queue slots available
/// - min_chunk_size: minimum chunk size (must be power of 2)
/// Returns a power-of-2 chunk size that maximizes queue utilization.
size_t calculate_chunk_size(size_t total_len, unsigned available_depth,
size_t min_chunk_size) const;
/// Allocate aligned buffer for O_DIRECT
void *alloc_aligned_buffer(size_t size) const;
/// Free aligned buffer
void free_aligned_buffer(void *ptr) const;
/// Mutex to serialize concurrent access to ring_
mutable Mutex ring_mutex_;
};
#endif // USE_URING
} // namespace mooncake
#ifdef USE_3FS

2
mooncake-store/include/file_storage.h
View File

@ -67,7 +67,7 @@ class FileStorage {
tl::expected<bool, ErrorCode> IsEnableOffloading();
tl::expected<void, ErrorCode> BatchLoad(
const std::unordered_map<std::string, Slice>& batch_object);
std::unordered_map<std::string, Slice>& batch_object);
tl::expected<void, ErrorCode> BatchQuerySegmentSlices(
const std::vector<std::string>& keys,

52
mooncake-store/include/storage_backend.h
View File

@ -200,6 +200,9 @@ struct FileStorageConfig {
uint32_t client_buffer_gc_interval_seconds = 1;
uint64_t client_buffer_gc_ttl_ms = 5000;
// Use io_uring for file I/O instead of POSIX pread/pwrite
bool use_uring = false;
// Validates the configuration for correctness and consistency
bool Validate() const;
@ -230,7 +233,7 @@ class StorageBackendInterface {
complete_handler) = 0;
virtual tl::expected<void, ErrorCode> BatchLoad(
const std::unordered_map<std::string, Slice>& batched_slices) = 0;
std::unordered_map<std::string, Slice>& batched_slices) = 0;
virtual tl::expected<bool, ErrorCode> IsExist(const std::string& key) = 0;
@ -438,6 +441,7 @@ class StorageBackend {
std::string fsdir_;
bool enable_eviction_{
true}; // User-configurable flag to enable/disable eviction
bool use_uring_{false}; // Use io_uring for file I/O
#ifdef USE_3FS
bool is_3fs_dir_{false}; // Flag to indicate if the storage is using 3FS
@ -614,7 +618,7 @@ class StorageBackendAdaptor : public StorageBackendInterface {
complete_handler) override;
tl::expected<void, ErrorCode> BatchLoad(
const std::unordered_map<std::string, Slice>& batched_slices) override;
std::unordered_map<std::string, Slice>& batched_slices) override;
tl::expected<bool, ErrorCode> IsExist(const std::string& key) override;
@ -674,6 +678,8 @@ class BucketStorageBackend : public StorageBackendInterface {
BucketStorageBackend(const FileStorageConfig& file_storage_config_,
const BucketBackendConfig& bucket_backend_config_);
~BucketStorageBackend();
/**
* @brief Offload objects in batches
* @param batch_object A map from object key to a list of data slices to be
@ -708,7 +714,7 @@ class BucketStorageBackend : public StorageBackendInterface {
* @return tl::expected<void, ErrorCode> indicating operation status.
*/
tl::expected<void, ErrorCode> BatchLoad(
const std::unordered_map<std::string, Slice>& batched_slices) override;
std::unordered_map<std::string, Slice>& batched_slices) override;
/**
* @brief Retrieves the list of object keys belonging to a specific bucket.
@ -858,11 +864,37 @@ class BucketStorageBackend : public StorageBackendInterface {
*/
void CleanupOrphanedBucket(int64_t bucket_id);
public:
/**
* @brief Get a file instance for external buffer registration
* Opens a temporary file to get access to the UringFile instance
* @return Shared pointer to StorageFile or error
*/
tl::expected<std::shared_ptr<StorageFile>, ErrorCode> GetFileInstance()
const;
private:
// Alignment helper functions for O_DIRECT I/O
static constexpr size_t kDirectIOAlignment = 4096;
static inline size_t align_up(size_t size, size_t alignment) {
return (size + alignment - 1) & ~(alignment - 1);
}
static inline int64_t align_down(int64_t offset, int64_t alignment) {
return offset & ~(alignment - 1);
}
std::atomic<bool> initialized_{false};
std::optional<BucketIdGenerator> bucket_id_generator_;
static constexpr const char* BUCKET_DATA_FILE_SUFFIX = ".bucket";
static constexpr const char* BUCKET_METADATA_FILE_SUFFIX = ".meta";
// Aligned buffer for O_DIRECT I/O operations
// We use a fixed-size buffer to avoid frequent allocations
static constexpr size_t kAlignedBufferSize = 16 * 1024 * 1024; // 16MB
std::unique_ptr<void, void (*)(void*)> aligned_io_buffer_{nullptr,
[](void*) {}};
/**
* @brief A shared mutex to protect concurrent access to metadata.
*
@ -886,6 +918,18 @@ class BucketStorageBackend : public StorageBackendInterface {
mutable Mutex offloading_mutex_;
std::unordered_map<std::string, int64_t> GUARDED_BY(offloading_mutex_)
ungrouped_offloading_objects_;
// File handle cache for UringFile to avoid repeated open/close overhead
mutable Mutex file_cache_mutex_;
mutable std::unordered_map<std::string, std::shared_ptr<StorageFile>>
file_cache_ GUARDED_BY(file_cache_mutex_);
// Get or open a file with caching support
tl::expected<std::shared_ptr<StorageFile>, ErrorCode> GetOrOpenFile(
const std::string& path, FileMode mode) const;
// Clear file cache (called on destruction or when needed)
void ClearFileCache();
};
class OffsetAllocatorStorageBackend : public StorageBackendInterface {
@ -922,7 +966,7 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface {
* @return tl::expected<void, ErrorCode> indicating operation status.
*/
tl::expected<void, ErrorCode> BatchLoad(
const std::unordered_map<std::string, Slice>& batched_slices) override;
std::unordered_map<std::string, Slice>& batched_slices) override;
/**
* @brief Checks whether an object with the specified key exists in the

17
mooncake-store/src/CMakeLists.txt
View File

@ -22,6 +22,7 @@ set(MOONCAKE_STORE_SOURCES
offset_allocator.cpp
posix_file.cpp
client_buffer.cpp
aligned_client_buffer.cpp
real_client.cpp
dummy_client.cpp
http_metadata_server.cpp
@ -42,6 +43,17 @@ if(USE_3FS)
set(EXTRA_LIBS ${HF3FS_API_LIB})
endif()
# io_uring support (auto-detected)
find_library(URING_LIB uring PATHS /usr/lib /usr/lib64 /usr/local/lib /usr/local/lib64)
find_path(URING_INCLUDE liburing.h PATHS /usr/include /usr/local/include)
if(URING_LIB AND URING_INCLUDE)
message(STATUS "io_uring: Enabled for mooncake_store")
list(APPEND MOONCAKE_STORE_SOURCES uring_file.cpp)
list(APPEND EXTRA_LIBS ${URING_LIB})
else()
message(STATUS "io_uring: Disabled (liburing not found)")
endif()
# The cache_allocator library
include_directories(${Python3_INCLUDE_DIRS})
add_library(mooncake_store ${MOONCAKE_STORE_SOURCES})
@ -63,6 +75,11 @@ if (STORE_USE_ETCD)
add_dependencies(mooncake_store build_etcd_wrapper)
endif()
if(URING_LIB AND URING_INCLUDE)
target_compile_definitions(mooncake_store PUBLIC USE_URING)
target_include_directories(mooncake_store PRIVATE ${URING_INCLUDE})
endif()
if (USE_ASCEND_DIRECT)
set(ACL_RUNTIME_HEADER_PATH "${ASCEND_INCLUDE_DIR}/acl/acl_rt.h")
if(EXISTS "${ACL_RUNTIME_HEADER_PATH}")

102
mooncake-store/src/aligned_client_buffer.cpp Normal file
View File

@ -0,0 +1,102 @@
#include "aligned_client_buffer.hpp"
#include <glog/logging.h>
#include <sys/mman.h>
#include <cstdlib>
#include <cstring>
#include "utils.h"
namespace mooncake {
std::shared_ptr<AlignedClientBufferAllocator>
AlignedClientBufferAllocator::create(size_t size, const std::string& protocol,
bool use_hugepage) {
if (size == 0) {
LOG(ERROR)
<< "AlignedClientBufferAllocator: size must be greater than 0";
return nullptr;
}
// Align size up to kDirectIOAlignment
size_t aligned_size = align_up(size, kDirectIOAlignment);
void* aligned_buffer = nullptr;
if (use_hugepage) {
// Use hugepage allocation (already aligned)
aligned_buffer =
allocate_buffer_mmap_memory(aligned_size, kDirectIOAlignment);
if (!aligned_buffer) {
LOG(ERROR) << "AlignedClientBufferAllocator: failed to allocate "
<< "hugepage memory of size " << aligned_size;
return nullptr;
}
} else {
// Use posix_memalign for 4096-byte alignment
int ret =
posix_memalign(&aligned_buffer, kDirectIOAlignment, aligned_size);
if (ret != 0) {
LOG(ERROR) << "AlignedClientBufferAllocator: posix_memalign failed "
<< "with error " << ret << " (" << strerror(ret) << ")";
return nullptr;
}
// Zero-initialize the allocated memory
memset(aligned_buffer, 0, aligned_size);
}
// Verify alignment
if (reinterpret_cast<uintptr_t>(aligned_buffer) % kDirectIOAlignment != 0) {
LOG(ERROR) << "AlignedClientBufferAllocator: allocated buffer is not "
<< "aligned to " << kDirectIOAlignment << " bytes";
if (use_hugepage) {
free_buffer_mmap_memory(aligned_buffer, aligned_size);
} else {
free(aligned_buffer);
}
return nullptr;
}
LOG(INFO) << "AlignedClientBufferAllocator: allocated " << aligned_size
<< " bytes at address " << aligned_buffer << " (aligned to "
<< kDirectIOAlignment << " bytes)";
// Use custom deleter to properly free the aligned memory
return std::shared_ptr<AlignedClientBufferAllocator>(
new AlignedClientBufferAllocator(aligned_buffer, aligned_size, protocol,
use_hugepage));
}
AlignedClientBufferAllocator::AlignedClientBufferAllocator(
void* aligned_buffer, size_t size, const std::string& protocol,
bool use_hugepage)
: ClientBufferAllocator(aligned_buffer, size, protocol),
owns_memory_(true),
allocated_size_(size) {
// Store hugepage flag in parent class's use_hugepage_ member
// We need this for proper cleanup
use_hugepage_ = use_hugepage;
}
AlignedClientBufferAllocator::~AlignedClientBufferAllocator() {
// Free the aligned memory we allocated
// The parent class destructor will not free it because is_external_memory_
// is set to true
if (owns_memory_ && buffer_) {
if (use_hugepage_) {
LOG(INFO)
<< "AlignedClientBufferAllocator: freeing hugepage memory "
<< "at " << buffer_ << " (" << allocated_size_ << " bytes)";
free_buffer_mmap_memory(buffer_, allocated_size_);
} else {
LOG(INFO) << "AlignedClientBufferAllocator: freeing aligned memory "
<< "at " << buffer_ << " (" << allocated_size_
<< " bytes)";
free(buffer_);
}
buffer_ = nullptr;
}
}
} // namespace mooncake

102
mooncake-store/src/file_storage.cpp
View File

@ -3,8 +3,13 @@
#include <memory>
#include <vector>
#include "aligned_client_buffer.hpp"
#include "storage_backend.h"
#include "utils.h"
#ifdef USE_URING
#include "file_interface.h"
#endif
namespace mooncake {
FileStorageConfig FileStorageConfig::FromEnvironment() {
@ -52,6 +57,9 @@ FileStorageConfig FileStorageConfig::FromEnvironment() {
GetEnvOr<uint64_t>("MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_TTL_MS",
config.client_buffer_gc_ttl_ms);
auto use_uring_str = GetEnvStringOr("MOONCAKE_USE_URING", "false");
config.use_uring = (use_uring_str == "true" || use_uring_str == "1");
return config;
}
@ -142,7 +150,7 @@ FileStorage::FileStorage(const FileStorageConfig& config,
client_(client),
local_rpc_addr_(local_rpc_addr),
client_buffer_allocator_(
ClientBufferAllocator::create(config.local_buffer_size, "")) {
AlignedClientBufferAllocator::create(config.local_buffer_size, "")) {
if (!config.Validate()) {
throw std::invalid_argument("Invalid FileStorage configuration");
}
@ -150,9 +158,48 @@ FileStorage::FileStorage(const FileStorageConfig& config,
auto create_storage_backend_result = CreateStorageBackend(config_);
if (!create_storage_backend_result) {
LOG(ERROR) << "Failed to create storage backend";
throw std::runtime_error("Failed to create storage backend");
}
storage_backend_ = create_storage_backend_result.value();
// Register buffer with UringFile if using BucketStorageBackend
#ifdef USE_URING
if (config.storage_backend_type == StorageBackendType::kBucket) {
auto bucket_backend =
std::dynamic_pointer_cast<BucketStorageBackend>(storage_backend_);
if (bucket_backend) {
auto file_result = bucket_backend->GetFileInstance();
if (file_result) {
auto file = file_result.value();
auto uring_file = std::dynamic_pointer_cast<UringFile>(file);
if (uring_file) {
auto aligned_allocator =
std::static_pointer_cast<AlignedClientBufferAllocator>(
client_buffer_allocator_);
if (aligned_allocator) {
void* base_ptr = aligned_allocator->get_base_pointer();
size_t size = aligned_allocator->get_total_size();
if (uring_file->register_buffer(base_ptr, size)) {
LOG(INFO)
<< "Successfully registered buffer with "
"UringFile: "
<< "base=" << base_ptr << ", size=" << size;
} else {
LOG(WARNING)
<< "Failed to register buffer with UringFile";
}
}
}
} else {
LOG(WARNING)
<< "Failed to get file instance for buffer registration: "
<< file_result.error();
}
}
}
#endif
}
FileStorage::~FileStorage() {
@ -251,6 +298,18 @@ tl::expected<std::vector<uint64_t>, ErrorCode> FileStorage::BatchGet(
LOG(ERROR) << "Batch load object failed,err_code = " << result.error();
return tl::make_unexpected(result.error());
}
// After BatchLoad, slice.ptr may have been adjusted by offset_in_buffer
// (for O_DIRECT aligned reads). Update pointers to reflect actual data
// positions.
for (size_t i = 0; i < keys.size(); ++i) {
auto it = allocated_batch->slices.find(keys[i]);
if (it != allocated_batch->slices.end()) {
allocated_batch->pointers[i] =
reinterpret_cast<uintptr_t>(it->second.ptr);
}
}
MutexLocker locker(&client_buffer_mutex_);
client_buffer_allocated_batches_.emplace_back(std::move(allocated_batch));
auto end_time = std::chrono::steady_clock::now();
@ -373,7 +432,7 @@ tl::expected<void, ErrorCode> FileStorage::Heartbeat() {
}
tl::expected<void, ErrorCode> FileStorage::BatchLoad(
const std::unordered_map<std::string, Slice>& batch_object) {
std::unordered_map<std::string, Slice>& batch_object) {
auto start_time = std::chrono::steady_clock::now();
auto result = storage_backend_->BatchLoad(batch_object);
auto end_time = std::chrono::steady_clock::now();
@ -447,11 +506,22 @@ FileStorage::AllocateBatch(const std::vector<std::string>& keys,
std::chrono::steady_clock::now();
auto lease_timeout =
now + std::chrono::milliseconds(config_.client_buffer_gc_ttl_ms);
static constexpr size_t kDirectIOAlignment = 4096;
u_int64_t total_size = 0;
bool gc_triggered = false;
for (size_t i = 0; i < keys.size(); ++i) {
assert(sizes[i] <= kMaxSliceSize);
auto alloc_result = client_buffer_allocator_->allocate(sizes[i]);
// Allocate oversized buffer for O_DIRECT alignment:
// +4096 for aligning the ptr to 4096 boundary
// +4096 for aligned read tail padding (actual_offset may not be
// aligned)
size_t data_size = static_cast<size_t>(sizes[i]);
size_t alloc_size =
align_up(data_size, kDirectIOAlignment) + 2 * kDirectIOAlignment;
auto alloc_result = client_buffer_allocator_->allocate(alloc_size);
if (!alloc_result && !gc_triggered) {
gc_triggered = true;
{
@ -466,18 +536,28 @@ FileStorage::AllocateBatch(const std::vector<std::string>& keys,
}
}
}
alloc_result = client_buffer_allocator_->allocate(sizes[i]);
alloc_result = client_buffer_allocator_->allocate(alloc_size);
}
if (!alloc_result) {
LOG(ERROR) << "Failed to allocate slice buffer, size = " << sizes[i]
<< ", key = " << keys[i];
LOG(ERROR) << "Failed to allocate slice buffer, size = "
<< alloc_size << " (data_size=" << data_size
<< "), key = " << keys[i];
return tl::make_unexpected(ErrorCode::BUFFER_OVERFLOW);
}
total_size += sizes[i];
result->slices.emplace(
keys[i], Slice{alloc_result->ptr(), static_cast<size_t>(sizes[i])});
result->pointers.emplace_back(
reinterpret_cast<uintptr_t>(alloc_result->ptr()));
// Align ptr to 4096 boundary for O_DIRECT
void* raw_ptr = alloc_result->ptr();
void* aligned_ptr = reinterpret_cast<void*>(
(reinterpret_cast<uintptr_t>(raw_ptr) + kDirectIOAlignment - 1) &
~(kDirectIOAlignment - 1));
total_size += data_size;
// Slice records data_size; the buffer behind aligned_ptr is oversized
// to accommodate aligned reads
result->slices.emplace(keys[i], Slice{aligned_ptr, data_size});
// pointers will be adjusted after BatchLoad (offset_in_buffer
// correction)
result->pointers.emplace_back(reinterpret_cast<uintptr_t>(aligned_ptr));
result->handles.emplace_back(std::move(alloc_result.value()));
result->lease_timeout = lease_timeout;
}

278
mooncake-store/src/storage_backend.cpp
View File

@ -767,6 +767,11 @@ std::unique_ptr<StorageFile> StorageBackend::create_file(
}
#endif
#ifdef USE_URING
if (use_uring_) {
return std::make_unique<UringFile>(path, fd, 32, true);
}
#endif
return std::make_unique<PosixFile>(path, fd);
}
@ -932,6 +937,7 @@ tl::expected<void, ErrorCode> StorageBackendAdaptor::Init() {
storage_backend_ = std::make_unique<StorageBackend>(
file_storage_config_.storage_filepath, file_per_key_config_.fsdir,
file_per_key_config_.enable_eviction);
storage_backend_->use_uring_ = file_storage_config_.use_uring;
auto init_result = storage_backend_->Init();
if (!init_result) {
LOG(ERROR) << "Failed to init storage backend";
@ -1078,7 +1084,7 @@ tl::expected<bool, ErrorCode> StorageBackendAdaptor::IsExist(
}
tl::expected<void, ErrorCode> StorageBackendAdaptor::BatchLoad(
const std::unordered_map<std::string, Slice>& batched_slices) {
std::unordered_map<std::string, Slice>& batched_slices) {
for (const auto& [key, slice] : batched_slices) {
KVEntry kv;
kv.key = key;
@ -1231,7 +1237,29 @@ BucketStorageBackend::BucketStorageBackend(
const BucketBackendConfig& bucket_backend_config_)
: StorageBackendInterface(file_storage_config_),
storage_path_(file_storage_config_.storage_filepath),
bucket_backend_config_(bucket_backend_config_) {}
bucket_backend_config_(bucket_backend_config_) {
// Allocate aligned buffer for O_DIRECT I/O operations
void* buf = nullptr;
int ret = posix_memalign(&buf, kDirectIOAlignment, kAlignedBufferSize);
if (ret != 0) {
LOG(ERROR)
<< "BucketStorageBackend: Failed to allocate aligned buffer: "
<< strerror(ret);
} else {
aligned_io_buffer_.reset(buf);
// Update the deleter to use free
aligned_io_buffer_ = std::unique_ptr<void, void (*)(void*)>(
buf, [](void* p) { free(p); });
LOG(INFO) << "BucketStorageBackend: Allocated " << kAlignedBufferSize
<< " bytes aligned buffer at " << buf;
}
}
BucketStorageBackend::~BucketStorageBackend() {
// Clear file cache to release UringFile instances before destruction
// This ensures orderly cleanup of io_uring resources
ClearFileCache();
}
tl::expected<int64_t, ErrorCode> BucketStorageBackend::BatchOffload(
const std::unordered_map<std::string, std::vector<Slice>>& batch_object,
@ -1331,7 +1359,7 @@ tl::expected<void, ErrorCode> BucketStorageBackend::BatchQuery(
}
tl::expected<void, ErrorCode> BucketStorageBackend::BatchLoad(
const std::unordered_map<std::string, Slice>& batch_object) {
std::unordered_map<std::string, Slice>& batch_object) {
// Step 1: Build read plan by copying metadata under lock
// BucketReadGuard increments inflight_reads_ to prevent deletion during IO.
// When the guard goes out of scope, it decrements the counter.
@ -1413,10 +1441,44 @@ tl::expected<void, ErrorCode> BucketStorageBackend::BatchLoad(
// Read each key's data
for (const auto& plan : read_plans) {
// Read value (skip key in file: offset + key_size)
iovec iov{plan.dest_slice.ptr, plan.dest_slice.size};
auto read_res =
file->vector_read(&iov, 1, plan.offset + plan.key_size);
int64_t actual_offset = plan.offset + plan.key_size;
tl::expected<size_t, ErrorCode> read_res;
#ifdef USE_URING
// Try to use read_aligned for O_DIRECT I/O if file is UringFile
UringFile* uring_file = dynamic_cast<UringFile*>(file.get());
if (uring_file != nullptr) {
// Calculate aligned read range
int64_t aligned_offset =
align_down(actual_offset, kDirectIOAlignment);
int64_t data_end =
actual_offset + static_cast<int64_t>(plan.dest_slice.size);
int64_t aligned_end = static_cast<int64_t>(align_up(
static_cast<size_t>(data_end), kDirectIOAlignment));
size_t aligned_size =
static_cast<size_t>(aligned_end - aligned_offset);
int64_t offset_in_buffer = actual_offset - aligned_offset;
// Zero-copy path: read directly into the slice buffer.
// dest_slice.ptr is 4096-aligned and oversized (from
// AllocateBatch) to accommodate the full aligned read range.
read_res = uring_file->read_aligned(
plan.dest_slice.ptr, aligned_size, aligned_offset);
if (read_res) {
// Adjust ptr to point to actual data start (no memcpy)
batch_object.at(plan.key).ptr =
static_cast<char*>(plan.dest_slice.ptr) +
offset_in_buffer;
read_res = plan.dest_slice.size;
}
} else
#endif
{
// Fallback to vector_read for non-UringFile
iovec iov{plan.dest_slice.ptr, plan.dest_slice.size};
read_res = file->vector_read(&iov, 1, actual_offset);
}
if (!read_res) {
LOG(ERROR) << "vector_read failed for key: " << plan.key
@ -1873,18 +1935,105 @@ tl::expected<void, ErrorCode> BucketStorageBackend::WriteBucket(
}
auto file = std::move(open_file_result.value());
auto write_result = file->vector_write(iovs.data(), iovs.size(), 0);
if (!write_result) {
LOG(ERROR) << "vector_write failed for: " << bucket_id
<< ", error: " << write_result.error();
return tl::make_unexpected(write_result.error());
}
if (static_cast<int64_t>(write_result.value()) !=
bucket_metadata->data_size) {
LOG(ERROR) << "Write size mismatch for: " << bucket_data_path
<< ", expected: " << bucket_metadata->data_size
<< ", got: " << write_result.value();
return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL);
#ifdef USE_URING
// Try to use write_aligned for O_DIRECT I/O if file is UringFile
UringFile* uring_file = dynamic_cast<UringFile*>(file.get());
if (uring_file != nullptr) {
size_t total_size = static_cast<size_t>(bucket_metadata->data_size);
size_t aligned_size = align_up(total_size, kDirectIOAlignment);
// Allocate aligned buffer if needed
void* write_buffer = nullptr;
std::unique_ptr<void, void (*)(void*)> temp_buffer{nullptr,
[](void*) {}};
if (aligned_size <= kAlignedBufferSize && aligned_io_buffer_) {
// Use the pre-allocated buffer
write_buffer = aligned_io_buffer_.get();
} else {
// Allocate a temporary larger buffer
void* buf = nullptr;
int ret = posix_memalign(&buf, kDirectIOAlignment, aligned_size);
if (ret != 0) {
LOG(ERROR)
<< "Failed to allocate aligned buffer for WriteBucket: "
<< strerror(ret);
return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
}
temp_buffer.reset(buf);
temp_buffer = std::unique_ptr<void, void (*)(void*)>(
buf, [](void* p) { free(p); });
write_buffer = buf;
LOG(WARNING) << "WriteBucket: bucket_id=" << bucket_id
<< " requires " << aligned_size
<< " bytes, exceeds buffer size " << kAlignedBufferSize
<< ", using temporary allocation";
}
// Aggregate all iovs data into the aligned buffer
char* dst = static_cast<char*>(write_buffer);
for (const auto& iov : iovs) {
memcpy(dst, iov.iov_base, iov.iov_len);
dst += iov.iov_len;
}
// Zero-pad the remaining bytes
if (aligned_size > total_size) {
memset(dst, 0, aligned_size - total_size);
}
// Write using write_aligned
auto write_result =
uring_file->write_aligned(write_buffer, aligned_size, 0);
if (!write_result) {
LOG(ERROR) << "write_aligned failed for: " << bucket_id
<< ", error: " << write_result.error();
return tl::make_unexpected(write_result.error());
}
if (write_result.value() != aligned_size) {
LOG(ERROR) << "Write size mismatch for: " << bucket_data_path
<< ", expected: " << aligned_size
<< ", got: " << write_result.value();
return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL);
}
// Flush bucket data to stable storage before writing metadata.
// This prevents a crash from leaving valid metadata pointing at
// incomplete data (write-ordering durability guarantee).
auto sync_result = uring_file->datasync();
if (!sync_result) {
LOG(ERROR) << "datasync failed for bucket: " << bucket_id;
return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL);
}
// Invalidate cache for this file since content changed
{
MutexLocker cache_locker(&file_cache_mutex_);
file_cache_.erase(bucket_data_path);
}
} else
#endif
{
// Fallback to vector_write for non-UringFile
auto write_result = file->vector_write(iovs.data(), iovs.size(), 0);
if (!write_result) {
LOG(ERROR) << "vector_write failed for: " << bucket_id
<< ", error: " << write_result.error();
return tl::make_unexpected(write_result.error());
}
if (static_cast<int64_t>(write_result.value()) !=
bucket_metadata->data_size) {
LOG(ERROR) << "Write size mismatch for: " << bucket_data_path
<< ", expected: " << bucket_metadata->data_size
<< ", got: " << write_result.value();
return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL);
}
// Invalidate cache for this file since content changed
{
MutexLocker cache_locker(&file_cache_mutex_);
file_cache_.erase(bucket_data_path);
}
}
auto store_bucket_metadata_result =
StoreBucketMetadata(bucket_id, bucket_metadata);
@ -2138,13 +2287,62 @@ BucketStorageBackend::OpenFile(const std::string& path, FileMode mode) const {
break;
}
#ifdef USE_URING
// Add O_DIRECT flag when using uring for direct I/O
if (file_storage_config_.use_uring) {
flags |= O_DIRECT;
}
#endif
int fd = open(path.c_str(), flags | access_mode, 0644);
if (fd < 0) {
LOG(ERROR) << "Failed to open file: " << path << ", errno=" << errno
<< " (" << strerror(errno) << ")";
return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL);
}
#ifdef USE_URING
if (file_storage_config_.use_uring) {
return std::make_unique<UringFile>(path, fd, 32, true);
}
#endif
return std::make_unique<PosixFile>(path, fd);
}
tl::expected<std::shared_ptr<StorageFile>, ErrorCode>
BucketStorageBackend::GetOrOpenFile(const std::string& path,
FileMode mode) const {
// Only cache read-mode files (write mode needs O_TRUNC which invalidates
// cache)
if (mode == FileMode::Read) {
MutexLocker locker(&file_cache_mutex_);
auto it = file_cache_.find(path);
if (it != file_cache_.end()) {
return it->second;
}
}
// Open new file
auto result = OpenFile(path, mode);
if (!result) {
return tl::make_unexpected(result.error());
}
auto file = std::shared_ptr<StorageFile>(std::move(result.value()));
// Cache read-mode files
if (mode == FileMode::Read) {
MutexLocker locker(&file_cache_mutex_);
file_cache_[path] = file;
}
return file;
}
void BucketStorageBackend::ClearFileCache() {
MutexLocker locker(&file_cache_mutex_);
file_cache_.clear();
}
tl::expected<void, ErrorCode> BucketStorageBackend::HandleNext(
const std::function<
ErrorCode(const std::vector<std::string>& keys,
@ -2175,6 +2373,32 @@ tl::expected<bool, ErrorCode> BucketStorageBackend::HasNext() {
return next_bucket_ != 0;
}
tl::expected<std::shared_ptr<StorageFile>, ErrorCode>
BucketStorageBackend::GetFileInstance() const {
// Create a temporary file to get access to the file instance
// This is used for external buffer registration with UringFile
namespace fs = std::filesystem;
std::string temp_path =
(fs::path(storage_path_) / "temp_for_registration").string();
auto open_result = OpenFile(temp_path, FileMode::Write);
if (!open_result) {
LOG(ERROR) << "Failed to open temporary file for GetFileInstance: "
<< temp_path;
return tl::make_unexpected(open_result.error());
}
auto file = std::move(open_result.value());
// Remove the temporary file from disk now that the fd is open.
// The fd remains valid (Unix semantics) until the StorageFile is destroyed.
fs::remove(temp_path);
// Convert unique_ptr to shared_ptr
return std::shared_ptr<StorageFile>(std::move(file));
}
// ============================================================================
// OffsetAllocatorStorageBackend Implementation
// ============================================================================
@ -2264,9 +2488,17 @@ tl::expected<void, ErrorCode> OffsetAllocatorStorageBackend::Init() {
}
}
// Release fd to PosixFile (PosixFile takes ownership and will close it)
data_file_ =
std::make_unique<PosixFile>(data_file_path_, fd_guard.release());
// Release fd to StorageFile (takes ownership and will close it)
#ifdef USE_URING
if (file_storage_config_.use_uring) {
data_file_ = std::make_unique<UringFile>(
data_file_path_, fd_guard.release(), 32, true);
} else
#endif
{
data_file_ = std::make_unique<PosixFile>(data_file_path_,
fd_guard.release());
}
// Create allocator with base=0, size=capacity
allocator_ = offset_allocator::OffsetAllocator::create(0, capacity_);
@ -2467,7 +2699,7 @@ tl::expected<int64_t, ErrorCode> OffsetAllocatorStorageBackend::BatchOffload(
//-----------------------------------------------------------------------------
tl::expected<void, ErrorCode> OffsetAllocatorStorageBackend::BatchLoad(
const std::unordered_map<std::string, Slice>& batched_slices) {
std::unordered_map<std::string, Slice>& batched_slices) {
if (!initialized_.load(std::memory_order_acquire)) {
LOG(ERROR)
<< "Storage backend is not initialized. Call Init() before use.";

871
mooncake-store/src/uring_file.cpp Normal file
View File

@ -0,0 +1,871 @@
#ifdef USE_URING
#include <algorithm>
#include <cerrno>
#include <chrono>
#include <cstring>
#include <string>
#include <sys/uio.h>
#include <unistd.h>
#include <cmath>
#include <glog/logging.h>
#include <liburing.h>
#include "file_interface.h"
namespace mooncake {
// ---------------------------------------------------------------------------
// Construction / Destruction
// ---------------------------------------------------------------------------
UringFile::UringFile(const std::string &filename, int fd, unsigned queue_depth,
bool use_direct_io)
: StorageFile(filename, fd),
ring_initialized_(false),
files_registered_(false),
buffer_registered_(false),
queue_depth_(queue_depth),
use_direct_io_(use_direct_io),
registered_buffer_(nullptr),
registered_buffer_size_(0) {
if (fd < 0) {
error_code_ = ErrorCode::FILE_INVALID_HANDLE;
return;
}
int ret = io_uring_queue_init(queue_depth, &ring_, 0);
if (ret < 0) {
LOG(ERROR) << "Failed to initialize io_uring: " << strerror(-ret);
error_code_ = ErrorCode::FILE_INVALID_HANDLE;
return;
}
ring_initialized_ = true;
// Register the file descriptor to avoid per-I/O fd lookup overhead.
ret = io_uring_register_files(&ring_, &fd_, 1);
if (ret >= 0) {
files_registered_ = true;
} else {
LOG(WARNING) << "[UringFile] io_uring_register_files failed: "
<< strerror(-ret) << " (continuing without)";
}
if (use_direct_io_) {
LOG(INFO) << "[UringFile] O_DIRECT mode enabled for " << filename;
}
}
UringFile::~UringFile() {
auto dtor_start = std::chrono::steady_clock::now();
if (ring_initialized_) {
if (buffer_registered_) {
io_uring_unregister_buffers(&ring_);
}
if (files_registered_) {
io_uring_unregister_files(&ring_);
}
io_uring_queue_exit(&ring_);
}
if (fd_ >= 0) {
if (close(fd_) != 0) {
LOG(WARNING) << "Failed to close file: " << filename_;
}
if (error_code_ == ErrorCode::FILE_WRITE_FAIL) {
if (::unlink(filename_.c_str()) == -1) {
LOG(ERROR) << "Failed to delete corrupted file: " << filename_;
} else {
LOG(INFO) << "Deleted corrupted file: " << filename_;
}
}
}
fd_ = -1;
auto dtor_end = std::chrono::steady_clock::now();
auto dtor_elapsed_ms =
std::chrono::duration_cast<std::chrono::milliseconds>(dtor_end -
dtor_start)
.count();
if (dtor_elapsed_ms > 1) {
LOG(WARNING) << "[UringFile::~UringFile] cleanup took "
<< dtor_elapsed_ms << "ms for " << filename_;
}
}
// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------
namespace {
inline size_t next_power_of_2(size_t n) {
if (n == 0) return 1;
n--;
n |= n >> 1;
n |= n >> 2;
n |= n >> 4;
n |= n >> 8;
n |= n >> 16;
n |= n >> 32;
return n + 1;
}
} // namespace
void *UringFile::alloc_aligned_buffer(size_t size) const {
// Align size up to ALIGNMENT_
size_t aligned_size = ((size + ALIGNMENT_ - 1) / ALIGNMENT_) * ALIGNMENT_;
void *ptr = nullptr;
if (posix_memalign(&ptr, ALIGNMENT_, aligned_size) != 0) {
LOG(ERROR) << "[UringFile] Failed to allocate aligned buffer of size "
<< aligned_size;
return nullptr;
}
return ptr;
}
void UringFile::free_aligned_buffer(void *ptr) const {
if (ptr) {
free(ptr);
}
}
/**
* @brief Calculate optimal chunk size for parallel I/O
*
* Strategy:
* 1. Chunk size must be a power of 2
* 2. Chunk size = max(min_chunk_size, optimal_size)
* 3. optimal_size = the size that can fully utilize remaining queue depth
*
* @param total_len Total length to transfer
* @param available_depth Available queue depth slots
* @param min_chunk_size Minimum chunk size (must be power of 2)
* @return Optimal chunk size (power of 2)
*/
size_t UringFile::calculate_chunk_size(size_t total_len,
unsigned available_depth,
size_t min_chunk_size) const {
if (total_len == 0 || available_depth == 0) {
return min_chunk_size;
}
// Calculate the size that would fully utilize available queue depth
size_t optimal_size = (total_len + available_depth - 1) / available_depth;
// Round up to next power of 2
optimal_size = next_power_of_2(optimal_size);
// Return max(min_chunk_size, optimal_size)
return std::max(min_chunk_size, optimal_size);
}
// ---------------------------------------------------------------------------
// Internal helpers - I/O submission
// ---------------------------------------------------------------------------
tl::expected<size_t, ErrorCode> UringFile::submit_and_wait_n(int n) {
int ret = io_uring_submit_and_wait(&ring_, n);
if (ret < 0) {
LOG(ERROR) << "[UringFile] io_uring_submit_and_wait failed: "
<< strerror(-ret);
return make_error<size_t>(ErrorCode::INTERNAL_ERROR);
}
size_t total_bytes = 0;
bool has_error = false;
unsigned head;
unsigned count = 0;
struct io_uring_cqe *cqe;
io_uring_for_each_cqe(&ring_, head, cqe) {
if (cqe->res < 0) {
LOG(ERROR) << "[UringFile] I/O failed: " << strerror(-cqe->res);
has_error = true;
} else {
total_bytes += static_cast<size_t>(cqe->res);
}
count++;
}
io_uring_cq_advance(&ring_, count);
if (has_error) {
return make_error<size_t>(ErrorCode::INTERNAL_ERROR);
}
return total_bytes;
}
// ---------------------------------------------------------------------------
// Write implementation with intelligent chunking
// ---------------------------------------------------------------------------
tl::expected<size_t, ErrorCode> UringFile::write(const std::string &buffer,
size_t length) {
return write(std::span<const char>(buffer.data(), length), length);
}
tl::expected<size_t, ErrorCode> UringFile::write(std::span<const char> data,
size_t length) {
if (fd_ < 0 || !ring_initialized_) {
return make_error<size_t>(ErrorCode::FILE_NOT_FOUND);
}
if (length == 0) {
return make_error<size_t>(ErrorCode::FILE_INVALID_BUFFER);
}
constexpr size_t MIN_CHUNK_SIZE = 4096;
// If using O_DIRECT, allocate aligned buffer and copy data
void *aligned_buffer = nullptr;
const char *source_ptr = data.data();
size_t actual_length = length;
if (use_direct_io_) {
// Align length up to ALIGNMENT_
actual_length = ((length + ALIGNMENT_ - 1) / ALIGNMENT_) * ALIGNMENT_;
aligned_buffer = alloc_aligned_buffer(actual_length);
if (!aligned_buffer) {
return make_error<size_t>(ErrorCode::FILE_WRITE_FAIL);
}
// Copy data to aligned buffer and zero-pad if necessary
std::memcpy(aligned_buffer, data.data(), length);
if (actual_length > length) {
std::memset(static_cast<char *>(aligned_buffer) + length, 0,
actual_length - length);
}
source_ptr = static_cast<const char *>(aligned_buffer);
}
size_t total_written = 0;
const char *ptr = source_ptr;
size_t remaining = actual_length;
off_t current_offset = 0;
int target_fd = files_registered_ ? 0 : fd_;
while (remaining > 0) {
size_t chunk_size =
calculate_chunk_size(remaining, queue_depth_, MIN_CHUNK_SIZE);
unsigned num_chunks = std::min(
static_cast<unsigned>((remaining + chunk_size - 1) / chunk_size),
queue_depth_);
for (unsigned i = 0; i < num_chunks; i++) {
size_t this_chunk = std::min(chunk_size, remaining);
struct io_uring_sqe *sqe = io_uring_get_sqe(&ring_);
if (!sqe) {
LOG(ERROR) << "[UringFile::write] Failed to get SQE";
if (aligned_buffer) free_aligned_buffer(aligned_buffer);
return make_error<size_t>(ErrorCode::FILE_WRITE_FAIL);
}
io_uring_prep_write(sqe, target_fd, ptr, this_chunk,
current_offset);
if (files_registered_) {
io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE);
}
ptr += this_chunk;
current_offset += this_chunk;
remaining -= this_chunk;
if (remaining == 0) break;
}
auto result = submit_and_wait_n(num_chunks);
if (!result) {
if (aligned_buffer) free_aligned_buffer(aligned_buffer);
return make_error<size_t>(ErrorCode::FILE_WRITE_FAIL);
}
size_t bytes_written = result.value();
if (bytes_written == 0) {
LOG(ERROR) << "[UringFile::write] Zero bytes written";
if (aligned_buffer) free_aligned_buffer(aligned_buffer);
return make_error<size_t>(ErrorCode::FILE_WRITE_FAIL);
}
total_written += bytes_written;
}
if (aligned_buffer) {
free_aligned_buffer(aligned_buffer);
}
// For O_DIRECT, we may have written more than requested (due to alignment)
// but return the original length
if (total_written < length) {
LOG(WARNING) << "[UringFile::write] Incomplete write: " << total_written
<< " / " << length;
return make_error<size_t>(ErrorCode::FILE_WRITE_FAIL);
}
return length; // Return original length, not aligned length
}
// ---------------------------------------------------------------------------
// Read implementation with intelligent chunking
// ---------------------------------------------------------------------------
tl::expected<size_t, ErrorCode> UringFile::read(std::string &buffer,
size_t length) {
if (fd_ < 0 || !ring_initialized_) {
return make_error<size_t>(ErrorCode::FILE_NOT_FOUND);
}
if (length == 0) {
return make_error<size_t>(ErrorCode::FILE_INVALID_BUFFER);
}
constexpr size_t MIN_CHUNK_SIZE = 4096;
// If using O_DIRECT, allocate aligned buffer
void *aligned_buffer = nullptr;
char *read_ptr = nullptr;
size_t actual_length = length;
if (use_direct_io_) {
// Align length up to ALIGNMENT_
actual_length = ((length + ALIGNMENT_ - 1) / ALIGNMENT_) * ALIGNMENT_;
aligned_buffer = alloc_aligned_buffer(actual_length);
if (!aligned_buffer) {
return make_error<size_t>(ErrorCode::FILE_READ_FAIL);
}
read_ptr = static_cast<char *>(aligned_buffer);
} else {
buffer.resize(length);
read_ptr = buffer.data();
}
char *ptr = read_ptr;
size_t remaining = actual_length;
size_t total_read = 0;
off_t current_offset = 0;
int target_fd = files_registered_ ? 0 : fd_;
while (remaining > 0) {
size_t chunk_size =
calculate_chunk_size(remaining, queue_depth_, MIN_CHUNK_SIZE);
unsigned num_chunks = std::min(
static_cast<unsigned>((remaining + chunk_size - 1) / chunk_size),
queue_depth_);
size_t batch_size = 0;
for (unsigned i = 0; i < num_chunks; i++) {
size_t this_chunk = std::min(chunk_size, remaining);
struct io_uring_sqe *sqe = io_uring_get_sqe(&ring_);
if (!sqe) {
LOG(ERROR) << "[UringFile::read] Failed to get SQE";
if (aligned_buffer) {
free_aligned_buffer(aligned_buffer);
} else {
buffer.clear();
}
return make_error<size_t>(ErrorCode::FILE_READ_FAIL);
}
io_uring_prep_read(sqe, target_fd, ptr, this_chunk, current_offset);
if (files_registered_) {
io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE);
}
ptr += this_chunk;
current_offset += this_chunk;
batch_size += this_chunk;
remaining -= this_chunk;
if (remaining == 0) break;
}
auto result = submit_and_wait_n(num_chunks);
if (!result) {
if (aligned_buffer) {
free_aligned_buffer(aligned_buffer);
} else {
buffer.clear();
}
return make_error<size_t>(ErrorCode::FILE_READ_FAIL);
}
size_t bytes_read = result.value();
if (bytes_read == 0) {
break; // EOF
}
total_read += bytes_read;
if (bytes_read < batch_size) {
break; // EOF
}
}
// Copy from aligned buffer to std::string if using O_DIRECT
if (use_direct_io_) {
size_t actual_read = std::min(total_read, length);
buffer.assign(static_cast<const char *>(aligned_buffer), actual_read);
free_aligned_buffer(aligned_buffer);
total_read = actual_read;
} else {
buffer.resize(total_read);
}
if (total_read != length && total_read == 0) {
return make_error<size_t>(ErrorCode::FILE_READ_FAIL);
}
return total_read;
}
// ---------------------------------------------------------------------------
// Zero-copy aligned I/O interface for O_DIRECT
// ---------------------------------------------------------------------------
tl::expected<size_t, ErrorCode> UringFile::read_aligned(void *buffer,
size_t length,
off_t offset) {
MutexLocker lock(&ring_mutex_);
if (fd_ < 0 || !ring_initialized_) {
return make_error<size_t>(ErrorCode::FILE_NOT_FOUND);
}
if (length == 0 || buffer == nullptr) {
return make_error<size_t>(ErrorCode::FILE_INVALID_BUFFER);
}
// Verify alignment when using O_DIRECT
if (use_direct_io_) {
if (reinterpret_cast<uintptr_t>(buffer) % ALIGNMENT_ != 0) {
LOG(ERROR) << "[UringFile::read_aligned] Buffer not aligned to "
<< ALIGNMENT_;
return make_error<size_t>(ErrorCode::FILE_INVALID_BUFFER);
}
if (length % ALIGNMENT_ != 0) {
LOG(ERROR) << "[UringFile::read_aligned] Length not aligned to "
<< ALIGNMENT_;
return make_error<size_t>(ErrorCode::FILE_INVALID_BUFFER);
}
if (offset % ALIGNMENT_ != 0) {
LOG(ERROR) << "[UringFile::read_aligned] Offset not aligned to "
<< ALIGNMENT_;
return make_error<size_t>(ErrorCode::FILE_INVALID_BUFFER);
}
}
constexpr size_t MIN_CHUNK_SIZE = 4096;
char *ptr = static_cast<char *>(buffer);
size_t remaining = length;
size_t total_read = 0;
off_t current_offset = offset;
int target_fd = files_registered_ ? 0 : fd_;
// Check if this buffer falls within the registered buffer range
bool use_fixed_buffer =
(buffer_registered_ &&
reinterpret_cast<uintptr_t>(buffer) >=
reinterpret_cast<uintptr_t>(registered_buffer_) &&
reinterpret_cast<uintptr_t>(buffer) + length <=
reinterpret_cast<uintptr_t>(registered_buffer_) +
registered_buffer_size_);
while (remaining > 0) {
size_t chunk_size =
calculate_chunk_size(remaining, queue_depth_, MIN_CHUNK_SIZE);
unsigned num_chunks = std::min(
static_cast<unsigned>((remaining + chunk_size - 1) / chunk_size),
queue_depth_);
size_t batch_size = 0;
for (unsigned i = 0; i < num_chunks; i++) {
size_t this_chunk = std::min(chunk_size, remaining);
struct io_uring_sqe *sqe = io_uring_get_sqe(&ring_);
if (!sqe) {
LOG(ERROR) << "[UringFile::read_aligned] Failed to get SQE";
return make_error<size_t>(ErrorCode::FILE_READ_FAIL);
}
if (use_fixed_buffer) {
// Use registered fixed buffer - avoids get_user_pages()
// overhead
io_uring_prep_read_fixed(sqe, target_fd, ptr, this_chunk,
current_offset, 0);
} else {
// Use regular buffer
io_uring_prep_read(sqe, target_fd, ptr, this_chunk,
current_offset);
}
if (files_registered_) {
io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE);
}
ptr += this_chunk;
current_offset += this_chunk;
batch_size += this_chunk;
remaining -= this_chunk;
if (remaining == 0) break;
}
auto result = submit_and_wait_n(num_chunks);
if (!result) {
return make_error<size_t>(ErrorCode::FILE_READ_FAIL);
}
size_t bytes_read = result.value();
if (bytes_read == 0) {
break; // EOF
}
total_read += bytes_read;
if (bytes_read < batch_size) {
break; // EOF
}
}
return total_read;
}
tl::expected<size_t, ErrorCode> UringFile::write_aligned(const void *buffer,
size_t length,
off_t offset) {
if (fd_ < 0 || !ring_initialized_) {
return make_error<size_t>(ErrorCode::FILE_NOT_FOUND);
}
if (length == 0 || buffer == nullptr) {
return make_error<size_t>(ErrorCode::FILE_INVALID_BUFFER);
}
// Verify alignment when using O_DIRECT
if (use_direct_io_) {
if (reinterpret_cast<uintptr_t>(buffer) % ALIGNMENT_ != 0) {
LOG(ERROR) << "[UringFile::write_aligned] Buffer not aligned to "
<< ALIGNMENT_;
return make_error<size_t>(ErrorCode::FILE_INVALID_BUFFER);
}
if (length % ALIGNMENT_ != 0) {
LOG(ERROR) << "[UringFile::write_aligned] Length not aligned to "
<< ALIGNMENT_;
return make_error<size_t>(ErrorCode::FILE_INVALID_BUFFER);
}
if (offset % ALIGNMENT_ != 0) {
LOG(ERROR) << "[UringFile::write_aligned] Offset not aligned to "
<< ALIGNMENT_;
return make_error<size_t>(ErrorCode::FILE_INVALID_BUFFER);
}
}
constexpr size_t MIN_CHUNK_SIZE = 4096;
const char *ptr = static_cast<const char *>(buffer);
size_t remaining = length;
size_t total_written = 0;
off_t current_offset = offset;
int target_fd = files_registered_ ? 0 : fd_;
// Check if this buffer falls within the registered buffer range
bool use_fixed_buffer =
(buffer_registered_ &&
reinterpret_cast<uintptr_t>(buffer) >=
reinterpret_cast<uintptr_t>(registered_buffer_) &&
reinterpret_cast<uintptr_t>(buffer) + length <=
reinterpret_cast<uintptr_t>(registered_buffer_) +
registered_buffer_size_);
while (remaining > 0) {
size_t chunk_size =
calculate_chunk_size(remaining, queue_depth_, MIN_CHUNK_SIZE);
unsigned num_chunks = std::min(
static_cast<unsigned>((remaining + chunk_size - 1) / chunk_size),
queue_depth_);
for (unsigned i = 0; i < num_chunks; i++) {
size_t this_chunk = std::min(chunk_size, remaining);
struct io_uring_sqe *sqe = io_uring_get_sqe(&ring_);
if (!sqe) {
LOG(ERROR) << "[UringFile::write_aligned] Failed to get SQE";
return make_error<size_t>(ErrorCode::FILE_WRITE_FAIL);
}
if (use_fixed_buffer) {
// Use registered fixed buffer - avoids get_user_pages()
// overhead
io_uring_prep_write_fixed(sqe, target_fd, ptr, this_chunk,
current_offset, 0);
} else {
// Use regular buffer
io_uring_prep_write(sqe, target_fd, ptr, this_chunk,
current_offset);
}
if (files_registered_) {
io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE);
}
ptr += this_chunk;
current_offset += this_chunk;
remaining -= this_chunk;
if (remaining == 0) break;
}
auto result = submit_and_wait_n(num_chunks);
if (!result) {
return make_error<size_t>(ErrorCode::FILE_WRITE_FAIL);
}
size_t bytes_written = result.value();
if (bytes_written == 0) {
LOG(ERROR) << "[UringFile::write_aligned] Zero bytes written";
return make_error<size_t>(ErrorCode::FILE_WRITE_FAIL);
}
total_written += bytes_written;
}
if (total_written != length) {
LOG(WARNING) << "[UringFile::write_aligned] Incomplete write: "
<< total_written << " / " << length;
return make_error<size_t>(ErrorCode::FILE_WRITE_FAIL);
}
return total_written;
}
// ---------------------------------------------------------------------------
// Vectored I/O — multi-SQE parallel submission
//
// Each iovec is submitted as an independent SQE so the NVMe device can
// serve them concurrently, matching the pipelining approach used in the
// benchmark. When iovcnt > queue_depth_, requests are batched.
// ---------------------------------------------------------------------------
tl::expected<size_t, ErrorCode> UringFile::vector_write(const iovec *iov,
int iovcnt,
off_t offset) {
if (fd_ < 0 || !ring_initialized_) {
return make_error<size_t>(ErrorCode::FILE_NOT_FOUND);
}
int target_fd = files_registered_ ? 0 : fd_;
size_t total_written = 0;
off_t cur_offset = offset;
int remaining = iovcnt;
int idx = 0;
while (remaining > 0) {
int batch = std::min(remaining, static_cast<int>(queue_depth_));
for (int i = 0; i < batch; i++) {
struct io_uring_sqe *sqe = io_uring_get_sqe(&ring_);
if (!sqe) {
LOG(ERROR) << "[UringFile::vector_write] Failed to get SQE";
return make_error<size_t>(ErrorCode::FILE_WRITE_FAIL);
}
io_uring_prep_write(sqe, target_fd, iov[idx].iov_base,
iov[idx].iov_len, cur_offset);
if (files_registered_) {
io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE);
}
cur_offset += static_cast<off_t>(iov[idx].iov_len);
idx++;
}
auto result = submit_and_wait_n(batch);
if (!result) {
return make_error<size_t>(ErrorCode::FILE_WRITE_FAIL);
}
total_written += result.value();
remaining -= batch;
}
return total_written;
}
tl::expected<size_t, ErrorCode> UringFile::vector_read(const iovec *iov,
int iovcnt,
off_t offset) {
if (fd_ < 0 || !ring_initialized_) {
return make_error<size_t>(ErrorCode::FILE_NOT_FOUND);
}
size_t expected_bytes = 0;
for (int i = 0; i < iovcnt; ++i) {
expected_bytes += iov[i].iov_len;
}
auto start = std::chrono::steady_clock::now();
int target_fd = files_registered_ ? 0 : fd_;
size_t total_read = 0;
off_t cur_offset = offset;
int remaining = iovcnt;
int idx = 0;
while (remaining > 0) {
int batch = std::min(remaining, static_cast<int>(queue_depth_));
for (int i = 0; i < batch; i++) {
struct io_uring_sqe *sqe = io_uring_get_sqe(&ring_);
if (!sqe) {
LOG(ERROR) << "[UringFile::vector_read] Failed to get SQE";
return make_error<size_t>(ErrorCode::FILE_READ_FAIL);
}
io_uring_prep_read(sqe, target_fd, iov[idx].iov_base,
iov[idx].iov_len, cur_offset);
if (files_registered_) {
io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE);
}
cur_offset += static_cast<off_t>(iov[idx].iov_len);
idx++;
}
auto result = submit_and_wait_n(batch);
if (!result) {
auto end = std::chrono::steady_clock::now();
auto elapsed_us =
std::chrono::duration_cast<std::chrono::microseconds>(end -
start)
.count();
LOG(ERROR) << "[UringFile::vector_read] FAILED: fd=" << fd_
<< ", offset=" << offset << ", iovcnt=" << iovcnt
<< ", expected_bytes=" << expected_bytes
<< ", time=" << elapsed_us << "us";
return make_error<size_t>(ErrorCode::FILE_READ_FAIL);
}
total_read += result.value();
remaining -= batch;
}
auto end = std::chrono::steady_clock::now();
auto elapsed_us =
std::chrono::duration_cast<std::chrono::microseconds>(end - start)
.count();
if (elapsed_us > 1000 || expected_bytes > 1024 * 1024) {
double throughput_mbps =
(elapsed_us > 0)
? (static_cast<double>(total_read) / (1024.0 * 1024.0)) /
(static_cast<double>(elapsed_us) / 1000000.0)
: 0;
LOG(INFO) << "[UringFile::vector_read] fd=" << fd_
<< ", offset=" << offset << ", iovcnt=" << iovcnt
<< ", bytes=" << total_read << ", time=" << elapsed_us
<< "us (" << (elapsed_us / 1000.0) << "ms)"
<< ", throughput=" << throughput_mbps << "MB/s";
}
return total_read;
}
// ---------------------------------------------------------------------------
// Buffer registration for high-performance I/O
// ---------------------------------------------------------------------------
bool UringFile::register_buffer(void *buffer, size_t length) {
if (!ring_initialized_) {
LOG(ERROR) << "[UringFile::register_buffer] io_uring not initialized";
return false;
}
if (buffer_registered_) {
LOG(WARNING) << "[UringFile::register_buffer] Buffer already "
"registered, unregistering first";
unregister_buffer();
}
if (!buffer || length == 0) {
LOG(ERROR) << "[UringFile::register_buffer] Invalid buffer or length";
return false;
}
// Verify alignment when using O_DIRECT
if (use_direct_io_) {
if (reinterpret_cast<uintptr_t>(buffer) % ALIGNMENT_ != 0) {
LOG(ERROR) << "[UringFile::register_buffer] Buffer not aligned to "
<< ALIGNMENT_;
return false;
}
if (length % ALIGNMENT_ != 0) {
LOG(ERROR) << "[UringFile::register_buffer] Length not aligned to "
<< ALIGNMENT_;
return false;
}
}
registered_iovec_.iov_base = buffer;
registered_iovec_.iov_len = length;
int ret = io_uring_register_buffers(&ring_, &registered_iovec_, 1);
if (ret < 0) {
LOG(ERROR)
<< "[UringFile::register_buffer] io_uring_register_buffers failed: "
<< strerror(-ret);
return false;
}
registered_buffer_ = buffer;
registered_buffer_size_ = length;
buffer_registered_ = true;
LOG(INFO) << "[UringFile::register_buffer] Successfully registered buffer: "
<< "addr=" << buffer << ", size=" << length;
return true;
}
void UringFile::unregister_buffer() {
if (!buffer_registered_) {
return;
}
if (ring_initialized_) {
int ret = io_uring_unregister_buffers(&ring_);
if (ret < 0) {
LOG(ERROR) << "[UringFile::unregister_buffer] "
"io_uring_unregister_buffers failed: "
<< strerror(-ret);
} else {
LOG(INFO) << "[UringFile::unregister_buffer] Successfully "
"unregistered buffer";
}
}
registered_buffer_ = nullptr;
registered_buffer_size_ = 0;
buffer_registered_ = false;
}
// ---------------------------------------------------------------------------
// datasync — flush bucket data to stable storage before metadata write
// ---------------------------------------------------------------------------
tl::expected<void, ErrorCode> UringFile::datasync() {
if (fd_ < 0 || !ring_initialized_) {
return make_error<void>(ErrorCode::FILE_NOT_FOUND);
}
int target_fd = files_registered_ ? 0 : fd_;
struct io_uring_sqe *sqe = io_uring_get_sqe(&ring_);
if (!sqe) {
LOG(ERROR) << "[UringFile::datasync] Failed to get SQE";
return make_error<void>(ErrorCode::INTERNAL_ERROR);
}
io_uring_prep_fsync(sqe, target_fd, IORING_FSYNC_DATASYNC);
if (files_registered_) {
io_uring_sqe_set_flags(sqe, IOSQE_FIXED_FILE);
}
auto result = submit_and_wait_n(1);
if (!result) {
LOG(ERROR) << "[UringFile::datasync] fsync failed for: " << filename_;
return make_error<void>(ErrorCode::FILE_WRITE_FAIL);
}
return {};
}
} // namespace mooncake
#endif // USE_URING

2
mooncake-store/tests/file_storage_test.cpp
View File

@ -58,7 +58,7 @@ class FileStorageTest : public ::testing::Test {
tl::expected<void, ErrorCode> FileStorageBatchLoad(
FileStorage& fileStorage,
const std::unordered_map<std::string, Slice>& batch_object) {
std::unordered_map<std::string, Slice>& batch_object) {
return fileStorage.BatchLoad(batch_object);
}