IOCP Sample

/*
==========================================================================

Purpose:

This is a sample code that demonstrates for the following:

* Use of the I/O Completion ports with WinSock.  The idea is to create a 
  simple application that will use IOCP, highlight how to use IOCP, 
  and will server as baseline code for a much complicated, robust, 
  and scalable application. This code written with the intension to let the 
  reader understand IOCP.  Please ignore what server and clients are doing, 
  there operations doesn't warrant the use of IOCP.

Notes:

* The server will create IOCP, Worker threads, all incoming client sockets 
  will be associated with IOCP, the server will accept the client sent 
  message will display it and then send a message back as an acknowledgement.

Author:

* Swarajya Pendharkar

Date:

* 10th March 2006

Updates:

* Implemented IOCP with Overlapped I/O - 24th March 2006
* More updates pertaining to Overlapped I/O - 19th April 2006
                      
==========================================================================
*/

#include "stdafx.h"

#include <stdio.h>
#include <stdlib.h>
#include <conio.h>
#include <string.h>
#include <winsock2.h>

#define ACK_MESG_RECV "Message received successfully"

//Op codes for IOCP
#define OP_READ     0
#define OP_WRITE    1
#define OP_ACCEPT   2

//The norm is to create two worker threads per processor
//Many programs will find out how many processors are there on the host
//and then multiply it by two to decide on the number of worker threads
#define MAX_WORKER_THREADS 2

//Data structures required for implementation of IOCP
CRITICAL_SECTION g_csConsole; //When thread write to console we need mutual exclusion

class COverlappedPlus //Extended Overlapped class
{
public:
     OVERLAPPED        m_ol;
     
     //Get/Set calls
     void SetOpCode(int n)
     {
          m_nOpCode = n;
     }
     
     int GetOpCode()
     {
          return m_nOpCode;
     }
     
     //Constructors
     COverlappedPlus()
     {
          m_nOpCode = 0;
          ZeroMemory(&m_ol, sizeof(OVERLAPPED));
     }
     
     COverlappedPlus(int nOpCode)
     {
          m_nOpCode = nOpCode;
          ZeroMemory(&m_ol, sizeof(OVERLAPPED));
     }
     
private:
     int               m_nOpCode; //will be used by the worker thread to decide what operation to perform
};

class CClientContext  //To store and manage client related information
{
public:
     //Get/Set calls
     void SetSocket(SOCKET s)
     {
          m_Socket = s;
     }
     
     SOCKET GetSocket()
     {
          return m_Socket;
     }
     
     void SetBuffer(char *szBuffer)
     {
          strcpy(m_szBuffer, szBuffer);
     }
     
     void GetBuffer(char *szBuffer)
     {
          strcpy(szBuffer, m_szBuffer);
     }
     
     //Constructor
     CClientContext()
     {
          m_Socket =  SOCKET_ERROR;
          ZeroMemory(m_szBuffer, 256);
     }
     
private:
     SOCKET            m_Socket;  //accepted socket
     char              m_szBuffer[256]; //Used in passing messages thru IOCP
};

//Global I/O completion port handle
HANDLE g_hIOCompletionPort = NULL;

//global functions
bool InitializeIOCP();
void AcceptConnections(SOCKET ListenSocket);
DWORD WINAPI WorkerThread(LPVOID lpParam);
void WriteToConsole(char *szBuffer);

int main(void)
{
     //Initialize the Console Critical Section
     InitializeCriticalSection(&g_csConsole);
     
     // Initialize Winsock
     WSADATA wsaData;
     
     int nResult;
     nResult = WSAStartup(MAKEWORD(2,2), &wsaData);
     
     if (NO_ERROR != nResult)
     {
          printf("\nError occurred while executing WSAStartup().");
          return 1; //error
     }
     else
     {
          printf("\nWSAStartup() successful.");
     }
     
     if (false == InitializeIOCP())
     {
          printf("\nError occurred while initializing IOCP");
          goto error;
     }
     else
     {
          printf("\nIOCP initialization successful.");
     }
     
     SOCKET ListenSocket; 
     
     struct sockaddr_in ServerAddress;
     
     //Overlapped I/O follows the model established in Windows and can be performed only on 
     //sockets created through the WSASocket function 
     ListenSocket = WSASocket(AF_INET, SOCK_STREAM, 0, NULL, 0, WSA_FLAG_OVERLAPPED);
     
     if (INVALID_SOCKET == ListenSocket) 
     {
          printf("\nError occurred while opening socket: %ld.", WSAGetLastError());
          goto error;
     }
     else
     {
          printf("\nWSASocket() successful.");
     }
     
     //Cleanup and Init with 0 the ServerAddress
     ZeroMemory((char *)&ServerAddress, sizeof(ServerAddress));
     
     //Port number will be supplied as a commandline argument
     int nPortNo;
     nPortNo = 10000;
     
     //Fill up the address structure
     ServerAddress.sin_family = AF_INET;
     ServerAddress.sin_addr.s_addr = INADDR_ANY; //WinSock will supply address
     ServerAddress.sin_port = htons(nPortNo);    //comes from commandline
     
     //Assign local address and port number
     if (SOCKET_ERROR == bind(ListenSocket, (struct sockaddr *) &ServerAddress, sizeof(ServerAddress))) 
     {
          closesocket(ListenSocket);
          printf("\nError occurred while binding.");
          goto error;
     }
     else
     {
          printf("\nbind() successful.");
     }
     
     //Make the socket a listening socket
     if (SOCKET_ERROR == listen(ListenSocket,SOMAXCONN))
     {
          closesocket(ListenSocket);
          printf("\nError occurred while listening.");
          goto error;
     }
     else
     {
          printf("\nlisten() successful.");
     }
     
     //This function will take care of multiple clients using IOCP
     AcceptConnections(ListenSocket);
     
     //Close open sockets
     closesocket(ListenSocket);
     
     //Delete the Console Critical Section
     DeleteCriticalSection(&g_csConsole);
     
     //Cleanup Winsock
     WSACleanup();
     return 0; //success
     
error:
     
     //Delete the Console Critical Section
     DeleteCriticalSection(&g_csConsole);
     
     // Cleanup Winsock
     WSACleanup();
     return 1; //error
}

