diff --git a/mooncake-store/benchmarks/CMakeLists.txt b/mooncake-store/benchmarks/CMakeLists.txt index b3a22253..ac8994f0 100644 --- a/mooncake-store/benchmarks/CMakeLists.txt +++ b/mooncake-store/benchmarks/CMakeLists.txt @@ -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 diff --git a/mooncake-store/benchmarks/file_interface_bench.cpp b/mooncake-store/benchmarks/file_interface_bench.cpp new file mode 100644 index 00000000..8af9e8a6 --- /dev/null +++ b/mooncake-store/benchmarks/file_interface_bench.cpp @@ -0,0 +1,864 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#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::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 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 GenerateRandomData(size_t size) { + std::vector data(size); + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_int_distribution dist(0, 255); + + for (size_t i = 0; i < size; ++i) { + data[i] = static_cast(dist(gen)); + } + return data; + } + + void GenerateRandomDataAligned(void* ptr, size_t size) { + char* data = static_cast(ptr); + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_int_distribution dist(0, 255); + + for (size_t i = 0; i < size; ++i) { + data[i] = static_cast(dist(gen)); + } + } + + std::unique_ptr CreateFile(int fd) { +#ifdef USE_URING + if (config_.use_uring) { + return std::make_unique(config_.file_path, fd, + config_.uring_queue_depth, + config_.use_uring_direct_io); + } +#endif + return std::make_unique(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 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(static_cast(data), size), size); + + auto end = std::chrono::high_resolution_clock::now(); + + if (!result) { + std::cerr << "Write operation failed with error code: " + << static_cast(result.error()) << std::endl; + return std::nullopt; + } + + size_t bytes_written = result.value(); + auto duration_us = + std::chrono::duration_cast(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> 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(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( + 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 result; + if (is_write) { + auto* uring_file = dynamic_cast(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(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( + 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(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(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> 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(buffer) + offset, + this_chunk, offset); + } else { + ret = pread(fd, static_cast(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(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> 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(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(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(expected); + const char* act_data = static_cast(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(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 to test file (default: " + "/tmp/file_bench_test.dat)" + << std::endl; + std::cout << " --size Size of test data in bytes (default: " + "104857600 = 100MB)" + << std::endl; + std::cout + << " --size-mb Size of test data in MB (overrides --size)" + << std::endl; + std::cout << " --chunk-size 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 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 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; +} diff --git a/mooncake-store/include/aligned_client_buffer.hpp b/mooncake-store/include/aligned_client_buffer.hpp new file mode 100644 index 00000000..57f81e92 --- /dev/null +++ b/mooncake-store/include/aligned_client_buffer.hpp @@ -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 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 diff --git a/mooncake-store/include/client_buffer.hpp b/mooncake-store/include/client_buffer.hpp index 6a4720d3..1758047d 100644 --- a/mooncake-store/include/client_buffer.hpp +++ b/mooncake-store/include/client_buffer.hpp @@ -51,19 +51,22 @@ class ClientBufferAllocator [[nodiscard]] std::optional 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 allocator_; std::string protocol; - void* buffer_; size_t buffer_size_; bool is_external_memory_ = false; - bool use_hugepage_ = false; }; /** diff --git a/mooncake-store/include/file_interface.h b/mooncake-store/include/file_interface.h index 48f67dc2..550d5670 100644 --- a/mooncake-store/include/file_interface.h +++ b/mooncake-store/include/file_interface.h @@ -5,9 +5,13 @@ #include #include #include "types.h" +#include "mutex.h" #include #include #include +#ifdef USE_URING +#include +#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 write(const std::string &buffer, + size_t length) override; + tl::expected write(std::span data, + size_t length) override; + tl::expected read(std::string &buffer, + size_t length) override; + tl::expected vector_write(const iovec *iov, int iovcnt, + off_t offset) override; + tl::expected vector_read(const iovec *iov, int iovcnt, + off_t offset) override; + + // Zero-copy interface for O_DIRECT: caller must provide aligned buffer + tl::expected read_aligned(void *buffer, size_t length, + off_t offset = 0); + tl::expected 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 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 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 diff --git a/mooncake-store/include/file_storage.h b/mooncake-store/include/file_storage.h index b153cc06..8399b8fc 100644 --- a/mooncake-store/include/file_storage.h +++ b/mooncake-store/include/file_storage.h @@ -67,7 +67,7 @@ class FileStorage { tl::expected IsEnableOffloading(); tl::expected BatchLoad( - const std::unordered_map& batch_object); + std::unordered_map& batch_object); tl::expected BatchQuerySegmentSlices( const std::vector& keys, diff --git a/mooncake-store/include/storage_backend.h b/mooncake-store/include/storage_backend.h index 19616101..993cef55 100644 --- a/mooncake-store/include/storage_backend.h +++ b/mooncake-store/include/storage_backend.h @@ -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 BatchLoad( - const std::unordered_map& batched_slices) = 0; + std::unordered_map& batched_slices) = 0; virtual tl::expected 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 BatchLoad( - const std::unordered_map& batched_slices) override; + std::unordered_map& batched_slices) override; tl::expected 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 indicating operation status. */ tl::expected BatchLoad( - const std::unordered_map& batched_slices) override; + std::unordered_map& 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, 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 initialized_{false}; std::optional 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 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 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> + file_cache_ GUARDED_BY(file_cache_mutex_); + + // Get or open a file with caching support + tl::expected, 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 indicating operation status. */ tl::expected BatchLoad( - const std::unordered_map& batched_slices) override; + std::unordered_map& batched_slices) override; /** * @brief Checks whether an object with the specified key exists in the diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 7d817c0e..bbf4e958 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -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}") diff --git a/mooncake-store/src/aligned_client_buffer.cpp b/mooncake-store/src/aligned_client_buffer.cpp new file mode 100644 index 00000000..6ac7ea90 --- /dev/null +++ b/mooncake-store/src/aligned_client_buffer.cpp @@ -0,0 +1,102 @@ +#include "aligned_client_buffer.hpp" + +#include +#include +#include +#include + +#include "utils.h" + +namespace mooncake { + +std::shared_ptr +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(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( + 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 diff --git a/mooncake-store/src/file_storage.cpp b/mooncake-store/src/file_storage.cpp index 12ad2731..1bcc8e0e 100644 --- a/mooncake-store/src/file_storage.cpp +++ b/mooncake-store/src/file_storage.cpp @@ -3,8 +3,13 @@ #include #include +#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("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(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(file); + if (uring_file) { + auto aligned_allocator = + std::static_pointer_cast( + 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, 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(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 FileStorage::Heartbeat() { } tl::expected FileStorage::BatchLoad( - const std::unordered_map& batch_object) { + std::unordered_map& 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& 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(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& 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(sizes[i])}); - result->pointers.emplace_back( - reinterpret_cast(alloc_result->ptr())); + + // Align ptr to 4096 boundary for O_DIRECT + void* raw_ptr = alloc_result->ptr(); + void* aligned_ptr = reinterpret_cast( + (reinterpret_cast(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(aligned_ptr)); result->handles.emplace_back(std::move(alloc_result.value())); result->lease_timeout = lease_timeout; } diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index 9714e9b4..64e2f63d 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -767,6 +767,11 @@ std::unique_ptr StorageBackend::create_file( } #endif +#ifdef USE_URING + if (use_uring_) { + return std::make_unique(path, fd, 32, true); + } +#endif return std::make_unique(path, fd); } @@ -932,6 +937,7 @@ tl::expected StorageBackendAdaptor::Init() { storage_backend_ = std::make_unique( 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 StorageBackendAdaptor::IsExist( } tl::expected StorageBackendAdaptor::BatchLoad( - const std::unordered_map& batched_slices) { + std::unordered_map& 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( + 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 BucketStorageBackend::BatchOffload( const std::unordered_map>& batch_object, @@ -1331,7 +1359,7 @@ tl::expected BucketStorageBackend::BatchQuery( } tl::expected BucketStorageBackend::BatchLoad( - const std::unordered_map& batch_object) { + std::unordered_map& 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 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 read_res; + +#ifdef USE_URING + // Try to use read_aligned for O_DIRECT I/O if file is UringFile + UringFile* uring_file = dynamic_cast(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(plan.dest_slice.size); + int64_t aligned_end = static_cast(align_up( + static_cast(data_end), kDirectIOAlignment)); + size_t aligned_size = + static_cast(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(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 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(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(file.get()); + if (uring_file != nullptr) { + size_t total_size = static_cast(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 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( + 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(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(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(path, fd, 32, true); + } +#endif return std::make_unique(path, fd); } +tl::expected, 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(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 BucketStorageBackend::HandleNext( const std::function< ErrorCode(const std::vector& keys, @@ -2175,6 +2373,32 @@ tl::expected BucketStorageBackend::HasNext() { return next_bucket_ != 0; } +tl::expected, 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(std::move(file)); +} + // ============================================================================ // OffsetAllocatorStorageBackend Implementation // ============================================================================ @@ -2264,9 +2488,17 @@ tl::expected OffsetAllocatorStorageBackend::Init() { } } - // Release fd to PosixFile (PosixFile takes ownership and will close it) - data_file_ = - std::make_unique(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( + data_file_path_, fd_guard.release(), 32, true); + } else +#endif + { + data_file_ = std::make_unique(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 OffsetAllocatorStorageBackend::BatchOffload( //----------------------------------------------------------------------------- tl::expected OffsetAllocatorStorageBackend::BatchLoad( - const std::unordered_map& batched_slices) { + std::unordered_map& batched_slices) { if (!initialized_.load(std::memory_order_acquire)) { LOG(ERROR) << "Storage backend is not initialized. Call Init() before use."; diff --git a/mooncake-store/src/uring_file.cpp b/mooncake-store/src/uring_file.cpp new file mode 100644 index 00000000..3a53976a --- /dev/null +++ b/mooncake-store/src/uring_file.cpp @@ -0,0 +1,871 @@ +#ifdef USE_URING + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#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(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 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(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(cqe->res); + } + count++; + } + io_uring_cq_advance(&ring_, count); + + if (has_error) { + return make_error(ErrorCode::INTERNAL_ERROR); + } + return total_bytes; +} + +// --------------------------------------------------------------------------- +// Write implementation with intelligent chunking +// --------------------------------------------------------------------------- + +tl::expected UringFile::write(const std::string &buffer, + size_t length) { + return write(std::span(buffer.data(), length), length); +} + +tl::expected UringFile::write(std::span data, + size_t length) { + if (fd_ < 0 || !ring_initialized_) { + return make_error(ErrorCode::FILE_NOT_FOUND); + } + if (length == 0) { + return make_error(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(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(aligned_buffer) + length, 0, + actual_length - length); + } + source_ptr = static_cast(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((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(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(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(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(ErrorCode::FILE_WRITE_FAIL); + } + + return length; // Return original length, not aligned length +} + +// --------------------------------------------------------------------------- +// Read implementation with intelligent chunking +// --------------------------------------------------------------------------- + +tl::expected UringFile::read(std::string &buffer, + size_t length) { + if (fd_ < 0 || !ring_initialized_) { + return make_error(ErrorCode::FILE_NOT_FOUND); + } + if (length == 0) { + return make_error(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(ErrorCode::FILE_READ_FAIL); + } + read_ptr = static_cast(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((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(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(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(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(ErrorCode::FILE_READ_FAIL); + } + + return total_read; +} + +// --------------------------------------------------------------------------- +// Zero-copy aligned I/O interface for O_DIRECT +// --------------------------------------------------------------------------- + +tl::expected UringFile::read_aligned(void *buffer, + size_t length, + off_t offset) { + MutexLocker lock(&ring_mutex_); + if (fd_ < 0 || !ring_initialized_) { + return make_error(ErrorCode::FILE_NOT_FOUND); + } + if (length == 0 || buffer == nullptr) { + return make_error(ErrorCode::FILE_INVALID_BUFFER); + } + + // Verify alignment when using O_DIRECT + if (use_direct_io_) { + if (reinterpret_cast(buffer) % ALIGNMENT_ != 0) { + LOG(ERROR) << "[UringFile::read_aligned] Buffer not aligned to " + << ALIGNMENT_; + return make_error(ErrorCode::FILE_INVALID_BUFFER); + } + if (length % ALIGNMENT_ != 0) { + LOG(ERROR) << "[UringFile::read_aligned] Length not aligned to " + << ALIGNMENT_; + return make_error(ErrorCode::FILE_INVALID_BUFFER); + } + if (offset % ALIGNMENT_ != 0) { + LOG(ERROR) << "[UringFile::read_aligned] Offset not aligned to " + << ALIGNMENT_; + return make_error(ErrorCode::FILE_INVALID_BUFFER); + } + } + + constexpr size_t MIN_CHUNK_SIZE = 4096; + + char *ptr = static_cast(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(buffer) >= + reinterpret_cast(registered_buffer_) && + reinterpret_cast(buffer) + length <= + reinterpret_cast(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((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(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(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 UringFile::write_aligned(const void *buffer, + size_t length, + off_t offset) { + if (fd_ < 0 || !ring_initialized_) { + return make_error(ErrorCode::FILE_NOT_FOUND); + } + if (length == 0 || buffer == nullptr) { + return make_error(ErrorCode::FILE_INVALID_BUFFER); + } + + // Verify alignment when using O_DIRECT + if (use_direct_io_) { + if (reinterpret_cast(buffer) % ALIGNMENT_ != 0) { + LOG(ERROR) << "[UringFile::write_aligned] Buffer not aligned to " + << ALIGNMENT_; + return make_error(ErrorCode::FILE_INVALID_BUFFER); + } + if (length % ALIGNMENT_ != 0) { + LOG(ERROR) << "[UringFile::write_aligned] Length not aligned to " + << ALIGNMENT_; + return make_error(ErrorCode::FILE_INVALID_BUFFER); + } + if (offset % ALIGNMENT_ != 0) { + LOG(ERROR) << "[UringFile::write_aligned] Offset not aligned to " + << ALIGNMENT_; + return make_error(ErrorCode::FILE_INVALID_BUFFER); + } + } + + constexpr size_t MIN_CHUNK_SIZE = 4096; + + const char *ptr = static_cast(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(buffer) >= + reinterpret_cast(registered_buffer_) && + reinterpret_cast(buffer) + length <= + reinterpret_cast(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((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(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(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(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(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 UringFile::vector_write(const iovec *iov, + int iovcnt, + off_t offset) { + if (fd_ < 0 || !ring_initialized_) { + return make_error(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(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(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(iov[idx].iov_len); + idx++; + } + + auto result = submit_and_wait_n(batch); + if (!result) { + return make_error(ErrorCode::FILE_WRITE_FAIL); + } + total_written += result.value(); + remaining -= batch; + } + + return total_written; +} + +tl::expected UringFile::vector_read(const iovec *iov, + int iovcnt, + off_t offset) { + if (fd_ < 0 || !ring_initialized_) { + return make_error(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(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(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(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(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(ErrorCode::FILE_READ_FAIL); + } + total_read += result.value(); + remaining -= batch; + } + + auto end = std::chrono::steady_clock::now(); + auto elapsed_us = + std::chrono::duration_cast(end - start) + .count(); + + if (elapsed_us > 1000 || expected_bytes > 1024 * 1024) { + double throughput_mbps = + (elapsed_us > 0) + ? (static_cast(total_read) / (1024.0 * 1024.0)) / + (static_cast(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(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_, ®istered_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 UringFile::datasync() { + if (fd_ < 0 || !ring_initialized_) { + return make_error(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(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(ErrorCode::FILE_WRITE_FAIL); + } + return {}; +} + +} // namespace mooncake + +#endif // USE_URING diff --git a/mooncake-store/tests/file_storage_test.cpp b/mooncake-store/tests/file_storage_test.cpp index c145e49d..28cdbb2b 100644 --- a/mooncake-store/tests/file_storage_test.cpp +++ b/mooncake-store/tests/file_storage_test.cpp @@ -58,7 +58,7 @@ class FileStorageTest : public ::testing::Test { tl::expected FileStorageBatchLoad( FileStorage& fileStorage, - const std::unordered_map& batch_object) { + std::unordered_map& batch_object) { return fileStorage.BatchLoad(batch_object); }