Commit 0ff8d5c5 authored by luxq's avatar luxq

update test

parent fd96d122
...@@ -10,7 +10,7 @@ import ( ...@@ -10,7 +10,7 @@ import (
type DbWorkerRunningInfo struct { type DbWorkerRunningInfo struct {
ID string `bson:"_id,omitempty" json:"id"` ID string `bson:"_id,omitempty" json:"id"`
WorkerId string `bson:"worker_id" json:"worker_id"` WorkerId string `bson:"worker_id" json:"worker_id"`
ModelId string `bson:"model_id" json:"model_id"` ModelId int `bson:"model_id" json:"model_id"`
WaitTime int `bson:"wait_time" json:"wait_time"` WaitTime int `bson:"wait_time" json:"wait_time"`
} }
...@@ -36,7 +36,7 @@ func (d *dbWorkerRunning) UpdateWaitTime(ctx context.Context, id string, waitTim ...@@ -36,7 +36,7 @@ func (d *dbWorkerRunning) UpdateWaitTime(ctx context.Context, id string, waitTim
return err return err
} }
func (d *dbWorkerRunning) FindWorkerByModelId(ctx context.Context, modelId string, limit int) ([]*DbWorkerRunningInfo, error) { func (d *dbWorkerRunning) FindWorkerByModelId(ctx context.Context, modelId int, limit int) ([]*DbWorkerRunningInfo, error) {
// find all worker that at least one running model's mode_id is equal modelId // find all worker that at least one running model's mode_id is equal modelId
// sort by wait time // sort by wait time
findOptions := options.Find() findOptions := options.Find()
......
...@@ -47,9 +47,10 @@ func initdata(client *mongo.Client) []string { ...@@ -47,9 +47,10 @@ func initdata(client *mongo.Client) []string {
{ {
// add worker running info to dbRunning // add worker running info to dbRunning
for _, model := range worker.Models.RunningModels { for _, model := range worker.Models.RunningModels {
id, _ := strconv.Atoi(model.ModelID)
runningInfo := &DbWorkerRunningInfo{ runningInfo := &DbWorkerRunningInfo{
WorkerId: worker.WorkerId, WorkerId: worker.WorkerId,
ModelId: model.ModelID, ModelId: id,
WaitTime: model.WaitTime, WaitTime: model.WaitTime,
} }
_, err := dbRunning.InsertWorker(context.Background(), runningInfo) _, err := dbRunning.InsertWorker(context.Background(), runningInfo)
...@@ -75,6 +76,10 @@ func getRandId(max int) string { ...@@ -75,6 +76,10 @@ func getRandId(max int) string {
return strconv.Itoa(rand.Intn(max) + 1) return strconv.Itoa(rand.Intn(max) + 1)
} }
func getRandIdInt(max int) int {
return rand.Intn(max) + 1
}
func generateAWroker() *DbWorkerInfo { func generateAWroker() *DbWorkerInfo {
return &DbWorkerInfo{ return &DbWorkerInfo{
WorkerId: uuid.NewString(), WorkerId: uuid.NewString(),
...@@ -462,11 +467,11 @@ func BenchmarkDbWorkerRunning_FindWorkerByModelId(b *testing.B) { ...@@ -462,11 +467,11 @@ func BenchmarkDbWorkerRunning_FindWorkerByModelId(b *testing.B) {
if err != nil { if err != nil {
log.Fatal(err) log.Fatal(err)
} }
db := NewDBWorkerRunning(client, database, collection) db := NewDBWorkerRunning(client, database, workerRunningCollection)
defer db.client.Disconnect(context.Background()) defer db.client.Disconnect(context.Background())
b.ResetTimer() b.ResetTimer()
for i := 0; i < b.N; i++ { for i := 0; i < b.N; i++ {
modelId := getRandId(100) modelId := getRandIdInt(100)
if w, err := db.FindWorkerByModelId(context.Background(), modelId, 10); err != nil { if w, err := db.FindWorkerByModelId(context.Background(), modelId, 10); err != nil {
panic(fmt.Sprintf("find worker failed with err:%s", err)) panic(fmt.Sprintf("find worker failed with err:%s", err))
} else if len(w) == 0 { } else if len(w) == 0 {
...@@ -484,7 +489,7 @@ func BenchmarkDbWorkerRunning_FindWorkerByModelId_Parallel(b *testing.B) { ...@@ -484,7 +489,7 @@ func BenchmarkDbWorkerRunning_FindWorkerByModelId_Parallel(b *testing.B) {
defer db.client.Disconnect(context.Background()) defer db.client.Disconnect(context.Background())
b.RunParallel(func(pb *testing.PB) { b.RunParallel(func(pb *testing.PB) {
for pb.Next() { for pb.Next() {
modelId := getRandId(100) modelId := getRandIdInt(100)
if w, err := db.FindWorkerByModelId(context.Background(), modelId, 10); err != nil { if w, err := db.FindWorkerByModelId(context.Background(), modelId, 10); err != nil {
panic(fmt.Sprintf("find worker failed with err:%s", err)) panic(fmt.Sprintf("find worker failed with err:%s", err))
} else if len(w) == 0 { } else if len(w) == 0 {
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment