Windows Socket I/O模型之 重叠I/O事件通知模式

#include <winsock2.h>
#include <ws2tcpip.h>
#include <mswsock.h>
#include <windows.h>

#include "resolve.h"
#include "public.h"

typedef SINGLE_LIST_HEADER SocketObjHeader;
typedef SINGLE_LIST SocketObj;
typedef SINGLE_LIST_HEADER ThreadObjHeader;
typedef SINGLE_LIST ThreadObj;
typedef DOUBLE_LIST_HEADER BufferObjHeader;
typedef DOUBLE_LIST BufferObj;

#define DEFAULT_OVERLAPPED_COUNT 5

struct _SOCKET_OBJ;
struct _THREAD_OBJ;

int gOverlappedCount = DEFAULT_OVERLAPPED_COUNT;

typedef struct _BUFFER_OBJ_EX {
    WSAOVERLAPPED        ol;            // Overlapped structure

    SOCKET               sclient;       // Used for AcceptEx client socket

    char                *buf;           // Buffer for send/recv/AcceptEx
    int                  buflen;        // Length of the buffer

    int                  operation;     // Type of operation submitted
#define OP_ACCEPT       0                   // AcceptEx
#define OP_READ         1                   // WSARecv/WSARecvFrom
#define OP_WRITE        2                   // WSASend?WSASendTo

    struct _SOCKET_OBJ *Socket;         // SOCKET_OBJ that this I/O belongs to
    struct _THREAD_OBJ *Thread;         // THREAD_OBJ this I/O is assigned to

    SOCKADDR_STORAGE     addr;          // Remote address (UDP)
    int                  addrlen;       // Remote address length

    BufferObj           bufDList;
}BUFFER_OBJ_EX;

typedef struct _SOCKET_OBJ
{
    SOCKET               s;             // Socket handle for client connection

    int                  af,            // Address family of socket (AF_INET or AF_INET6)
                         bClosing;      // Indicates socket is closing

    volatile LONG        OutstandingOps;    // Number of outstanding overlapped ops

    BUFFER_OBJ_EX        **PendingAccepts;    // Array of pending AcceptEx calls (listening socket only)

    // Pointers to Microsoft specific extensions (listening socket only)
    LPFN_ACCEPTEX        lpfnAcceptEx;
    LPFN_GETACCEPTEXSOCKADDRS lpfnGetAcceptExSockaddrs;

    CRITICAL_SECTION     SockCritSec;   // Synchronize access to this SOCKET_OBJ

    SocketObj            next;          // Used to chain SOCKET_OBJ together
} SOCKET_OBJ;

typedef struct _THREAD_OBJ
{
    DOUBLE_LIST_HEADER BufferList;            // Linked list of all sockets allocated

    int         EventCount;            // How many events are in the array to wait on?

    HANDLE      Event;                 // Used to signal new clients assigned
                                       //  to this thread
    HANDLE      Thread;                // Handle to the curren thread

    HANDLE      Handles[MAXIMUM_WAIT_OBJECTS]; // Array of socket's event handles

    CRITICAL_SECTION ThreadCritSec;    // Protect access to SOCKET_OBJ lists

    ThreadObj   next;          // Next thread object in list
} THREAD_OBJ;

CRITICAL_SECTION gThreadListCritSec;
ThreadObjHeader threadobjHeader;

void AssignIoToThread(BUFFER_OBJ_EX *buf);

SOCKET_OBJ *GetSocketObj(SOCKET s, int af) {
    SOCKET_OBJ *sockobj = NULL;

    sockobj = (SOCKET_OBJ *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(SOCKET_OBJ));
    if (sockobj == NULL)
    {
        fprintf(stderr, "GetSocketObj: HeapAlloc failed: %d\n", GetLastError());
        ExitProcess(-1);
    }

    sockobj->s = s;
    sockobj->af = af;

    InitializeCriticalSection(&sockobj->SockCritSec);

    return sockobj;
}

