Add ExecuteWithTimeout to graph clients

example/nebula-test-timeout.js

@@ -0,0 +1,49 @@
+import nebulaPool from 'k6/x/nebulagraph';
+import { check } from 'k6';
+import { Trend } from 'k6/metrics';
+
+var latencyTrend = new Trend('latency', true);
+var responseTrend = new Trend('responseTime', true);
+
+var graph_option = {
+	address: "192.168.8.6:10010",
+	space: "sf1",
+	csv_path: "person.csv",
+	csv_delimiter: "|",
+	csv_with_header: true,
+	output: "output.csv"
+};
+
+nebulaPool.setOption(graph_option);
+var pool = nebulaPool.init();
+var session = pool.getSession()
+
+String.prototype.format = function() {
+	var formatted = this;
+	var data = arguments[0]
+
+	formatted = formatted.replace(/\{(\d+)\}/g, function(match, key) {
+		return data[key]
+	})
+	return formatted
+};
+
+export default function() {
+	let d = session.getData()
+	let ngql = 'go 2 steps from {0} over KNOWS yield dst(edge)'.format(d)
+	let timeoutMs = 1000
+	let response = session.executeWithTimeout(ngql, timeoutMs)
+
+	check(response, {
+		"IsSucceed": (r) => r !== null && r.isSucceed() === true
+	});
+
+	if (response !== null) {
+		latencyTrend.add(response.getLatency() / 1000);
+		responseTrend.add(response.getResponseTime() / 1000);
+	}
+};
+
+export function teardown() {
+	pool.close()
+}

pkg/common/types.go

@@ -28,6 +28,7 @@ type (
 		GetData() (Data, error)
 		GetFileData(sourceFile string) (Data, error)
 		Execute(stmt string) (IGraphResponse, error)
+		ExecuteWithTimeout(stmt string, timeoutMs int) (IGraphResponse, error)
 	}
 
 	// IGraphResponse graph response, just support some functions to user.

pkg/nebulagraph/client.go

@@ -437,6 +437,10 @@ func (gc *GraphClient) Execute(stmt string) (common.IGraphResponse, error) {
 	return result, nil
 }
 
+func (gc *GraphClient) ExecuteWithTimeout(stmt string, timeoutMs int) (common.IGraphResponse, error) {
+	return nil, fmt.Errorf("ExecuteWithTimeout is not supported by nebulagraph client")
+}
+
 // GetResponseTime GetResponseTime
 func (r *Response) GetResponseTime() int32 {
 	return r.ResponseTime

pkg/nebulagraph5/client.go

@@ -1,6 +1,7 @@
 package nebulagraph5
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"math"
@@ -295,6 +296,14 @@ func (gc *GraphClient) GetFileData(sourceFile string) (common.Data, error) {
 
 // Execute executes nebula query
 func (gc *GraphClient) Execute(stmt string) (common.IGraphResponse, error) {
+	return gc.execute(stmt, 0)
+}
+
+func (gc *GraphClient) ExecuteWithTimeout(stmt string, timeoutMs int) (common.IGraphResponse, error) {
+	return gc.execute(stmt, timeoutMs)
+}
+
+func (gc *GraphClient) execute(stmt string, timeoutMs int) (common.IGraphResponse, error) {
 	var (
 		isSucceed  bool = true
 		errMessage string
@@ -320,7 +329,7 @@ func (gc *GraphClient) Execute(stmt string) (common.IGraphResponse, error) {
 		}
 		gc.since = time.Now()
 	}
-	resp, err = gc.executeWithRetry(stmt)
+	resp, err = gc.executeWithRetry(stmt, timeoutMs)
 
 	if err != nil {
 		isSucceed = false
@@ -386,7 +395,7 @@ func (gc *GraphClient) Execute(stmt string) (common.IGraphResponse, error) {
 	return &Response{ResultSet: resp, ResponseTime: responseTime, err: err}, nil
 }
 
-func (gc *GraphClient) executeWithRetry(stmt string) (types.Result, error) {
+func (gc *GraphClient) executeWithRetry(stmt string, timeoutMs int) (types.Result, error) {
 	var (
 		err  error
 		resp types.Result
@@ -403,7 +412,7 @@ func (gc *GraphClient) executeWithRetry(stmt string) (types.Result, error) {
 		if i > 0 {
 			gc.Pool.logger.Warnf("execute statement failed, retry %d time, error: %s\n", i, err.Error())
 		}
-		resp, err = gc.execute(stmt)
+		resp, err = gc.executeOnce(stmt, timeoutMs)
 		if err == nil {
 			return resp, nil
 		} else {
@@ -415,15 +424,25 @@ func (gc *GraphClient) executeWithRetry(stmt string) (types.Result, error) {
 	return nil, err
 }
 
-func (gc *GraphClient) execute(stmt string) (types.Result, error) {
+func (gc *GraphClient) executeOnce(stmt string, timeoutMs int) (types.Result, error) {
 	if gc.Session == nil || gc.Session.IsClosed() {
 		sess, err := gc.Pool.pool.GetClient()
 		if err != nil {
 			return nil, err
 		}
 		gc.Session = sess
 	}
-	resp, err := gc.Session.Execute(stmt)
+	var (
+		resp types.Result
+		err  error
+	)
+	if timeoutMs > 0 {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeoutMs)*time.Millisecond)
+		defer cancel()
+		resp, err = gc.Session.ExecuteContext(ctx, stmt)
+	} else {
+		resp, err = gc.Session.Execute(stmt)
+	}
 	if err != nil {
 		return nil, fmt.Errorf("execute statement failed: %s, error: %w", stmt, err)
 	}