Files
weaseldb/src/main.cpp
2025-08-18 13:33:58 -04:00

345 lines
9.1 KiB
C++

#include "config.hpp"
#include <atomic>
#include <cassert>
#include <csignal>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <inttypes.h>
#include <iostream>
#include <netdb.h>
#include <netinet/tcp.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/un.h>
#include <thread>
#include <unistd.h>
#ifndef __has_feature
#define __has_feature(x) 0
#endif
// Adapted from getaddrinfo man page
int getListenFd(const char *node, const char *service) {
struct addrinfo hints;
struct addrinfo *result, *rp;
int sfd, s;
memset(&hints, 0, sizeof(hints));
hints.ai_family = AF_UNSPEC; /* Allow IPv4 or IPv6 */
hints.ai_socktype = SOCK_STREAM; /* stream socket */
hints.ai_flags = AI_PASSIVE; /* For wildcard IP address */
hints.ai_protocol = 0; /* Any protocol */
hints.ai_canonname = nullptr;
hints.ai_addr = nullptr;
hints.ai_next = nullptr;
s = getaddrinfo(node, service, &hints, &result);
if (s != 0) {
fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(s));
abort();
}
/* getaddrinfo() returns a list of address structures.
Try each address until we successfully bind(2).
If socket(2) (or bind(2)) fails, we (close the socket
and) try the next address. */
for (rp = result; rp != nullptr; rp = rp->ai_next) {
sfd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
if (sfd == -1) {
continue;
}
int val = 1;
setsockopt(sfd, SOL_SOCKET, SO_REUSEADDR, &val, sizeof(val));
if (bind(sfd, rp->ai_addr, rp->ai_addrlen) == 0) {
break; /* Success */
}
close(sfd);
}
freeaddrinfo(result); /* No longer needed */
if (rp == nullptr) { /* No address succeeded */
fprintf(stderr, "Could not bind\n");
abort();
}
int rv = listen(sfd, SOMAXCONN);
if (rv) {
perror("listen");
abort();
}
return sfd;
}
int getAcceptFd(int listenFd, struct sockaddr *addr) {
socklen_t addrlen = sizeof(sockaddr);
int fd = accept4(listenFd, addr, &addrlen, SOCK_NONBLOCK);
return fd;
}
// Connection lifecycle. Only one of these is the case at a time
// - Created on an accept thread from a call to accept
// - Waiting on connection fd to be readable/writable
// - Owned by a network thread, which drains readable and writable bytes
// - Owned by a thread in the request processing pipeline
// - Closed by a network thread according to http protocol
//
// Since only one thread owns a connection at a time, no synchronization is
// necessary
struct Connection {
const int fd;
const int64_t id;
struct sockaddr addr;
Connection(struct sockaddr addr, int fd, int64_t id)
: fd(fd), id(id), addr(addr) {}
~Connection() {
int e = close(fd);
if (e == -1) {
perror("close");
abort();
}
}
struct Task {
std::string s;
bool closeConnection{false};
int written = 0;
};
std::deque<Task> tasks;
void readBytes() {
for (;;) {
// TODO make size configurable
char buf[1024];
int r = read(fd, buf, sizeof(buf));
if (r == -1) {
if (errno == EINTR) {
continue;
}
if (errno == EAGAIN) {
return;
}
perror("read");
goto close_connection;
}
if (r == 0) {
goto close_connection;
}
// pump parser
tasks.emplace_back(std::string{buf, size_t(r)});
}
close_connection:
tasks.emplace_back(std::string{}, true);
}
bool writeBytes() {
while (!tasks.empty()) {
auto &front = tasks.front();
if (front.closeConnection) {
return true;
}
int w;
for (;;) {
w = write(fd, front.s.data() + front.written,
front.s.size() - front.written);
if (w == -1) {
if (errno == EINTR) {
continue;
}
if (errno == EAGAIN) {
return false;
}
perror("write");
return true;
}
break;
}
assert(w != 0);
front.written += w;
if (front.written == int(front.s.size())) {
tasks.pop_front();
}
}
return false;
}
#if __has_feature(thread_sanitizer)
void tsan_acquire() { tsan_sync.load(std::memory_order_acquire); }
void tsan_release() { tsan_sync.store(0, std::memory_order_release); }
std::atomic<int> tsan_sync;
#else
void tsan_acquire() {}
void tsan_release() {}
#endif
};
int main(int argc, char *argv[]) {
std::string config_file = "config.toml";
if (argc > 1) {
config_file = argv[1];
}
auto config = weaseldb::ConfigParser::load_from_file(config_file);
if (!config) {
std::cerr << "Failed to load config from: " << config_file << std::endl;
std::cerr << "Using default configuration..." << std::endl;
config = weaseldb::Config{};
}
std::cout << "Configuration loaded successfully:" << std::endl;
std::cout << "Server bind address: " << config->server.bind_address
<< std::endl;
std::cout << "Server port: " << config->server.port << std::endl;
std::cout << "Max request size: " << config->server.max_request_size_bytes
<< " bytes" << std::endl;
std::cout << "Min request ID length: " << config->commit.min_request_id_length
<< std::endl;
std::cout << "Request ID retention: "
<< config->commit.request_id_retention_hours.count() << " hours"
<< std::endl;
std::cout << "Subscription buffer size: "
<< config->subscription.max_buffer_size_bytes << " bytes"
<< std::endl;
std::cout << "Keepalive interval: "
<< config->subscription.keepalive_interval.count() << " seconds"
<< std::endl;
signal(SIGPIPE, SIG_IGN);
int sockfd = getListenFd(config->server.bind_address.c_str(),
std::to_string(config->server.port).c_str());
std::vector<std::thread> threads;
int epollfd = epoll_create(/*ignored*/ 1);
if (epollfd == -1) {
perror("epoll_create");
abort();
}
// Network threads
// TODO make configurable
int networkThreads = 1;
// TODO make configurable
constexpr int kEventBatchSize = 10;
for (int i = 0; i < networkThreads; ++i) {
threads.emplace_back([epollfd, i]() {
pthread_setname_np(pthread_self(),
("network-" + std::to_string(i)).c_str());
for (;;) {
struct epoll_event events[kEventBatchSize]{};
int eventCount;
for (;;) {
eventCount =
epoll_wait(epollfd, events, kEventBatchSize, /*no timeout*/ -1);
if (eventCount == -1) {
if (errno == EINTR) {
continue;
}
perror("epoll_wait");
abort();
}
break;
}
for (int i = 0; i < eventCount; ++i) {
std::unique_ptr<Connection> conn{
static_cast<Connection *>(events[i].data.ptr)};
conn->tsan_acquire();
events[i].data.ptr = nullptr;
const int fd = conn->fd;
if (events[i].events & EPOLLERR) {
// Done with connection
continue;
}
if (events[i].events & EPOLLOUT) {
// Write bytes, maybe close connection
bool finished = conn->writeBytes();
if (finished) {
// Done with connection
continue;
}
}
if (events[i].events & EPOLLIN) {
conn->readBytes();
}
if (events[i].events & EPOLLOUT) {
bool done = conn->writeBytes();
if (done) {
continue;
}
}
if (conn->tasks.empty()) {
// Transfer back to epoll instance. This thread or another thread
// will wake when fd is ready
events[i].events = EPOLLIN | EPOLLONESHOT;
} else {
events[i].events = EPOLLOUT | EPOLLONESHOT;
}
conn->tsan_release();
events[i].data.ptr = conn.release();
int e = epoll_ctl(epollfd, EPOLL_CTL_MOD, fd, &events[i]);
if (e == -1) {
perror("epoll_ctl");
abort();
}
}
}
});
}
std::atomic<int64_t> connectionId{0};
// TODO make configurable
int acceptThreads = 1;
for (int i = 0; i < acceptThreads; ++i) {
threads.emplace_back([epollfd, i, sockfd, &connectionId]() {
pthread_setname_np(pthread_self(),
("accept-" + std::to_string(i)).c_str());
// Call accept in a loop
for (;;) {
struct sockaddr addr;
int fd = getAcceptFd(sockfd, &addr);
if (fd == -1) {
perror("accept4");
continue;
}
auto conn = std::make_unique<Connection>(
addr, fd, connectionId.fetch_add(1, std::memory_order_relaxed));
// Post to epoll instance
struct epoll_event event{};
event.events = EPOLLIN | EPOLLONESHOT;
conn->tsan_release();
event.data.ptr = conn.release();
int e = epoll_ctl(epollfd, EPOLL_CTL_ADD, fd, &event);
if (e == -1) {
perror("epoll_ctl");
abort();
}
}
});
}
for (auto &t : threads) {
t.join();
}
return 0;
}