IOCP应用

//---------------------------------------------------------------------------

#ifndef IOCP_H
#define IOCP_H
//---------------------------------------------------------------------------
#include <winsock2.h>
#include <windows.h>
#include <process.h>

#pragma comment(lib,"ws2_32.lib")

#define BUFFER_SIZE 1024

typedef struct _PER_HANDLE_DATA		// per-handle数据
{
	SOCKET s;			// 对应的套节字句柄
	sockaddr_in addr;	// 客户方地址
} PER_HANDLE_DATA, *PPER_HANDLE_DATA;


typedef struct _PER_IO_DATA			// per-I/O数据
{
	OVERLAPPED ol;			// 重叠结构
	char buf[BUFFER_SIZE];	// 数据缓冲区
	int nOperationType;		// 操作类型
        _PER_IO_DATA *perSend; 
#define OP_READ   1
#define OP_WRITE  2
#define OP_ACCEPT 3
} PER_IO_DATA, *PPER_IO_DATA;



class InterLock
{
CRITICAL_SECTION cs;
public:
  InterLock()
  {
    InitializeCriticalSectionAndSpinCount(&cs,1024);
  }
  ~InterLock()
  {
    DeleteCriticalSection(&cs);
  }
  void Enter()
  {
    EnterCriticalSection(&cs);
  }
  void Leave()
  {
    LeaveCriticalSection(&cs);
  }
};
class LockManager
{
InterLock *pLock;
public:
   LockManager(InterLock* lock):pLock(lock)
   {
     if(pLock)
     pLock->Enter();
   }
   ~LockManager()
   {
     if(pLock)
     pLock->Leave();
   }
};
struct THREAD_DATA
{
  InterLock lock;
  SOCKET sListen;
  bool bExit;
};
class InitSock
{
  WSADATA wsaData;
public:
  InitSock()
  {
    if(WSAStartup(MAKEWORD(2,2),&wsaData)!=0)
    exit(1);
  }
  ~InitSock()
  {
    WSACleanup();
  }
};

UINT WINAPI BeginThread(LPVOID param);

bool PostRecv(PPER_IO_DATA pPerIO,SOCKET s);
bool PostSend(PPER_IO_DATA pPerIO,SOCKET s,const char* buffer,int bufferlen);
#endif
 

//---------------------------------------------------------------------------


#pragma hdrstop

#include "IOCP.h"
#include <list>
using std::list;
#include <cassert>
//---------------------------------------------------------------------------

#pragma package(smart_init)




//-------------
// 初始化Winsock库
InitSock theSock;

class PerDataManager
{
   static list<PPER_HANDLE_DATA> handlelist;
   static list<PPER_IO_DATA>    iolist;
public:
  static InterLock lock;
  
  static void Add(const PPER_HANDLE_DATA &pd)
  {
     LockManager manager(&lock);
     handlelist.push_back(pd);
  }
  static void Delete(const PPER_HANDLE_DATA &pd)
  {
    LockManager manager(&lock);
    list<PPER_HANDLE_DATA>::iterator it=find(handlelist.begin(),handlelist.end(),pd);
    if(it!=handlelist.end())
    {
        ::closesocket((*it)->s);
        ::GlobalFree((*it));
        handlelist.erase(it);
    }
  }
  static void Clear()
  {
    for(list<PPER_HANDLE_DATA>::iterator it=handlelist.begin();it!=handlelist.end();++it)
    {
        ::closesocket((*it)->s);
        ::GlobalFree((*it));
    }
    handlelist.clear();

    for(list<PPER_IO_DATA>::iterator it=iolist.begin();it!=iolist.end();++it)
    {
        ::GlobalFree((*it));
    }
    iolist.clear();
  }
  static void Add(const PPER_IO_DATA &pd)
  {
     LockManager manager(&lock);
     iolist.push_back(pd);
  }
  static void Delete(const PPER_IO_DATA &pd)
  {
    LockManager manager(&lock);
    list<PPER_IO_DATA>::iterator it=find(iolist.begin(),iolist.end(),pd);
    if(it!=iolist.end())
    {
        ::GlobalFree((*it));
        iolist.erase(it);
    }
  }
};
list<PPER_HANDLE_DATA> PerDataManager::handlelist;
list<PPER_IO_DATA> PerDataManager::iolist;
InterLock PerDataManager::lock;
PerDataManager pdatamanager;

bool PostRecv(PPER_IO_DATA pPerIO,SOCKET s)
{
    WSABUF buf;
    buf.buf = pPerIO->buf ;
    buf.len = BUFFER_SIZE;
    pPerIO->nOperationType = OP_READ;

    DWORD nFlags = 0;
    DWORD dwTrans = 0;
    int ret=::WSARecv(s, &buf, 1, &dwTrans, &nFlags, &pPerIO->ol, NULL);
    if((ret==SOCKET_ERROR)&&(WSAGetLastError()!=WSA_IO_PENDING))
    return false;
    return true;
}
bool PostSend(PPER_IO_DATA pPerIO,SOCKET s,const char* buffer,int bufferlen)
{
    memset(pPerIO->buf,0,BUFFER_SIZE);
    if(bufferlen>BUFFER_SIZE)
    bufferlen=BUFFER_SIZE;
    
    memcpy(pPerIO->buf,buffer,bufferlen);
    
    pPerIO->nOperationType = OP_WRITE;
    WSABUF buf_w;
    buf_w.buf = pPerIO->buf;
    buf_w.len = bufferlen;	
    DWORD dwRecv_w;
    DWORD dwFlags_w = 0;
    int ret=::WSASend(s, &buf_w, 1, &dwRecv_w, dwFlags_w, &pPerIO->ol, NULL);
    if((ret==SOCKET_ERROR)&&(WSAGetLastError()!=WSA_IO_PENDING))
    return false;
    return true;
}

