wehub-resource-sync
378 lines
11 KiB
C++
378 lines
11 KiB
C++
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
#include "connector.h"
|
|
#include <cerrno>
|
|
#include <stdexcept>
|
|
#include <string_view>
|
|
|
|
namespace lmcache {
|
|
namespace connector {
|
|
|
|
void WorkerConn::connect(const std::string& h, int p) {
|
|
host = h;
|
|
port = p;
|
|
|
|
// 1. create socket
|
|
fd = socket(AF_INET, SOCK_STREAM, 0);
|
|
if (fd < 0) {
|
|
throw std::runtime_error("failed to create socket");
|
|
}
|
|
|
|
// 2. resolve host
|
|
struct addrinfo hints = {}, *result = nullptr;
|
|
hints.ai_family = AF_INET;
|
|
hints.ai_socktype = SOCK_STREAM;
|
|
|
|
std::string port_str = std::to_string(p);
|
|
int err = getaddrinfo(h.c_str(), port_str.c_str(), &hints, &result);
|
|
|
|
if (err != 0) {
|
|
::close(fd);
|
|
throw std::runtime_error(std::string("getaddrinfo failed: ") +
|
|
gai_strerror(err));
|
|
}
|
|
|
|
// 3. connect to host
|
|
if (::connect(fd, result->ai_addr, result->ai_addrlen) < 0) {
|
|
freeaddrinfo(result);
|
|
::close(fd);
|
|
throw std::runtime_error("connection failed");
|
|
}
|
|
|
|
freeaddrinfo(result);
|
|
}
|
|
|
|
// call send multiple times until all the data is sent
|
|
void WorkerConn::send_all(const void* data, size_t len) {
|
|
size_t sent_so_far = 0;
|
|
const char* ptr = static_cast<const char*>(data);
|
|
while (sent_so_far < len) {
|
|
ssize_t n = ::send(fd, ptr + sent_so_far, len - sent_so_far, 0);
|
|
if (n < 0) {
|
|
if (errno == EINTR) {
|
|
continue; // retry on EINTR
|
|
}
|
|
throw std::runtime_error("socket send failed");
|
|
}
|
|
if (n == 0) {
|
|
throw std::runtime_error("socket send failed: connection closed");
|
|
}
|
|
sent_so_far += n;
|
|
}
|
|
}
|
|
|
|
// scatter gather send
|
|
void WorkerConn::send_multipart(
|
|
const std::vector<std::pair<const void*, size_t>>& parts) {
|
|
if (parts.empty()) return;
|
|
|
|
std::vector<struct iovec> iov;
|
|
iov.reserve(parts.size());
|
|
for (const auto& part : parts) {
|
|
iov.push_back({const_cast<void*>(part.first), part.second});
|
|
}
|
|
|
|
size_t total_to_send = 0;
|
|
for (const auto& part : parts) {
|
|
total_to_send += part.second;
|
|
}
|
|
|
|
size_t sent_so_far = 0;
|
|
size_t iov_idx = 0;
|
|
|
|
while (sent_so_far < total_to_send) {
|
|
ssize_t n = ::writev(fd, &iov[iov_idx], iov.size() - iov_idx);
|
|
if (n < 0) {
|
|
if (errno == EINTR) {
|
|
continue;
|
|
}
|
|
throw std::runtime_error("socket writev failed");
|
|
}
|
|
if (n == 0) {
|
|
throw std::runtime_error("socket writev failed: connection closed");
|
|
}
|
|
|
|
sent_so_far += n;
|
|
|
|
// adjust iovec for partial writes
|
|
size_t remaining = n;
|
|
while (remaining > 0 && iov_idx < iov.size()) {
|
|
if (remaining >= iov[iov_idx].iov_len) {
|
|
remaining -= iov[iov_idx].iov_len;
|
|
iov_idx++;
|
|
} else {
|
|
iov[iov_idx].iov_base =
|
|
static_cast<char*>(iov[iov_idx].iov_base) + remaining;
|
|
iov[iov_idx].iov_len -= remaining;
|
|
remaining = 0;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
void WorkerConn::recv_exactly(void* buf, size_t len) {
|
|
size_t recv_so_far = 0;
|
|
char* ptr = static_cast<char*>(buf);
|
|
while (recv_so_far < len) {
|
|
ssize_t n = ::recv(fd, ptr + recv_so_far, len - recv_so_far, 0);
|
|
if (n < 0) {
|
|
if (errno == EINTR) {
|
|
continue;
|
|
}
|
|
throw std::runtime_error("socket recv failed");
|
|
}
|
|
if (n == 0) {
|
|
throw std::runtime_error("socket recv failed: connection closed");
|
|
}
|
|
recv_so_far += n;
|
|
}
|
|
}
|
|
|
|
std::string WorkerConn::recv_line() {
|
|
std::string line;
|
|
line.reserve(128);
|
|
for (;;) {
|
|
char c;
|
|
recv_exactly(&c, 1);
|
|
line.push_back(c);
|
|
size_t n = line.size();
|
|
if (n >= 2 && line[n - 2] == '\r' && line[n - 1] == '\n') {
|
|
return line;
|
|
}
|
|
}
|
|
}
|
|
|
|
void WorkerConn::authenticate_if_needed() {
|
|
if (authed) return;
|
|
if (password.empty()) return;
|
|
|
|
if (!username.empty()) {
|
|
// auth username password
|
|
std::string u =
|
|
"$" + std::to_string(username.size()) + "\r\n" + username + "\r\n";
|
|
std::string p =
|
|
"$" + std::to_string(password.size()) + "\r\n" + password + "\r\n";
|
|
static constexpr std::string_view auth_prefix = "*3\r\n$4\r\nAUTH\r\n";
|
|
|
|
send_multipart({
|
|
{auth_prefix.data(), auth_prefix.size()},
|
|
{u.data(), u.size()},
|
|
{p.data(), p.size()},
|
|
});
|
|
} else {
|
|
// auth password
|
|
std::string p =
|
|
"$" + std::to_string(password.size()) + "\r\n" + password + "\r\n";
|
|
static constexpr std::string_view auth_prefix = "*2\r\n$4\r\nAUTH\r\n";
|
|
|
|
send_multipart({
|
|
{auth_prefix.data(), auth_prefix.size()},
|
|
{p.data(), p.size()},
|
|
});
|
|
}
|
|
|
|
std::string line = recv_line();
|
|
if (line.rfind("+OK\r\n", 0) == 0) {
|
|
authed = true;
|
|
return;
|
|
}
|
|
if (!line.empty() && line[0] == '-') {
|
|
throw std::runtime_error("AUTH failed: " + line);
|
|
}
|
|
throw std::runtime_error("AUTH failed: unexpected reply: " + line);
|
|
}
|
|
|
|
const std::string& WorkerConn::make_key_header(const std::string& key) {
|
|
key_header_buf.clear();
|
|
|
|
const size_t needed = key.size() + 16;
|
|
if (needed > key_header_buf.capacity()) {
|
|
key_header_buf.reserve(needed);
|
|
}
|
|
|
|
key_header_buf += '$';
|
|
key_header_buf += std::to_string(key.size());
|
|
key_header_buf += crlf;
|
|
key_header_buf += key;
|
|
key_header_buf += crlf;
|
|
|
|
return key_header_buf;
|
|
}
|
|
|
|
const std::string& WorkerConn::make_size_header(size_t batch_chunk_num_bytes) {
|
|
size_header_buf.clear();
|
|
|
|
size_header_buf += '$';
|
|
size_header_buf += std::to_string(batch_chunk_num_bytes);
|
|
size_header_buf += crlf;
|
|
|
|
return size_header_buf;
|
|
}
|
|
|
|
/*
|
|
RESP set, get, exists are very fragile here since we make hard assumptions about
|
|
the RESP responses. a single error could break our assumptions. we can mitigate
|
|
this by actually parsing the headers and trailers.
|
|
*/
|
|
|
|
RedisConnector::RedisConnector(std::string host, int port, int num_workers,
|
|
std::string username, std::string password)
|
|
: ConnectorBase(num_workers),
|
|
host_(std::move(host)),
|
|
port_(port),
|
|
username_(std::move(username)),
|
|
password_(std::move(password)) {
|
|
worker_fds_.reserve(num_workers);
|
|
start_workers(); // start after derived class is fully constructed
|
|
}
|
|
|
|
RedisConnector::~RedisConnector() { close(); }
|
|
|
|
WorkerConn RedisConnector::create_connection() {
|
|
WorkerConn conn;
|
|
conn.connect(host_, port_);
|
|
|
|
conn.username = username_;
|
|
conn.password = password_;
|
|
conn.authenticate_if_needed();
|
|
|
|
// track socket fd for shutdown
|
|
{
|
|
std::lock_guard<std::mutex> lk(worker_fds_mu_);
|
|
worker_fds_.push_back(conn.fd);
|
|
}
|
|
|
|
return conn;
|
|
}
|
|
|
|
void RedisConnector::do_single_get(WorkerConn& conn, const std::string& key,
|
|
void* buf, size_t len, size_t chunk_size) {
|
|
if (len != chunk_size) {
|
|
throw std::runtime_error("buffer size mismatch");
|
|
}
|
|
|
|
const std::string& size_header = conn.make_size_header(chunk_size);
|
|
const std::string& key_header = conn.make_key_header(key);
|
|
|
|
conn.send_multipart({{conn.get_prefix.data(), conn.get_prefix.size()},
|
|
{key_header.data(), key_header.size()}});
|
|
|
|
// parse response in 3 steps
|
|
|
|
// 1. recv size header
|
|
std::vector<char> recv_size_header_buf(size_header.size());
|
|
conn.recv_exactly(recv_size_header_buf.data(), size_header.size());
|
|
if (std::memcmp(recv_size_header_buf.data(), size_header.data(),
|
|
size_header.size()) != 0) {
|
|
throw std::runtime_error("GET: size header mismatch");
|
|
}
|
|
|
|
// 2. recv KV Cache (payload) without parsing
|
|
conn.recv_exactly(buf, len);
|
|
|
|
// 3. recv trailer
|
|
char trailer[WorkerConn::crlf_len];
|
|
conn.recv_exactly(trailer, WorkerConn::crlf_len);
|
|
if (std::memcmp(trailer, WorkerConn::crlf.data(), WorkerConn::crlf_len) !=
|
|
0) {
|
|
throw std::runtime_error("GET: trailer mismatch");
|
|
}
|
|
}
|
|
|
|
// RESP SET
|
|
void RedisConnector::do_single_set(WorkerConn& conn, const std::string& key,
|
|
const void* buf, size_t len,
|
|
size_t chunk_size) {
|
|
// we only write exactly batch_chunk_num_bytes bytes (save_unfull_chunk must
|
|
// be off)
|
|
if (len != chunk_size) {
|
|
throw std::runtime_error("buffer size mismatch");
|
|
}
|
|
|
|
// build headers using reusable buffers
|
|
const std::string& size_header = conn.make_size_header(chunk_size);
|
|
const std::string& key_header = conn.make_key_header(key);
|
|
|
|
// send SET cmd
|
|
// iovec let's us combine pre-built parts and dynamic strings
|
|
conn.send_multipart({{conn.set_prefix.data(), conn.set_prefix.size()},
|
|
{key_header.data(), key_header.size()},
|
|
{size_header.data(), size_header.size()},
|
|
{buf, len},
|
|
{WorkerConn::crlf.data(), WorkerConn::crlf_len}});
|
|
|
|
// parse response which should be exactly +OK\r\n
|
|
char response[WorkerConn::ok_response_len];
|
|
conn.recv_exactly(response, WorkerConn::ok_response_len);
|
|
|
|
if (std::memcmp(response, WorkerConn::ok_response.data(),
|
|
WorkerConn::ok_response_len) != 0) {
|
|
throw std::runtime_error("SET: response was not OK");
|
|
}
|
|
}
|
|
|
|
// RESP EXISTS
|
|
bool RedisConnector::do_single_exists(WorkerConn& conn,
|
|
const std::string& key) {
|
|
// build key header using reusable buffer
|
|
const std::string& key_header = conn.make_key_header(key);
|
|
|
|
// send EXISTS cmd
|
|
// iovec let's us combine pre-built parts and dynamic strings
|
|
conn.send_multipart({{conn.exists_prefix.data(), conn.exists_prefix.size()},
|
|
{key_header.data(), key_header.size()}});
|
|
|
|
// parse response (either :0\r\n or :1\r\n)
|
|
char response[WorkerConn::exists_response_len];
|
|
conn.recv_exactly(response, WorkerConn::exists_response_len);
|
|
|
|
if (std::memcmp(response, WorkerConn::exists_one.data(),
|
|
WorkerConn::exists_response_len) == 0) {
|
|
return true;
|
|
} else if (std::memcmp(response, WorkerConn::exists_zero.data(),
|
|
WorkerConn::exists_response_len) == 0) {
|
|
return false;
|
|
} else {
|
|
throw std::runtime_error(
|
|
"EXISTS returned invalid response that wasn't :0\r\n or :1\r\n");
|
|
}
|
|
}
|
|
|
|
// RESP DEL
|
|
bool RedisConnector::do_single_delete(WorkerConn& conn,
|
|
const std::string& key) {
|
|
// build key header using reusable buffer
|
|
const std::string& key_header = conn.make_key_header(key);
|
|
|
|
// send DEL cmd
|
|
conn.send_multipart({{conn.del_prefix.data(), conn.del_prefix.size()},
|
|
{key_header.data(), key_header.size()}});
|
|
|
|
// parse response (either :0\r\n or :1\r\n, same format as EXISTS)
|
|
char response[WorkerConn::exists_response_len];
|
|
conn.recv_exactly(response, WorkerConn::exists_response_len);
|
|
|
|
if (std::memcmp(response, WorkerConn::exists_one.data(),
|
|
WorkerConn::exists_response_len) == 0) {
|
|
return true; // key was deleted
|
|
} else if (std::memcmp(response, WorkerConn::exists_zero.data(),
|
|
WorkerConn::exists_response_len) == 0) {
|
|
return false; // key did not exist
|
|
} else {
|
|
throw std::runtime_error(
|
|
"DEL returned invalid response that wasn't :0\r\n or :1\r\n");
|
|
}
|
|
}
|
|
|
|
void RedisConnector::shutdown_connections() {
|
|
std::lock_guard<std::mutex> lk(worker_fds_mu_);
|
|
for (int fd : worker_fds_) {
|
|
if (fd >= 0) {
|
|
::shutdown(fd, SHUT_RDWR);
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace connector
|
|
} // namespace lmcache
|