WinSock IOCP 模型总结(附一个带缓存池的IOCP类)

 1 BOOL IOCPBase::InitializeListenSocket()
 2 {
 3     // 生成用于监听的socket的Context
 4     listenSockContext = new SocketContext;
 5     listenSockContext->connSocket = WSASocket(AF_INET, SOCK_STREAM, 0, NULL, 0, WSA_FLAG_OVERLAPPED);
 6     if (INVALID_SOCKET == listenSockContext->connSocket)
 7         return false;
 8 
 9     // 将socket绑定到完成端口中
10     if (NULL == CreateIoCompletionPort((HANDLE)listenSockContext->connSocket, completionPort, (DWORD)listenSockContext, 0))
11     {
12         RELEASE_SOCKET(listenSockContext->connSocket);
13         return false;
14     }
15 
16     //服务器地址信息，用于绑定socket
17     sockaddr_in serverAddr;
18 
19     // 填充地址信息
20     ZeroMemory((char *)&serverAddr, sizeof(serverAddr));
21     serverAddr.sin_family = AF_INET;
22     serverAddr.sin_addr.s_addr = htonl(INADDR_ANY);
23     serverAddr.sin_port = htons(port);
24 
25     // 绑定地址和端口
26     if (SOCKET_ERROR == bind(listenSockContext->connSocket, (sockaddr *)&serverAddr, sizeof(serverAddr)))
27     {
28         return false;
29     }
30 
31     // 开始监听
32     if (SOCKET_ERROR == listen(listenSockContext->connSocket, SOMAXCONN))
33     {
34         return false;
35     }
36 
37     GUID guidAcceptEx = WSAID_ACCEPTEX;
38     GUID guidGetAcceptSockAddrs = WSAID_GETACCEPTEXSOCKADDRS;
39     // 提取扩展函数指针
40     DWORD dwBytes = 0;
41     if (SOCKET_ERROR == WSAIoctl(
42         listenSockContext->connSocket,
43         SIO_GET_EXTENSION_FUNCTION_POINTER,
44         &guidAcceptEx,
45         sizeof(guidAcceptEx),
46         &fnAcceptEx,
47         sizeof(fnAcceptEx),
48         &dwBytes,
49         NULL,
50         NULL))
51     {
52         DeInitialize();
53         return false;
54     }
55 
56     if (SOCKET_ERROR == WSAIoctl(
57         listenSockContext->connSocket,
58         SIO_GET_EXTENSION_FUNCTION_POINTER,
59         &guidGetAcceptSockAddrs,
60         sizeof(guidGetAcceptSockAddrs),
61         &fnGetAcceptExSockAddrs,
62         sizeof(fnGetAcceptExSockAddrs),
63         &dwBytes,
64         NULL,
65         NULL))
66     {
67         DeInitialize();
68         return false;
69     }
70 
71     for (size_t i = 0; i < MAX_POST_ACCEPT; i++)
72     {
73         IOContext *ioContext = listenSockContext->GetNewIOContext();
74         if (false == PostAccept(listenSockContext, ioContext))
75         {
76             listenSockContext->RemoveContext(ioContext);
77             return false;
78         }
79     }
80     return true;
81 }

 1 class SocketContext
 2 {
 3 public:
 4     SOCKET connSocket;                                // 连接的socket
 5     SOCKADDR_IN clientAddr;                            // 连接的远程地址
 6 
 7 private:
 8     vector<IOContext*> arrIoContext;                // 同一个socket上的多个IO请求
 9     static IOContextPool ioContextPool;                // 空闲的IOContext池
10     CRITICAL_SECTION csLock;
11 
12 public:
13     SocketContext()
14     {
15         InitializeCriticalSection(&csLock);
16         arrIoContext.clear();
17         connSocket = INVALID_SOCKET;
18         ZeroMemory(&clientAddr, sizeof(clientAddr));
19     }
20 
21     ~SocketContext()
22     {
23         RELEASE_SOCKET(connSocket);
24 
25         // 回收所有的IOContext
26         for (vector<IOContext*>::iterator it = arrIoContext.begin(); it != arrIoContext.end(); it++)
27         {
28             ioContextPool.ReleaseIOContext(*it);
29         }
30 
31         EnterCriticalSection(&csLock);
32         arrIoContext.clear();
33         LeaveCriticalSection(&csLock);
34 
35         DeleteCriticalSection(&csLock);
36     }
37 
38     // 获取一个新的IoContext
39     IOContext *GetNewIOContext()
40     {
41         IOContext *context = ioContextPool.AllocateIoContext();
42         if (context != NULL)
43         {
44             EnterCriticalSection(&csLock);
45             arrIoContext.push_back(context);
46             LeaveCriticalSection(&csLock);
47         }
48         return context;
49     }
50 
51     // 从数组中移除一个指定的IoContext
52     void RemoveContext(IOContext* pContext)
53     {
54         for (vector<IOContext*>::iterator it = arrIoContext.begin(); it != arrIoContext.end(); it++)
55         {
56             if (pContext == *it)
57             {
58                 ioContextPool.ReleaseIOContext(*it);
59 
60                 EnterCriticalSection(&csLock);
61                 arrIoContext.erase(it);
62                 LeaveCriticalSection(&csLock);
63 
64                 break;
65             }
66         }
67     }
68 
69     // 
70 };

 1 // 空闲的IOContext管理类(IOContext池)
 2 class IOContextPool
 3 {
 4 private:
 5     list<IOContext *> contextList;
 6     CRITICAL_SECTION csLock;
 7 
 8 public:
 9     IOContextPool()
10     {
11         InitializeCriticalSection(&csLock);
12         contextList.clear();
13 
14         EnterCriticalSection(&csLock);
15         for (size_t i = 0; i < INIT_IOCONTEXT_NUM; i++)
16         {
17             IOContext *context = new IOContext;
18             contextList.push_back(context);
19         }
20         LeaveCriticalSection(&csLock);
21         
22     }
23 
24     ~IOContextPool()
25     {
26         EnterCriticalSection(&csLock);
27         for (list<IOContext *>::iterator it = contextList.begin(); it != contextList.end(); it++)
28         {
29             delete (*it);
30         }
31         contextList.clear();
32         LeaveCriticalSection(&csLock);
33 
34         DeleteCriticalSection(&csLock);
35     }
36 
37     // 分配一个IOContxt
38     IOContext *AllocateIoContext()
39     {
40         IOContext *context = NULL;
41 
42         EnterCriticalSection(&csLock);
43         if (contextList.size() > 0) //list不为空，从list中取一个
44         {
45             context = contextList.back();
46             contextList.pop_back();
47         }
48         else    //list为空，新建一个
49         {
50             context = new IOContext;
51         }
52         LeaveCriticalSection(&csLock);
53 
54         return context;
55     }
56 
57     // 回收一个IOContxt
58     void ReleaseIOContext(IOContext *pContext)
59     {
60         pContext->Reset();
61         EnterCriticalSection(&csLock);
62         contextList.push_front(pContext);
63         LeaveCriticalSection(&csLock);
64     }
65 };

