iocpserverex.cpp

应用完成端口进行tcp的连接

// THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF
// ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A
// PARTICULAR PURPOSE.
//
// Copyright (C) 1998 - 2000  Microsoft Corporation.  All Rights Reserved.
//
// Module:
//      iocpserverex.cpp
//
// Abstract:
//      This program is a Winsock echo server program that demonstrates the usage
//      of AcceptEx with IOCP. The AcceptEx function accepts a new connection, 
//      returns the local and remote address, and receives the first block of data 
//      sent by the client application. The design of this program is based on that 
//      in the iocpserver.cpp. But it uses overlapped AcceptEx on the IOCP also. 
//      AcceptEx allows data to be "returned" from an accepted connection.
//
//      Another point worth noting is that the Win32 API CreateThread() does not 
//      initialize the C Runtime and therefore, C runtime functions such as 
//      printf() have been avoid or rewritten (see myprintf()) to use just Win32 APIs.
//
//
//  Usage:
//      Start the server and wait for connections on port 6001
//          iocpserverex -e:6001
//
//  Build:
//      Use the headers and libs from the April98 Platform SDK or later.
//      Link with ws2_32.lib and mswsock.lib
//      
//  Author: Wei Hua, Barry Butterklee - Microsoft Developer Support
//
//

#ifdef _IA64_
	#pragma warning(disable:4127 4267)
#endif 

#ifndef WIN32_LEAN_AND_MEAN
	#define WIN32_LEAN_AND_MEAN
#endif

#define xmalloc(s) HeapAlloc(GetProcessHeap(),HEAP_ZERO_MEMORY,(s))
#define xfree(p)   HeapFree(GetProcessHeap(),0,(p))

#include <winsock2.h>
#include <mswsock.h>
#include <Ws2tcpip.h>
#include <stdio.h>
#include <stdlib.h>

#include "iocpserver.h"

char *g_Port = DEFAULT_PORT;
BOOL g_bEndServer = FALSE;			// set to TRUE on CTRL-C
BOOL g_bRestart = TRUE;				// set to TRUE to CTRL-BRK
BOOL g_bVerbose = FALSE;
HANDLE g_hIOCP = INVALID_HANDLE_VALUE;
SOCKET g_sdListen = INVALID_SOCKET;
HANDLE g_ThreadHandles[MAX_WORKER_THREAD];
WSAEVENT g_hCleanupEvent[1];
PPER_SOCKET_CONTEXT g_pCtxtListenSocket = NULL;
PPER_SOCKET_CONTEXT g_pCtxtList = NULL;		// linked list of context info structures
											// maintained to allow the the cleanup 
											// handler to cleanly close all sockets and 
											// free resources.

CRITICAL_SECTION g_CriticalSection;		// guard access to the global context list

int myprintf(const char *lpFormat, ...);

