add more test

db/workerinfo.go

@@ -81,3 +81,15 @@ func (d *dbWorker) UpdateModel(ctx context.Context, id string, models *types.Mod
 	_, err := d.col.UpdateOne(ctx, bson.M{"_id": id}, update)
 	return err
 }
+
+func (d *dbWorker) UpdateHardware(ctx context.Context, id string, hardware *types.HardwareInfo) error {
+	update := bson.M{"$set": bson.M{"hardware": hardware}}
+	_, err := d.col.UpdateOne(ctx, bson.M{"_id": id}, update)
+	return err
+}
+
+func (d *dbWorker) UpdateNodeInfo(ctx context.Context, id string, nodeInfo *types.NodeInfo) error {
+	update := bson.M{"$set": bson.M{"node_info": nodeInfo}}
+	_, err := d.col.UpdateOne(ctx, bson.M{"_id": id}, update)
+	return err
+}

db/workerinfo_test.go

@@ -2,6 +2,8 @@ package db
 
 import (
 	"context"
+	crand "crypto/rand"
+	"encoding/hex"
 	"fmt"
 	"github.com/google/uuid"
 	"go.mongodb.org/mongo-driver/bson/primitive"
@@ -15,66 +17,27 @@ import (
 )
 
 var (
-	idlist []interface{}
+	idlist = make([]string, 0, 1000000)
+	//workers    = make([]*DbWorkerInfo, 0, 1000000)
+	database   = "test"
+	collection = "workers"
 )
 
-func getRandId(max int) string {
-	return strconv.Itoa(rand.Intn(max) + 1)
+func init() {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
+	}
+	idlist = initdata(client)
 }
 
-func initdata(client *mongo.Client) []interface{} {
+func initdata(client *mongo.Client) []string {
 	t1 := time.Now()
-	// insert 1000000 DbWorkerInfo to db
-
-	idlist := make([]interface{}, 0, 1000000)
-	db := NewDBWorker(client, "test", "workers")
+	db := NewDBWorker(client, database, collection)
 
 	// Insert 1,000,000 DbWorkerInfo to db
 	for i := 0; i < 1000; i++ {
-		worker := &DbWorkerInfo{
-			WorkerId: uuid.NewString(),
-			NodeInfo: &types.NodeInfo{
-				MinerPubkey:    "0x3434332432432231213213213423142132132132132132132132132132132133",
-				BenefitAddress: "0x8888888333338383833838332432432454545444",
-				DeviceIp:       "192.168.1.200",
-			},
-			Models: &types.ModelInfo{
-				InstalledModels: []*types.InstalledModelInfo{
-					&types.InstalledModelInfo{
-						ModelId:  getRandId(100),
-						DiskSize: 100,
-					},
-				},
-				RunningModels: []*types.RunningModelInfo{
-					&types.RunningModelInfo{
-						ModelId:     getRandId(100),
-						StartedTime: uint64((time.Now().Add(-time.Hour)).Unix()),
-						LatestTime:  uint64(time.Now().Add(-time.Minute * 10).Unix()),
-						RunCount:    100,
-						GpuRamUsed:  80,
-					},
-				},
-			},
-			Hardware: &types.HardwareInfo{
-				Devices: []*types.DeviceInfo{
-					&types.DeviceInfo{
-						DeviceType:  "cpu",
-						DeviceModel: "xxxxxxxxx",
-						DeviceParam: "kkkkkkkk",
-						DevicePower: 0,
-					},
-					&types.DeviceInfo{
-						DeviceType:  "gpu",
-						DeviceModel: "Nvidia 3090",
-						DeviceParam: "kkkkkkkk",
-						DevicePower: 100,
-					},
-				},
-				Usages: []*types.DeviceUsage{
-					&types.DeviceUsage{},
-				},
-			},
-		}
+		worker := generateAWroker()
 		result, err := db.InsertWorker(context.Background(), worker)
 		if err != nil {
 			panic(fmt.Sprintf("insert worker failed with err:%s", err))
@@ -92,21 +55,58 @@ func initdata(client *mongo.Client) []interface{} {
 	return idlist
 }
 
-func init() {
-	client, err := ConnectMongoDB("mongodb://localhost:27017")
-	if err != nil {
-		log.Fatal(err)
+func getRandId(max int) string {
+	return strconv.Itoa(rand.Intn(max) + 1)
+}
+
+func generateAWroker() *DbWorkerInfo {
+	return &DbWorkerInfo{
+		WorkerId: uuid.NewString(),
+		NodeInfo: generateANodeInfo(),
+		Models:   generateAmodel(),
+		Hardware: generateAHardware(),
 	}
-	idlist = initdata(client)
 }
 
-func BenchmarkDbWorkerInfo_UpdateResource(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017")
-	if err != nil {
-		log.Fatal(err)
+func generateANodeInfo() *types.NodeInfo {
+	x := make([]byte, 32)
+	y := make([]byte, 32)
+	z := make([]byte, 20)
+	crand.Read(x)
+	crand.Read(y)
+	crand.Read(z)
+
+	return &types.NodeInfo{
+		MinerPubkey:    hex.EncodeToString(x) + hex.EncodeToString(y),
+		BenefitAddress: hex.EncodeToString(z),
+		DeviceIp:       fmt.Sprintf("192.168.1.%d", rand.Intn(255)),
 	}
-	b.ResetTimer()
-	nresource := &types.ModelInfo{
+}
+
+func generateAHardware() *types.HardwareInfo {
+	return &types.HardwareInfo{
+		Devices: []*types.DeviceInfo{
+			&types.DeviceInfo{
+				DeviceType:  "cpu",
+				DeviceModel: "xxxxxxxxx",
+				DeviceParam: "kkkkkkkk",
+				DevicePower: 0,
+			},
+			&types.DeviceInfo{
+				DeviceType:  "gpu",
+				DeviceModel: "Nvidia 3090",
+				DeviceParam: "kkkkkkkk",
+				DevicePower: 100,
+			},
+		},
+		Usages: []*types.DeviceUsage{
+			&types.DeviceUsage{},
+		},
+	}
+}
+
+func generateAmodel() *types.ModelInfo {
+	return &types.ModelInfo{
 		InstalledModels: []*types.InstalledModelInfo{
 			&types.InstalledModelInfo{
 				ModelId:  getRandId(100),
@@ -123,54 +123,150 @@ func BenchmarkDbWorkerInfo_UpdateResource(b *testing.B) {
 				StartedTime: uint64((time.Now().Add(-time.Hour)).Unix()),
 				LatestTime:  uint64(time.Now().Add(-time.Minute * 20).Unix()),
 				RunCount:    101,
-				GpuRamUsed:  80,
+				GpuRamUsed:  uint64(rand.Intn(20) + 10),
 			},
 		},
 	}
 
-	db := NewDBWorker(client, "test", "workers")
+}
+
+func BenchmarkGenerateWorker(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		generateAWroker()
+	}
+}
+
+func BenchmarkDbWorker_InsertWorker(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
+	}
+	workers := make([]*DbWorkerInfo, 0, b.N)
+	for i := 0; i < b.N; i++ {
+		workers = append(workers, generateAWroker())
+	}
+	db := NewDBWorker(client, database, collection)
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		worker := workers[i]
+		if _, err := db.InsertWorker(context.Background(), worker); err != nil {
+			panic(fmt.Sprintf("insert worker failed with err:%s", err))
+		}
+	}
+}
+
+func BenchmarkDbWorker_InsertWorker_Parallel(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorker(client, database, collection)
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			worker := generateAWroker()
+			if _, err := db.InsertWorker(context.Background(), worker); err != nil {
+				panic(fmt.Sprintf("insert worker failed with err:%s", err))
+			}
+		}
+	})
+}
+
+func BenchmarkDbWorker_UpdateHardware(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
+	}
+	b.ResetTimer()
+	nhardware := generateAHardware()
+
+	db := NewDBWorker(client, database, collection)
 	for i := 0; i < b.N; i++ {
 		idx := rand.Intn(len(idlist))
-		if err := db.UpdateModel(context.Background(), idlist[idx].(string), nresource); err != nil {
+		if err := db.UpdateHardware(context.Background(), idlist[idx], nhardware); err != nil {
 			panic(fmt.Sprintf("update worker failed with err:%s", err))
-			//ufailed++
 		}
 	}
-	//fmt.Println("update failed:", ufailed)
 }
 
-func BenchmarkDbWorkerInfo_UpdateResource_Parallel(b *testing.B) {
+func BenchmarkDbWorker_UpdateHardware_Parallel(b *testing.B) {
 	client, err := ConnectMongoDB("mongodb://localhost:27017")
 	if err != nil {
 		log.Fatal(err)
 	}
-	nresource := &types.ModelInfo{
-		InstalledModels: []*types.InstalledModelInfo{
-			&types.InstalledModelInfo{
-				ModelId:  getRandId(100),
-				DiskSize: 101,
-			},
-			&types.InstalledModelInfo{
-				ModelId:  getRandId(100),
-				DiskSize: 44,
-			},
-		},
-		RunningModels: []*types.RunningModelInfo{
-			&types.RunningModelInfo{
-				ModelId:     getRandId(100),
-				StartedTime: uint64((time.Now().Add(-time.Hour)).Unix()),
-				LatestTime:  uint64(time.Now().Add(-time.Minute * 20).Unix()),
-				RunCount:    101,
-				GpuRamUsed:  80,
-			},
-		},
+	nhardware := generateAHardware()
+	db := NewDBWorker(client, database, collection)
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			idx := rand.Intn(len(idlist))
+			if err := db.UpdateHardware(context.Background(), idlist[idx], nhardware); err != nil {
+				panic(fmt.Sprintf("update worker failed with err:%s", err))
+			}
+		}
+	})
+}
+
+func BenchmarkDbWorker_UpdateModel(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
+	}
+	b.ResetTimer()
+	nresource := generateAmodel()
+
+	db := NewDBWorker(client, database, collection)
+	for i := 0; i < b.N; i++ {
+		idx := rand.Intn(len(idlist))
+		if err := db.UpdateModel(context.Background(), idlist[idx], nresource); err != nil {
+			panic(fmt.Sprintf("update worker failed with err:%s", err))
+		}
+	}
+}
+
+func BenchmarkDbWorker_UpdateModel_Parallel(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
 	}
-	db := NewDBWorker(client, "test", "workers")
+	nresource := generateAmodel()
+	db := NewDBWorker(client, database, collection)
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			idx := rand.Intn(len(idlist))
+			if err := db.UpdateModel(context.Background(), idlist[idx], nresource); err != nil {
+				panic(fmt.Sprintf("update worker failed with err:%s", err))
+			}
+		}
+	})
+}
 
+func BenchmarkDbWorker_UpdateNodeInfo(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
+	}
+	b.ResetTimer()
+	nnodeinfo := generateANodeInfo()
+
+	db := NewDBWorker(client, database, collection)
+	for i := 0; i < b.N; i++ {
+		idx := rand.Intn(len(idlist))
+		if err := db.UpdateNodeInfo(context.Background(), idlist[idx], nnodeinfo); err != nil {
+			panic(fmt.Sprintf("update worker failed with err:%s", err))
+		}
+	}
+}
+
+func BenchmarkDbWorker_UpdateNodeInfo_Parallel(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
+	}
+	nnodeinfo := generateANodeInfo()
+	db := NewDBWorker(client, database, collection)
 	b.RunParallel(func(pb *testing.PB) {
 		for pb.Next() {
 			idx := rand.Intn(len(idlist))
-			if err := db.UpdateModel(context.Background(), idlist[idx].(string), nresource); err != nil {
+			if err := db.UpdateNodeInfo(context.Background(), idlist[idx], nnodeinfo); err != nil {
 				panic(fmt.Sprintf("update worker failed with err:%s", err))
 			}
 		}