DWORD IOCPBase::WorkerThreadProc(LPVOID lpParam)
{
    IOCPBase *iocp = (IOCPBase*)lpParam;
    OVERLAPPED *ol = NULL;
    SocketContext *sockContext;
    DWORD dwBytes = 0;
    IOContext *ioContext = NULL;

    while (WAIT_OBJECT_0 != WaitForSingleObject(iocp->stopEvent, 0))
    {
        BOOL bRet = GetQueuedCompletionStatus(iocp->completionPort, &dwBytes, (PULONG_PTR)&sockContext, &ol, INFINITE);

        // 读取传入的参数
        ioContext = CONTAINING_RECORD(ol, IOContext, overLapped);

        // 收到退出标志
        if (EXIT_CODE == (DWORD)sockContext)
        {
            break;
        }

        if (!bRet)
        {
            DWORD dwErr = GetLastError();

            // 如果是超时了，就再继续等吧  
            if (WAIT_TIMEOUT == dwErr)
            {
                // 确认客户端是否还活着...
                if (!iocp->IsSocketAlive(sockContext->connSocket))
                {
                    iocp->OnConnectionClosed(sockContext);

                    // 回收socket
                    iocp->DoClose(sockContext);
                    continue;
                }
                else
                {
                    continue;
                }
            }
            // 可能是客户端异常退出了(64)
            else if (ERROR_NETNAME_DELETED == dwErr)
            {
                iocp->OnConnectionError(sockContext, dwErr);

                // 回收socket
                iocp->DoClose(sockContext);
                continue;
            }
            else
            {
                iocp->OnConnectionError(sockContext, dwErr);

                // 回收socket
                iocp->DoClose(sockContext);
                continue;
            }
        }
        else
        {
            // 判断是否有客户端断开
            if ((0 == dwBytes) && (RECV_POSTED == ioContext->ioType || SEND_POSTED == ioContext->ioType))
            {
                iocp->OnConnectionClosed(sockContext);

                // 回收socket
                iocp->DoClose(sockContext);
                continue;
            }
            else
            {
                switch (ioContext->ioType)
                {
                case ACCEPT_POSTED:
                    iocp->DoAccpet(sockContext, ioContext);
                    break;
                case RECV_POSTED:
                    iocp->DoRecv(sockContext, ioContext);
                    break;
                case SEND_POSTED:
                    iocp->DoSend(sockContext, ioContext);
                    break;
                default:
                    break;
                }
            }
        }
    }

    // 释放线程参数
    RELEASE(lpParam);
    return 0;
}

 1 BOOL IOCPBase::DoAccpet(SocketContext * sockContext, IOContext * ioContext)
 2 {
 3     
 4     InterlockedIncrement(&connectCnt);
 5     InterlockedDecrement(&acceptPostCnt);
 6     SOCKADDR_IN *clientAddr = NULL;
 7     SOCKADDR_IN *localAddr = NULL;
 8     int clientAddrLen, localAddrLen;
 9     clientAddrLen = localAddrLen = sizeof(SOCKADDR_IN);
10 
11     // 1. 获取地址信息 （GetAcceptExSockAddrs函数不仅可以获取地址信息，还可以顺便取出第一组数据）
12     fnGetAcceptExSockAddrs(ioContext->wsaBuf.buf, 0, localAddrLen, clientAddrLen, (LPSOCKADDR *)&localAddr, &localAddrLen, (LPSOCKADDR *)&clientAddr, &clientAddrLen);
13 
14     // 2. 为新连接建立一个SocketContext 
15     SocketContext *newSockContext = new SocketContext;
16     newSockContext->connSocket = ioContext->ioSocket;
17     memcpy_s(&(newSockContext->clientAddr), sizeof(SOCKADDR_IN), clientAddr, sizeof(SOCKADDR_IN));
18 
19     // 3. 将listenSocketContext的IOContext 重置后继续投递AcceptEx
20     ioContext->Reset();
21     if (false == PostAccept(listenSockContext, ioContext))
22     {
23         listenSockContext->RemoveContext(ioContext);
24     }
25 
26     // 4. 将新socket和完成端口绑定
27     if (NULL == CreateIoCompletionPort((HANDLE)newSockContext->connSocket, completionPort, (DWORD)newSockContext, 0))
28     {
29         DWORD dwErr = WSAGetLastError();
30         if (dwErr != ERROR_INVALID_PARAMETER)
31         {
32             DoClose(newSockContext);
33             return false;
34         }
35     }
36 
37     // 并设置tcp_keepalive
38     tcp_keepalive alive_in;
39     tcp_keepalive alive_out;
40     alive_in.onoff = TRUE;
41     alive_in.keepalivetime = 1000 * 60;  // 60s  多长时间（ ms ）没有数据就开始 send 心跳包
42     alive_in.keepaliveinterval = 1000 * 10; //10s  每隔多长时间（ ms ） send 一个心跳包
43     unsigned long ulBytesReturn = 0;
44     if (SOCKET_ERROR == WSAIoctl(newSockContext->connSocket, SIO_KEEPALIVE_VALS, &alive_in, sizeof(alive_in), &alive_out, sizeof(alive_out), &ulBytesReturn, NULL, NULL))
45     {
46         TRACE(L"WSAIoctl failed: %d/n", WSAGetLastError());
47     }
48 
49 
50     OnConnectionEstablished(newSockContext);
51 
52     // 5. 建立recv操作所需的ioContext，在新连接的socket上投递recv请求
53     IOContext *newIoContext = newSockContext->GetNewIOContext();
54     newIoContext->ioType = RECV_POSTED;
55     newIoContext->ioSocket = newSockContext->connSocket;
56     // 投递recv请求
57     if (false == PostRecv(newSockContext, newIoContext))
58     {
59         DoClose(sockContext);
60         return false;
61     }
62 
63     return true;
64 }

 1 BOOL IOCPBase::PostAccept(SocketContext * sockContext, IOContext * ioContext)
 2 {
 3     DWORD dwBytes = 0;
 4     ioContext->ioType = ACCEPT_POSTED;
 5     ioContext->ioSocket = WSASocket(AF_INET, SOCK_STREAM, 0, NULL, 0, WSA_FLAG_OVERLAPPED);
 6     if (INVALID_SOCKET == ioContext->ioSocket)
 7     {
 8         return false;
 9     }
10     
11     // 将接收缓冲置为0,令AcceptEx直接返回,防止拒绝服务攻击
12     if (false == fnAcceptEx(listenSockContext->connSocket, ioContext->ioSocket, ioContext->wsaBuf.buf, 0, sizeof(sockaddr_in) + 16, sizeof(sockaddr_in) + 16, &dwBytes, &ioContext->overLapped))
13     {
14         if (WSA_IO_PENDING != WSAGetLastError())
15         {
16             return false;
17         }
18     }
19 
20     InterlockedIncrement(&acceptPostCnt);
21     return true;
22 }

 1 BOOL IOCPBase::PostRecv(SocketContext * sockContext, IOContext *ioContext)
 2 {
 3     DWORD dwFlags = 0, dwBytes = 0;
 4     ioContext->Reset();
 5     ioContext->ioType = RECV_POSTED;
 6 
 7     int nBytesRecv = WSARecv(ioContext->ioSocket, &ioContext->wsaBuf, 1, &dwBytes, &dwFlags, &ioContext->overLapped, NULL);
 8     // 如果返回值错误，并且错误的代码并非是Pending的话，那就说明这个重叠请求失败了
 9     if ((SOCKET_ERROR == nBytesRecv) && (WSA_IO_PENDING != WSAGetLastError()))
10     {
11         DoClose(sockContext);
12         return false;
13     }
14     return true;
15 }