116 lines
3.5 KiB
C++
116 lines
3.5 KiB
C++
#pragma once
|
|
#include "Utils/ThreadPool.hpp"
|
|
#include "Socket/WSAManager.hpp"
|
|
#include "Socket/TCPSocket.hpp"
|
|
#include "Socket/Log.hpp"
|
|
#include <functional>
|
|
|
|
#ifndef _WIN32
|
|
|
|
typedef struct _OVERLAPPED {
|
|
char dummy;
|
|
} OVERLAPPED;
|
|
|
|
typedef struct __WSABUF {
|
|
std::uint32_t len;
|
|
char *buf;
|
|
} WSABUF;
|
|
|
|
#endif
|
|
|
|
namespace Chattr {
|
|
|
|
class IOCP;
|
|
|
|
struct IOCPPASSINDATA {
|
|
OVERLAPPED overlapped;
|
|
TCPSocket socket;
|
|
char buf[1501];
|
|
int recvbytes;
|
|
int sendbytes;
|
|
WSABUF wsabuf;
|
|
IOCP* IOCPInstance;
|
|
};
|
|
|
|
class IOCP {
|
|
public:
|
|
#ifdef _WIN32
|
|
static void iocpWather(ThreadPool* threadPool, HANDLE completionPort_, std::function<void(ThreadPool*, IOCPPASSINDATA*)> callback) {
|
|
DWORD tid = GetCurrentThreadId();
|
|
spdlog::debug("Waiting IO to complete on TID: {}.", tid);
|
|
IOCPPASSINDATA* data;
|
|
SOCKET sock;
|
|
DWORD cbTransfrred;
|
|
int retVal = GetQueuedCompletionStatus(completionPort_, &cbTransfrred, (PULONG_PTR)&sock, (LPOVERLAPPED*)&data, INFINITE);
|
|
if (retVal == 0 || cbTransfrred == 0) {
|
|
spdlog::info("Client disconnected. [{}]", (std::string)(data->socket.remoteAddr));
|
|
delete data;
|
|
threadPool->enqueueJob(iocpWather, completionPort_, callback);
|
|
return;
|
|
}
|
|
threadPool->enqueueJob(callback, data);
|
|
threadPool->enqueueJob(iocpWather, completionPort_, callback);
|
|
};
|
|
#elif __linux__
|
|
static void iocpWather(ThreadPool* threadPool, std::function<void(ThreadPool*, IOCPPASSINDATA*)> callback) {
|
|
// DWORD tid = GetCurrentThreadId();
|
|
// spdlog::debug("Waiting IO to complete on TID: {}.", tid);
|
|
// IOCPPASSINDATA* data;
|
|
// SOCKET sock;
|
|
// DWORD cbTransfrred;
|
|
// int retVal = GetQueuedCompletionStatus(completionPort_, &cbTransfrred, (PULONG_PTR)&sock, (LPOVERLAPPED*)&data, INFINITE);
|
|
// if (retVal == 0 || cbTransfrred == 0) {
|
|
// spdlog::info("Client disconnected. [{}]", (std::string)(data->socket.remoteAddr));
|
|
// delete data;
|
|
// threadPool->enqueueJob(iocpWather, completionPort_, callback);
|
|
// return;
|
|
// }
|
|
// threadPool->enqueueJob(callback, data);
|
|
// threadPool->enqueueJob(iocpWather, completionPort_, callback);
|
|
};
|
|
#endif
|
|
|
|
template<typename _Callable>
|
|
void init(ThreadPool* __IOCPThread, _Callable&& callback) {
|
|
IOCPThread_ = __IOCPThread;
|
|
#ifdef _WIN32
|
|
completionPort_ = ::CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
|
|
if (completionPort_ == NULL)
|
|
log::critical("CreateIoCompletionPort()");
|
|
#elif __linux__
|
|
|
|
#endif
|
|
auto boundFunc = [callback = std::move(callback)](ThreadPool* __IOCPThread, IOCPPASSINDATA* data) mutable {
|
|
callback(__IOCPThread, data);
|
|
};
|
|
|
|
int tCount = __IOCPThread->threadCount;
|
|
|
|
spdlog::info("Resizing threadpool size to: {}", tCount * 2);
|
|
|
|
__IOCPThread->respawnWorker(tCount * 2);
|
|
|
|
spdlog::info("Set IOCP Worker count to: {}", tCount);
|
|
for (int i = 0; i < tCount; i++) {
|
|
std::function<void(ThreadPool*, IOCPPASSINDATA*)> task(boundFunc);
|
|
__IOCPThread->enqueueJob(iocpWather, task);
|
|
}
|
|
}
|
|
|
|
void registerSocket(SOCKET sock);
|
|
|
|
int recv(Chattr::IOCPPASSINDATA* data);
|
|
int send(Chattr::IOCPPASSINDATA* data, int __flags);
|
|
|
|
private:
|
|
struct Chattr::WSAManager wsaManager;
|
|
ThreadPool* IOCPThread_;
|
|
|
|
#ifdef _WIN32
|
|
HANDLE completionPort_ = INVALID_HANDLE_VALUE;
|
|
#elif __linux__
|
|
int epollfd = -1;
|
|
#endif
|
|
};
|
|
|
|
} |