Np_Term/include/socket/iocp.h

#pragma once
#include <openssl/bio.h>
#include <openssl/err.h>
#include <openssl/ssl.h>

#include <functional>
#include <queue>
#include <vector>

#include "socket.h"
#include "utils/thread_pool.h"
#include "wsa_manager.h"

#ifdef __linux__

typedef struct _OVERLAPPED {
  char dummy;
} OVERLAPPED;

typedef struct __WSABUF {
  std::uint32_t len;
  char* buf;
} WSABUF;

#endif

namespace Network {

class IOCP;

enum class IOCPEVENT { QUIT, READ, WRITE };

struct IOCPPASSINDATA {
  OVERLAPPED overlapped;
  IOCPEVENT event;
  std::shared_ptr<Socket> socket;
  SSL* ssl;
  std::uint32_t transferredbytes;
  WSABUF wsabuf;
  std::uint32_t bufsize;
  IOCP* IOCPInstance;
#ifdef __linux__
  std::shared_ptr<std::queue<IOCPPASSINDATA*>> sendQueue;
#endif

  IOCPPASSINDATA(std::uint32_t bufsize) {
    std::memset(&overlapped, 0, sizeof(overlapped));
    event = IOCPEVENT::QUIT;
    socket = nullptr;
    transferredbytes = 0;
    this->bufsize = bufsize;
    IOCPInstance = nullptr;

    wsabuf.buf = new char[bufsize];
  }
  IOCPPASSINDATA(const IOCPPASSINDATA& other)
      : event(other.event),
        socket(other.socket),
        transferredbytes(other.transferredbytes),
        bufsize(other.bufsize),
        IOCPInstance(other.IOCPInstance)
#ifdef __linux__
        ,
        sendQueue(other.sendQueue)
#endif
  {
    std::memset(&overlapped, 0, sizeof(overlapped));
    wsabuf.buf = new char[bufsize];
    std::memcpy(wsabuf.buf, other.wsabuf.buf, other.wsabuf.len);
  }

  ~IOCPPASSINDATA() {
    if (wsabuf.buf != nullptr) delete wsabuf.buf;
  }

  IOCPPASSINDATA& operator=(const IOCPPASSINDATA& other) {
    if (this != &other) {
      std::memset(&overlapped, 0, sizeof(overlapped));
      event = other.event;
      socket = other.socket;
      transferredbytes = other.transferredbytes;
      bufsize = other.bufsize;
      IOCPInstance = other.IOCPInstance;
#ifdef __linux__
      sendQueue = other.sendQueue;
#endif
      wsabuf.buf = new char[bufsize];
      std::memcpy(wsabuf.buf, other.wsabuf.buf, other.wsabuf.len);
    }
    return *this;
  }
};

class IOCP {
 public:
  ~IOCP();

  template <typename... _args>
  void init(utils::ThreadPool* __IOCPThread, SessionProtocol proto,
            std::function<void(utils::ThreadPool*, IOCPPASSINDATA*, _args&&...)>
                callback,
            _args&&... args) {
    IOCPThread_ = __IOCPThread;
    proto_ = proto;

    if (proto == SessionProtocol::TLS || proto == SessionProtocol::QUIC) {
      rbio_ = ::BIO_new(::BIO_s_mem());
      wbio_ = ::BIO_new(::BIO_s_mem());
    }
#ifdef _WIN32
    completionPort_ =
        ::CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
    if (completionPort_ == NULL) {
      spdlog::critical("CreateIoCompletionPort()");
      std::exit(EXIT_FAILURE);
    }
    int tCount = __IOCPThread->threadCount;

    spdlog::info("Resizing threadpool size to: {}", tCount * 2);

    __IOCPThread->respawnWorker(tCount * 2);

    for (int i = 0; i < tCount; i++)
      __IOCPThread->enqueueJob(iocpWatcher_, IOCPThread_, callback, args...);
#endif
  }

  void destruct();

  void registerSocket(IOCPPASSINDATA* data);

  int send(std::vector<IOCPPASSINDATA> data);
  int recv(IOCPPASSINDATA& data);

