Merge pull request #7309 from taosdata/feature/TD-5838

[TD-5838]<feature> add function of speed test

src/kit/shell/inc/shell.h

@@ -56,6 +56,8 @@ typedef struct SShellArguments {
   int   abort;
   int   port;
   int   pktLen;
+  int   pktNum;
+  char* pktType;
   char* netTestRole;
 } SShellArguments;
 

src/kit/shell/src/shellLinux.c

@@ -50,6 +50,8 @@ static struct argp_option options[] = {
   {"timezone",   't', "TIMEZONE",   0,                   "Time zone of the shell, default is local."},
   {"netrole",    'n', "NETROLE",    0,                   "Net role when network connectivity test, default is startup, options: client|server|rpc|startup|sync."},
   {"pktlen",     'l', "PKTLEN",     0,                   "Packet length used for net test, default is 1000 bytes."},
+  {"pktnum",     'N', "PKTNUM",     0,                   "Packet numbers used for net test, default is 100."},
+  {"pkttype",    'S', "PKTTYPE",    0,                   "Packet type used for net test, default is TCP."},
   {0}};
 
 static error_t parse_opt(int key, char *arg, struct argp_state *state) {
@@ -146,6 +148,17 @@ static error_t parse_opt(int key, char *arg, struct argp_state *state) {
         return -1;
       }
       break;
+    case 'N':
+      if (arg) {
+        arguments->pktNum = atoi(arg);
+      } else {
+        fprintf(stderr, "Invalid packet number\n");
+        return -1;
+      }
+      break;
+    case 'S':
+      arguments->pktType = arg;
+      break;
     case OPT_ABORT:
       arguments->abort = 1;
       break;

src/kit/shell/src/shellMain.c

@@ -85,6 +85,8 @@ SShellArguments args = {
   .threadNum = 5,
   .commands = NULL,
   .pktLen = 1000,
+  .pktNum = 100,
+  .pktType = "TCP",
   .netTestRole = NULL
 };
 
@@ -118,7 +120,7 @@ int main(int argc, char* argv[]) {
       printf("Failed to init taos");
       exit(EXIT_FAILURE);
     }
-    taosNetTest(args.netTestRole, args.host, args.port, args.pktLen);
+    taosNetTest(args.netTestRole, args.host, args.port, args.pktLen, args.pktNum, args.pktType);
     exit(0);
   }
 

src/kit/shell/src/shellWindows.c

@@ -58,6 +58,10 @@ void printHelp() {
   printf("%s%s%s\n", indent, indent, "Net role when network connectivity test, default is startup, options: client|server|rpc|startup|sync.");
   printf("%s%s\n", indent, "-l");
   printf("%s%s%s\n", indent, indent, "Packet length used for net test, default is 1000 bytes.");
+  printf("%s%s\n", indent, "-N");
+  printf("%s%s%s\n", indent, indent, "Packet numbers used for net test, default is 100.");
+  printf("%s%s\n", indent, "-S");
+  printf("%s%s%s\n", indent, indent, "Packet type used for net test, default is TCP.");
   printf("%s%s\n", indent, "-V");
   printf("%s%s%s\n", indent, indent, "Print program version.");
 
@@ -184,6 +188,22 @@ void shellParseArgument(int argc, char *argv[], SShellArguments *arguments) {
         exit(EXIT_FAILURE);
       }
     }
