update test code

operator/workerInstalledInfo.go

@@ -22,7 +22,7 @@ type WorkerInstalledOperator struct {
 func NewDBWorkerInstalled(client *mongo.Client, database string) *WorkerInstalledOperator {
 	return &WorkerInstalledOperator{
 		client: client,
-		col:    client.Database(database).Collection("worker_running_info"),
+		col:    client.Database(database).Collection("worker_installed_info"),
 	}
 }
 
@@ -30,9 +30,9 @@ func (d *WorkerInstalledOperator) Insert(ctx context.Context, worker *WorkerInst
 	return d.col.InsertOne(ctx, worker)
 }
 
-func (d *WorkerInstalledOperator) UpdateGpuFree(ctx context.Context, id string, gpuFree int64) error {
+func (d *WorkerInstalledOperator) UpdateGpuFree(ctx context.Context, workerid string, gpuFree int64) error {
 	update := bson.M{"$set": bson.M{"gpu_free": gpuFree}}
-	_, err := d.col.UpdateOne(ctx, bson.M{"_id": id}, update)
+	_, err := d.col.UpdateMany(ctx, bson.M{"worker_id": workerid}, update)
 	return err
 }
 

operator/workerInstalledInfo_test.go

+package operator
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"math/rand"
+	"strconv"
+	"testing"
+)
+
+func TestWorkerInstalledOperator_Insert(t *testing.T) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	if err != nil {
+		log.Fatal(err)
+	}
+	count := 1
+	infodb := NewDBWorker(client, database)
+	defer infodb.client.Disconnect(context.Background())
+	db := NewDBWorkerInstalled(client, database)
+	defer db.client.Disconnect(context.Background())
+	rundb := NewDBWorkerRunning(client, database)
+	defer rundb.client.Disconnect(context.Background())
+
+	for i := 0; i < count; i++ {
+		worker := generateAWroker()
+		infodb.InsertWorker(context.Background(), worker)
+		for _, installed := range worker.Models.InstalledModels {
+			id, _ := strconv.Atoi(installed.ModelID)
+			_, err := db.Insert(context.Background(), &WorkerInstalledInfo{
+				WorkerId: worker.WorkerId,
+				ModelId:  id,
+				GpuFree:  10011,
+			})
+			if err != nil {
+				panic(fmt.Sprintf("insert worker failed with err:%s", err))
+			}
+		}
+
+		for _, running := range worker.Models.RunningModels {
+			id, _ := strconv.Atoi(running.ModelID)
+			_, err := rundb.Insert(context.Background(), &WorkerRunningInfo{
+				WorkerId: worker.WorkerId,
+				ModelId:  id,
+				ExecTime: 100,
+			})
+			if err != nil {
+				panic(fmt.Sprintf("insert worker failed with err:%s", err))
+			}
+		}
+	}
+}
+
+func BenchmarkWorkerInstalledOperator_UpdateGpuFree(b *testing.B) {
+
+	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	if err != nil {
+		log.Fatal(err)
+	}
+	infodb := NewDBWorker(client, database)
+	defer infodb.client.Disconnect(context.Background())
+	db := NewDBWorkerInstalled(client, database)
+	defer db.client.Disconnect(context.Background())
+
+	allWorker, err := infodb.GetAllWorkerId(context.Background())
+	if err != nil {
+		panic(fmt.Sprintf("get all worker failed with err:%s", err))
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		idx := rand.Intn(len(allWorker))
+		workerid := allWorker[idx]
+		gpufree := rand.Intn(100) + 12000
+		if err := db.UpdateGpuFree(context.Background(), workerid, int64(gpufree)); err != nil {
+			panic(fmt.Sprintf("update worker failed with err:%s", err))
+		}
+	}
+}
+
+func BenchmarkWorkerInstalledOperator_UpdateGpuFree_Parallel(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	if err != nil {
+		log.Fatal(err)
+	}
+	infodb := NewDBWorker(client, database)
+	defer infodb.client.Disconnect(context.Background())
+	db := NewDBWorkerInstalled(client, database)
+	defer db.client.Disconnect(context.Background())
+
+	allWorker, err := infodb.GetAllWorkerId(context.Background())
+	if err != nil {
+		panic(fmt.Sprintf("get all worker failed with err:%s", err))
+
+	}
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			idx := rand.Intn(len(allWorker))
+			workerid := allWorker[idx]
+			gpufree := rand.Intn(100) + 12000
+			if err := db.UpdateGpuFree(context.Background(), workerid, int64(gpufree)); err != nil {
+				panic(fmt.Sprintf("update worker failed with err:%s", err))
+			}
+		}
+	})
+}
+
+func TestWorkerInstalledOperator_UpdateGpuFree(t *testing.T) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerInstalled(client, database)
+	defer db.client.Disconnect(context.Background())
+	workerid := "ebb6a2df-9e91-4b3c-ace8-748aaeca718c"
+	gpufree := 102222
+	if err := db.UpdateGpuFree(context.Background(), workerid, int64(gpufree)); err != nil {
+		panic(fmt.Sprintf("update worker failed with err:%s", err))
+	}
+}

