update workerInstallinfo and test case

operator/workerInstalledInfo.go

@@ -27,6 +27,13 @@ func NewDBWorkerInstalled(client *mongo.Client, database string) *WorkerInstalle
 		col:    client.Database(database).Collection("worker_installed_info"),
 	}
 }
+func (d WorkerInstalledOperator) CreateIndex(ctx context.Context) error {
+	_, err := d.col.Indexes().CreateOne(ctx, mongo.IndexModel{
+		Keys: bson.D{{"model_id", 1}, {"gpu_free", 1}},
+		//Options: options.Index().SetUnique(true),
+	})
+	return err
+}
 
 func (d *WorkerInstalledOperator) InsertMany(ctx context.Context, workers []*WorkerInstalledInfo) (*mongo.InsertManyResult, error) {
 	var docs []interface{}

operator/workerInstalledInfo_test.go

@@ -331,16 +331,28 @@ func BenchmarkWorkerInstalledOperator_FindWorkerByModelIdAndGpuMem(b *testing.B)
 	if err != nil {
 		log.Fatal(err)
 	}
+	b.StopTimer()
 	db := NewDBWorkerInstalled(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, true, true)
+	})
+	b.StartTimer()
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		id := rand.Intn(maxModelId)
-		mem := rand.Intn(100) + 12000
-		_, err := db.FindWorkerByModelIdAndGpuMem(context.Background(), id, int64(mem), 10)
+		id := rand.Intn(maxModelId) + 1
+		mem := 100
+		find, err := db.FindWorkerByModelIdAndGpuMem(context.Background(), id, int64(mem), 10)
 		if err != nil {
 			panic(fmt.Sprintf("find worker failed with err:%s", err))
 		}
+		if len(find) == 0 {
+			b.Logf("find worker failed with err: not found with id %d", id)
+		}
 	}
 }
 
@@ -349,17 +361,29 @@ func BenchmarkWorkerInstalledOperator_FindWorkerByModelIdAndGpuMem_Parallel(b *t
 	if err != nil {
 		log.Fatal(err)
 	}
+	b.StopTimer()
 	db := NewDBWorkerInstalled(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, true, true)
+	})
+	b.StartTimer()
 	b.ResetTimer()
 	b.RunParallel(func(pb *testing.PB) {
 		for pb.Next() {
-			id := rand.Intn(maxModelId)
-			mem := rand.Intn(100) + 12000
-			_, err := db.FindWorkerByModelIdAndGpuMem(context.Background(), id, int64(mem), 10)
+			id := rand.Intn(maxModelId) + 1
+			mem := 100
+			find, err := db.FindWorkerByModelIdAndGpuMem(context.Background(), id, int64(mem), 10)
 			if err != nil {
 				panic(fmt.Sprintf("find worker failed with err:%s", err))
 			}
+			if len(find) == 0 {
+				b.Logf("find worker failed with err: not found with id %d", id)
+			}
 		}
 	})