enhance csv reader

Makefile

@@ -5,7 +5,13 @@ pairs := darwin/amd64 linux/amd64 linux/arm64
 GOPATH ?= ~/go
 export GO111MODULE=on
 VERSION ?= v1.0.1
-K6_VERSION ?= v0.33.0
+K6_VERSION ?= v0.40.0
+
+fmt:
+	find . -name '*.go' -exec gofmt -s -w {} +
+
+lint :
+	golangci-lint run --out-format=tab ./...
 
 build: 
 	go install github.com/k6io/xk6/cmd/xk6@v0.4.1

example/nebula-test.js

@@ -9,7 +9,7 @@ var responseTrend = new Trend('responseTime');
 var pool = nebulaPool.init("192.168.8.152:9669", 400);
 // initial session for every vu
 var session = pool.getSession("root", "nebula")
-session.execute("USE ldbc")
+session.execute("USE sf1")
 
 export function setup() {
 	// config csv file
@@ -23,7 +23,7 @@ export default function (data) {
 	// get csv data from csv file
 	let d = session.getData()
 	// d[0] means the first column data in the csv file
-	let ngql = 'go 2 steps from ' + d[0] + ' over KNOWS '
+	let ngql = 'go 2 steps from ' + d[0] + ' over KNOWS yield $$.Person.firstName as firstName'
 	let response = session.execute(ngql)
 	check(response, {
 		"IsSucceed": (r) => r.isSucceed() === true

go.mod

@@ -4,6 +4,5 @@ go 1.16
 
 require (
 	github.com/vesoft-inc/nebula-http-gateway/ccore v0.0.0-20220413113447-a3f4c56287d8
-	go.k6.io/k6 v0.33.0
-	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
+	go.k6.io/k6 v0.40.0
 )

pkg/aggcsv/output.go

@@ -53,7 +53,10 @@ func (o *Output) Start() error {
 	if err != nil {
 		return err
 	}
-	o.outputFile.Write([]byte("#timestamp,vus,requestCount,errorCount,latencyAvg,latencyP90,latencyP95,latencyP99,responseTimeAvg,responseTimeP90,responseTimeP95,responseTimeP99,rowSizePerReq\n"))
+	_, err = o.outputFile.Write([]byte("#timestamp,vus,requestCount,errorCount,latencyAvg,latencyP90,latencyP95,latencyP99,responseTimeAvg,responseTimeP90,responseTimeP95,responseTimeP99,rowSizePerReq\n"))
+	if err != nil {
+		return err
+	}
 
 	pf, err := output.NewPeriodicFlusher(time.Duration(o.config.aggregationInterval)*time.Second, o.aggregateAndFlush)
 	if err != nil {
@@ -143,7 +146,7 @@ func (o *Output) aggregateAndFlush() {
 		latencyAvg, latencyP90, latencyP95, latencyP99,
 		rtAvg, rtP90, rtP95, rtP99,
 		rowSizePerReq)
-	o.outputFile.Write([]byte(line))
+	_, _ = o.outputFile.Write([]byte(line))
 }
 
 func average(xs []float64) float64 {

pkg/common/csv.go

@@ -11,9 +11,7 @@ type (
 		Path       string
 		Delimiter  string
 		WithHeader bool
-		divisor    int
-		remainder  int
-		DataCh     chan<- Data
+		limit      int
 	}
 
 	CSVWriter struct {
@@ -24,12 +22,12 @@ type (
 	}
 )
 
-func NewCsvReader(path, delimiter string, withHeader bool, dataCh chan<- Data) *CSVReader {
+func NewCsvReader(path, delimiter string, withHeader bool, limit int) *CSVReader {
 	return &CSVReader{
 		Path:       path,
 		Delimiter:  delimiter,
 		WithHeader: withHeader,
-		DataCh:     dataCh,
+		limit:      limit,
 	}
 }
 
@@ -42,59 +40,50 @@ func NewCsvWriter(path, delimiter string, header []string, dataCh <-chan []strin
 	}
 }
 
-func (c *CSVReader) SetDivisor(divisor int) {
-	if divisor > 0 {
-		c.divisor = divisor
-	}
-}
-
-func (c *CSVReader) SetRemainder(remainder int) {
-	if remainder >= 0 {
-		c.remainder = remainder
-	}
-}
-
-func (c *CSVReader) ReadForever() error {
-	line := 0
+// ReadForever read the csv in slice first, and send to the data channel forever.
+func (c *CSVReader) ReadForever(dataCh chan<- Data) error {
+	lines := make([]Data, 0, c.limit)
 	file, err := os.Open(c.Path)
-	defer file.Close()
 	if err != nil {
 		return err
 	}
-	go func() {
-		file, err := os.Open(c.Path)
-		defer file.Close()
+	defer func() {
+		_ = file.Close()
+	}()
+	reader := csv.NewReader(file)
+	comma := []rune(c.Delimiter)
+	if len(comma) > 0 {
+		reader.Comma = comma[0]
+	}
+	if c.WithHeader {
+		_, err := reader.Read()
 		if err != nil {
-			return
+			return err
 		}
-		reader := csv.NewReader(file)
-		comma := []rune(c.Delimiter)
-		if len(comma) > 0 {
-			reader.Comma = comma[0]
+	}
+	for {
+		row, err := reader.Read()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return err
 		}
-		var offset int64 = 0
-		if c.WithHeader {
-			offset = 1
+
+		lines = append(lines, row)
+		if len(lines) == c.limit {
+			break
 		}
-		file.Seek(offset, 0)
+	}
 
+	go func() {
+		index := 0
 		for {
-			row, err := reader.Read()
-			if err == io.EOF {
-				line = 0
-				file.Seek(offset, 0)
-				row, err = reader.Read()
-			}
-			if err != nil {
-				return
-			}
-			line++
-			if c.divisor > 0 && c.remainder >= 0 {
-				if line%c.divisor != c.remainder {
-					continue
-				}
+			if index == len(lines) {
+				index = 0
 			}
-			c.DataCh <- row
+			dataCh <- lines[index]
+			index++
 		}
 	}()
 	return nil
@@ -103,7 +92,9 @@ func (c *CSVReader) ReadForever() error {
 
 func (c *CSVWriter) WriteForever() error {
 	file, err := os.OpenFile(c.Path, os.O_CREATE|os.O_TRUNC|os.O_RDWR, 0644)
-	defer file.Close()
+	defer func() {
+		_ = file.Close()
+	}()
 	if err != nil {
 		return err
 	}
@@ -112,11 +103,15 @@ func (c *CSVWriter) WriteForever() error {
 	if len(comma) > 0 {
 		w.Comma = comma[0]
 	}
-	w.Write(c.Header)
+	if err := w.Write(c.Header); err != nil {
+		return err
+	}
 	w.Flush()
 	go func() {
 		file, err := os.OpenFile(c.Path, os.O_APPEND|os.O_RDWR, 0644)
-		defer file.Close()
+		defer func() {
+			_ = file.Close()
+		}()
 		if err != nil {
 			return
 		}
@@ -127,7 +122,7 @@ func (c *CSVWriter) WriteForever() error {
 		}
 
 		for {
-			w.Write(<-c.DataCh)
+			_ = w.Write(<-c.DataCh)
 			w.Flush()
 		}
 

pkg/common/types.go

@@ -37,11 +37,14 @@ type (
 		// Init initialize the poop with default channel bufferSize
 		Init(address string, concurrent int) (IGraphClientPool, error)
 		InitWithSize(address string, concurrent int, size int) (IGraphClientPool, error)
-		ConfigCSV(path, delimiter string, withHeader bool) error
+		ConfigCSV(path, delimiter string, withHeader bool, opts ...interface{}) error
 		ConfigOutput(path string) error
-		// ConfigCsvStrategy, csv strategy
-		// 0 means all vus have the same csv reader.
-		// 1 means each vu has a separate csv reader.
+
+		// Deprecated
 		ConfigCsvStrategy(strategy int)
 	}
+
+	ICsvReader interface {
+		ReadForever(dataCh chan<- Data) error
+	}
 )

pkg/nebulagraph/client.go

@@ -27,6 +27,7 @@ type (
 		hosts             []string
 		mutex             sync.Mutex
 		clientGetter      graphClientGetter
+		csvReader         common.ICsvReader
 	}
 
 	graphClientGetter func(endpoint, username, password string) (nebula.GraphClient, error)
@@ -117,9 +118,8 @@ func (gp *GraphPool) InitWithSize(address string, concurrent int, chanSize int)
 	if gp.initialized {
 		return gp, nil
 	}
-	var err error
 
-	err = gp.initAndVerifyPool(address, concurrent, chanSize)
+	err := gp.initAndVerifyPool(address, concurrent, chanSize)
 	if err != nil {
 		return nil, err
 	}
@@ -148,31 +148,30 @@ func (gp *GraphPool) initAndVerifyPool(address string, concurrent int, chanSize
 	return nil
 }
 
-// ConfigCsvStrategy sets csv reader strategy
+// Deprecated ConfigCsvStrategy sets csv reader strategy
 func (gp *GraphPool) ConfigCsvStrategy(strategy int) {
-	gp.csvStrategy = csvReaderStrategy(strategy)
+	return
 }
 
 // ConfigCSV makes the read csv file configuration
-func (gp *GraphPool) ConfigCSV(path, delimiter string, withHeader bool) error {
-	dataCh := gp.DataCh
-	if gp.csvStrategy == AllInOne {
-		reader := common.NewCsvReader(path, delimiter, withHeader, dataCh)
-		if err := reader.ReadForever(); err != nil {
-			return err
-		}
-	} else {
-		// read the csv concurrently
-		l := len(gp.clients)
-		for c := 0; c < l; c++ {
-			reader := common.NewCsvReader(path, delimiter, withHeader, dataCh)
-			reader.SetDivisor(l)
-			reader.SetRemainder(c)
-			if err := reader.ReadForever(); err != nil {
-				return err
-			}
+func (gp *GraphPool) ConfigCSV(path, delimiter string, withHeader bool, opts ...interface{}) error {
+	var (
+		limit int = 500 * 10000
+	)
+	if gp.csvReader != nil {
+		return nil
+	}
+	if len(opts) > 0 {
+		l, ok := opts[0].(int)
+		if ok {
+			limit = l
 		}
 	}
+	gp.csvReader = common.NewCsvReader(path, delimiter, withHeader, limit)
+
+	if err := gp.csvReader.ReadForever(gp.DataCh); err != nil {
+		return err
+	}
 
 	return nil
 }
@@ -280,11 +279,9 @@ func (gc *GraphClient) Execute(stmt string) (common.IGraphResponse, error) {
 	if gc.Pool.OutputCh != nil {
 		var fr []string
 		if rows != 0 {
-			for _, r := range rs.GetRows() {
-				for _, c := range r.GetValues() {
-					fr = append(fr, ValueToString(c))
-				}
-				break
+			r := rs.GetRows()[0]
+			for _, c := range r.GetValues() {
+				fr = append(fr, ValueToString(c))
 			}
 		}
 		o := &output{