测试Winsock的select

说明

实现了一个回显一行字符串的服务器：客户端发送一行字符串，一'\n'结尾，服务器接受完一行后就原封不动地发回给客户端。

windows下对select的能监控的Socket数量是有限制的，若超过，一种方案是再开一个线程。

#ifndef FD_SETSIZE
#define FD_SETSIZE      64
#endif /* FD_SETSIZE */

代码

#include <iostream>
#define _WINSOCK_DEPRECATED_NO_WARNINGS
#include <WinSock2.h>
#include <vector>
#include <memory>

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


struct ClientSocketItem
{
	ClientSocketItem()
	{
		hSocket = NULL;
		memset(szRecv, 0, sizeof(szRecv));
		nRecvSize = 0;
		bNeedWrite = false;
		nWriteOffset = 0;
	}

	SOCKET hSocket;
	char szRecv[1024];
	unsigned int nRecvSize;
	std::string strIp;

	bool bNeedWrite;//接受完成，可以发送
	unsigned int nWriteOffset = 0;
};

std::vector<std::shared_ptr<ClientSocketItem>> g_Clients;


void do_accept(SOCKET hListenSocket)
{
	sockaddr_in mPeerAddr = { 0 };
	int nAddrLen = sizeof(sockaddr);
	SOCKET hClientSocket = accept(hListenSocket, (sockaddr*)(&mPeerAddr), &nAddrLen);
	if (INVALID_SOCKET == hClientSocket)
	{
		std::cout << "accept failed with error "
			<< WSAGetLastError() << std::endl;
	}
	else
	{
		unsigned long nNoBlock = 0;
		ioctlsocket(hClientSocket, FIONBIO, &nNoBlock);

		std::string strIpAddr = inet_ntoa(mPeerAddr.sin_addr);
		std::cout << "accept success, peer ip is " << strIpAddr.c_str() << std::endl;

		auto pClient = std::make_shared<ClientSocketItem>();
		pClient->hSocket = hClientSocket;
		pClient->strIp = strIpAddr;
		g_Clients.push_back(pClient);
	}
}

bool do_read(const std::shared_ptr<ClientSocketItem>& pClient )
{
	if (!pClient)
	{
		return false;
	}

	char c = 0; //测试用，每次只读一个字符
	int nRecvValue = recv(pClient->hSocket, &c, 1, 0);
	if (nRecvValue > 0)
	{
		pClient->szRecv[pClient->nRecvSize] = c;
		pClient->nRecvSize += 1;
		std::cout << "read one char: " << c << std::endl;
		if (c == '\n')
		{
			std::cout << "read finished" << std::endl;
			pClient->bNeedWrite = true;
		}

		return true;
	}
	else if (0 == nRecvValue)
	{
		std::cout << "peer client closed" << std::endl;
		closesocket(pClient->hSocket);
		return false;
	}
	else
	{
		int nError = WSAGetLastError();
		if (WSAEWOULDBLOCK != nError)
		{
			std::cerr << "recv failed with error " << nError << std::endl;
			closesocket(pClient->hSocket);
			return false;
		}

		std::cout << "next recv" << std::endl;
		return true;
	}
}


bool do_write(const std::shared_ptr<ClientSocketItem>& pClient)
{
	if (!pClient)
	{
		return false;
	}

	if (!pClient->bNeedWrite)
	{
		return true;
	}

	//测试用，每次只发送一个字符
	int nRet = send(pClient->hSocket, pClient->szRecv + pClient->nWriteOffset, 1, 0);
	if (nRet >= 0)
	{
		std::cout << "send one char: " << pClient->szRecv[pClient->nWriteOffset] << std::endl;

		pClient->nWriteOffset += 1;
		if (pClient->nWriteOffset == pClient->nRecvSize)
		{
			std::cout << "send finished, close client(" << pClient->strIp.c_str()
				<< ")" << std::endl;
			pClient->bNeedWrite = false;
			closesocket(pClient->hSocket);
			return false;
		}

		return true;
	}
	else
	{
		int nError = WSAGetLastError();
		if (WSAEWOULDBLOCK != nError)
		{
			std::cerr << "send failed with error " << nError << std::endl;
			closesocket(pClient->hSocket);
			return false;
		}

		std::cout << "next send" << std::endl;
		return true;
	}
}