void FreeSocketObj(SOCKET_OBJ *obj)
{
    BUFFER_OBJ  *ptr=NULL,
                *tmp=NULL;


    if (obj->OutstandingOps != 0)
    {
        return;
    }
    if (obj->s != INVALID_SOCKET)
    {
        closesocket(obj->s);
        obj->s = INVALID_SOCKET;
    }

    DeleteCriticalSection(&obj->SockCritSec);

    HeapFree(GetProcessHeap(), 0, obj);
}

BUFFER_OBJ_EX *GetBufferObjEx(SOCKET_OBJ *sock, int buflen) {
    BUFFER_OBJ_EX *newobj = NULL;

    newobj = (BUFFER_OBJ_EX *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(BUFFER_OBJ_EX));
    if (newobj == NULL)
    {
        fprintf(stderr, "GetBufferObj: HeapAlloc failed: %d\n", GetLastError());
        ExitProcess(-1);
    }

    newobj->buf = (char *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(BYTE) *buflen);
    if (newobj->buf == NULL)
    {
        fprintf(stderr, "GetBufferObj: HeapAlloc failed: %d\n", GetLastError());
        ExitProcess(-1);
    }
    newobj->buflen = buflen;

    newobj->addrlen = sizeof(newobj->addr);
    newobj->Socket = sock;

    newobj->ol.hEvent = WSACreateEvent();
    if (newobj->ol.hEvent == NULL) {
        fprintf(stderr, "WSACreateEvent failed.\n");
        ExitProcess(-1);
    }

    return newobj;
}

BUFFER_OBJ_EX *FindBufferObjEx(THREAD_OBJ *thread, WSAEVENT hEvent) {
    BufferObj *ptr=NULL;
    BUFFER_OBJ_EX *bufobj = NULL;

    EnterCriticalSection(&thread->ThreadCritSec);

    ptr = (BufferObj *)GotoNextDoubleList(&(thread->BufferList), &(thread->BufferList.head));
    while (ptr) {
        bufobj = (BUFFER_OBJ_EX *)container_of(BUFFER_OBJ_EX, bufDList, ptr);

        if (bufobj->ol.hEvent == hEvent)
            break;

        ptr = (BufferObj *)GotoNextDoubleList(&(thread->BufferList), ptr);
    }

    LeaveCriticalSection(&thread->ThreadCritSec);

    return bufobj;
}

void FreeBufferObjEx(BUFFER_OBJ_EX *obj)
{
    // Close the event
    WSACloseEvent(obj->ol.hEvent);
    obj->ol.hEvent = NULL;
    // Free the buffers
    HeapFree(GetProcessHeap(), 0, obj->buf);
    HeapFree(GetProcessHeap(), 0, obj);
}

int InsertBufferObjExToThread(THREAD_OBJ *thread, BUFFER_OBJ_EX *buf) {
    int ret;

    EnterCriticalSection(&thread->ThreadCritSec);

    if (thread->EventCount < MAXIMUM_WAIT_OBJECTS - 1) {
        EnqueueDoubleList(&(thread->BufferList), &(buf->bufDList));
        thread->Handles[thread->EventCount++] = buf->ol.hEvent;

        ret = NO_ERROR;
    } else {
        ret = SOCKET_ERROR;
    }

    LeaveCriticalSection(&thread->ThreadCritSec);

    return ret;
}

void RemoveBufferFromThread(SOCKET_OBJ *sock, BUFFER_OBJ_EX *buf)
{
    EnterCriticalSection(&buf->Thread->ThreadCritSec);

    // Remove buffer from the list
    RemoveDoubleList(&buf->Thread->BufferList, &buf->bufDList);
    // Decrement the event count for the thread
    buf->Thread->EventCount--;
    // Set the thread's event
    WSASetEvent(buf->Thread->Event);

    LeaveCriticalSection(&buf->Thread->ThreadCritSec);

}

THREAD_OBJ *GetThreadObj() {
    THREAD_OBJ *thread = NULL;

    thread = (THREAD_OBJ *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(THREAD_OBJ));
    if (thread == NULL) {
        fprintf(stderr, "GetThreadObj:HeapAlloc failed.\n");
        ExitProcess(-1);
    }

    thread->Event = WSACreateEvent();
    if (thread->Event == NULL) {
        fprintf(stderr, "GetThreadObj: WSACreateEvent failed.\n");
        ExitProcess(-1);
    }

    thread->Handles[0] = thread->Event;
    thread->EventCount = 1;

    InitializeDoubleHead(&thread->BufferList);

    InitializeCriticalSection(&thread->ThreadCritSec);

    return thread;
}

