firecgi/server.cc

#include <arpa/inet.h>
#include <iomanip>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <signal.h>
#include <sys/epoll.h>
#include <sys/eventfd.h>
#include <sys/resource.h>
#include <sys/socket.h>
#include <thread>

#include "fireusage/usage.h"

#include "server.h"
#include "connection.h"

namespace firecgi {

Server::Server(int port, const std::function<void(Request*)>& callback, int threads, const std::unordered_set<std::string_view>& headers, int max_request_len)
		: port_(port),
		  callback_(callback),
		  threads_(threads),
		  headers_(headers),
		  max_request_len_(max_request_len),
		  close_fd_(eventfd(0, 0)) {
	CHECK_GE(close_fd_, 0);

	LOG(INFO) << "listening on [::1]:" << port_;
	signal(SIGPIPE, SIG_IGN);
}

void Server::Serve() {
	std::vector<std::thread> threads;
	for (int i = 0; i < threads_ - 1; ++i) {
		threads.emplace_back([this]() { ServeInt(); });
	}
	ServeInt();
	for (auto& thread : threads) {
		thread.join();
	}
	LOG(INFO) << "all threads shut down";
}

void Server::Shutdown() {
	uint64_t shutdown = 1;
	PCHECK(write(close_fd_, &shutdown, sizeof(shutdown)) == sizeof(shutdown));
}

namespace {
Server* shutdown_server = nullptr;
} // namespace

void Server::RegisterSignalHandlers() {
	shutdown_server = this;
	for (auto sig : {SIGINT, SIGTERM}) {
		signal(sig, [](int signum) {
			LOG(INFO) << "received " << strsignal(signum);
			shutdown_server->Shutdown();
		});
	}
}

void Server::ServeInt() {
	auto epoll_fd = epoll_create1(0);
	PCHECK(epoll_fd >= 0) << "epoll_create()";

	auto listen_sock = NewListenSock();

	char new_conn;
	{
		struct epoll_event ev{
			.events = EPOLLIN,
			.data = {
				.ptr = &new_conn,
			},
		};
		PCHECK(epoll_ctl(epoll_fd, EPOLL_CTL_ADD, listen_sock, &ev) == 0);
	}

	char shutdown;
	{
		struct epoll_event ev{
			.events = EPOLLIN,
			.data = {
				.ptr = &shutdown,
			},
		};
		PCHECK(epoll_ctl(epoll_fd, EPOLL_CTL_ADD, close_fd_, &ev) == 0);
	}

	std::unordered_set<Connection*> connections;

	fireusage::UsageTracker usage_tracker;
	usage_tracker.Start();

	while (true) {
		constexpr auto max_events = 256;
		struct epoll_event events[max_events];
		auto num_fd = epoll_wait(epoll_fd, events, max_events, -1);
		if (num_fd == -1 && errno == EINTR) {
			continue;
		}
		PCHECK(num_fd > 0) << "epoll_wait()";

		for (auto i = 0; i < num_fd; ++i) {
			if (events[i].data.ptr == &new_conn) {
				connections.insert(CHECK_NOTNULL(NewConn(listen_sock, epoll_fd)));
			} else if (events[i].data.ptr == &shutdown) {
				for (auto& conn : connections) {
					usage_tracker.AddEvents(conn->Requests());
					delete conn;
				}
				usage_tracker.Stop();
				PCHECK(close(listen_sock) == 0);
				PCHECK(close(epoll_fd) == 0);
				usage_tracker.Log();
				return;
			} else {
				auto conn = static_cast<Connection*>(events[i].data.ptr);
				auto fd = conn->Read();
				if (fd != -1) {
					PCHECK(epoll_ctl(epoll_fd, EPOLL_CTL_DEL, fd, nullptr) == 0);
					usage_tracker.AddEvents(conn->Requests());
					connections.erase(conn);
					delete conn;
				}
			}
		}
	}
}

Connection* Server::NewConn(int listen_sock, int epoll_fd) {
	sockaddr_in6 client_addr;
	socklen_t client_addr_len = sizeof(client_addr);

	auto client_sock = accept(listen_sock, (sockaddr*) &client_addr, &client_addr_len);
	PCHECK(client_sock >= 0) << "accept()";
	CHECK_EQ(client_addr.sin6_family, AF_INET6);

	int flags = 1;
	PCHECK(setsockopt(client_sock, SOL_TCP, TCP_NODELAY, &flags, sizeof(flags)) == 0);

	auto *conn = new Connection(client_sock, client_addr, callback_, headers_, max_request_len_);
	{
		struct epoll_event ev{
			.events = EPOLLIN,
			.data = {
				.ptr = conn,
			},
		};
		PCHECK(epoll_ctl(epoll_fd, EPOLL_CTL_ADD, client_sock, &ev) == 0);
	}

	return conn;
}

int Server::NewListenSock() {
	auto sock = socket(AF_INET6, SOCK_STREAM, 0);
	PCHECK(sock >= 0) << "socket()";

	{
		int optval = 1;
		PCHECK(setsockopt(sock, SOL_SOCKET, SO_REUSEPORT, &optval, sizeof(optval)) == 0);
	}

	{
		sockaddr_in6 bind_addr = {
			.sin6_family = AF_INET6,
			.sin6_port = htons(port_),
			.sin6_addr = IN6ADDR_LOOPBACK_INIT,
		};
		PCHECK(bind(sock, (sockaddr*) &bind_addr, sizeof(bind_addr)) == 0);
	}

	PCHECK(listen(sock, 128) == 0);
	return sock;
}

} // namespace firecgi