void __cdecl main (int argc, char *argv[])	{

    SYSTEM_INFO systemInfo;
	WSADATA wsaData;
	DWORD dwThreadCount = 0;
	int nRet = 0;

    g_ThreadHandles[0] = (HANDLE)WSA_INVALID_EVENT;

	for( int i = 0; i < MAX_WORKER_THREAD; i++ ) {
		g_ThreadHandles[i] = INVALID_HANDLE_VALUE;
	}

	if( !ValidOptions(argc, argv) )
		return;

	if( !SetConsoleCtrlHandler(CtrlHandler, TRUE) ) {
		myprintf("SetConsoleCtrlHandler() failed to install console handler: %d\n", 
				GetLastError());
		return;
	}

	GetSystemInfo(&systemInfo);
	dwThreadCount = systemInfo.dwNumberOfProcessors * 2;

	g_hCleanupEvent[0] = WSACreateEvent();
	if( g_hCleanupEvent[0] == WSA_INVALID_EVENT ) {
		myprintf("WSACreateEvent() failed: %d\n", WSAGetLastError());
		return;
	}

	if( (nRet = WSAStartup(0x202, &wsaData)) != 0 ) {
		myprintf("WSAStartup() failed: %d\n",nRet);
		SetConsoleCtrlHandler(CtrlHandler, FALSE);
		if(g_hCleanupEvent[0] != WSA_INVALID_EVENT) {
			WSACloseEvent(g_hCleanupEvent[0]);
			g_hCleanupEvent[0] = WSA_INVALID_EVENT;
		}
		return;
	}

	InitializeCriticalSection(&g_CriticalSection);

	while( g_bRestart ) {
		g_bRestart = FALSE;
		g_bEndServer = FALSE;
		WSAResetEvent(g_hCleanupEvent[0]);

		__try	{

            //
			// notice that we will create more worker threads (dwThreadCount) than 
			// the thread concurrency limit on the IOCP.
			//
			g_hIOCP = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
            if( g_hIOCP == NULL ) {
				myprintf("CreateIoCompletionPort() failed to create I/O completion port: %d\n", 
						GetLastError());
				__leave;
			}

			for( DWORD dwCPU=0; dwCPU<dwThreadCount; dwCPU++ ) {

				//
				// Create worker threads to service the overlapped I/O requests.  The decision
				// to create 2 worker threads per CPU in the system is a heuristic.  Also,
				// note that thread handles are closed right away, because we will not need them
				// and the worker threads will continue to execute.
				//
				HANDLE  hThread;
				DWORD   dwThreadId;

				hThread = CreateThread(NULL, 0, WorkerThread, g_hIOCP, 0, &dwThreadId);
				if( hThread == NULL ) {
					myprintf("CreateThread() failed to create worker thread: %d\n", 
						   GetLastError());
					__leave;
				}
				g_ThreadHandles[dwCPU] = hThread;
				hThread = INVALID_HANDLE_VALUE;
			}

			if( !CreateListenSocket() )
				__leave;

			if( !CreateAcceptSocket(TRUE) )
				__leave;

			WSAWaitForMultipleEvents(1, g_hCleanupEvent, TRUE, WSA_INFINITE, FALSE);
		}

		__finally	{

			g_bEndServer = TRUE;

			//
			// Cause worker threads to exit
			//
			if( g_hIOCP ) {
				for( DWORD i = 0; i < dwThreadCount; i++ )
					PostQueuedCompletionStatus(g_hIOCP, 0, 0, NULL);
			}

            //
			// Make sure worker threads exits.
			//
			if( WAIT_OBJECT_0 != WaitForMultipleObjects(dwThreadCount,  g_ThreadHandles, TRUE, 1000) )
				myprintf("WaitForMultipleObjects() failed: %d\n", GetLastError());
			else
				for( DWORD i=0; i<dwThreadCount; i++ ) {
					if( g_ThreadHandles[i] != INVALID_HANDLE_VALUE )
						CloseHandle(g_ThreadHandles[i]);
					g_ThreadHandles[i] = INVALID_HANDLE_VALUE;
				}

			if( g_sdListen != INVALID_SOCKET ) {
				closesocket(g_sdListen);                                
				g_sdListen = INVALID_SOCKET;
			}

			if( g_pCtxtListenSocket ) {
				while( !HasOverlappedIoCompleted((LPOVERLAPPED)&g_pCtxtListenSocket->pIOContext->Overlapped) )
					Sleep(0);

				if( g_pCtxtListenSocket->pIOContext->SocketAccept != INVALID_SOCKET )
					closesocket(g_pCtxtListenSocket->pIOContext->SocketAccept);
				g_pCtxtListenSocket->pIOContext->SocketAccept = INVALID_SOCKET;

                //
				// We know there is only one overlapped I/O on the listening socket
				//
				if( g_pCtxtListenSocket->pIOContext )
					xfree(g_pCtxtListenSocket->pIOContext);

				if( g_pCtxtListenSocket )
					xfree(g_pCtxtListenSocket);
				g_pCtxtListenSocket = NULL;
			}

			CtxtListFree();

			if( g_hIOCP ) {
				CloseHandle(g_hIOCP);
				g_hIOCP = NULL;
			}
		} //finally

		if( g_bRestart ) {
			myprintf("\niocpserverex is restarting...\n");
		} else
			myprintf("\niocpserverex is exiting...\n");

	} //while (g_bRestart)

	DeleteCriticalSection(&g_CriticalSection);
	if(g_hCleanupEvent[0] != WSA_INVALID_EVENT) {
		WSACloseEvent(g_hCleanupEvent[0]);
		g_hCleanupEvent[0] = WSA_INVALID_EVENT;
	}
	WSACleanup();
	SetConsoleCtrlHandler(CtrlHandler, FALSE);
} //main