void RenumberEvents(THREAD_OBJ *thread) {
    EnterCriticalSection(&thread->ThreadCritSec);

    BUFFER_OBJ_EX *obj = NULL;
    int         i = 0;

    BufferObj *sptr = NULL;
    sptr = (BufferObj *)GotoNextDoubleList(&(thread->BufferList), &(thread->BufferList.head));
    thread->EventCount = 1;
    while (sptr) {
        obj = (BUFFER_OBJ_EX *)container_of(BUFFER_OBJ_EX, bufDList, sptr);

        thread->Handles[thread->EventCount++] = obj->ol.hEvent;

        sptr = (BufferObj *)GotoNextDoubleList(&(thread->BufferList), sptr);
    }

    LeaveCriticalSection(&thread->ThreadCritSec);
}

int PostRecv(BUFFER_OBJ_EX *recvobj) {
    WSABUF  wbuf;
    DWORD   bytes,
            flags;
    int     rc=NO_ERROR;

    EnterCriticalSection(&recvobj->Socket->SockCritSec);

    recvobj->operation = OP_READ;

    wbuf.buf = recvobj->buf;
    wbuf.len = recvobj->buflen;

    flags = 0;

    if (gProtocol == IPPROTO_TCP) {
        rc = WSARecv(recvobj->Socket->s, &wbuf, 1, &bytes, &flags, &recvobj->ol, NULL);
    }
    else {
        ExitProcess(-1);
    }
    if (rc == SOCKET_ERROR)
    {
        rc = NO_ERROR;
        if (WSAGetLastError() != WSA_IO_PENDING)
        {
            fprintf(stderr, "PostRecv: WSARecv* failed: %d\n", WSAGetLastError());
            rc = SOCKET_ERROR;
        }
    }

    // Increment outstanding overlapped operations
    InterlockedIncrement(&recvobj->Socket->OutstandingOps);

    LeaveCriticalSection(&recvobj->Socket->SockCritSec);

    return NO_ERROR;
}

int PostSend(BUFFER_OBJ_EX *sendobj) {
    WSABUF  wbuf;
    DWORD   bytes;
    int     rc;

    rc = NO_ERROR;

    sendobj->operation = OP_WRITE;

    wbuf.buf = sendobj->buf;
    wbuf.len = sendobj->buflen;

    EnterCriticalSection(&sendobj->Socket->SockCritSec);
    if (gProtocol == IPPROTO_TCP) {
        rc = WSASend(sendobj->Socket->s, &wbuf, 1, &bytes, 0, &sendobj->ol, NULL);
    }
    else {
        ExitProcess(-1);
    }
    if (rc == SOCKET_ERROR)
    {
        rc = NO_ERROR;
        if (WSAGetLastError() != WSA_IO_PENDING)
        {
            fprintf(stderr, "PostSend: WSASend* failed: %d\n", WSAGetLastError());
            rc = SOCKET_ERROR;
        }
    }

    // Increment the outstanding operation count
    InterlockedIncrement(&sendobj->Socket->OutstandingOps);

    LeaveCriticalSection(&sendobj->Socket->SockCritSec);

    return rc;
}

int PostAccept(BUFFER_OBJ_EX *acceptobj) {
    DWORD   bytes;
    int     rc = NO_ERROR;

    acceptobj->operation = OP_ACCEPT;

    EnterCriticalSection(&acceptobj->Socket->SockCritSec);

    acceptobj->sclient = socket(acceptobj->Socket->af, SOCK_STREAM, IPPROTO_TCP);
    if (acceptobj->sclient == INVALID_SOCKET) {
        fprintf(stderr, "PostAccept: create socket failed.\n");
        return -1;
    }

    rc = acceptobj->Socket->lpfnAcceptEx(acceptobj->Socket->s, acceptobj->sclient, acceptobj->buf,
        acceptobj->buflen - ((sizeof(SOCKADDR_STORAGE) + 16) * 2), sizeof(SOCKADDR_STORAGE) + 16, sizeof(SOCKADDR_STORAGE) + 16,
           &bytes, &acceptobj->ol);
    if (rc == FALSE) {
        if (WSAGetLastError() != WSA_IO_PENDING)
        {
            fprintf(stderr, "PostAccept: AcceptEx failed: %d\n", WSAGetLastError());
            ExitProcess(-1);
        }
    }

    InterlockedIncrement(&acceptobj->Socket->OutstandingOps);

    LeaveCriticalSection(&acceptobj->Socket->SockCritSec);

    return rc;
}

void HandleIo(BUFFER_OBJ_EX *buf) {
    SOCKET_OBJ      *sock=NULL,
                    *clientobj=NULL;     // New client object for accepted connections
    BUFFER_OBJ_EX   *recvobj=NULL,       // Used to post new receives on accepted connections
                    *sendobj=NULL;       // Used to post sends for data received
    DWORD            bytes,
                     flags;
    BOOL             bFreeSocketObj;
    int              error,
                     rc;

    sock = buf->Socket;
    error = NO_ERROR;

    bFreeSocketObj = FALSE;

    InterlockedDecrement(&sock->OutstandingOps);

    rc = WSAGetOverlappedResult(sock->s, &buf->ol, &bytes, FALSE, &flags);
    if (rc == FALSE) {
        fprintf(stderr, "HandleIo: WSAGetOverlappedResult failed.\n");
        ExitProcess(-1);
    }

    if (buf->operation == OP_ACCEPT) {
        SOCKADDR_STORAGE *LocalSockaddr=NULL,
                         *RemoteSockaddr=NULL;
        int               LocalSockaddrLen,
                          RemoteSockaddrLen;

        sock->lpfnGetAcceptExSockaddrs(buf->buf, buf->buflen - ((sizeof(SOCKADDR_STORAGE) + 16) * 2),
                sizeof(SOCKADDR_STORAGE) + 16, sizeof(SOCKADDR_STORAGE) + 16,
                (SOCKADDR **)&LocalSockaddr, &LocalSockaddrLen,
                (SOCKADDR **)&RemoteSockaddr, &RemoteSockaddrLen);

        clientobj = GetSocketObj(buf->sclient, buf->Socket->af);
        sendobj = GetBufferObjEx(clientobj, gBufferSize);

        sendobj->buflen = bytes;
        memcpy(sendobj->buf, buf->buf, bytes);

        AssignIoToThread(sendobj);

        if (PostSend(sendobj) != NO_ERROR) {
            RemoveBufferFromThread(clientobj, sendobj);
            FreeBufferObjEx(sendobj);

            closesocket(clientobj->s);
            clientobj->s = INVALID_SOCKET;

            FreeSocketObj(clientobj);
        }

        PostAccept(buf);
    } else if ((buf->operation == OP_READ) && (error == NO_ERROR)) {
        if ((bytes > 0) || (gProtocol == IPPROTO_UDP))
        {
            // Create a buffer to send
            sendobj = buf;

            sendobj->buflen = bytes;

            // Initiate the send
            if (PostSend(sendobj) != NO_ERROR)
            {
                // In the event of an error, clean up the socket object
                RemoveBufferFromThread(sock, sendobj);
                FreeBufferObjEx(sendobj);

                closesocket(sock->s);
                sock->s = INVALID_SOCKET;

                bFreeSocketObj = TRUE;
            }
        }
        else
        {
            // Graceful close
            sock->bClosing = TRUE;

            // Free the completed operation
            RemoveBufferFromThread(sock, buf);
            FreeBufferObjEx(buf);

            // Check to see if there are more outstanding operations. If so, wait
            //    for them to complete; otherwise, clean up the socket object.
            EnterCriticalSection(&sock->SockCritSec);
            if (sock->OutstandingOps == 0)
            {
                closesocket(sock->s);

                bFreeSocketObj = TRUE;
            }
            LeaveCriticalSection(&sock->SockCritSec);
        }
    } else if (buf->operation == OP_WRITE) {
        EnterCriticalSection(&sock->SockCritSec);
        if (sock->bClosing && (sock->OutstandingOps == 0))
        {
            RemoveBufferFromThread(sock, buf);
            closesocket(sock->s);
            FreeBufferObjEx(buf);

            bFreeSocketObj = TRUE;
        }
        else
        {
            buf->buflen = gBufferSize;

            // Free the send op that just completed
            if (PostRecv(buf) != NO_ERROR)
            {
                RemoveBufferFromThread(sock, buf);
                FreeBufferObjEx(buf);
            }
        }
        LeaveCriticalSection(&sock->SockCritSec);
    }
}

