wehub-resource-sync
250 lines
8.1 KiB
C++
250 lines
8.1 KiB
C++
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
#include "connector.h"
|
|
#include <fcntl.h>
|
|
#include <cerrno>
|
|
#include <cstdio>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
|
|
namespace example_connector {
|
|
|
|
// ---------------------------------------------------------------
|
|
// Shared I/O helpers
|
|
// ---------------------------------------------------------------
|
|
|
|
static void write_all(int fd, const void* data, size_t len) {
|
|
size_t written = 0;
|
|
const char* ptr = static_cast<const char*>(data);
|
|
while (written < len) {
|
|
ssize_t n = ::write(fd, ptr + written, len - written);
|
|
if (n < 0) {
|
|
if (errno == EINTR) continue;
|
|
throw std::runtime_error("write failed: " + std::string(strerror(errno)));
|
|
}
|
|
if (n == 0) throw std::runtime_error("write returned 0");
|
|
written += static_cast<size_t>(n);
|
|
}
|
|
}
|
|
|
|
static size_t read_all(int fd, void* buf, size_t len) {
|
|
size_t total = 0;
|
|
char* ptr = static_cast<char*>(buf);
|
|
while (total < len) {
|
|
ssize_t n = ::read(fd, ptr + total, len - total);
|
|
if (n < 0) {
|
|
if (errno == EINTR) continue;
|
|
throw std::runtime_error("read failed: " + std::string(strerror(errno)));
|
|
}
|
|
if (n == 0) break;
|
|
total += static_cast<size_t>(n);
|
|
}
|
|
return total;
|
|
}
|
|
|
|
// ---------------------------------------------------------------
|
|
// ExampleFSConnector
|
|
// ---------------------------------------------------------------
|
|
|
|
// Key encoding constants — must match fs_l2_adapter.py
|
|
static constexpr char KEY_SEP = '@';
|
|
static constexpr const char* PATH_SLASH_REPLACEMENT = "-SEP-";
|
|
static constexpr const char* FILE_EXT = ".data";
|
|
|
|
std::string ExampleFSConnector::replace_all(const std::string& str,
|
|
const std::string& from,
|
|
const std::string& to) {
|
|
std::string result = str;
|
|
size_t pos = 0;
|
|
while ((pos = result.find(from, pos)) != std::string::npos) {
|
|
result.replace(pos, from.size(), to);
|
|
pos += to.size();
|
|
}
|
|
return result;
|
|
}
|
|
|
|
std::string ExampleFSConnector::safe_filename(const std::string& key) {
|
|
// Input key format (from _object_key_to_string):
|
|
// Unsalted: <model_name>@<kv_rank_hex>@<chunk_hash_hex>
|
|
// Salted : <model_name>@<kv_rank_hex>@<chunk_hash_hex>@<cache_salt>
|
|
//
|
|
// Output filename:
|
|
// Unsalted: <model_name_safe>@0x<kv_rank_hex>@<chunk_hash_hex>.data
|
|
// Salted :
|
|
// <model_name_safe>@0x<kv_rank_hex>@<chunk_hash_hex>@<cache_salt>.data
|
|
//
|
|
// Both model_name and cache_salt are forbidden from containing '@'
|
|
// (invariant enforced on the Python side), so splitting on '@' is
|
|
// unambiguous.
|
|
|
|
std::vector<std::string> parts;
|
|
size_t start = 0;
|
|
for (size_t pos = 0; pos <= key.size(); ++pos) {
|
|
if (pos == key.size() || key[pos] == KEY_SEP) {
|
|
parts.emplace_back(key.substr(start, pos - start));
|
|
start = pos + 1;
|
|
}
|
|
}
|
|
if (parts.size() != 3 && parts.size() != 4) {
|
|
throw std::runtime_error(
|
|
"ExampleFSConnector: malformed key "
|
|
"(expected 3 or 4 '@'-separated fields): " +
|
|
key);
|
|
}
|
|
|
|
const std::string& model = parts[0];
|
|
const std::string& kv_rank_hex = parts[1];
|
|
const std::string& chunk_hash = parts[2];
|
|
const std::string cache_salt = parts.size() == 4 ? parts[3] : std::string();
|
|
|
|
// Replace '/' with '-SEP-' for filesystem safety
|
|
std::string safe_model = replace_all(model, "/", PATH_SLASH_REPLACEMENT);
|
|
|
|
std::string result;
|
|
result.reserve(safe_model.size() + kv_rank_hex.size() + chunk_hash.size() +
|
|
cache_salt.size() + 32);
|
|
result += safe_model;
|
|
result += KEY_SEP;
|
|
result += "0x";
|
|
result += kv_rank_hex;
|
|
result += KEY_SEP;
|
|
result += chunk_hash;
|
|
if (!cache_salt.empty()) {
|
|
result += KEY_SEP;
|
|
result += cache_salt;
|
|
}
|
|
result += FILE_EXT;
|
|
return result;
|
|
}
|
|
|
|
ExampleFSConnector::ExampleFSConnector(std::string base_path, int num_workers)
|
|
: ConnectorBase(num_workers), base_path_(std::move(base_path)) {
|
|
std::filesystem::create_directories(base_path_);
|
|
start_workers();
|
|
}
|
|
|
|
ExampleFSConnector::~ExampleFSConnector() { close(); }
|
|
|
|
WorkerFSConn ExampleFSConnector::create_connection() {
|
|
WorkerFSConn conn;
|
|
conn.base_path = base_path_;
|
|
return conn;
|
|
}
|
|
|
|
void ExampleFSConnector::do_single_get(WorkerFSConn& conn,
|
|
const std::string& key, void* buf,
|
|
size_t len, size_t /*chunk_size*/) {
|
|
auto path = conn.base_path / safe_filename(key);
|
|
int fd = ::open(path.c_str(), O_RDONLY);
|
|
if (fd < 0) {
|
|
throw std::runtime_error("open for read failed: " + path.string() + ": " +
|
|
strerror(errno));
|
|
}
|
|
size_t n = read_all(fd, buf, len);
|
|
::close(fd);
|
|
if (n != len) {
|
|
throw std::runtime_error("incomplete read for " + path.string());
|
|
}
|
|
}
|
|
|
|
void ExampleFSConnector::do_single_set(WorkerFSConn& conn,
|
|
const std::string& key, const void* buf,
|
|
size_t len, size_t /*chunk_size*/) {
|
|
auto file_path = conn.base_path / safe_filename(key);
|
|
if (std::filesystem::exists(file_path)) return;
|
|
|
|
auto tmp_path = file_path;
|
|
tmp_path.replace_extension(".tmp");
|
|
|
|
int fd = ::open(tmp_path.c_str(), O_CREAT | O_WRONLY | O_TRUNC, 0644);
|
|
if (fd < 0) {
|
|
throw std::runtime_error("open for write failed: " + tmp_path.string() +
|
|
": " + strerror(errno));
|
|
}
|
|
try {
|
|
write_all(fd, buf, len);
|
|
} catch (...) {
|
|
::close(fd);
|
|
std::filesystem::remove(tmp_path);
|
|
throw;
|
|
}
|
|
::close(fd);
|
|
|
|
std::error_code ec;
|
|
std::filesystem::rename(tmp_path, file_path, ec);
|
|
if (ec) {
|
|
std::error_code remove_ec;
|
|
std::filesystem::remove(tmp_path, remove_ec);
|
|
throw std::runtime_error("rename failed: " + ec.message());
|
|
}
|
|
}
|
|
|
|
bool ExampleFSConnector::do_single_exists(WorkerFSConn& conn,
|
|
const std::string& key) {
|
|
auto path = conn.base_path / safe_filename(key);
|
|
return std::filesystem::exists(path);
|
|
}
|
|
|
|
bool ExampleFSConnector::do_single_delete(WorkerFSConn& conn,
|
|
const std::string& key) {
|
|
auto path = conn.base_path / safe_filename(key);
|
|
std::error_code ec;
|
|
return std::filesystem::remove(path, ec);
|
|
}
|
|
|
|
// ---------------------------------------------------------------
|
|
// ExampleMemoryConnector
|
|
// ---------------------------------------------------------------
|
|
|
|
ExampleMemoryConnector::ExampleMemoryConnector(int num_workers)
|
|
: ConnectorBase(num_workers),
|
|
store_(std::make_shared<SharedMemoryStore>()) {
|
|
start_workers();
|
|
}
|
|
|
|
ExampleMemoryConnector::~ExampleMemoryConnector() { close(); }
|
|
|
|
WorkerMemConn ExampleMemoryConnector::create_connection() {
|
|
WorkerMemConn conn;
|
|
conn.store = store_;
|
|
return conn;
|
|
}
|
|
|
|
void ExampleMemoryConnector::do_single_get(WorkerMemConn& conn,
|
|
const std::string& key, void* buf,
|
|
size_t len, size_t /*chunk_size*/) {
|
|
std::lock_guard<std::mutex> lk(conn.store->mu);
|
|
auto it = conn.store->data.find(key);
|
|
if (it == conn.store->data.end()) {
|
|
throw std::runtime_error("key not found: " + key);
|
|
}
|
|
if (it->second.size() != len) {
|
|
throw std::runtime_error("size mismatch for key: " + key);
|
|
}
|
|
std::memcpy(buf, it->second.data(), len);
|
|
}
|
|
|
|
void ExampleMemoryConnector::do_single_set(WorkerMemConn& conn,
|
|
const std::string& key,
|
|
const void* buf, size_t len,
|
|
size_t /*chunk_size*/) {
|
|
std::lock_guard<std::mutex> lk(conn.store->mu);
|
|
auto& vec = conn.store->data[key];
|
|
vec.resize(len);
|
|
std::memcpy(vec.data(), buf, len);
|
|
}
|
|
|
|
bool ExampleMemoryConnector::do_single_exists(WorkerMemConn& conn,
|
|
const std::string& key) {
|
|
std::lock_guard<std::mutex> lk(conn.store->mu);
|
|
return conn.store->data.count(key) > 0;
|
|
}
|
|
|
|
bool ExampleMemoryConnector::do_single_delete(WorkerMemConn& conn,
|
|
const std::string& key) {
|
|
std::lock_guard<std::mutex> lk(conn.store->mu);
|
|
return conn.store->data.erase(key) > 0;
|
|
}
|
|
|
|
} // namespace example_connector
|