IOCP例子二

#include <winsock2.h>
#include <Mswsock.h>
#include <stdio.h>
#include <windows.h>
#include <assert.h>
#include <list>
#pragma comment(lib, "ws2_32.lib")
#pragma comment(lib,"Mswsock.lib")
using namespace std;

LONG volatile				g_ThreadNum;															//线程数目
HANDLE							g_hCompletion;														//												 
SOCKET							g_socket;																	//sock	 


#define						BUFFER_SIZE								20480						//最大缓冲
enum IO_TYPE 
{
	IO_TYPE_ACCEPT,					//网络连接
	IO_TYPE_READ,						//网络读取	
	IO_TYPE_WRITE,						//网络写入
	IO_TYPE_UNKNOWN				//无效类型
};

#pragma pack(1)

struct OverLappedEx 
{
	OVERLAPPED				m_OLap;
	IO_TYPE						m_IOType;									//IO请求类型
	char								m_szBuf[BUFFER_SIZE];				//接收缓冲
	OverLappedEx(IO_TYPE ioType)
	{
		ZeroMemory(&m_OLap, sizeof(OVERLAPPED));
		ZeroMemory(m_szBuf, BUFFER_SIZE);
		m_IOType = ioType;
	}
};


struct PerSocketData 
{
	SOCKET   m_Socket;
	SOCKET m_AccSocket;
	PerSocketData()
	{
		m_Socket = INVALID_SOCKET;
		m_AccSocket = INVALID_SOCKET;
	}
};


#pragma pack()

list <PerSocketData *> g_ArrSocketData;
list <OverLappedEx *> g_ArrOverLapEx;

PerSocketData *AssignSockToCompletionPort(SOCKET tSocket)
{
	assert(tSocket != INVALID_SOCKET);
	PerSocketData *pSockData = new PerSocketData();
	pSockData->m_Socket = tSocket;
	g_ArrSocketData.push_back(pSockData);
	CreateIoCompletionPort((HANDLE)tSocket, g_hCompletion, (ULONG_PTR)pSockData, 0);
	return pSockData;
}


BOOL PostAccept(PerSocketData *pSockData)
{
	DWORD dwBytesRecv = 0;
	assert(pSockData != NULL);
	OverLappedEx *m_pOverLap = new OverLappedEx(IO_TYPE_ACCEPT);
	g_ArrOverLapEx.push_back(m_pOverLap);
	pSockData->m_AccSocket = WSASocket(AF_INET , SOCK_STREAM , IPPROTO_TCP, NULL, 0, WSA_FLAG_OVERLAPPED);
	BOOL bRet = AcceptEx(pSockData->m_Socket, pSockData->m_AccSocket, m_pOverLap->m_szBuf, 0, sizeof(sockaddr_in)+16, sizeof(sockaddr_in)+16, &dwBytesRecv, &m_pOverLap->m_OLap);
	if (!bRet)
	{
		if (WSAGetLastError() != WSA_IO_PENDING) 
		{
			return FALSE;
		}
	}
	return TRUE;
}
BOOL PostRecv(PerSocketData *pSockData)
{

	assert(pSockData != NULL);
	WSABUF wsaBuf = {0};
	OverLappedEx *m_pOverLap = new OverLappedEx(IO_TYPE_READ);
	g_ArrOverLapEx.push_back(m_pOverLap);
	wsaBuf.buf = m_pOverLap->m_szBuf;
	wsaBuf.len = BUFFER_SIZE;
	DWORD dwBytesRecv = 0, dwFlags = 0;
	int iRet = WSARecv(pSockData->m_Socket, &wsaBuf, 1, &dwBytesRecv, &dwFlags, &(m_pOverLap->m_OLap), NULL);
	if (iRet != NO_ERROR)
	{
		if (WSAGetLastError() != WSA_IO_PENDING)
		{
			return FALSE;
		}
	}
	return TRUE;
}






