簡單的IOCP程序(TCP)

// tcp_iocp_serve.cpp

#include <WinSock2.h>
#include <Windows.h>
#include <vector>
#include <iostream>

using namespace std;

#pragma comment(lib, "Ws2_32.lib")      // Socket編程需用的動態鏈接庫
#pragma comment(lib, "Kernel32.lib")    // IOCP需要用到的動態鏈接庫

/**
* 結構體名稱：PER_IO_DATA
* 結構體功能：重疊I/O需要用到的結構體，臨時記錄IO數據
**/
const int g_nDataBuffSize = 2 * 1024;
typedef struct
{
	OVERLAPPED sOverLapped;
	WSABUF sDatabuff;
	char szBuffer[g_nDataBuffSize];
	int nBufferLen;
	int nOperationType;
}PER_IO_OPERATEION_DATA, *LPPER_IO_OPERATION_DATA, *LPPER_IO_DATA, PER_IO_DATA;

/**
* 結構體名稱：PER_HANDLE_DATA
* 結構體存儲：記錄單個套接字的數據，包括了套接字的變量及套接字的對應的客戶端的地址。
* 結構體作用：當服務器連接上客戶端時，信息存儲到該結構體中，知道客戶端的地址以便於回訪。
**/
typedef struct
{
	SOCKET socket;
	SOCKADDR_STORAGE clientaddr;
}PER_HANDLE_DATA, *LPPER_HANDLE_DATA;

// 定義全局變量
const int g_nDefaultPort = 6000;
vector < PER_HANDLE_DATA* > g_vctClientGroup;      // 記錄客戶端的向量組

HANDLE g_hMutex = CreateMutex(NULL, FALSE, NULL);
DWORD WINAPI ServerWorkThread(LPVOID lpCompletionPortID);
DWORD WINAPI ServerSendThread(LPVOID lpParam);

