epoll.cc

epoll.cc

#define THROW_IF(arg, msg) \
  if (arg) throw runtime_error(msg);

using namespace std;

int main() {
  sockaddr_in addr{
          .sin_family = AF_INET,
          .sin_port = htons(PORT),
          .sin_addr = {.s_addr = INADDR_ANY},
  };
  socklen_t addrLen = sizeof(addr);

  int opt = 1;
  int socketFd = socket(AF_INET, SOCK_STREAM, 0);
  assert(socketFd != 0);
  assert(setsockopt(socketFd, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt)) == 0);
  assert(bind(socketFd, (sockaddr *) &addr, addrLen) >= 0);
  assert(listen(socketFd, MAX_SOCKET) >= 0);
  cout << "Listen at port " << PORT << endl;


  epoll_event acceptEvents{
          .events = EPOLLIN,
          .data{
                  .fd = socketFd}},
          events[MAX_EPOLL_EVENT];
  const int epollFd = epoll_create1(0);
  assert(epollFd >= 0);
  assert(epoll_ctl(epollFd, EPOLL_CTL_ADD, socketFd, &acceptEvents) >= 0);

  while (true) {
    const int ready = epoll_wait(epollFd, events, MAX_EPOLL_EVENT, EPOLL_TIMEOUT);
    for (int i = 0; i < ready; i++) {
      THROW_IF(events[i].events & EPOLLERR, "epoll_wait returns error")

      epoll_event event{.events = 0};
      // listener socket is ready, a new client is trying to connect
      if (events[i].data.fd == socketFd) {
        sockaddr_in clientAddr{};
        socklen_t clientAddrSize = sizeof clientAddr;

        const int fd = accept4(socketFd, (sockaddr *) &clientAddr, &clientAddrSize, SOCK_NONBLOCK);
        THROW_IF(fd < 0, "Failed to accept a client connection")
        THROW_IF(fd > MAX_SOCKET, "Maximum socket connection exceeded")

        event.data.fd = fd;
        event.events |= EPOLLOUT;
        THROW_IF(epoll_ctl(epollFd, EPOLL_CTL_ADD, fd, &event) < 0, "Failed to add socket to the interesting list")
      } else {
        const int fd = events[i].data.fd;
        char buffer[BUFFER_SIZE] = {0};
        int flag = 0;
        event.data.fd = fd;

        if (events[i].events & EPOLLIN) {
          const u_int byteSent = recv(fd, buffer, sizeof buffer, 0);
          if (byteSent < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
              flag |= EPOLLIN;
            } else {
              throw runtime_error("Receive data failed");
            }
          } else if (byteSent == 0) {
            cout << "Client " << fd << " is disconnected" << endl;
          } else {
            const long t = chrono::duration_cast<chrono::milliseconds>(chrono::system_clock::now().time_since_epoch()).count();
            cout << "[" << t << "] (" << fd << "): " << string_view(buffer) << endl;
          }
        } else if (events[i].events & EPOLLOUT) {
          const u_int byteSent = send(fd, buffer, sizeof buffer, 0);
          if (byteSent == -1) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
              flag |= EPOLLOUT;
            } else {
              throw runtime_error("Sent data failed");
            }
          } else {
            flag |= EPOLLIN;
          }
        }
        event.events = flag;

        if (event.events == 0) {
          cout << "closing " << fd << " socket..." << endl;
          THROW_IF(epoll_ctl(epollFd, EPOLL_CTL_DEL, fd, nullptr) < 0, "Failed to remove socket from the interesting list")
          close(fd);
        } else {
          THROW_IF(epoll_ctl(epollFd, EPOLL_CTL_MOD, fd, &event) < 0, "Failed to modify socket in the interesting list")
        }
      }
    }
  }
}