Net-IOCP/Net-Iocp/Core/IocpCore.cpp at develop · GeonRyoung/Net-IOCP · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
#include "IocpCore.h"
#include <iostream>

#include "SessionManager.h"
#include "RudpSession.h"
#include "MetricManager.h"

void IocpCore::WorkerThreadMain()
{
    while (IsRunning())
    {
        DWORD bytesTransferred = 0;
        ULONG_PTR completionKey = 0;
        LPOVERLAPPED pOverlapped = nullptr;

        BOOL bSuccess = GetQueuedCompletionStatus(
            m_hIocp,
            &bytesTransferred,
            &completionKey,
            &pOverlapped,
            INFINITE
        );

        if (!m_bIsRunning.load(std::memory_order_acquire))
            break;

        if (pOverlapped == nullptr)
            continue;

        if (m_sessionManager == nullptr)
            continue;

        const uint64_t completionSessionId = static_cast<uint64_t>(completionKey);
        IocpContext* context = reinterpret_cast<IocpContext*>(pOverlapped);

        if (context->type == IocpOpType::Recv)
        {
            RecvContext* recvCtx = static_cast<RecvContext*>(context);
            if (bSuccess && bytesTransferred > 0)
            {
                m_sessionManager->OnReceiveRouted(recvCtx, static_cast<DWORD>(bytesTransferred));
            }
            m_sessionManager->PostRecvFrom(recvCtx);
            continue;
        }
        else if (context->type == IocpOpType::Send)
        {
            uint64_t sessionId = context->sessionId;
            std::shared_ptr<RudpSession> session = m_sessionManager->FindSession(sessionId);
            if (!session)
                continue;

            if (!bSuccess)
            {
                session->DisConnect();
                m_sessionManager->RemoveSession(sessionId);
                continue;
            }

            session->OnSendCompleted(static_cast<size_t>(bytesTransferred));
            if (session->GetState() == SessionState::Closed)
                m_sessionManager->RemoveSession(sessionId);
        }
    }
}

bool IocpCore::CreateIocpPort(DWORD numThreads)
{
    if (m_bIsRunning.load(std::memory_order_acquire))
        return true;

    m_hIocp = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, numThreads);
    if (m_hIocp == NULL)
        return false;

    m_bIsRunning.store(true, std::memory_order_release);

    if (numThreads == 0)
    {
        SYSTEM_INFO sysInfo;
        GetSystemInfo(&sysInfo);
        numThreads = sysInfo.dwNumberOfProcessors * 2;
    }

    for (DWORD i = 0; i < numThreads; i++)
        m_workerThreads.emplace_back([this]() { WorkerThreadMain(); });

    return true;
}

bool IocpCore::RegisterSocket(SOCKET socket, ULONG_PTR completionKey)
{
    HANDLE hResult = CreateIoCompletionPort(
        reinterpret_cast<HANDLE>(socket),
        m_hIocp,
        completionKey,
        0
    );
    return (hResult != NULL);
}

void IocpCore::Shutdown()
{
    bool expected = true;
    if (!m_bIsRunning.compare_exchange_strong(
        expected, false, std::memory_order_acq_rel))
        return;

    for (size_t i = 0; i < m_workerThreads.size(); i++)
        PostQueuedCompletionStatus(m_hIocp, 0, 0, nullptr);

    for (std::thread& thread : m_workerThreads)
    {
        if (thread.joinable())
            thread.join();
    }
    m_workerThreads.clear();

    if (m_hIocp != INVALID_HANDLE_VALUE)
    {
        CloseHandle(m_hIocp);
        m_hIocp = INVALID_HANDLE_VALUE;
    }
}