UINT WINAPI ServerThread(LPVOID lpParam)
{
	// 得到完成端口对象句柄
	HANDLE hCompletion = (HANDLE)lpParam;

        InitCOM initcom;
	DWORD dwTrans;
	PPER_HANDLE_DATA pPerHandle;
	PPER_IO_DATA pPerIO;
	while(TRUE)
	{
		// 在关联到此完成端口的所有套节字上等待I/O完成
		BOOL bOK = ::GetQueuedCompletionStatus(hCompletion, 
			&dwTrans, (LPDWORD)&pPerHandle, (LPOVERLAPPED*)&pPerIO, WSA_INFINITE);

                if(!pPerHandle)//结束
                break;

		if(!bOK)						// 在此套节字上有错误发生
		{
                        pdatamanager.Delete(pPerHandle);
                        pdatamanager.Delete(pPerIO);
			continue;
		}
		
		if(dwTrans == 0 &&				// 套节字被对方关闭
			(pPerIO->nOperationType == OP_READ || pPerIO->nOperationType == OP_WRITE))	
			
		{
                        pdatamanager.Delete(pPerHandle);
                        pdatamanager.Delete(pPerIO);
			continue;
		}

		switch(pPerIO->nOperationType)	// 通过per-I/O数据中的nOperationType域查看什么I/O请求完成了
		{
		case OP_READ:	// 完成一个接收请求
			{
				pPerIO->buf[dwTrans] = '\0';

                                
                                memset(pPerIO->buf,0,BUFFER_SIZE);
                                PostRecv(pPerIO,pPerHandle->s);
			}
			break;
		case OP_WRITE: // 应该等待通知再继续Post（因为使用同一块缓存区），我这里没有作处理，因为我的数据量很少（不超过128字节)。
                        {
                           memset(pPerIO->buf,0,BUFFER_SIZE);
                        }
                        break;
		case OP_ACCEPT:
			break;
		}
	}
	return 0;
}

UINT WINAPI BeginThread(LPVOID param)
{
        THREAD_DATA *pThreaddata=((THREAD_DATA*)param);
	SOCKET sListen = pThreaddata->sListen;

	// 创建完成端口对象，创建工作线程处理完成端口对象中事件
	HANDLE hCompletion = ::CreateIoCompletionPort(INVALID_HANDLE_VALUE, 0, 0, 0);
        assert(hCompletion);
	HANDLE hWorker=(HANDLE)_beginthreadex(NULL, 0, ServerThread, (LPVOID)hCompletion, 0, 0);
        assert(hWorker);
	// 循环处理到来的连接

	while(TRUE)
	{
                {
                   LockManager manager(&(pThreaddata->lock));
                   if(pThreaddata->bExit)
                   break;
                }
		// 等待接受未决的连接请求
		SOCKADDR_IN saRemote;
		int nRemoteLen = sizeof(saRemote);
		SOCKET sNew = ::accept(sListen, (sockaddr*)&saRemote, &nRemoteLen);
                

		// 接受到新连接之后，为它创建一个per-handle数据，并将它们关联到完成端口对象。
		PPER_HANDLE_DATA pPerHandle =(PPER_HANDLE_DATA)::GlobalAlloc(GPTR, sizeof(PER_HANDLE_DATA));
                assert(pPerHandle);
		pPerHandle->s = sNew;
		memcpy(&pPerHandle->addr, &saRemote, nRemoteLen);

		::CreateIoCompletionPort((HANDLE)pPerHandle->s, hCompletion, (DWORD)pPerHandle, 0);

                pdatamanager.Add(pPerHandle);
                
		// 投递一个接收请求
		PPER_IO_DATA pPerIO = (PPER_IO_DATA)::GlobalAlloc(GPTR, sizeof(PER_IO_DATA));
                assert(pPerIO);
                
		pPerIO->nOperationType = OP_READ;
		WSABUF buf;
		buf.buf = pPerIO->buf;
		buf.len = BUFFER_SIZE;	
		DWORD dwRecv;
		DWORD dwFlags = 0;
		::WSARecv(pPerHandle->s, &buf, 1, &dwRecv, &dwFlags, &pPerIO->ol, NULL);

                pdatamanager.Add(pPerIO);

		pPerIO->perSend = (PPER_IO_DATA)::GlobalAlloc(GPTR, sizeof(PER_IO_DATA));
		assert(pPerIO->perSend);
                pdatamanager.Add(pPerIO->perSend);
	}
        PostQueuedCompletionStatus(hCompletion, 0, 0, NULL);
        WaitForSingleObject(hWorker,INFINITE);
        CloseHandle(hWorker);
        CloseHandle(hCompletion);
        
        pdatamanager.Clear();
        return 0;
}

HANDLE hBegin=NULL;
SOCKET sListen=INVALID_SOCKET;
THREAD_DATA tdata;

bool Listen(int nPort)
{
  if(sListen==INVALID_SOCKET)
  sListen=::socket(AF_INET,SOCK_STREAM,IPPROTO_TCP);
  if(sListen==INVALID_SOCKET)
  return false;
  sockaddr_in addr;
  addr.sin_family=AF_INET;
  addr.sin_addr.S_un.S_addr=INADDR_ANY;
  addr.sin_port=::htons(nPort);

  if(::bind(sListen,(PSOCKADDR)&addr,sizeof(addr))==SOCKET_ERROR)
  return false;

  return ::listen(sListen,5)!=SOCKET_ERROR;
}
void Close()
{
  ::closesocket(sListen);
  sListen=INVALID_SOCKET;
}