// 開始主函數
int main()
{
	// 加載socket動態鏈接庫
	WORD wVersionRequested = MAKEWORD(2, 2); // 請求2.2版本的WinSock庫
	WSADATA wsaData;    // 接收Windows Socket的結構信息
	DWORD dwErr = WSAStartup(wVersionRequested, &wsaData);

	if (0 != dwErr)
	{
		// 檢查套接字庫是否申請成功
		cerr << "Request Windows Socket Library Error!\n";
		system("pause");
		return -1;
	}
	if (LOBYTE(wsaData.wVersion) != 2
		|| HIBYTE(wsaData.wVersion) != 2)
	{
		// 檢查是否申請了所需版本的套接字庫
		WSACleanup();
		cerr << "Request Windows Socket Version 2.2 Error!\n";
		system("pause");
		return -1;
	}

	// 創建IOCP的內核對象
	/**
	* 需要用到的函數的原型：
	* HANDLE WINAPI CreateIoCompletionPort(
	*    __in   HANDLE FileHandle,     // 已經打開的文件句柄或者空句柄，一般是客戶端的句柄
	*    __in   HANDLE ExistingCompletionPort, // 已經存在的IOCP句柄
	*    __in   ULONG_PTR CompletionKey,   // 完成鍵，包含了指定I/O完成包的指定文件
	*    __in   DWORD NumberOfConcurrentThreads // 真正併發同時執行最大線程數，一般推介是CPU核心數*2
	* );
	**/
	HANDLE hCompletionPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
	if (NULL == hCompletionPort)
	{
		// 創建IO內核對象失敗
		cerr << "CreateIoCompletionPort failed. Error:" << GetLastError() << endl;
		system("pause");
		return -1;
	}

	// 創建IOCP線程--線程裏面創建線程池

	// 確定處理器的核心數量
	SYSTEM_INFO sSysInfo;
	GetSystemInfo(&sSysInfo);

	// 基於處理器的核心數量創建線程
	for (DWORD i = 0; i < (sSysInfo.dwNumberOfProcessors * 2); ++i)
	{
		// 創建服務器工作器線程，並將完成端口傳遞到該線程
		HANDLE hThreadHandle = CreateThread(NULL, 0, ServerWorkThread, hCompletionPort, 0, NULL);
		if (NULL == hThreadHandle)
		{
			cerr << "Create Thread Handle failed. Error:" << GetLastError() << endl;
			system("pause");
			return -1;
		}
		CloseHandle(hThreadHandle);
	}

	// 建立流式套接字
	SOCKET srvSocket = socket(AF_INET, SOCK_STREAM, 0);

	// 綁定SOCKET到本機
	SOCKADDR_IN srvAddr;
	srvAddr.sin_addr.S_un.S_addr = htonl(INADDR_ANY);
	srvAddr.sin_family = AF_INET;
	srvAddr.sin_port = htons(g_nDefaultPort);
	int nBindResult = bind(srvSocket, (SOCKADDR*)&srvAddr, sizeof(SOCKADDR));
	if (SOCKET_ERROR == nBindResult)
	{
		cerr << "Bind failed. Error:" << GetLastError() << endl;
		system("pause");
		return -1;
	}

	// 將SOCKET設置爲監聽模式
	int nListenResult = listen(srvSocket, 10);
	if (SOCKET_ERROR == nListenResult)
	{
		cerr << "Listen failed. Error: " << GetLastError() << endl;
		system("pause");
		return -1;
	}

	// 開始處理IO數據
	cout << "本服務器已準備就緒，正在等待客戶端的接入...\n";

	// 創建用於發送數據的線程
	HANDLE hSendThread = CreateThread(NULL, 0, ServerSendThread, 0, 0, NULL);

	while (true)
	{
		PER_HANDLE_DATA * pPerHandleData = NULL;
		SOCKADDR_IN saRemote;
		int nRemoteLen = 0;
		SOCKET acceptSocket;

		// 接收連接，並分配完成端，這兒可以用AcceptEx()
		nRemoteLen = sizeof(saRemote);
		acceptSocket = accept(srvSocket, (SOCKADDR*)&saRemote, &nRemoteLen);
		if (SOCKET_ERROR == acceptSocket)
		{
			// 接收客戶端失敗
			cerr << "Accept Socket Error: " << GetLastError() << endl;
			system("pause");
			return -1;
		}

		// 創建用來和套接字關聯的單句柄數據信息結構
		pPerHandleData = (LPPER_HANDLE_DATA)GlobalAlloc(GPTR, sizeof(PER_HANDLE_DATA));  // 在堆中爲這個PerHandleData申請指定大小的內存
		pPerHandleData->socket = acceptSocket;
		memcpy(&pPerHandleData->clientaddr, &saRemote, nRemoteLen);
		g_vctClientGroup.push_back(pPerHandleData);       // 將單個客戶端數據指針放到客戶端組中

		// 將接受套接字和完成端口關聯
		CreateIoCompletionPort((HANDLE)(pPerHandleData->socket), hCompletionPort, (DWORD)pPerHandleData, 0);


		// 開始在接受套接字上處理I/O使用重疊I/O機制
		// 在新建的套接字上投遞一個或多個異步
		// WSARecv或WSASend請求，這些I/O請求完成後，工作者線程會爲I/O請求提供服務
		// 單I/O操作數據(I/O重疊)
		LPPER_IO_OPERATION_DATA perIoData = NULL;
		perIoData = (LPPER_IO_OPERATION_DATA)GlobalAlloc(GPTR, sizeof(PER_IO_OPERATEION_DATA));
		ZeroMemory(&(perIoData->sOverLapped), sizeof(OVERLAPPED));
		perIoData->sDatabuff.len = 1024;
		perIoData->sDatabuff.buf = perIoData->szBuffer;
		perIoData->nOperationType = 0;    // read

		DWORD dwRecvBytes;
		DWORD dwFlags = 0;
		WSARecv(pPerHandleData->socket, &(perIoData->sDatabuff), 1, &dwRecvBytes, &dwFlags, &(perIoData->sOverLapped), NULL);
	}

	system("pause");
	return 0;
}