DWORD WINAPI IoThread(LPVOID lpParam) {
    THREAD_OBJ      *thread=NULL;
    int              index,
                     count,
                     rc,
                     i;

    thread = (THREAD_OBJ *)lpParam;

    RenumberEvents(thread);

    while (true) {
        rc = WaitForMultipleObjects(thread->EventCount, thread->Handles, FALSE, INFINITE);
        if (rc == WAIT_FAILED || rc == WAIT_TIMEOUT) {
            if (GetLastError() == ERROR_INVALID_HANDLE)
            {
                RenumberEvents(thread);
                continue;
            }
            else
            {
                fprintf(stderr, "IoThread: WaitForMultipleObjects failed: %d\n",
                        GetLastError());
                break;
            }
        }

        count = thread->EventCount;
        for (i = 0; i < count; ++i) {
            rc = WaitForSingleObject(thread->Handles[i], 0);
            if (rc == WAIT_TIMEOUT)
                continue;

            index = i;
            WSAResetEvent(thread->Handles[index]);

            if (index == 0) {
                RenumberEvents(thread);
                break;
            } else {
                HandleIo(FindBufferObjEx(thread, thread->Handles[index]));
            }
        }
    }

    ExitThread(0);
    return 0;
}

void AssignIoToThread(BUFFER_OBJ_EX *buf) {
    ThreadObj *threadObj;
    THREAD_OBJ *thread;

    EnterCriticalSection(&gThreadListCritSec);

    threadObj = (ThreadObj *)GotoNextSingleList(&threadobjHeader, threadobjHeader.head);
    while (threadObj) {
        thread = (THREAD_OBJ *)container_of(THREAD_OBJ, next, threadObj);

        if (InsertBufferObjExToThread(thread, buf) == NO_ERROR) {
            break;
        }

        threadObj = (ThreadObj *)GotoNextSingleList(&threadobjHeader, threadObj);
    }

    if (threadObj == NULL) {
        thread = GetThreadObj();

        thread->Thread = CreateThread(NULL, 0, IoThread, (LPVOID)thread, 0, NULL);
        if (thread->Thread == NULL) {
            fprintf(stderr, "AssignIoToThread: CreateThread failed.\n");
            ExitProcess(-1);
        }

        InsertBufferObjExToThread(thread, buf);
        EnqueueSingleListHead(&threadobjHeader, &(thread->next));
    }

    buf->Thread = thread;

    WSASetEvent(thread->Event);

    LeaveCriticalSection(&gThreadListCritSec);
}

