Reactor/lib/TcpServer.hpp

#pragma once

#include "Core.hpp"
#include "Socket.hpp"
#include "TcpConnection.hpp"
#include "EventLoopThread.hpp"
#include "Utilities.hpp"
#include <unordered_map>
#include <string>
#include <atomic>
#include <functional>

namespace reactor {

using NewConnectionCallback = std::function<void(int, const InetAddress&)>;

class Acceptor : public NonCopyable
{
private:
	EventLoop* loop_;
	Socket acceptSocket_;
	std::unique_ptr<Channel> acceptChannel_;
	NewConnectionCallback newConnectionCallback_;
	bool listening_;
	int idleFd_;

	void handleRead()
	{
		loop_->assertInLoopThread();
		InetAddress peerAddr;
		int connfd = acceptSocket_.accept(peerAddr);

		if (connfd >= 0) {
			if (newConnectionCallback_) {
				newConnectionCallback_(connfd, peerAddr);
			} else {
				close(connfd);
				LOG_WARN << "Acceptor no callback for new connection, closing fd=" << connfd;
			}
		} else {
			LOG_ERROR << "Acceptor accept failed: " << strerror(errno);
			if (errno == EMFILE) {
				close(idleFd_);
				idleFd_ = ::accept(acceptSocket_.fd(), nullptr, nullptr);
				close(idleFd_);
				idleFd_ = ::open("/dev/null", O_RDONLY | O_CLOEXEC);
			}
		}
	}

public:
	Acceptor(EventLoop* loop, const InetAddress& listenAddr, bool reusePort = true)
		: loop_(loop), acceptSocket_(Socket::createTcp(listenAddr.isIpV6())),
		  acceptChannel_(std::make_unique<Channel>(loop, acceptSocket_.fd())),
		  listening_(false), idleFd_(::open("/dev/null", O_RDONLY | O_CLOEXEC))
	{
		acceptSocket_.setReuseAddr(true);
		if (reusePort) {
			acceptSocket_.setReusePort(true);
		}
		acceptSocket_.bind(listenAddr);
		acceptChannel_->setReadCallback([this]() { handleRead(); });
		LOG_INFO << "Acceptor created for " << listenAddr.toIpPort();
	}

	~Acceptor()
	{
		acceptChannel_->disableAll();
		acceptChannel_->remove();
		close(idleFd_);
		LOG_INFO << "Acceptor destroyed";
	}

	void listen()
	{
		loop_->assertInLoopThread();
		listening_ = true;
		acceptSocket_.listen();
		acceptChannel_->enableReading();
		LOG_INFO << "Acceptor listening";
	}

	bool listening() const { return listening_; }

	void setNewConnectionCallback(NewConnectionCallback cb)
	{
		newConnectionCallback_ = std::move(cb);
	}
};

class TcpServer : public NonCopyable
{
private:
	EventLoop* loop_;
	std::string name_;
	std::unique_ptr<Acceptor> acceptor_;
	std::unique_ptr<EventLoopThreadPool> threadPool_;
	MessageCallback messageCallback_;
	ConnectionCallback connectionCallback_;
	WriteCompleteCallback writeCompleteCallback_;

	std::unordered_map<std::string, TcpConnectionPtr> connections_;
	std::atomic<int> nextConnId_;
	bool started_;

	void newConnection(int sockfd, const InetAddress& peerAddr)
	{
		loop_->assertInLoopThread();
		EventLoop* ioLoop = threadPool_->getNextLoop();
		if (!ioLoop) ioLoop = loop_;

		std::string connName = name_ + "-" + peerAddr.toIpPort() + "#" + std::to_string(nextConnId_++);
		InetAddress localAddr = Socket::getLocalAddr(sockfd);

		LOG_INFO << "TcpServer new connection " << connName << " from " << peerAddr.toIpPort();

		auto conn = std::make_shared<TcpConnection>(ioLoop, connName, sockfd, localAddr, peerAddr);
		connections_[connName] = conn;

		conn->setMessageCallback(messageCallback_);
		conn->setConnectionCallback(connectionCallback_);
		conn->setWriteCompleteCallback(writeCompleteCallback_);
		conn->setCloseCallback([this](const TcpConnectionPtr& conn) {
			removeConnection(conn);
		});

		ioLoop->runInLoop([conn]() { conn->connectEstablished(); });
	}