+    else if (strcmp(argv[i], "-N") == 0) {
+      if (i < argc - 1) {
+        arguments->pktNum = atoi(argv[++i]);
+      } else {
+        fprintf(stderr, "option -N requires an argument\n");
+        exit(EXIT_FAILURE);
+      }
+    }
+    else if (strcmp(argv[i], "-S") == 0) {
+      if (i < argc - 1) {
+        arguments->pktType = argv[++i];
+      } else {
+        fprintf(stderr, "option -S requires an argument\n");
+        exit(EXIT_FAILURE);
+      }
+    }
     else if (strcmp(argv[i], "-V") == 0) {
       printVersion();
       exit(EXIT_SUCCESS);

src/util/inc/tnettest.h

@@ -20,7 +20,7 @@
 extern "C" {
 #endif
 
-void taosNetTest(char *role, char *host, int port, int pkgLen);
+void taosNetTest(char *role, char *host, int32_t port, int32_t pkgLen, int32_t pkgNum, char *pkgType);
 
 #ifdef __cplusplus
 }

src/util/src/tnettest.c

@@ -27,6 +27,10 @@
 #include "syncMsg.h"
 
 #define MAX_PKG_LEN (64 * 1000)
+#define MAX_SPEED_PKG_LEN (1024 * 1024 * 1024)
+#define MIN_SPEED_PKG_LEN 1024
+#define MAX_SPEED_PKG_NUM 10000
+#define MIN_SPEED_PKG_NUM 1
 #define BUFFER_SIZE (MAX_PKG_LEN + 1024)
 
 extern int32_t tsRpcMaxUdpSize;
@@ -466,6 +470,7 @@ static void taosNetTestRpc(char *host, int32_t startPort, int32_t pkgLen) {
       sendpkgLen = pkgLen;
     }
 
+    tsRpcForceTcp = 1;
     int32_t ret = taosNetCheckRpc(host, port, sendpkgLen, spi, NULL);
     if (ret < 0) {
       printf("failed to test TCP port:%d\n", port);
@@ -479,6 +484,7 @@ static void taosNetTestRpc(char *host, int32_t startPort, int32_t pkgLen) {
       sendpkgLen = pkgLen;
     }
 
+    tsRpcForceTcp = 0;
     ret = taosNetCheckRpc(host, port, pkgLen, spi, NULL);
     if (ret < 0) {
       printf("failed to test UDP port:%d\n", port);
@@ -542,12 +548,122 @@ static void taosNetTestServer(char *host, int32_t startPort, int32_t pkgLen) {
   }
 }
 
-void taosNetTest(char *role, char *host, int32_t port, int32_t pkgLen) {
+static void taosNetTestFqdn(char *host) {
+  int code = 0;
+  uint64_t startTime = taosGetTimestampUs();
+  uint32_t ip = taosGetIpv4FromFqdn(host);
+  if (ip == 0xffffffff) {
+    uError("failed to get IP address from %s since %s", host, strerror(errno));
+    code = -1;
+  }
+  uint64_t endTime = taosGetTimestampUs();
+  uint64_t el = endTime - startTime;
+  printf("check convert fqdn spend, status: %d\tcost: %" PRIu64 " us\n", code, el);
+  return;
+}
+
+static void taosNetCheckSpeed(char *host, int32_t port, int32_t pkgLen,
+                              int32_t pkgNum, char *pkgType) {
+  // record config
+  int32_t compressTmp = tsCompressMsgSize;
+  int32_t maxUdpSize  = tsRpcMaxUdpSize;
+  int32_t forceTcp  = tsRpcForceTcp;
+
+  if (0 == strcmp("tcp", pkgType)){
+    tsRpcForceTcp = 1;
+    tsRpcMaxUdpSize = 0;            // force tcp
+  } else {
+    tsRpcForceTcp = 0;
+    tsRpcMaxUdpSize = INT_MAX;
+  }
+  tsCompressMsgSize = -1;
+
+  SRpcEpSet epSet;
+  SRpcMsg   reqMsg;
+  SRpcMsg   rspMsg;
+  void *    pRpcConn;
+  char secretEncrypt[32] = {0};
+  char    spi = 0;
+  pRpcConn = taosNetInitRpc(secretEncrypt, spi);
+  if (NULL == pRpcConn) {
+    uError("failed to init client rpc");
+    return;
+  }
+
+  printf("check net spend, host:%s port:%d pkgLen:%d pkgNum:%d pkgType:%s\n\n", host, port, pkgLen, pkgNum, pkgType);
+  int32_t totalSucc = 0;
+  uint64_t startT = taosGetTimestampUs();
+  for (int32_t i = 1; i <= pkgNum; i++) {
+    uint64_t startTime = taosGetTimestampUs();
+
+    memset(&epSet, 0, sizeof(SRpcEpSet));
+    epSet.inUse = 0;
+    epSet.numOfEps = 1;
+    epSet.port[0] = port;
+    strcpy(epSet.fqdn[0], host);
+
+    reqMsg.msgType = TSDB_MSG_TYPE_NETWORK_TEST;
+    reqMsg.pCont = rpcMallocCont(pkgLen);
+    reqMsg.contLen = pkgLen;
+    reqMsg.code = 0;
+    reqMsg.handle = NULL;   // rpc handle returned to app
+    reqMsg.ahandle = NULL;  // app handle set by client
+    strcpy(reqMsg.pCont, "nettest speed");
+
+    rpcSendRecv(pRpcConn, &epSet, &reqMsg, &rspMsg);
+
+    int code = 0;
+    if ((rspMsg.code != 0) || (rspMsg.msgType != TSDB_MSG_TYPE_NETWORK_TEST + 1)) {
+      uError("ret code 0x%x %s", rspMsg.code, tstrerror(rspMsg.code));
+      code = -1;
+    }else{
+      totalSucc ++;
+    }
+
+    rpcFreeCont(rspMsg.pCont);
+
+    uint64_t endTime = taosGetTimestampUs();
+    uint64_t el = endTime - startTime;
+    printf("progress:%5d/%d\tstatus:%d\tcost:%8.2lf ms\tspeed:%8.2lf MB/s\n", i, pkgNum, code, el/1000.0, pkgLen/(el/1000000.0)/1024.0/1024.0);
+  }
+  int64_t endT = taosGetTimestampUs();
+  uint64_t elT = endT - startT;
+  printf("\ntotal succ:%5d/%d\tcost:%8.2lf ms\tspeed:%8.2lf MB/s\n", totalSucc, pkgNum, elT/1000.0, pkgLen/(elT/1000000.0)/1024.0/1024.0*totalSucc);
+
+  rpcClose(pRpcConn);
+
+  // return config
+  tsCompressMsgSize = compressTmp;
+  tsRpcMaxUdpSize = maxUdpSize;
+  tsRpcForceTcp = forceTcp;
+  return;
+}
+
+static void taosNetTestSpeed(char *host, int32_t port, int32_t pkgLen,
+                             int32_t pkgNum, char *pkgType) {
+  if (0 == strcmp("fqdn", pkgType)){
+    taosNetTestFqdn(host);
+    return;
+  }
+
+  taosNetCheckSpeed(host, port, pkgLen, pkgNum, pkgType);
+  return;
+}
+
+void taosNetTest(char *role, char *host, int32_t port, int32_t pkgLen,
+                 int32_t pkgNum, char *pkgType) {
   tscEmbedded = 1;
   if (host == NULL) host = tsLocalFqdn;
   if (port == 0) port = tsServerPort;
-  if (pkgLen <= 10) pkgLen = 1000;
-  if (pkgLen > MAX_PKG_LEN) pkgLen = MAX_PKG_LEN;
+  if (0 == strcmp("speed", role)){
+    if (pkgLen <= MIN_SPEED_PKG_LEN) pkgLen = MIN_SPEED_PKG_LEN;
+    if (pkgLen > MAX_SPEED_PKG_LEN) pkgLen = MAX_SPEED_PKG_LEN;
+    if (pkgNum <= MIN_SPEED_PKG_NUM) pkgNum = MIN_SPEED_PKG_NUM;
+    if (pkgNum > MAX_SPEED_PKG_NUM) pkgNum = MAX_SPEED_PKG_NUM;
+  }else{
+    if (pkgLen <= 10) pkgLen = 1000;
+    if (pkgLen > MAX_PKG_LEN) pkgLen = MAX_PKG_LEN;
+  }
 
   if (0 == strcmp("client", role)) {
     taosNetTestClient(host, port, pkgLen);
@@ -560,7 +676,11 @@ void taosNetTest(char *role, char *host, int32_t port, int32_t pkgLen) {
     taosNetCheckSync(host, port);
   } else if (0 == strcmp("startup", role)) {
     taosNetTestStartup(host, port);
-  } else {
+  } else if (0 == strcmp("speed", role)) {
+    tscEmbedded = 0;
+    char type[10] = {0};
+    taosNetTestSpeed(host, port, pkgLen, pkgNum, strtolower(type, pkgType));
+  }else {
     taosNetTestStartup(host, port);
   }