// 開始服務工作線程函數
DWORD WINAPI ServerWorkThread(LPVOID lpParam)
{
	HANDLE hCompletionPort = (HANDLE)lpParam;
	DWORD dwBytesTransferred = 0;
	LPOVERLAPPED lpOverlapped;
	LPPER_HANDLE_DATA perHandleData = NULL;
	LPPER_IO_DATA perIoData = NULL;
	DWORD dwRecvBytes = 0;
	DWORD dwFlags = 0;
	BOOL bRet = FALSE;

	while (true)
	{
		bRet = GetQueuedCompletionStatus(hCompletionPort, &dwBytesTransferred, (PULONG_PTR)&perHandleData, (LPOVERLAPPED*)&lpOverlapped, INFINITE);
		if (bRet == 0)
		{
			cerr << "GetQueuedCompletionStatus Error: " << GetLastError() << endl;
			return -1;
		}
		perIoData = (LPPER_IO_DATA)CONTAINING_RECORD(lpOverlapped, PER_IO_DATA, sOverLapped);

		// 檢查在套接字上是否有錯誤發生
		if (0 == dwBytesTransferred)
		{
			closesocket(perHandleData->socket);
			GlobalFree(perHandleData);
			GlobalFree(perIoData);
			continue;
		}

		// 開始數據處理，接收來自客戶端的數據
		WaitForSingleObject(g_hMutex, INFINITE);
		cout << "A Client says: " << perIoData->sDatabuff.buf << endl;
		ReleaseMutex(g_hMutex);

		// 爲下一個重疊調用建立單I/O操作數據
		ZeroMemory(&(perIoData->sOverLapped), sizeof(OVERLAPPED)); // 清空內存
		perIoData->sDatabuff.len = 1024;
		perIoData->sDatabuff.buf = perIoData->szBuffer;
		perIoData->nOperationType = 0;    // read
		WSARecv(perHandleData->socket, &(perIoData->sDatabuff), 1, &dwRecvBytes, &dwFlags, &(perIoData->sOverLapped), NULL);
	}

	return 0;
}


// 發送信息的線程執行函數
DWORD WINAPI ServerSendThread(LPVOID /*lpParam*/)
{
	while (true)
	{
		char szTalk[200] = {0};
		gets_s(szTalk);
		int nLen = 0;
		for (nLen = 0; szTalk[nLen] != '\0'; ++nLen)
		{
			// 找出這個字符組的長度
		}
		szTalk[nLen] = '\n';
		szTalk[++nLen] = '\0';
		cout << "I Say:" << szTalk;
		WaitForSingleObject(g_hMutex, INFINITE);
		for (int i = 0; i < (int)g_vctClientGroup.size(); ++i)
		{
			send(g_vctClientGroup[i]->socket, szTalk, 200, 0);  // 發送信息
		}
		ReleaseMutex(g_hMutex);
	}

	return 0;
}

// tcp_iocp_client.cpp

#include <iostream>
#include <cstdio>
#include <string>
#include <cstring>
#include <winsock2.h>
#include <Windows.h>

using namespace std;

#pragma comment(lib, "Ws2_32.lib")      // Socket編程需用的動態鏈接庫

SOCKET g_sockClient;                    // 連接成功後的套接字
HANDLE g_hBufferMutex = NULL;           // 令其能互斥成功正常通信的信號量句柄
const int g_nDefaultPort = 6000;

DWORD WINAPI SendMessageThread(LPVOID lpParam);
DWORD WINAPI ReceiveMessageThread(LPVOID lpParam);