//Function to Initialize IOCP
bool InitializeIOCP()
{
     //Create I/O completion port
     g_hIOCompletionPort = CreateIoCompletionPort(INVALID_HANDLE_value, NULL, 0, 0 );
     
     if ( NULL == g_hIOCompletionPort)
     {
          printf("\nError occurred while creating IOCP: %ld.", WSAGetLastError());
          return false;
     }
     
     //Create worker threads
     for (int ii = 0; ii < MAX_WORKER_THREADS; ii++)
     {
          DWORD nThreadID;
          CreateThread(0, 0, WorkerThread, NULL, 0, &nThreadID);
     }
     
     return true;
}

//This function will loop on while and will associate incoming sockets to IOCP
void AcceptConnections(SOCKET ListenSocket)
{
     sockaddr_in ClientAddress;
     int nClientLength = sizeof(ClientAddress);
     char szConsole[256];
     
     //Infinite, no graceful shutdown of server implemented, 
     //preferably server should be implemented as a service
     //Events can also be used for graceful shutdown
     while (1) 
     {
          //Accept remote connection attempt from the client
          SOCKET Socket = accept(ListenSocket, (sockaddr*)&ClientAddress, &nClientLength);
          
          if (INVALID_SOCKET == Socket)
          {
               sprintf(szConsole, "Error occurred while accepting socket: %ld.", WSAGetLastError());
               WriteToConsole(szConsole);
          }
          
          //Display Client's IP
          sprintf(szConsole, "Client connected from: %s", inet_ntoa(ClientAddress.sin_addr)); 
          WriteToConsole(szConsole);
          
          //Disable Nagling
          int nFlag = 1;
          int nResult = setsockopt(Socket, IPPROTO_TCP, TCP_NODELAY, (char *)&nFlag, sizeof(char));
          if (-1 == nResult)
          {
               sprintf(szConsole, "Error occurred while executing setsockopt()");
               WriteToConsole(szConsole);
               continue; //Keep at it
          }
          
          //Create a new OverlappedPlus and ClientContext for this newly accepted client
          COverlappedPlus  *pOverlappedPlus = new COverlappedPlus;
          CClientContext   *pClientContext  = new CClientContext;
          
          pOverlappedPlus->SetOpCode(OP_ACCEPT);
          pClientContext->SetSocket(Socket);
          
          //Ask worker thread to associate the incoming socket connection to IOCP
          PostQueuedCompletionStatus(g_hIOCompletionPort, 0, (DWORD)pClientContext, &pOverlappedPlus->m_ol);
     }
}