DWORD WINAPI ServerThread(LPVOID lpParam)
{
	//参数效验
	assert(lpParam!=NULL);
	if (lpParam==NULL)return FALSE;

	//定义变量 
	HANDLE hCompletion = (HANDLE)lpParam;
	OverLappedEx *pOverLaps = NULL;
	PerSocketData *pPerSockData = NULL;
	DWORD dwTrans = 0;
	BOOL bOK=FALSE; 
	InterlockedIncrement(&g_ThreadNum);

	while(TRUE)
	{
		// 在关联到此完成端口的所有套节字上等待I/O完成
		bOK= GetQueuedCompletionStatus(hCompletion, &dwTrans, (LPDWORD)&pPerSockData, (LPOVERLAPPED*)&pOverLaps, INFINITE);

		//收到关闭通知
		if (pPerSockData==NULL || pOverLaps==NULL) break;

		// 在此套节字上有错误发生
		if(bOK==FALSE)continue;

		// 套节字被对方关闭
		if(dwTrans == 0 &&  pOverLaps!=NULL &&(pOverLaps->m_IOType == IO_TYPE_READ || pOverLaps->m_IOType == IO_TYPE_WRITE))	
		{
			closesocket(pPerSockData->m_Socket);
			g_ArrSocketData.remove(pPerSockData);
			g_ArrOverLapEx.remove(pOverLaps);
			delete pPerSockData;
			delete pOverLaps;
			pPerSockData=NULL;
			pOverLaps=NULL;
			continue;
		}

		switch(pOverLaps->m_IOType)
		{
		case IO_TYPE_READ:		//网络读取
			{


				PostRecv(pPerSockData);
				printf("收到客户端数据:%s\n",pOverLaps->m_szBuf);


				
				DWORD dwSend = 0;
				OverLappedEx *pOverLapEx = NULL;
				pOverLapEx = new OverLappedEx(IO_TYPE_WRITE);
				memcpy_s(pOverLapEx->m_szBuf, BUFFER_SIZE, "这是服务端发送的消息", ARRAYSIZE("这是服务端发送的消息"));
				WSABUF wsaBuf = {0};
				wsaBuf.buf = pOverLapEx->m_szBuf;
				wsaBuf.len = ARRAYSIZE("这是服务端发送的消息");
				g_ArrOverLapEx.push_back(pOverLapEx);
				//针对单个客户端发
				//WSASend(pPerSockData->m_Socket, &wsaBuf, 1, &dwSend, 0, &(pOverLapEx->m_OLap), NULL);
				//所有客户端群发
				list<PerSocketData*>::const_iterator cIterSocke;  
				for (cIterSocke = g_ArrSocketData.begin(); cIterSocke != g_ArrSocketData.end(); cIterSocke++)  
				{
					//过滤自己
					if ((*cIterSocke)->m_Socket==g_socket)continue;
					WSASend((*cIterSocke)->m_Socket, &wsaBuf, 1, &dwSend, 0, &(pOverLapEx->m_OLap), NULL);
				}



				g_ArrOverLapEx.remove(pOverLaps);
				if (pOverLaps)
				{
					delete pOverLaps;
					pOverLaps =NULL;
				}
			}
			break;
		case IO_TYPE_WRITE:
			{
				g_ArrOverLapEx.remove(pOverLaps);
				if (pOverLaps)
				{
					delete pOverLaps;
					pOverLaps =NULL;
				}
			}
			break;
		case IO_TYPE_ACCEPT:
			{
				if (pPerSockData->m_AccSocket == INVALID_SOCKET)continue;

				PerSocketData *pData = AssignSockToCompletionPort(pPerSockData->m_AccSocket);
				PostRecv(pData);
				PostAccept(pPerSockData);

				g_ArrOverLapEx.remove(pOverLaps);
				if (pOverLaps)
				{
					delete pOverLaps;
					pOverLaps =NULL;
				}

			}
			break;
		default:break;
		}
	}
	InterlockedDecrement(&g_ThreadNum);
	return 0;
}

//启动服务
DWORD ThreadProc(  LPVOID lpParameter)
{
	//设置变量
	g_ThreadNum=0;
	int nPort = 4567;

	//创建完成端口对象
	g_hCompletion = ::CreateIoCompletionPort(INVALID_HANDLE_VALUE, 0, 0, 0);

	// 创建监听套节字
	g_socket = WSASocket(AF_INET , SOCK_STREAM , IPPROTO_TCP, NULL, 0, WSA_FLAG_OVERLAPPED);
	assert(g_socket!=INVALID_SOCKET);
	if (g_socket==INVALID_SOCKET)return FALSE;


	//绑定端口
	SOCKADDR_IN  ServerAddr;
	ServerAddr.sin_family = AF_INET;
	ServerAddr.sin_port = ntohs(nPort);
	ServerAddr.sin_addr.S_un.S_addr = INADDR_ANY;
	if (bind(g_socket, (sockaddr*)&ServerAddr, sizeof(ServerAddr))==SOCKET_ERROR)
	{
		assert(FALSE);
		return FALSE;
	}

	//监听
	listen(g_socket, SOMAXCONN);

	//关联套接字
	PerSocketData *pSockData = AssignSockToCompletionPort(g_socket);

	//队列工作项在线程池中的工作线程
	QueueUserWorkItem(ServerThread, g_hCompletion, WT_EXECUTELONGFUNCTION);

	PostAccept(pSockData);

	return TRUE;
}

int main(void)
{

	// 初始化WS2_32.dll
	WSADATA wsaData;
	WORD sockVersion = MAKEWORD(2, 0);
	WSAStartup(sockVersion, &wsaData);

	//启动服务
	CreateThread(NULL,NULL,(LPTHREAD_START_ROUTINE)ThreadProc,NULL,NULL,NULL);

	getchar();

	//清理线程工作
	while(g_ThreadNum > 0)
	{
		PostQueuedCompletionStatus(g_hCompletion, 0, 0, NULL);
		Sleep(100);
	}

	 list<PerSocketData*>::const_iterator cIterSocke;  
	 for (cIterSocke = g_ArrSocketData.begin(); cIterSocke != g_ArrSocketData.end(); cIterSocke++)  
	 {
		 closesocket((*cIterSocke)->m_Socket);
		 delete *cIterSocke;
	 }
	 g_ArrSocketData.clear();

	 list<OverLappedEx*>::const_iterator cIter;  
	 for (cIter = g_ArrOverLapEx.begin(); cIter != g_ArrOverLapEx.end(); cIter++)  
	 {
		 delete *cIter;
	 }
	 g_ArrOverLapEx.clear();

	CloseHandle(g_hCompletion);	 

	WSACleanup();	

	return 0;
}