operator/workerinfo.go

@@ -28,6 +28,26 @@ func NewDBWorker(client *mongo.Client, database string) *WorkerInfoOperator {
 	}
 }
 
+func (d *WorkerInfoOperator) GetAllWorkerId(ctx context.Context) ([]string, error) {
+	filter := bson.D{}
+	opts := options.Find().SetProjection(bson.D{{"worker_id", 1}})
+
+	cursor, err := d.col.Find(ctx, filter, opts)
+	if err != nil {
+		return nil, err
+	}
+	defer cursor.Close(ctx)
+	var results []WorkerInfo
+	if err = cursor.All(context.TODO(), &results); err != nil {
+		return nil, err
+	}
+	var workerIds []string
+	for _, result := range results {
+		workerIds = append(workerIds, result.WorkerId)
+	}
+	return workerIds, err
+}
+
 func (d *WorkerInfoOperator) InsertWorker(ctx context.Context, worker *WorkerInfo) (*mongo.InsertOneResult, error) {
 	return d.col.InsertOne(ctx, worker)
 }

operator/workerinfo_test.go

@@ -6,12 +6,12 @@ import (
 	"encoding/hex"
 	"fmt"
 	"github.com/google/uuid"
+	"github.com/odysseus/mogo/types"
 	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
 	"log"
 	"math/rand"
-	"mogo/types"
 	"strconv"
 	"testing"
 	"time"
@@ -21,9 +21,10 @@ var (
 	maxModelId = 10000
 	idlist     = make([]string, 0, 1000000)
 	//workers    = make([]*DbWorkerInfo, 0, 1000000)
-	database                = "test"
-	collection              = "workers"
-	workerRunningCollection = "worker_running"
+	database                  = "test"
+	collection                = "workers"
+	workerRunningCollection   = "worker_running"
+	workerInstalledCollection = "worker_installed"
 )
 
 func ConnectMongoDB(uri string, username, passwd string) (*mongo.Client, error) {
@@ -39,13 +40,13 @@ func ConnectMongoDB(uri string, username, passwd string) (*mongo.Client, error)
 	return client, nil
 }
 
-func init() {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
-	if err != nil {
-		log.Fatal(err)
-	}
-	idlist = initdata(client)
-}
+//func init() {
+//	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+//	if err != nil {
+//		log.Fatal(err)
+//	}
+//	idlist = initdata(client)
+//}
 
 func initdata(client *mongo.Client) []string {
 	t1 := time.Now()
@@ -202,6 +203,11 @@ func generateAHardware() *types.HardwareInfo {
 			generateAIdleGpu(0),
 			generateAIdleGpu(1),
 			generateAUsageGpu(2),
+			generateAUsageGpu(3),
+			generateAUsageGpu(4),
+			generateAUsageGpu(5),
+			generateAUsageGpu(6),
+			generateAUsageGpu(7),
 		},
 		RAM:  generateARam(),
 		DISK: generateADisk(),
@@ -235,9 +241,9 @@ func generateAModel() *types.ModelInfo {
 		InstalledModels: make([]*types.InstalledModel, 0, 1000),
 		RunningModels:   make([]*types.RunningModel, 0, 1000),
 	}
-	for i := 0; i < 500; i++ {
+	for i := 0; i < 100; i++ {
 		m.InstalledModels = append(m.InstalledModels, generateAInstallModel())
-		if len(m.RunningModels) < 500 {
+		if len(m.RunningModels) < 100 {
 			m.RunningModels = append(m.RunningModels, generateARunningModel())
 		}
 	}