diff --git a/example/nebula-test-timeout.js b/example/nebula-test-timeout.js new file mode 100644 index 0000000..f801b29 --- /dev/null +++ b/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() +} diff --git a/pkg/common/types.go b/pkg/common/types.go index 11f08ce..8aa8891 100644 --- a/pkg/common/types.go +++ b/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. diff --git a/pkg/nebulagraph/client.go b/pkg/nebulagraph/client.go index ea9b00d..3cf662b 100644 --- a/pkg/nebulagraph/client.go +++ b/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 diff --git a/pkg/nebulagraph5/client.go b/pkg/nebulagraph5/client.go index 64bc941..3d30dd3 100644 --- a/pkg/nebulagraph5/client.go +++ b/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,7 +424,7 @@ 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 { @@ -423,7 +432,17 @@ func (gc *GraphClient) execute(stmt string) (types.Result, error) { } 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) }