fix: refactor code

taosdata · Mar 9, 2023 · 55622b0 · 55622b0
1 parent 100718d
commit 55622b0
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 44 deletions.
diff --git a/inc/bench.h b/inc/bench.h
@@ -41,6 +41,8 @@
 #include <sys/ioctl.h>
 #include <signal.h>
 
+#define SOCKET_ERROR  -1
+
 #elif DARWIN
 #include <argp.h>
 #include <unistd.h>
@@ -49,9 +51,19 @@
 #include <arpa/inet.h>
 #include <sys/time.h>
 #include <netdb.h>
-#else
+#else  // WINDOWS
+#define _CRT_RAND_S
+#include <windows.h>
 #include <winsock2.h>
-#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 <limits.h>
 #include <regex.h>

diff --git 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 <windows.h>
-#include <winsock2.h>
-
-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,52 +1148,48 @@ int createSockFd() {
     return sockfd;
 }
 
-void destroySockFd(int sockfd) {
-    // shutdown the connection since no more data will be sent
-    int result, how;
+static void printSocketMsg(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) {
+        printSocketMsg("Socket shutdown failed with error: %d\n", result);
+        closeSockFd(sockfd);
         return;
     }
     // Receive until the peer closes the connection
     do {
         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 {
+            printSocketMsg("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) {