//Worker thread will service IOCP requests
DWORD WINAPI WorkerThread(LPVOID lpParam)
{    
     DWORD            *lpContext = NULL;
     OVERLAPPED       *pOverlapped = NULL;
     COverlappedPlus  *pOverlappedPlus = NULL; 
     CClientContext   *pClientContext = NULL;
     DWORD            dwBytesXfered = 0;
     char szBuffer[256];
     char szConsole[256];
     int nBytesSent = 0;
     int nBytesRecv = 0;
     
     WSABUF            wbuf;
     DWORD             dwBytes = 0, dwFlags = 0;
     
     HANDLE hTemp = NULL;
     
     DWORD dwThreadId = GetCurrentThreadId();
     
     //Infinite, worker thread will be around to process requests
     while (1)
     {
          BOOL bReturn = GetQueuedCompletionStatus(
               g_hIOCompletionPort,
               &dwBytesXfered,
               (LPDWORD)&lpContext,
               &pOverlapped,
               INFINITE);
          
          if ( (FALSE == bReturn) || (NULL == pOverlapped) || (NULL == lpContext) )
          {
               //Operation failed
               continue;
          }
          
          pOverlappedPlus = CONTAINING_RECORD(pOverlapped, COverlappedPlus, m_ol); //Get OverlappedPlus
          pClientContext = (CClientContext *)lpContext;  //Get the client context
          
          switch (pOverlappedPlus->GetOpCode())
          {
          case OP_ACCEPT:
               
               //Associate the socket with IOCP
               hTemp = CreateIoCompletionPort((HANDLE)pClientContext->GetSocket(), g_hIOCompletionPort, NULL, 0);
               
               if (NULL == hTemp)
               {
                    sprintf(szConsole, "Thread Id: %d, Error occurred while executing CreateIoCompletionPort().", dwThreadId);
                    WriteToConsole(szConsole);
                    
                    //Let's not work with this client
                    delete pOverlappedPlus;
                    delete pClientContext;
                    
                    pOverlappedPlus = NULL;
                    pClientContext = NULL;
                    
                    break; //error
               }
               
               pOverlappedPlus->SetOpCode(OP_READ);
               
               //Client will send data ask worker thread to receive and display
               PostQueuedCompletionStatus(g_hIOCompletionPort, 0, (DWORD)pClientContext, &pOverlappedPlus->m_ol);
               
               break;
               
          case OP_READ:
               
               //Clear szBuffer, init with 0
               ZeroMemory(szBuffer, sizeof(szBuffer));
               
               wbuf.buf = szBuffer;
               wbuf.len = 256;
               
               dwFlags = MSG_PARTIAL;
               
               //Overlapped recv 
               nBytesRecv = WSARecv(pClientContext->GetSocket(), &wbuf, 1, 
                    &dwBytes, &dwFlags, &pOverlappedPlus->m_ol, NULL);
               
               if ((SOCKET_ERROR == nBytesRecv) && (WSA_IO_PENDING != WSAGetLastError()))
               {
                    sprintf(szConsole, "Thread Id: %d, Error occurred while executing WSARecv().", dwThreadId);
                    WriteToConsole(szConsole);
                    
                    //Let's not work with this client
                    delete pOverlappedPlus;
                    delete pClientContext;
                    
                    pOverlappedPlus = NULL;
                    pClientContext = NULL;
                    
                    break; //error
               }
               
               //Wait on this Overlapped I/O till we have successfully received the data
               //We want to display on console, what we receive
               while (!HasOverlappedIoCompleted((LPOVERLAPPED)&pOverlappedPlus->m_ol))
               {
                    Sleep(0); //Switch to some other thread
               }
               
               //Display the message received on console
               sprintf(szConsole, "Thread Id: %d, The following message was received: %s", dwThreadId, szBuffer);
               WriteToConsole(szConsole);
               
               pClientContext->SetBuffer(ACK_MESG_RECV);
               pOverlappedPlus->SetOpCode(OP_WRITE);
               
               //Ask worker thread to send a message back to the client acknowledging that his message was received
               PostQueuedCompletionStatus(g_hIOCompletionPort, 0, (DWORD)pClientContext, &pOverlappedPlus->m_ol);
               
               break;
               
          case OP_WRITE:
               
               //Clear szBuffer, init with 0
               ZeroMemory(szBuffer, sizeof(szBuffer));
               
               //Get the buffer sent for writing
               pClientContext->GetBuffer(szBuffer);
               
               wbuf.buf = szBuffer;
               wbuf.len = strlen(szBuffer);
               
               dwFlags = MSG_PARTIAL;
               
               //Overlapped send
               nBytesSent = WSASend(pClientContext->GetSocket(), &wbuf, 1, 
                    &dwBytes, dwFlags, &pOverlappedPlus->m_ol, NULL);
               
               if ((SOCKET_ERROR == nBytesSent) && (WSA_IO_PENDING != WSAGetLastError()))
               {
                    sprintf(szConsole, "Thread Id: %d, Error occurred while executing WSASend().", dwThreadId);
                    WriteToConsole(szConsole);
                    
                    //Let's not work with this client
                    delete pOverlappedPlus;
                    delete pClientContext;
                    
                    pOverlappedPlus = NULL;
                    pClientContext = NULL;
                    
                    break; //error
                    
               }
               
               //If you want to wait for completion of WSASend use HasOverlappedIoCompleted(), 
               //other options are - GetOverlappedResult(), WSAGetOverlappedResult()
               //One can also wait on the event handle in the overlapped structure.
               
               //Here I am moving on, not bothered with whether the overlapped send was successful.
               
               //Display the message received on console
               sprintf(szConsole, "Thread Id: %d, The message was sent successfully.", dwThreadId);
               WriteToConsole(szConsole);
               
               //We are done with this client, do a cleanup.
               closesocket(pClientContext->GetSocket());  //We are done with this client
               
               delete pOverlappedPlus;
               delete pClientContext;
               
               pOverlappedPlus = NULL;
               pClientContext = NULL;
               
               break;
               
          default:
               //We should never be reaching here, under normal circumstances.
               break;
          } // switch
     } // while
}

//Function to synchronize console output
//Threads need to be synchronized while they write to console.
//WriteConsole() API can be used, it is thread-safe, I think.
void WriteToConsole(char *szBuffer)
{
     EnterCriticalSection(&g_csConsole);
     printf("\n%s", szBuffer);
     LeaveCriticalSection(&g_csConsole);
}

你可能感兴趣的:(IOCP Sample)