 private:
#ifdef _WIN32
  template <typename... _args>
  void iocpWatcher_(
      utils::ThreadPool* IOCPThread,
      std::function<void(utils::ThreadPool*, IOCPPASSINDATA*, _args&&...)>
          callback,
      _args&&... args) {
    IOCPPASSINDATA* data;
    SOCKET sock;
    DWORD cbTransfrred;
    int retVal = GetQueuedCompletionStatus(completionPort_, &cbTransfrred,
                                           (PULONG_PTR)&sock,
                                           (LPOVERLAPPED*)&data, INFINITE);
    if (retVal == 0 || cbTransfrred == 0) {
      data->event = IOCPEVENT::QUIT;
      spdlog::debug("Disconnected. [{}]",
                    (std::string)(data->socket->remoteAddr));
    } else {
      if (data->event == IOCPEVENT::READ &&
          (proto_ == SessionProtocol::TLS || proto_ == SessionProtocol::QUIC)) {
        ::BIO_write(rbio_, data->wsabuf.buf, cbTransfrred);
        ::SSL_read_ex(data->ssl, data->wsabuf.buf, data->bufsize);
      }
      data->transferredbytes = cbTransfrred;
    }

    IOCPThread->enqueueJob(callback, IOCPThread, data, args...);
    IOCPThread->enqueueJob(iocpWatcher_, callback, args...);
  };
#elif __linux__

#endif

  void packet_sender_(utils::ThreadPool* IOCPThread) {
#ifdef _WIN32
    std::lock_guard<std::mutex> lock(send_queue_mutex_);
#elif __linux__

#endif
  }

  struct WSAManager* wsaManager = WSAManager::GetInstance();
  utils::ThreadPool* IOCPThread_;

  BIO* rbio_ = nullptr;
  BIO* wbio_ = nullptr;

  SessionProtocol proto_;

  // 밑의 unordered_map들에 키를 추가/제거 하려는 스레드는 이 뮤텍스를 잡아야 함.
  std::mutex socket_mod_mutex_;

  // 각 소켓별 뮤텍스. 다른 스레드가 읽는 중이라면 수신 순서 보장을 위해 다른
  // 스레드는 수신을 포기하고 이전 스레드에 전송을 위임해야 함. (ONESHOT으로
  // 인해 발생하지 않을 것으로 기대되는 행동이기는 하나 보장되지 않을 수
  // 있으므로) 수신 스레드는 epoll이나 iocp를 통해 적당히 수신받아 이 큐를
  // 채워야 함. (항시 socket에 대한 큐에 대해 읽기 시도가 행해질 수 있어야
  // 한다는 뜻임) EPOLLIN에 대해서는 ONESHOT으로 등록해 놓고 읽는 도중에 버퍼에
  // 새 값이 채워질 수 있으므로 읽기가 끝나고 나서 재등록한다
  std::unordered_map<SOCKET, std::mutex> recv_queue_mutex_;
  // 각 소켓별 패킷, int는 그 vector의 시작 인덱스(vector의 끝까지 다 읽었으면
  // 그 vector는 list에서 삭제되어야 하며, 데이터는 평문으로 변환하여 저장한다)
  std::unordered_map<SOCKET,
                     std::list<std::pair<std::vector<char>, std::uint32_t>>>
      recv_queue_;

  // 각 소켓별 뮤텍스. 다른 스레드가 쓰는 중이라면 송신 순서 보장을 위해 다른
  // 스레드는 대기한다. condition variable을 쓰지 않는 이유는 소켓 갯수만큼의
  // 스레드가 대기할 수 없기 때문이며, 송신 등록 시에 적당히 송신 스레드를
  // 스폰해야 함. 이 변수는 항상 송신 스레드에 의해 관리되어야만 하며,
  // 윈도우에서는 송신을 iocp가 대행하기 때문에 이 큐가 필요 없는 것처럼 느껴질
  // 수 있으나, 송신 순서를 보장하기 위해 WSASend를 한 스레드가 연속해서
  // 호출해야만 하는데 이는 한번 쓰기 호출 시에 송신 중 데이터를 추가하려는
  // 스레드가 데이터를 추가하고 미전송된 경우를 대비하여 최대 1개까지의 스레드가
  // 대기하도록 한다. 리눅스에서도 send_queue에 데이터를 쌓고 최대 1개까지의
  // 스레드가 대기하도록 한다.
  std::unordered_map<SOCKET, std::atomic_flag> pending_try_empty_;
  std::unordered_map<SOCKET, std::atomic_flag> sending_;
  std::unordered_map<SOCKET, std::mutex> send_queue_mutex_;
  std::unordered_map<SOCKET, std::list<IOCPPASSINDATA>> send_queue_;

#ifdef _WIN32
  HANDLE completionPort_ = INVALID_HANDLE_VALUE;
#elif __linux__

#endif
};

}  // namespace Network