diff --git a/inc/bench.h b/inc/bench.h index 16ba7053..a9513c65 100644 --- a/inc/bench.h +++ b/inc/bench.h @@ -41,6 +41,8 @@ #include #include +#define SOCKET_ERROR -1 + #elif DARWIN #include #include @@ -49,9 +51,22 @@ #include #include #include -#else + +#define SOCKET_ERROR -1 + +#else // WINDOWS +#define _CRT_RAND_S +#include #include -#endif +#pragma comment(lib, "ws2_32.lib") +#define SHUT_WR SD_SEND + +typedef unsigned __int32 uint32_t; +// Some old MinGW/CYGWIN distributions don't define this: +#ifndef ENABLE_VIRTUAL_TERMINAL_PROCESSING + #define ENABLE_VIRTUAL_TERMINAL_PROCESSING 0x0004 +#endif // ENABLE_VIRTUAL_TERMINAL_PROCESSING +#endif // LINUX #include #include diff --git a/src/benchUtil.c b/src/benchUtil.c index 6e715be2..deac14b1 100644 --- a/src/benchUtil.c +++ b/src/benchUtil.c @@ -50,17 +50,6 @@ void ERROR_EXIT(const char *msg) { } #ifdef WINDOWS -#define _CRT_RAND_S -#include -#include - -typedef unsigned __int32 uint32_t; - -#pragma comment(lib, "ws2_32.lib") -// Some old MinGW/CYGWIN distributions don't define this: -#ifndef ENABLE_VIRTUAL_TERMINAL_PROCESSING -#define ENABLE_VIRTUAL_TERMINAL_PROCESSING 0x0004 -#endif // ENABLE_VIRTUAL_TERMINAL_PROCESSING HANDLE g_stdoutHandle; DWORD g_consoleMode; @@ -1159,26 +1148,31 @@ int createSockFd() { return sockfd; } -void destroySockFd(int sockfd) { - // shutdown the connection since no more data will be sent - int result, how; +static void errorPrintSocketMsg(char *msg, int result) { + errorPrint("Socket shutdown failed with error: %d\n", #ifdef WINDOWS - how = SD_SEND; + WSAGetLastError()); #else - how = SHUT_WR; - #define SOCKET_ERROR -1 + result); #endif - result = shutdown(sockfd, how); - if (SOCKET_ERROR == result) { - errorPrint("Socket shutdown failed with error: %d\n", +} + +static void closeSockFd(int sockfd) { #ifdef WINDOWS - WSAGetLastError()); - closesocket(sockfd); - WSACleanup(); + closesocket(sockfd); + WSACleanup(); #else - result); - close(sockfd); + close(sockfd); #endif +} + +void destroySockFd(int sockfd) { + // shutdown the connection since no more data will be sent + int result; + result = shutdown(sockfd, SHUT_WR); + if (SOCKET_ERROR == result) { + errorPrintSocketMsg("Socket shutdown failed with error: %d\n", result); + closeSockFd(sockfd); return; } // Receive until the peer closes the connection @@ -1186,25 +1180,16 @@ void destroySockFd(int sockfd) { int recvbuflen = LARGE_BUFF_LEN; char recvbuf[LARGE_BUFF_LEN]; result = recv(sockfd, recvbuf, recvbuflen, 0); - if ( result > 0 ) + if ( result > 0 ) { debugPrint("Socket bytes received: %d\n", result); - else if ( result == 0 ) + } else if (result == 0) { infoPrint("Connection closed with result %d\n", result); - else - errorPrint("Socket recv failed with error: %d\n", -#ifdef WINDOWS - WSAGetLastError() -#else - result -#endif - ); - } while( result > 0 ); -#ifdef WINDOWS - closesocket(sockfd); - WSACleanup(); -#else - close(sockfd); -#endif + } else { + errorPrintSocketMsg("Socket recv failed with error: %d\n", result); + } + } while (result > 0); + + closeSockFd(sockfd); } FORCE_INLINE void printErrCmdCodeStr(char *cmd, int32_t code, TAOS_RES *res) {