	void removeConnection(const TcpConnectionPtr& conn)
	{
		loop_->runInLoop([this, conn]() {
			LOG_INFO << "TcpServer removing connection " << conn->name();
			size_t n = connections_.erase(conn->name());
			assert(n == 1);

			EventLoop* ioLoop = conn->getLoop();
			ioLoop->queueInLoop([conn]() { conn->connectDestroyed(); });
		});
	}

	void removeConnectionInLoop(const TcpConnectionPtr& conn)
	{
		loop_->assertInLoopThread();
		LOG_INFO << "TcpServer removing connection " << conn->name();
		size_t n = connections_.erase(conn->name());
		assert(n == 1);

		EventLoop* ioLoop = conn->getLoop();
		ioLoop->queueInLoop([conn]() { conn->connectDestroyed(); });
	}

public:
	TcpServer(EventLoop* loop, const InetAddress& listenAddr, const std::string& name,
			  bool reusePort = true)
		: loop_(loop), name_(name),
		  acceptor_(std::make_unique<Acceptor>(loop, listenAddr, reusePort)),
		  threadPool_(std::make_unique<EventLoopThreadPool>(0, name + "-EventLoop")),
		  nextConnId_(1), started_(false)
	{
		acceptor_->setNewConnectionCallback([this](int sockfd, const InetAddress& addr) {
			newConnection(sockfd, addr);
		});
		LOG_INFO << "TcpServer " << name_ << " created for " << listenAddr.toIpPort();
	}

	~TcpServer()
	{
		loop_->assertInLoopThread();
		LOG_INFO << "TcpServer " << name_ << " destructing with " << connections_.size() << " connections";

		for (auto& item : connections_) {
			auto conn = item.second;
			auto ioLoop = conn->getLoop();
			ioLoop->runInLoop([conn]() { conn->forceClose(); });
		}
	}

	void setThreadNum(int numThreads)
	{
		assert(0 <= numThreads);
		threadPool_ = std::make_unique<EventLoopThreadPool>(numThreads, name_ + "-EventLoop");
		LOG_INFO << "TcpServer " << name_ << " set thread pool size to " << numThreads;
	}

	void start()
	{
		if (!started_) {
			started_ = true;
			if (!acceptor_->listening()) {
				loop_->runInLoop([this]() { acceptor_->listen(); });
			}
			LOG_INFO << "TcpServer " << name_ << " started with " << threadPool_->size() << " threads";
		}
	}

	void setMessageCallback(MessageCallback cb) { messageCallback_ = std::move(cb); }
	void setConnectionCallback(ConnectionCallback cb) { connectionCallback_ = std::move(cb); }
	void setWriteCompleteCallback(WriteCompleteCallback cb) { writeCompleteCallback_ = std::move(cb); }

	const std::string& name() const { return name_; }
	const char* ipPort() const { return acceptor_ ? "listening" : "not-listening"; }
	EventLoop* getLoop() const { return loop_; }

	size_t numConnections() const
	{
		return connections_.size();
	}

	std::vector<TcpConnectionPtr> getConnections() const
	{
		std::vector<TcpConnectionPtr> result;
		result.reserve(connections_.size());
		for (const auto& item : connections_) {
			result.push_back(item.second);
		}
		return result;
	}

	TcpConnectionPtr getConnection(const std::string& name) const
	{
		auto it = connections_.find(name);
		return it != connections_.end() ? it->second : TcpConnectionPtr();
	}

	void forceCloseAllConnections()
	{
		for (auto& item : connections_) {
			auto conn = item.second;
			auto ioLoop = conn->getLoop();
			ioLoop->runInLoop([conn]() { conn->forceClose(); });
		}
	}
};

} // namespace reactor