Remove header filtering, since it's fully configurable in nginx

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "firebuf"]
 	path = firebuf
 	url = ../firebuf.git
+[submodule "fireusage"]
+	path = fireusage
+	url = ../fireusage.git

Makefile

@@ -10,13 +10,16 @@ objects = server.o connection.o request.o parse.o
 firebuf/firebuf.o:
 	$(MAKE) --directory=firebuf firebuf.o
 
+fireusage/fireusage.o:
+	$(MAKE) --directory=fireusage fireusage.o
+
 firecgi.a: $(objects)
 	ar rcs $@ $^
 
-firecgi.o: $(objects) firebuf/firebuf.o
+firecgi.o: $(objects) firebuf/firebuf.o fireusage/fireusage.o
 	gold -z relro -z now -r --output=$@ $+
 
-firecgi.so: $(objects) firebuf/firebuf.o
+firecgi.so: $(objects) firebuf/firebuf.o fireusage/fireusage.o
 	$(FIRE_CXX) $(FIRE_CXXFLAGS) $(FIRE_LDFLAGS) -shared -o $@ $+ $(FIRE_LDFLIBS)
 
 example_simple: example_simple.o firecgi.o
@@ -27,6 +30,7 @@ example_simple: example_simple.o firecgi.o
 
 clean:
 	$(MAKE) --directory=firebuf clean
+	$(MAKE) --directory=fireusage clean
 	rm --force example_simple connection_afl *.so *.o *.a
 
 afl:

connection.cc

@@ -9,10 +9,9 @@
 
 namespace firecgi {
 
-Connection::Connection(int sock, const sockaddr_in6& client_addr, const std::function<void(Request*)>& callback, const std::unordered_set<std::string_view>& headers, int max_request_len)
+Connection::Connection(int sock, const sockaddr_in6& client_addr, const std::function<void(Request*)>& callback, int max_request_len)
 		: sock_(sock),
 		  callback_(callback),
-		  headers_(headers),
 		  buf_(sock, max_request_len),
 		  request_(this) {
 	char client_addr_str[INET6_ADDRSTRLEN];
@@ -104,11 +103,9 @@ int Connection::Read() {
 					}
 					std::string_view value(value_buf, param_header->value_length);
 
-					if (headers_.find(key) != headers_.end()) {
 					request_.AddParam(key, value);
 				}
 			}
-			}
 			break;
 
 		  case 5:
@@ -147,4 +144,8 @@ int Connection::Read() {
 	return -1;
 }
 
+uint64_t Connection::Requests() const {
+	return requests_;
+}
+
 } // namespace firecgi

connection.h

@@ -13,16 +13,17 @@ namespace firecgi {
 
 class Connection {
   public:
-	Connection(int sock, const sockaddr_in6& client_addr, const std::function<void(Request*)>& callback, const std::unordered_set<std::string_view>& headers, int max_request_len);
+	Connection(int sock, const sockaddr_in6& client_addr, const std::function<void(Request*)>& callback, int max_request_len);
 	~Connection();
 
 	[[nodiscard]] int Read();
 	[[nodiscard]] bool Write(const std::vector<iovec>& vecs);
 
+	[[nodiscard]] uint64_t Requests() const;
+
   private:
   	const int sock_;
 	const std::function<void(Request*)>& callback_;
-	const std::unordered_set<std::string_view>& headers_;
 	firebuf::StreamBuffer buf_;
 	Request request_;
 

example_simple.cc

@@ -15,5 +15,9 @@ int main(int argc, char *argv[]) {
 		request->WriteBody("Hello world");
 		request->End();
 	}, FLAGS_threads);
+	server.RegisterSignalHandlers();
 	server.Serve();
+
+	gflags::ShutDownCommandLineFlags();
+	google::ShutdownGoogleLogging();
 }

request.cc

@@ -28,7 +28,7 @@ void Request::NewRequest(uint16_t request_id) {
 	body_written_ = false;
 }
 
