Winsock IOCP模型

// IOCP.cpp : Defines the entry point for the console application.
//
//

#include "stdafx.h"

#include <WinSock2.h>
#include <Windows.h>
#include <process.h>
#pragma comment(lib, "WS2_32.lib")

#define MAX_BUFFER	256
#define MAX_TIMEOUT 1000
#define MAX_SOCKET  1024
#define MAX_THREAD	64

typedef enum _OPERATION_INFO_
{
	OP_NULL,
	OP_READ,
	OP_WRITE
}OPERATIONINFO;

typedef struct _PER_HANDLE_DATA_
{
public:
	_PER_HANDLE_DATA_()
	{
		clean();
	}
	~_PER_HANDLE_DATA_()
	{
		clean();
	}
protected:
	void clean()
	{
		sock = INVALID_SOCKET;
		memset(&addr, 0, sizeof(addr));
		addr.sin_addr.S_un.S_addr = INADDR_ANY;
		addr.sin_port = htons(0);
		addr.sin_family = AF_INET;
	}
public:	
	SOCKET sock;
	SOCKADDR_IN addr;
	
}PERHANDLEDATA, *PPERHANDLEDATA;

typedef struct _PER_IO_DTATA_
{
public: 
	_PER_IO_DTATA_()
	{
		clean();
	}
	~_PER_IO_DTATA_()
	{
		clean();
	}
private:
	void clean()
	{
		ZeroMemory(&ol, sizeof(ol));
		memset(buf, 0, sizeof(buf));
		wsaBuf.buf = buf;
		wsaBuf.len = MAX_BUFFER;
		opType =  OP_NULL;
	}
public:
	WSAOVERLAPPED ol;
	WSABUF wsaBuf;
	char buf[MAX_BUFFER];
	OPERATIONINFO opType;
}PERIODATA, *PPERIODATA;

HANDLE hThread[MAX_THREAD] = {0};
int g_nThread = 0;
BOOL g_bExitThread = FALSE;

unsigned __stdcall ThreadProc(LPVOID lParam);

int _tmain(int argc, _TCHAR* argv[])
{
	WSADATA wsaData;
	if(0 != WSAStartup(MAKEWORD(2, 2), &wsaData))
	{
		printf("WSAStartup failed with error code: %d/n", GetLastError());
		return EXIT_FAILURE;
	}
	if(2 != HIBYTE(wsaData.wVersion) || 2 != LOBYTE(wsaData.wVersion))
	{
		printf("Socket version not supported./n");
		WSACleanup();
		return EXIT_FAILURE;
	}

	// Create I/O Completion Port
	HANDLE hIOCP = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, NULL, 0);
	if(NULL == hIOCP)
	{
		printf("CreateIoCompletionPort failed with error code: %d/n", GetLastError());
		WSACleanup();
		return EXIT_FAILURE;
	}
	// Create worker thread
	SYSTEM_INFO si = {0};
	GetSystemInfo(&si);
	for(int i = 0; i < (int)si.dwNumberOfProcessors+2; i++)
	{
		hThread[g_nThread] = (HANDLE)_beginthreadex(NULL, 0, ThreadProc, (LPVOID)hIOCP, 0, NULL);
		if(NULL == hThread[g_nThread])
		{
			printf("_beginthreadex failed with error code: %d/n", GetLastError());
			continue;
		}
		++g_nThread;

		if(g_nThread > MAX_THREAD)
		{
			break;
		}
	}

	// Create socket
	SOCKET sListen = WSASocket(AF_INET, SOCK_STREAM, IPPROTO_TCP, NULL, 0, WSA_FLAG_OVERLAPPED);
	if(INVALID_SOCKET == sListen)
	{
		printf("WSASocket failed with error code: %d/n", WSAGetLastError());
		goto EXIT_CODE;
	}
	SOCKADDR_IN addr;
	memset(&addr, 0, sizeof(addr));
	addr.sin_addr.S_un.S_addr = inet_addr("127.0.0.1");
	addr.sin_family = AF_INET;
	addr.sin_port = htons(5050);
	if(SOCKET_ERROR == bind(sListen, (LPSOCKADDR)&addr, sizeof(addr)))
	{
		printf("bind failed with error code: %d/n", WSAGetLastError());
		closesocket(sListen);
		sListen = INVALID_SOCKET;
		goto EXIT_CODE;
	}
	if(SOCKET_ERROR == listen(sListen, 5))
	{
		printf("listen failed with error code: %d/n", WSAGetLastError());
		closesocket(sListen);
		sListen = INVALID_SOCKET;
		goto EXIT_CODE;
	}

	printf("Server start, wait for client to connect .../n");
	while(TRUE)
	{
		SOCKADDR_IN remote;
		memset(&remote, 0, sizeof(remote));
		int len = sizeof(remote);

		SOCKET sNew = WSAAccept(sListen, (LPSOCKADDR)&remote, &len, NULL, NULL);
		if(INVALID_SOCKET == sNew)
		{
			printf("WSAAccept failed with error code: %d/n", WSAGetLastError());
			continue;
		}
		printf("Client <%s : %d> come in./n", inet_ntoa(remote.sin_addr), ntohs(remote.sin_port));
		PERHANDLEDATA* pPerHandleData = new PERHANDLEDATA;
		pPerHandleData->sock = sNew;		
		memcpy(&(pPerHandleData->addr), &remote, sizeof(remote));
		// Associate with IOCP
		if(NULL == CreateIoCompletionPort((HANDLE)(pPerHandleData->sock), hIOCP, (ULONG_PTR)pPerHandleData, 0))
		{
			printf("CreateIoCompletionPort failed with error code: %d/n", GetLastError());
			closesocket(pPerHandleData->sock);
			delete pPerHandleData;
			continue;
		}
		// Post Receive
		PERIODATA* pPerIoData = new PERIODATA;
		pPerIoData->opType = OP_READ;
		DWORD dwTrans = pPerIoData->wsaBuf.len;
		DWORD dwFlags = 0;
		if(SOCKET_ERROR == WSARecv(pPerHandleData->sock, &(pPerIoData->wsaBuf), 1, 
			&dwTrans, &dwFlags, &(pPerIoData->ol), NULL))
		{
			if(WSA_IO_PENDING != WSAGetLastError())
			{
				printf("WSARecv failed with error code: %d/n", WSAGetLastError());
				closesocket(pPerHandleData->sock);
				delete pPerHandleData;
				delete pPerIoData;
				continue;
			}
		}
	}
	closesocket(sListen);
	sListen = INVALID_SOCKET;

