update benchmark test

operator/workerinfo_test.go

@@ -21,6 +21,7 @@ import (
 var (
 	maxModelId   = 1000
 	idlist       = make([]string, 0, 1000000)
+	workeridList = make([]string, 0, 1000000)
 	database     = "test"
 	once         = sync.Once{}
 )
@@ -41,15 +42,20 @@ func ConnectMongoDB() (*mongo.Client, error) {
 	return client, nil
 }
 
-func initdata(client *mongo.Client, count int, running bool, installed bool) []string {
+func initdata(client *mongo.Client, count int, running bool, installed bool, workerid bool) []string {
 	t1 := time.Now()
 	db := NewDBWorker(client, database)
 	dbRunning := NewDBWorkerRunning(client, database)
 	dbInstalled := NewDBWorkerInstalled(client, database)
 
+	ids := make([]string, 0, count)
+
 	// Insert 1,000,000 DbWorkerInfo to operator
 	for i := 0; i < count; i++ {
 		worker := generateAWorker()
+		if workerid {
+			ids = append(ids, worker.WorkerId)
+		}
 		result, err := db.InsertWorker(context.Background(), worker)
 		if err != nil {
 			panic(fmt.Sprintf("insert worker failed with err:%s", err))
@@ -88,16 +94,18 @@ func initdata(client *mongo.Client, count int, running bool, installed bool) []s
 			}
 		}
 
+		if !workerid {
 			id, ok := result.InsertedID.(primitive.ObjectID)
 			if !ok {
 				panic("inserted id is not primitive.ObjectID")
 			}
-		idlist = append(idlist, id.Hex())
+			ids = append(ids, id.Hex())
+		}
 		//fmt.Printf("insert worker %s: %v\n", id.Hex(), worker)
 	}
 	t2 := time.Now()
 	fmt.Printf("init data cost %s\n", t2.Sub(t1).String())
-	return idlist
+	return ids
 }
 
 func getRandId(max int) string {
@@ -568,7 +576,7 @@ func BenchmarkDbWorker_UpdateHardware(b *testing.B) {
 		if err := db.CreateIndex(context.Background()); err != nil {
 			panic(fmt.Sprintf("create index failed with err:%s", err))
 		}
-		idlist = initdata(client, 10000, false, false)
+		idlist = initdata(client, 10000, false, false, true)
 	})
 	b.StartTimer()
 
@@ -599,7 +607,7 @@ func BenchmarkDbWorker_UpdateHardware_Parallel(b *testing.B) {
 		if err := db.CreateIndex(context.Background()); err != nil {
 			panic(fmt.Sprintf("create index failed with err:%s", err))
 		}
-		idlist = initdata(client, 10000, false, false)
+		idlist = initdata(client, 10000, false, false, true)
 	})
 
 	b.StartTimer()
@@ -618,6 +626,45 @@ func BenchmarkDbWorker_UpdateHardware_Parallel(b *testing.B) {
 	})
 }
 
+func BenchmarkDbWorker_UpdateGPUUsage_Parallel(b *testing.B) {
+	client, err := ConnectMongoDB()
+	if err != nil {
+		log.Fatal(err)
+	}
+	b.StopTimer()
+	db := NewDBWorker(client, database)
+	defer db.client.Disconnect(context.Background())
+	once.Do(func() {
+		db.Clear()
+		if err := db.CreateIndex(context.Background()); err != nil {
+			panic(fmt.Sprintf("create index failed with err:%s", err))
+		}
+		idlist = initdata(client, 10000, false, false, true)
+	})
+	b.StartTimer()
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			b.StopTimer()
+			idx := rand.Intn(len(idlist))
+			usage := rand.Intn(10) + 50
+			memfree := int64(rand.Intn(10) + 10)
+			pwoer := rand.Intn(10) + 20
+			temp := rand.Intn(10) + 30
+			gpuusages := []types.GpuUsage{
+				{0, usage, memfree, pwoer, temp},
+				{1, usage + 1, memfree + 1, pwoer + 1, temp + 1},
+				{2, usage + 2, memfree + 2, pwoer + 2, temp + 2},
+			}
+			b.StartTimer()
+			if err := db.UpdateGPUUsage(context.Background(), idlist[idx], gpuusages); err != nil {
+				panic(fmt.Sprintf("update gpu usage failed with err:%s", err))
+			}
+		}
+	})
+
+}
+
 func BenchmarkDbWorker_FindWorkerByInstallModelAndSortByGpuRam(b *testing.B) {
 	client, err := ConnectMongoDB()
 	if err != nil {
@@ -632,7 +679,7 @@ func BenchmarkDbWorker_FindWorkerByInstallModelAndSortByGpuRam(b *testing.B) {
 		if err := db.CreateIndex(context.Background()); err != nil {
 			panic(fmt.Sprintf("create index failed with err:%s", err))
 		}
-		idlist = initdata(client, 10000, false, false)
+		idlist = initdata(client, 10000, false, false, true)
 	})
 	b.StartTimer()
 	b.ResetTimer()
@@ -664,7 +711,7 @@ func BenchmarkDbWorker_FindWorkerByInstallModelAndSortByGpuRam_Parallel(b *testi
 		if err := db.CreateIndex(context.Background()); err != nil {
 			panic(fmt.Sprintf("create index failed with err:%s", err))
 		}
-		idlist = initdata(client, 10000, false, false)
+		idlist = initdata(client, 10000, false, false, true)
 	})
 	b.StartTimer()
 	b.ResetTimer()
@@ -700,7 +747,7 @@ func BenchmarkDbWorker_FindWorkerByRunningModelAndSortByWaitTime(b *testing.B) {
 		if err := db.CreateIndex(context.Background()); err != nil {
 			panic(fmt.Sprintf("create index failed with err:%s", err))
 		}
-		idlist = initdata(client, 1000, false, false)
+		idlist = initdata(client, 1000, false, false, true)
 	})
 	b.StartTimer()
 	b.ResetTimer()
@@ -730,7 +777,7 @@ func BenchmarkDbWorker_FindWorkerByRunningModelAndSortByWaitTime_Parallel(b *tes
 		if err := db.CreateIndex(context.Background()); err != nil {
 			panic(fmt.Sprintf("create index failed with err:%s", err))
 		}
-		idlist = initdata(client, 1000, false, false)
+		idlist = initdata(client, 1000, false, false, true)
 	})
 	b.StartTimer()
 	b.ResetTimer()