//
// Function: AssignToFreeThread
//
// Description:
// This routine assigns a socket connection to an available child
// thread to handle any IO on it. If no threads are available, a
// new thread is spawned to handle the connection.
//
void AssignToFreeThread(SOCKET_OBJ *sock)
{
THREAD_OBJ *thread=NULL;
thread = gChildThreads;
while (thread)
{
// If this routine returns something other than SOCKET_ERROR
// that it was successfully assigned to a child thread.
if (InsertSocketObj(thread, sock) != SOCKET_ERROR)
break;
thread = thread->next;
}
if (thread == NULL)
{
// No thread was found to assign the client socket to, create a new thread
//
printf("Creating new thread object\n");
thread = GetThreadObj();
thread->Thread = CreateThread(NULL, 0, ChildThread, (LPVOID)thread, 0, NULL);
if (thread->Thread == NULL)
{
fprintf(stderr, "AssignToFreeThread: CreateThread failed: %d\n", GetLastError());
ExitProcess(-1);
}
InsertSocketObj(thread, sock);
// Insert the thread the list of threads
if (gChildThreads == NULL)
{
gChildThreads = thread;
}
else
{
thread->next = gChildThreads;
gChildThreads = thread;
}
gChildThreadsCount++;
}
// signal child thread to rebuild the event list
WSASetEvent(thread->Event);
return;
}
void MightyTCPCompletionPortServer::HandleIO(DWORD dwKey, PCompletionPort_BufObj pBuffer, DWORD dwTrans, int nError)
{
PCompletionPortIOObject pContext = (PCompletionPortIOObject)dwKey;
//首先减少套节字上的未决I/O计数
if(pContext != NULL)
{
if(pBuffer->nOperation == OP_READ)
pContext->nOutstandingRecv --;
else if(pBuffer->nOperation == OP_WRITE)
pContext->nOutstandingSend --;
// 检查套节字是否已经被我们关闭
if(pContext->nCondition)
{
if(pContext->nOutstandingRecv == 0 && pContext->nOutstandingSend == 0)
{
RemoveSocketObj(pContext);
}
// 释放已关闭套节字的未决I/O
ReleaseBuffer(pBuffer);
return;
}
}
else
{
RemovePendingAccept(pBuffer);
}
//检查套节字上发生的错误，如果有的话，通知用户，然后关闭套节字
if(nError != NO_ERROR)
{
if(pBuffer->nOperation != OP_ACCEPT)
{
//if(nError==ERROR_NETNAME_DELETED)
RemoveSocketObj(pContext);
if(pContext->nOutstandingRecv == 0 && pContext->nOutstandingSend == 0)
{
//ReleaseContext(pContext);
}
}
else // 在监听套节字上发生错误，也就是监听套节字处理的客户出错了
{
// 客户端出错，释放I/O缓冲区
if(pBuffer->sAccept!= INVALID_SOCKET)
{
::closesocket(pBuffer->sAccept);
pBuffer->sAccept= INVALID_SOCKET;
}
}
ReleaseBuffer(pBuffer);
return;
}
// 开始处理
if(pBuffer->nOperation == OP_ACCEPT)
{
if(dwTrans == 0)
{
if(pBuffer->sAccept!= INVALID_SOCKET)
{
::closesocket(pBuffer->sAccept);
pBuffer->sAccept= INVALID_SOCKET;
}
}
else
{
// 为新接受的连接申请客户上下文对象
PCompletionPortIOObject pClient =GetSocketObj(pBuffer->sAccept);
if(pClient != NULL)
{
if(InsertSocketObj(pClient)==0)
{
// 取得客户地址
int nLocalLen, nRmoteLen;
LPSOCKADDR pLocalAddr, pRemoteAddr;
m_lpfnGetAcceptExSockaddrs(
pBuffer->buff,
pBuffer->nLen - ((sizeof(sockaddr_in) + 16) * 2),
sizeof(sockaddr_in) + 16,
sizeof(sockaddr_in) + 16,
(SOCKADDR **)&pLocalAddr,
&nLocalLen,
(SOCKADDR **)&pRemoteAddr,
&nRmoteLen);
memcpy(&pClient->addrLocal, pLocalAddr, nLocalLen);
memcpy(&pClient->addrRemote, pRemoteAddr, nRmoteLen);
// 关联新连接到完成端口对象
::CreateIoCompletionPort((HANDLE)pClient->s, m_hCompletionPort, (DWORD)pClient, 0);
// 通知用户
pBuffer->nLen = dwTrans;
m_pNotifyCallBack->OnRecvData(pBuffer->buff,dwTrans,pClient->identityIDKey);
// 向新连接投递几个Read请求，这些空间在套节字关闭或出错时释放
for(int i=0; i<5; i++)
{
PCompletionPort_BufObj p = AllocateBuffer(pClient,2048);
if(p != NULL)
{
if(!PostRecv(p))
{
RemoveSocketObj(pClient);
break;
}
}
}
}
else // 连接数量已满，关闭连接
{
RemoveSocketObj(pClient);
}
}
else
{
// 资源不足，关闭与客户的连接即可
::closesocket(pBuffer->sAccept);
pBuffer->sAccept= INVALID_SOCKET;
}
}
// Accept请求完成，释放I/O缓冲区
ReleaseBuffer(pBuffer);
// 通知监听线程继续再投递一个Accept请求
::InterlockedIncrement(&m_nRepostCount);
::SetEvent(m_hRepostEvent);
}
else if(pBuffer->nOperation == OP_READ)
{
if(dwTrans == 0) // 对方关闭套节字
{
// 先通知用户
pBuffer->nLen = 0;
RemoveSocketObj(pContext);
// 释放客户上下文和缓冲区对象
if(pContext->nOutstandingRecv == 0 && pContext->nOutstandingSend == 0)
{
//ReleaseContext(pContext);
}
//ReleaseBuffer(pBuffer);
}
else
{
pBuffer->nLen = dwTrans;
// 按照I/O投递的顺序读取接收到的数据
PCompletionPort_BufObj pNextIO = GetNextReadBuffer(pContext, pBuffer);
while(pNextIO != NULL)
{
m_pNotifyCallBack->OnRecvData(pNextIO->buff,dwTrans,pContext->identityIDKey);
// 增加要读的序列号的值
::InterlockedIncrement((LONG*)&pContext->nCurrentReadSequence);
// 释放这个已完成的I/O
ReleaseBuffer(pNextIO);
pNextIO= GetNextReadBuffer(pContext, NULL);
}
// 继续投递一个新的接收请求
pBuffer = AllocateBuffer(pContext,4096);
if(pBuffer == NULL || !PostRecv(pBuffer))
{
//closeAConnection(pContext);
}
}
}
else if(pBuffer->nOperation == OP_WRITE)
{
if(dwTrans == 0) // 对方关闭套节字
{
// 先通知用户
pBuffer->nLen = 0;
RemoveSocketObj(pContext);
// 再关闭连接
//CloseAConnection(pContext);
// 释放客户上下文和缓冲区对象
if(pContext->nOutstandingRecv == 0 && pContext->nOutstandingSend == 0)
{
//ReleaseContext(pContext);
}
//ReleaseBuffer(pBuffer);
}
else
{
// 写操作完成，通知用户
pBuffer->nLen = dwTrans;
//OnWriteCompleted(pContext, pBuffer);
// 释放SendText函数申请的缓冲区
ReleaseBuffer(pBuffer);
}
}
}
//
// Function: main
//
// Description:
// This is the main program. It parses the command line and creates
// the main socket. For UDP this socket is used to receive datagrams.
// For TCP the socket is used to accept incoming client connections.
// Each client TCP connection is handed off to a worker thread which
// will receive any data on that connection until the connection is
// closed.
//
int __cdecl main(int argc, char **argv)
{
WSADATA wsd;
THREAD_OBJ *thread=NULL;
SOCKET_OBJ *sockobj=NULL,
*newsock=NULL;
int index,
rc;
struct addrinfo *res=NULL,
*ptr=NULL;
// Validate the command line
ValidateArgs(argc, argv);
// Load Winsock
if (WSAStartup(MAKEWORD(2,2), &wsd) != 0)
{
fprintf(stderr, "unable to load Winsock!\n");
return -1;
}
printf("Local address: %s; Port: %s; Family: %d\n",
gBindAddr, gBindPort, gAddressFamily);
res = ResolveAddress(gBindAddr, gBindPort, gAddressFamily, gSocketType, gProtocol);
if (res == NULL)
{
fprintf(stderr, "ResolveAddress failed to return any addresses!\n");
return -1;
}
thread = GetThreadObj();
// For each local address returned, create a listening/receiving socket
ptr = res;
while (ptr)
{
PrintAddress(ptr->ai_addr, ptr->ai_addrlen); printf("\n");
sockobj = GetSocketObj(INVALID_SOCKET, (gProtocol == IPPROTO_TCP) ? TRUE : FALSE);
// create the socket
sockobj->s = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
if (sockobj->s == INVALID_SOCKET)
{
fprintf(stderr,"socket failed: %d\n", WSAGetLastError());
return -1;
}
InsertSocketObj(thread, sockobj);
// bind the socket to a local address and port
rc = bind(sockobj->s, ptr->ai_addr, ptr->ai_addrlen);
if (rc == SOCKET_ERROR)
{
fprintf(stderr, "bind failed: %d\n", WSAGetLastError());
return -1;
}
if (gProtocol == IPPROTO_TCP)
{
rc = listen(sockobj->s, 200);
if (rc == SOCKET_ERROR)
{
fprintf(stderr, "listen failed: %d\n", WSAGetLastError());
return -1;
}
// Register events on the socket
rc = WSAEventSelect(
sockobj->s,
sockobj->event,
FD_ACCEPT | FD_CLOSE
);
if (rc == SOCKET_ERROR)
{
fprintf(stderr, "WSAEventSelect failed: %d\n", WSAGetLastError());
return -1;
}
}
else
{
// Register events on the socket
rc = WSAEventSelect(
sockobj->s,
sockobj->event,
FD_READ | FD_WRITE | FD_CLOSE
);
if (rc == SOCKET_ERROR)
{
fprintf(stderr, "WSAEventSelect failed: %d\n", WSAGetLastError());
return -1;
}
}
ptr = ptr->ai_next;
}
// free the addrinfo structure for the 'bind' address
freeaddrinfo(res);
gStartTime = gStartTimeLast = GetTickCount();
while (1)
{
rc = WaitForMultipleObjects(
thread->SocketCount + 1,
thread->Handles,
FALSE,
5000
);
if (rc == WAIT_FAILED)
{
fprintf(stderr, "WaitForMultipleObjects failed: %d\n", GetLastError());
break;
}
else if (rc == WAIT_TIMEOUT)
{
PrintStatistics();
}
else
{
index = rc - WAIT_OBJECT_0;
sockobj = FindSocketObj(thread, index-1);
if (gProtocol == IPPROTO_TCP)
{
SOCKADDR_STORAGE sa;
WSANETWORKEVENTS ne;
SOCKET sc;
int salen;
rc = WSAEnumNetworkEvents(
sockobj->s,
thread->Handles[index],
&ne
);
if (rc == SOCKET_ERROR)
{
fprintf(stderr, "WSAEnumNetworkEvents failed: %d\n", WSAGetLastError());
break;
}
while (1)
{
sc = INVALID_SOCKET;
salen = sizeof(sa);
//
// For TCP, accept the connection and hand off the client socket
// to a worker thread
//
sc = accept(
sockobj->s,
(SOCKADDR *)&sa,
&salen
);
if ((sc == INVALID_SOCKET) && (WSAGetLastError() != WSAEWOULDBLOCK))
{
fprintf(stderr, "accept failed: %d\n", WSAGetLastError());
break;
}
else if (sc != INVALID_SOCKET)
{
newsock = GetSocketObj(INVALID_SOCKET, FALSE);
// Copy address information
memcpy(&newsock->addr, &sa, salen);
newsock->addrlen = salen;
newsock->s = sc;
InterlockedIncrement(&gTotalConnections);
InterlockedIncrement(&gCurrentConnections);
/*
printf("Accepted connection from: ");
PrintAddress((SOCKADDR *)&newsock->addr, newsock->addrlen);
printf("\n");
*/
// Register for read, write and close on the client socket
rc = WSAEventSelect(
newsock->s,
newsock->event,
FD_READ | FD_WRITE | FD_CLOSE
);
if (rc == SOCKET_ERROR)
{
fprintf(stderr, "WSAEventSelect failed: %d\n", WSAGetLastError());
break;
}
AssignToFreeThread(newsock);
}
else
{
// Failed with WSAEWOULDBLOCK -- just continue
break;
}
}
}
else
{
// For UDP all we have to do is handle events on the main
// threads.
if (HandleIo(thread, sockobj) == SOCKET_ERROR)
{
RenumberThreadArray(thread);
}
}
}
}
WSACleanup();
return 0;
}
//
// Function: main
//
// Description:
// This is the main program. It parses the command line and creates
// the main socket. For UDP this socket is used to receive datagrams.
// For TCP the socket is used to accept incoming client connections.
// Each client TCP connection is handed off to a worker thread which
// will receive any data on that connection until the connection is
// closed.
//
int __cdecl main(int argc, char **argv)
{
WSADATA wsd;
SOCKET_OBJ *sockobj=NULL,
*sptr=NULL,
*tmp=NULL;
HANDLE hThread;
ULONG lastprint=0;
MSG msg;
int rc;
struct addrinfo *res=NULL,
*ptr=NULL;
ValidateArgs(argc, argv);
InitializeCriticalSection(&gSocketCritSec);
// Load winsock
if (WSAStartup(MAKEWORD(2,2), &wsd) != 0)
{
fprintf(stderr, "unable to load Winsock!\n");
return -1;
}
gWorkerWindow = MakeWorkerWindow();
printf("Local address: %s; Port: %s; Family: %d\n",
gBindAddr, gBindPort, gAddressFamily);
res = ResolveAddress(gBindAddr, gBindPort, gAddressFamily, gSocketType, gProtocol);
if (res == NULL)
{
fprintf(stderr, "ResolveAddress failed to return any addresses!\n");
return -1;
}
// For each local address returned, create a listening/receiving socket
ptr = res;
while (ptr)
{
PrintAddress(ptr->ai_addr, ptr->ai_addrlen); printf("\n");
sockobj = GetSocketObj(INVALID_SOCKET);
// create the socket
sockobj->s = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
if (sockobj->s == INVALID_SOCKET)
{
fprintf(stderr,"socket failed: %d\n", WSAGetLastError());
return -1;
}
InsertSocketObj(sockobj);
// bind the socket to a local address and port
rc = bind(sockobj->s, ptr->ai_addr, ptr->ai_addrlen);
if (rc == SOCKET_ERROR)
{
fprintf(stderr, "bind failed: %d\n", WSAGetLastError());
return -1;
}
if (gProtocol == IPPROTO_TCP)
{
rc = listen(sockobj->s, 200);
if (rc == SOCKET_ERROR)
{
fprintf(stderr, "listen failed: %d\n", WSAGetLastError());
return -1;
}
// Register for notification
rc = WSAAsyncSelect(
sockobj->s,
gWorkerWindow,
WM_SOCKET,
FD_ACCEPT | FD_CLOSE
);
if (rc == SOCKET_ERROR)
{
fprintf(stderr, "WSAAsyncSelect failed: %d\n", WSAGetLastError());
return -1;
}
}
else
{
// Register for notification
rc = WSAAsyncSelect(
sockobj->s,
gWorkerWindow,
WM_SOCKET,
FD_READ | FD_WRITE | FD_CLOSE
);
if (rc == SOCKET_ERROR)
{
fprintf(stderr, "WSAAsyncSelect failed: %d\n", WSAGetLastError());
return -1;
}
}
ptr = ptr->ai_next;
}
// free the addrinfo structure for the 'bind' address
freeaddrinfo(res);
gStartTime = gStartTimeLast = lastprint = GetTickCount();
// Start a thread to print statistics
hThread = CreateThread(NULL, 0, StatisticsThread, NULL, 0, NULL);
if (hThread == NULL)
{
fprintf(stderr, "CreateThread failed: %d\n", WSAGetLastError());
return -1;
}
CloseHandle(hThread);
while(rc = GetMessage(&msg, NULL, 0, 0))
{
if (rc == -1)
{
fprintf(stderr, "GetMessage() failed with error %d\n", GetLastError());
return -1;
}
TranslateMessage(&msg);
DispatchMessage(&msg);
}
WSACleanup();
DeleteCriticalSection(&gSocketCritSec);
return 0;
}
//
// Function: WindowProc
//
// Description:
// This is the window procedure which handles the window messages for
// our hidden window. It handles all the WM_SOCKET messages and performs
// the correct actions for each message type (FD_READ, FD_WRITE, etc.).
//
LRESULT CALLBACK WindowProc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
SOCKET_OBJ *sockobj=NULL,
*newsock=NULL;
int rc;
if (uMsg == WM_SOCKET)
{
// Check for an error on the socket
if (WSAGETSELECTERROR(lParam))
{
// An error occured on the socket, close it down
fprintf(stderr, "Socket failed with error %d\n", WSAGETSELECTERROR(lParam));
closesocket(wParam);
RemoveSocketObjByHandle(wParam);
}
else
{
// Find the socket object for this event
sockobj = FindSocketObj(wParam);
if (sockobj == NULL)
return 0;
switch (WSAGETSELECTEVENT(lParam))
{
case FD_ACCEPT:
// Get a new object for the client socket
newsock = GetSocketObj(INVALID_SOCKET);
newsock->s = accept(
wParam,
(SOCKADDR *)&newsock->addr,
&newsock->addrlen
);
if (newsock->s == INVALID_SOCKET)
{
fprintf(stderr, "accept failed: %d\n", WSAGetLastError());
break;
}
InterlockedIncrement(&gCurrentConnections);
// Create a socket information structure to associate with the
// socket for processing I/O.
InsertSocketObj(newsock);
/*
printf("Accepted connection from: ");
PrintAddress((SOCKADDR *)&newsock->addr, newsock->addrlen);
printf("\n");
*/
rc = WSAAsyncSelect(
newsock->s,
hwnd,
WM_SOCKET,
FD_READ | FD_WRITE | FD_CLOSE
);
if (rc == SOCKET_ERROR)
{
fprintf(stderr, "WSAAsyncSelect failed: %d\n", WSAGetLastError());
return -1;
}
break;
case FD_READ:
rc = ReceivePendingData(sockobj);
if (rc == -1)
{
RemoveSocketObj(sockobj);
break;
}
else if (rc != WSAEWOULDBLOCK)
{
PostMessage(hwnd, WM_SOCKET, wParam, FD_READ);
}
//
// Don't break fall through and attempt to send data
//
case FD_WRITE:
//
// Send routine automatically tries to send all queued buffers.
//
rc = SendPendingData(sockobj);
if (rc == -1)
{
RemoveSocketObj(sockobj);
}
break;
case FD_CLOSE:
sockobj->closing = TRUE;
//
// Post an FD_READ message to force another receive
// This is to ensure we recv() until 0 is returned.
//
PostMessage(hwnd, WM_SOCKET, wParam, FD_READ);
break;
default:
printf("Unknown message received: %d\n", WSAGETSELECTEVENT(lParam));
break;
}
}
return 0;
}
return DefWindowProc(hwnd, uMsg, wParam, lParam);
}