EXIT_CODE:
	g_bExitThread = TRUE;
	PostQueuedCompletionStatus(hIOCP, 0, NULL, NULL);
	WaitForMultipleObjects(g_nThread, hThread, TRUE, INFINITE);
	for(int i = 0; i < g_nThread; i++)
	{
		CloseHandle(hThread[g_nThread]);
	}
	CloseHandle(hIOCP); // Close IOCP
	WSACleanup();
	return 0;
}

unsigned __stdcall ThreadProc(LPVOID lParam)
{
	HANDLE hIOCP = (HANDLE)lParam;

	PERHANDLEDATA* pPerHandleData = NULL;
	PERIODATA* pPerIoData = NULL;
	WSAOVERLAPPED* lpOverlapped = NULL;
	DWORD dwTrans = 0;
	DWORD dwFlags = 0;
	while(!g_bExitThread)
	{
		BOOL bRet = GetQueuedCompletionStatus(hIOCP, &dwTrans, (PULONG_PTR)&pPerHandleData, &lpOverlapped, INFINITE);
		if(!bRet)
		{
			printf("GetQueuedCompletionStatus failed with error: %d/n", WSAGetLastError());
			continue;
		}
		else
		{
			pPerIoData = CONTAINING_RECORD(lpOverlapped, PERIODATA, ol);
			if(0 == dwTrans)
			{
				printf("Client: <%s : %d> leave./n", inet_ntoa(pPerHandleData->addr.sin_addr), ntohs(pPerHandleData->addr.sin_port));
				closesocket(pPerHandleData->sock);
				delete pPerHandleData;
				delete pPerIoData;
				continue;
			}
			else
			{
				switch(pPerIoData->opType)
				{
				case OP_READ:
					printf("recv client <%s : %d> data: %s/n", inet_ntoa(pPerHandleData->addr.sin_addr), ntohs(pPerHandleData->addr.sin_port), pPerIoData->buf);
					pPerIoData->opType = OP_WRITE;
					memset(&(pPerIoData->ol), 0, sizeof(pPerIoData->ol));
					if(SOCKET_ERROR == WSASend(pPerHandleData->sock, &(pPerIoData->wsaBuf), 1, &dwTrans, dwFlags, &(pPerIoData->ol), NULL))
					{
						if(WSA_IO_PENDING != WSAGetLastError())
						{
							printf("WSASend failed with error code: %d./n", WSAGetLastError());
							continue;
						}
					}
					break;

				case OP_WRITE:
					{
						pPerIoData->opType = OP_READ;
						dwFlags = 0;
						memset(&(pPerIoData->ol), 0, sizeof(pPerIoData->ol));
						memset(pPerIoData->buf, 0, sizeof(pPerIoData->buf));
						pPerIoData->wsaBuf.buf = pPerIoData->buf;
						dwTrans = pPerIoData->wsaBuf.len = MAX_BUFFER;
						if(SOCKET_ERROR == WSARecv(pPerHandleData->sock, &(pPerIoData->wsaBuf), 1, &dwTrans, &dwFlags, &(pPerIoData->ol), NULL))
						{
							if(WSA_IO_PENDING != WSAGetLastError())
							{
								printf("WSARecv failed with error code: %d./n", WSAGetLastError());
								continue;
							}
						}
					}
					break;

				default:
					break;
				}
			}
		}
	}
	return 0;
}