From 105b64decf2595a27326eaff793467766f9201cc Mon Sep 17 00:00:00 2001
From: bigsheeper <yihao.dai@zilliz.com>
Date: Sun, 27 Sep 2020 14:43:04 +0800
Subject: [PATCH] Add write node benchmark

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
---
 writer/main.go                   | 30 +++++++++++++++++-------------
 writer/write_node/writer_node.go |  7 ++++++-
 2 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/writer/main.go b/writer/main.go
index f862775c8..878f127d4 100644
--- a/writer/main.go
+++ b/writer/main.go
@@ -45,11 +45,12 @@ func main() {
 	}
 
 	const Debug = true
-	const CountMsgNum = 1000 * 1000
+	const MB = 1024 * 1024
+	const timeInterval = time.Second * 5
+	const CountMsgNum = 10000 * 10
 
 	if Debug {
-		var printFlag = true
-		var startTime = true
+		var shouldBenchmark = false
 		var start time.Time
 
 		for {
@@ -57,22 +58,25 @@ func main() {
 				break
 			}
 			msgLength := wn.MessageClient.PrepareBatchMsg()
+			// wait until first 100,000 rows are successfully wrote
+			if wn.MsgCounter.InsertCounter >= CountMsgNum {
+				shouldBenchmark = true
+				start = time.Now()
+			}
 			if msgLength > 0 {
-				if startTime {
-					fmt.Println("============> Start Test <============")
-					startTime = false
-					start = time.Now()
-				}
-
 				wn.DoWriteNode(ctx, &wg)
 				fmt.Println("write node do a batch message, storage len: ", msgLength)
 			}
-
 			// Test insert time
-			if printFlag && wn.MsgCounter.InsertCounter >= CountMsgNum {
-				printFlag = false
+			// ignore if less than 1000 records per time interval
+			if shouldBenchmark && wn.MsgCounter.InsertCounter > 1000{
 				timeSince := time.Since(start)
-				fmt.Println("============> Do", wn.MsgCounter.InsertCounter, "Insert in", timeSince, "<============")
+				if timeSince >= timeInterval {
+					speed := wn.MsgCounter.InsertedRecordSize / timeInterval.Seconds() / MB
+					fmt.Println("============> Insert", wn.MsgCounter.InsertCounter,"records, cost:" , timeSince, "speed:", speed, "M/s", "<============")
+					wn.MsgCounter.InsertCounter = 0
+					start = time.Now()
+				}
 			}
 		}
 	}
diff --git a/writer/write_node/writer_node.go b/writer/write_node/writer_node.go
index f74e9dbbb..c01058431 100644
--- a/writer/write_node/writer_node.go
+++ b/writer/write_node/writer_node.go
@@ -19,6 +19,7 @@ type SegmentIdInfo struct {
 
 type MsgCounter struct {
 	InsertCounter int64
+	InsertedRecordSize float64
 	DeleteCounter int64
 }
 
@@ -43,6 +44,7 @@ func NewWriteNode(ctx context.Context,
 	msgCounter := MsgCounter{
 		InsertCounter: 0,
 		DeleteCounter: 0,
+		InsertedRecordSize: 0,
 	}
 
 	return &WriteNode{
@@ -71,7 +73,10 @@ func (wn *WriteNode) InsertBatchData(ctx context.Context, data []*msgpb.InsertOr
 	}
 
 	wn.MsgCounter.InsertCounter += int64(len(timeStamp))
-
+	if len(timeStamp) > 0 {
+		// assume each record is same size
+		wn.MsgCounter.InsertedRecordSize += float64(len(timeStamp) * len(data[0].RowsData.Blob))
+	}
 	error := (*wn.KvStore).PutRows(ctx, prefixKeys, binaryData, suffixKeys, timeStamp)
 	if error != nil {
 		fmt.Println("Can't insert data!")
-- 
GitLab