//
//  Just validate the command line options.
//
BOOL ValidOptions(int argc, char *argv[]) {
	BOOL bRet = TRUE;

	for( int i=1; i<argc; i++ ) {
		if( (argv[i][0] =='-') || (argv[i][0] == '/') ) {
			switch( tolower(argv[i][1]) ) {
			case 'e':
				if( strlen(argv[i]) > 3 )
					g_Port = &argv[i][3];
				break;

			case 'v':
				g_bVerbose = TRUE;
				break;

			case '?':
				myprintf("Usage:\n  iocpserver [-p:port] [-v] [-?]\n");
				myprintf("  -e:port\tSpecify echoing port number\n");        
				myprintf("  -v\t\tVerbose\n");        
				myprintf("  -?\t\tDisplay this help\n");
				bRet = FALSE;
				break;

			default:
				myprintf("Unknown options flag %s\n", argv[i]);
				bRet = FALSE;
				break;
			}
		}
	}   

	return(bRet);
}

//
//  Intercept CTRL-C or CTRL-BRK events and cause the server to initiate shutdown.
//  CTRL-BRK resets the restart flag, and after cleanup the server restarts.
//
BOOL WINAPI CtrlHandler (DWORD dwEvent)	{

	switch( dwEvent ) {
	case CTRL_BREAK_EVENT: 
		g_bRestart = TRUE;
	case CTRL_C_EVENT:
	case CTRL_LOGOFF_EVENT:
	case CTRL_SHUTDOWN_EVENT:
	case CTRL_CLOSE_EVENT:
		if( g_bVerbose )
			myprintf("CtrlHandler: closing listening socket\n");

		g_bEndServer = TRUE;

		WSASetEvent(g_hCleanupEvent[0]);
		break;

	default:
		//
		// unknown type--better pass it on.
		//

		return(FALSE);
	}
	return(TRUE);
}

//
// Create a socket with all the socket options we need, namely disable buffering
// and set linger.
//
SOCKET CreateSocket(void) {
	int nRet = 0;
	int nZero = 0;
	SOCKET sdSocket = INVALID_SOCKET;

	sdSocket = WSASocket(AF_INET, SOCK_STREAM, IPPROTO_IP, NULL, 0, WSA_FLAG_OVERLAPPED); 
	if( sdSocket == INVALID_SOCKET ) {
		myprintf("WSASocket(sdSocket) failed: %d\n", WSAGetLastError());
		return(sdSocket);
	}

    //
	// Disable send buffering on the socket.  Setting SO_SNDBUF
	// to 0 causes winsock to stop buffering sends and perform
	// sends directly from our buffers, thereby save one memory copy.
	//
    // However, this does prevent the socket from ever filling the
    // send pipeline. This can lead to packets being sent that are
    // not full (i.e. the overhead of the IP and TCP headers is 
    // great compared to the amount of data being carried).
    //
    // Disabling the send buffer has less serious repercussions 
    // than disabling the receive buffer.
	//
	nZero = 0;
	nRet = setsockopt(sdSocket, SOL_SOCKET, SO_SNDBUF, (char *)&nZero, sizeof(nZero));
	if( nRet == SOCKET_ERROR) {
		myprintf("setsockopt(SNDBUF) failed: %d\n", WSAGetLastError());
		return(sdSocket);
	}

    //
    // Don't disable receive buffering. This will cause poor network
    // performance since if no receive is posted and no receive buffers,
    // the TCP stack will set the window size to zero and the peer will
    // no longer be allowed to send data.
    //

    // 
    // Do not set a linger value...especially don't set it to an abortive
    // close. If you set abortive close and there happens to be a bit of
    // data remaining to be transfered (or data that has not been 
    // acknowledged by the peer), the connection will be forcefully reset
    // and will lead to a loss of data (i.e. the peer won't get the last
    // bit of data). This is BAD. If you are worried about malicious
    // clients connecting and then not sending or receiving, the server
    // should maintain a timer on each connection. If after some point,
    // the server deems a connection is "stale" it can then set linger