int _tmain(int argc, _TCHAR* argv[])
{
    WSADATA          wsd;
    THREAD_OBJ      *thread=NULL;
    SOCKET_OBJ      *sockobj=NULL;
    SocketObjHeader  ListenSockets;
    int              endpointcount=0,
                     interval = 0,
                     rc,
                     i;
    struct addrinfo *res=NULL,
                    *ptr=NULL;

    if (WSAStartup(MAKEWORD(2,2), &wsd) != 0)
    {
        fprintf(stderr, "unable to load Winsock!\n");
        return -1;
    }

    InitializeCriticalSection(&gThreadListCritSec);
    InitializeSingleHead(&ListenSockets);
    InitializeSingleHead(&threadobjHeader);

    res = ResolveAddress(gSrvAddr, gPort, gAddressFamily, gSocketType, gProtocol);
    if (res == NULL)
    {
        fprintf(stderr, "ResolveAddress failed to return any addresses!\n");
        return -1;
    }

    ptr = res;
    while (ptr) {
        sockobj = GetSocketObj(INVALID_SOCKET, ptr->ai_family);

        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;
        }

        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) {
            BUFFER_OBJ_EX *acceptobj = NULL;
            GUID guidAcceptEx = WSAID_ACCEPTEX,
                 guidGetAcceptExSockaddrs = WSAID_GETACCEPTEXSOCKADDRS;
            DWORD       bytes;

            rc = WSAIoctl(sockobj->s, SIO_GET_EXTENSION_FUNCTION_POINTER, &guidAcceptEx, sizeof(guidAcceptEx),
                &sockobj->lpfnAcceptEx, sizeof(sockobj->lpfnAcceptEx), &bytes, NULL, NULL);
            if (rc == SOCKET_ERROR)
            {
                fprintf(stderr, "WSAIoctl: SIO_GET_EXTENSION_FUNCTION_POINTER failed: %d\n",
                        WSAGetLastError());
                return -1;
            }

            rc = WSAIoctl(sockobj->s, SIO_GET_EXTENSION_FUNCTION_POINTER, &guidGetAcceptExSockaddrs, sizeof(guidGetAcceptExSockaddrs),
                &sockobj->lpfnGetAcceptExSockaddrs, sizeof(sockobj->lpfnGetAcceptExSockaddrs), &bytes, NULL, NULL);
            if (rc == SOCKET_ERROR)
            {
                fprintf(stderr, "WSAIoctl: SIO_GET_EXTENSION_FUNCTION_POINTER failed: %d\n",
                        WSAGetLastError());
                return -1;
            }

            rc = listen(sockobj->s, 200);
            if (rc == SOCKET_ERROR)
            {
                fprintf(stderr, "listen failed: %d\n", WSAGetLastError());
                return -1;
            }

            sockobj->PendingAccepts = (BUFFER_OBJ_EX **)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, (sizeof(BUFFER_OBJ *) * gOverlappedCount));
            if (sockobj->PendingAccepts == NULL)
            {
                fprintf(stderr, "PendingAccepts HeapAlloc failed: %d\n", GetLastError());
                ExitProcess(-1);
            }
            for (i = 0; i < gOverlappedCount; ++i) {
                sockobj->PendingAccepts[i] = acceptobj = GetBufferObjEx(sockobj, gBufferSize);
                AssignIoToThread(acceptobj);

                if (PostAccept(acceptobj) != NO_ERROR) {
                    ExitProcess(-1);
                }
            }

            EnqueueSingleList(&ListenSockets, &(sockobj->next));
        }


        ptr = ptr->ai_next;
    }

    freeaddrinfo(res);

    while (true) {
        Sleep(5000);

        interval++;

        if (interval == 12) {
            SocketObj   *listenptr=NULL;
            int          optval,
                         optlen;

            // Walk the list of outstanding accepts
            listenptr = (SocketObj *)GotoNextSingleList(&ListenSockets, ListenSockets.head);
            while (listenptr)
            {
                sockobj = (SOCKET_OBJ *)container_of(SOCKET_OBJ, next, listenptr);
                for(i=0; i < gOverlappedCount ;i++)
                {
                    optlen = sizeof(optval);
                    rc = getsockopt(
                            sockobj->PendingAccepts[i]->sclient,
                            SOL_SOCKET,
                            SO_CONNECT_TIME,
                            (char *)&optval,
                           &optlen
                           );
                    if (rc == SOCKET_ERROR)
                    {
                        fprintf(stderr, "getsockopt: SO_CONNECT_TIME failed: %d\n", WSAGetLastError());
                        return -1;
                    }
                    if ((optval != 0xFFFFFFFF) && (optval > 300))
                    {
                        closesocket(sockobj->PendingAccepts[i]->sclient);
                    }
                }
                listenptr = (SocketObj *)GotoNextSingleList(&ListenSockets, listenptr);
            }
            interval = 0;
        }
    }

    WSACleanup();

    return 0;
}