int main()
{
	// 加載socket動態鏈接庫(dll)
	WORD wVersionRequested;
	WSADATA wsaData;    // 這結構是用於接收Wjndows Socket的結構信息的
	wVersionRequested = MAKEWORD(2, 2);   // 請求2.2版本的WinSock庫
	int nErr = WSAStartup(wVersionRequested, &wsaData);
	if (nErr != 0) 
	{
		// 返回值爲零的時候是表示成功申請WSAStartup
		return -1;
	}
	if (LOBYTE(wsaData.wVersion) != 2
		|| HIBYTE(wsaData.wVersion) != 2)
	{
		// 檢查版本號是否正確
		WSACleanup();
		return -1;
	}

	// 創建socket操作，建立流式套接字，返回套接字號sockClient
	g_sockClient = socket(AF_INET, SOCK_STREAM, 0);
	if (g_sockClient == INVALID_SOCKET)
	{
		printf("Error at socket():%ld\n", WSAGetLastError());
		WSACleanup();
		return -1;
	}

	// 將套接字sockClient與遠程主機相連
	// int connect( SOCKET s,  const struct sockaddr* name,  int namelen);
	// 第一個參數：需要進行連接操作的套接字
	// 第二個參數：設定所需要連接的地址信息
	// 第三個參數：地址的長度
	SOCKADDR_IN addrSrv;
	addrSrv.sin_addr.S_un.S_addr = inet_addr("127.0.0.1");      // 本地迴路地址是127.0.0.1;
	addrSrv.sin_family = AF_INET;
	addrSrv.sin_port = htons(g_nDefaultPort);
	while (SOCKET_ERROR == connect(g_sockClient, (SOCKADDR*)&addrSrv, sizeof(SOCKADDR)))
	{
		// 如果還沒連接上服務器則要求重連
		cout << "服務器連接失敗，是否重新連接？（Y/N):";
		char choice;
		while (cin >> choice && (!((choice != 'Y' && choice == 'N') || (choice == 'Y' && choice != 'N'))))
		{
			cout << "輸入錯誤，請重新輸入:";
			cin.sync();
			cin.clear();
		}
		if (choice == 'Y')
		{
			continue;
		}
		else
		{
			cout << "退出系統中...";
			system("pause");
			return 0;
		}
	}
	cin.sync();
	cout << "本客戶端已準備就緒，用戶可直接輸入文字向服務器反饋信息。\n";

	send(g_sockClient, "\nAttention: A Client has enter...\n", 200, 0);

	g_hBufferMutex = CreateSemaphore(NULL, 1, 1, NULL);


	HANDLE hSendThread = CreateThread(NULL, 0, SendMessageThread, NULL, 0, NULL);
	HANDLE hReceiveThread = CreateThread(NULL, 0, ReceiveMessageThread, NULL, 0, NULL);


	WaitForSingleObject(hSendThread, INFINITE);  // 等待線程結束
	closesocket(g_sockClient);
	CloseHandle(hSendThread);
	CloseHandle(hReceiveThread);
	CloseHandle(g_hBufferMutex);
	WSACleanup();   // 終止對套接字庫的使用

	printf("End linking...\n");
	printf("\n");

	system("pause");

	return 0;
}


DWORD WINAPI SendMessageThread(LPVOID /*lpParam*/)
{
	while (true)
	{
		string strTalk;
		getline(cin, strTalk);
		WaitForSingleObject(g_hBufferMutex, INFINITE);     // P（資源未被佔用）
		if ("quit" == strTalk)
		{
			strTalk.push_back('\0');
			send(g_sockClient, strTalk.c_str(), 200, 0);
			break;
		}
		else
		{
			strTalk.append("\n");
		}
		printf("\nI Say:(\"quit\"to exit):");
		cout << strTalk;
		int nSendResult = send(g_sockClient, strTalk.c_str(), /*200*/strTalk.length() + 1, 0); // 發送信息
		ReleaseSemaphore(g_hBufferMutex, 1, NULL);     // V（資源佔用完畢）

		if (SOCKET_ERROR == nSendResult)
		{
			cout << "Send failed.Error:" << WSAGetLastError() << endl;
			break ;
		}
	}

	return 0;
}


DWORD WINAPI ReceiveMessageThread(LPVOID /*lpParam*/)
{
	while (true)
	{
		char szRecvBuf[300] = {0};

		int nRecvResult = recv(g_sockClient, szRecvBuf, 200, 0);

		if (SOCKET_ERROR == nRecvResult)
		{
			cout << "Recv failed.Error:" << WSAGetLastError() << endl;
			break ;
		}

		WaitForSingleObject(g_hBufferMutex, INFINITE);     // P（資源未被佔用）

		printf("%s Says: %s", "Server", szRecvBuf);        // 接收信息

		ReleaseSemaphore(g_hBufferMutex, 1, NULL);         // V（資源佔用完畢）
	}

	return 0;
}