-uint16_t Request::RequestId() {
+uint16_t Request::RequestId() const {
 	return request_id_;
 }
 
@@ -40,7 +40,7 @@ void Request::SetBody(const std::string_view& body) {
 	body_ = body;
 }
 
-const std::string_view& Request::GetParam(const std::string_view& key) {
+const std::string_view& Request::GetParam(const std::string_view& key) const {
 	auto iter = params_.find(key);
 	if (iter == params_.end()) {
 		static const std::string_view none;
@@ -49,7 +49,7 @@ const std::string_view& Request::GetParam(const std::string_view& key) {
 	return iter->second;
 }
 
-const std::string_view& Request::GetBody() {
+const std::string_view& Request::GetBody() const {
 	return body_;
 }
 

request.h

@@ -18,13 +18,13 @@ class Request {
 
 	void NewRequest(uint16_t request_id);
 
-	uint16_t RequestId();
+	uint16_t RequestId() const;
 
 	void AddParam(const std::string_view& key, const std::string_view& value);
 	void SetBody(const std::string_view& in);
 
-	const std::string_view& GetParam(const std::string_view& key);
-	const std::string_view& GetBody();
+	const std::string_view& GetParam(const std::string_view& key) const;
+	const std::string_view& GetBody() const;
 
 	void WriteHeader(const std::string_view& name, const std::string_view& value);
 	void WriteBody(const std::string_view& body);

server.cc

@@ -1,24 +1,30 @@
 #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)
+Server::Server(int port, const std::function<void(Request*)>& callback, int threads, int max_request_len)
 		: port_(port),
 		  callback_(callback),
 		  threads_(threads),
-		  headers_(headers),
-		  max_request_len_(max_request_len) {
-	LOG(INFO) << "listening on [::1]:" << port_;
+		  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);
 }
 
@@ -28,6 +34,29 @@ void Server::Serve() {
 		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() {
@@ -36,16 +65,33 @@ void Server::ServeInt() {
 
 	auto listen_sock = NewListenSock();
 
+	char new_conn;
 	{
 		struct epoll_event ev{
 			.events = EPOLLIN,
 			.data = {
-				.ptr = nullptr,
+				.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];
@@ -56,13 +102,25 @@ void Server::ServeInt() {
 		PCHECK(num_fd > 0) << "epoll_wait()";
 
 		for (auto i = 0; i < num_fd; ++i) {
-			if (events[i].data.ptr == nullptr) {
-				NewConn(listen_sock, epoll_fd);
+			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;
 				}
 			}
@@ -70,7 +128,7 @@ void Server::ServeInt() {
 	}
 }
 
-void Server::NewConn(int listen_sock, int epoll_fd) {
+Connection* Server::NewConn(int listen_sock, int epoll_fd) {
 	sockaddr_in6 client_addr;
 	socklen_t client_addr_len = sizeof(client_addr);
 
@@ -81,8 +139,8 @@ void Server::NewConn(int listen_sock, int epoll_fd) {
 	int flags = 1;
 	PCHECK(setsockopt(client_sock, SOL_TCP, TCP_NODELAY, &flags, sizeof(flags)) == 0);
 
+	auto *conn = new Connection(client_sock, client_addr, callback_, max_request_len_);
 	{
-		auto *conn = new Connection(client_sock, client_addr, callback_, headers_, max_request_len_);
 		struct epoll_event ev{
 			.events = EPOLLIN,
 			.data = {
@@ -91,6 +149,8 @@ void Server::NewConn(int listen_sock, int epoll_fd) {
 		};
 		PCHECK(epoll_ctl(epoll_fd, EPOLL_CTL_ADD, client_sock, &ev) == 0);
 	}
+
+	return conn;
 }
 
 int Server::NewListenSock() {

server.h

@@ -10,19 +10,22 @@ namespace firecgi {
 
 class Server {
   public:
-	Server(int port, const std::function<void(Request*)>& callback, int threads=1, const std::unordered_set<std::string_view>& headers={}, int max_request_len=(16*1024));
+	Server(int port, const std::function<void(Request*)>& callback, int threads=1, int max_request_len=(16*1024));
 	void Serve();
+	void Shutdown();
+	void RegisterSignalHandlers();
 
   private:
-	void NewConn(int listen_sock, int epoll_fd);
+	Connection *NewConn(int listen_sock, int epoll_fd);
 	int NewListenSock();
 	void ServeInt();
 
 	const int port_;
 	const std::function<void(Request*)> callback_;
 	const int threads_;
-	const std::unordered_set<std::string_view> headers_;
 	const int max_request_len_;
+
+	int close_fd_;
 };
 
 } // firecgi