int main(int argc, char* argv)
{
	WORD wVersionRequested = MAKEWORD(2, 2);
	WSADATA wsaData = { 0 };
	int err = WSAStartup(wVersionRequested, &wsaData);
	if (err != 0)
	{
		return -1;
	}

	if (LOBYTE(wsaData.wVersion) != 2 ||
		HIBYTE(wsaData.wVersion) != 2)
	{
		WSACleanup();
		return -1;
	}

	SOCKET hListenSocket = socket(AF_INET, SOCK_STREAM, 0);
	if (INVALID_SOCKET == hListenSocket)
	{
		std::cerr << "create socket failed with error " << WSAGetLastError()
			<< std::endl;
		return -1;
	}

	sockaddr_in mSockAddrIn = { 0 };
	mSockAddrIn.sin_family = AF_INET;
	mSockAddrIn.sin_port = htons((u_short)8878);
	mSockAddrIn.sin_addr.S_un.S_addr = inet_addr("0.0.0.0");
	if (SOCKET_ERROR == bind(hListenSocket, (sockaddr*)(&mSockAddrIn),
		sizeof(sockaddr)))
	{
		std::cerr << "bind failed with error " << WSAGetLastError() << std::endl;
		return -1;
	}

	if (SOCKET_ERROR == listen(hListenSocket, SOMAXCONN))
	{
		std::cerr << "listen failed with error " << WSAGetLastError() << std::endl;
		return -1;
	}


	std::vector<SOCKET> mAllClients;
	while (true)
	{
		fd_set readfds;
		fd_set writefds;
		fd_set exceptfds;
		FD_ZERO(&readfds);
		FD_ZERO(&writefds);
		FD_ZERO(&exceptfds);

		FD_SET(hListenSocket, &readfds);
		//FD_SET(hListenSocket, &writefds);
		//FD_SET(hListenSocket, &exceptfds);

		for (auto it = g_Clients.begin(); it != g_Clients.end(); ++it)
		{
			std::shared_ptr<ClientSocketItem> pClientItem = *it;
			if (!pClientItem) continue;
			FD_SET(pClientItem->hSocket, &readfds);

			if (pClientItem->bNeedWrite)//否则select会一直有事件
			{
				FD_SET(pClientItem->hSocket, &writefds);
			}
			
			FD_SET(pClientItem->hSocket, &exceptfds);
		}

		int nRet = select(0, &readfds, &writefds, &exceptfds, nullptr);
		std::cout << "select return with " << nRet << std::endl;
		if (nRet > 0)
		{
			//read
			if (FD_ISSET(hListenSocket, &readfds))
			{
				do_accept(hListenSocket);
			}
			for (auto it = g_Clients.begin(); it != g_Clients.end();)
			{
				std::shared_ptr<ClientSocketItem> pClientItem = *it;
				if (pClientItem && FD_ISSET(pClientItem->hSocket, &readfds))
				{
					if (!do_read(pClientItem))
					{
						it = g_Clients.erase(it);
						continue;
					}
				}

				++it;
			}

			//write
			for (auto it = g_Clients.begin(); it != g_Clients.end();)
			{
				std::shared_ptr<ClientSocketItem> pClientItem = *it;
				if (pClientItem && FD_ISSET(pClientItem->hSocket, &writefds))
				{
					if (!do_write(pClientItem))
					{
						it = g_Clients.erase(it);
						continue;
					}
				}

				++it;
			}//end of for

			//error 
			for (auto it = g_Clients.begin(); it != g_Clients.end();)
			{
				std::shared_ptr<ClientSocketItem> pClientItem = *it;
				if (pClientItem && FD_ISSET(pClientItem->hSocket, &exceptfds))
				{
					std::cerr << "client socket except, close and remove it" << std::endl;
					closesocket(pClientItem->hSocket);
					it = g_Clients.erase(it);
					continue;
				}

				++it;
			}//end of for
		}
		else
		{
			std::cerr << "select failed with error " << WSAGetLastError() << std::endl;
		}

	}


	return 0;
}