add sortset

go.mod

@@ -3,12 +3,14 @@ module mogo
 go 1.20
 
 require (
-	github.com/docker/docker v26.1.1+incompatible
 	github.com/google/uuid v1.6.0
+	github.com/redis/go-redis/v9 v9.5.1
 	go.mongodb.org/mongo-driver v1.15.0
 )
 
 require (
+	github.com/cespare/xxhash/v2 v2.2.0 // indirect
+	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/golang/snappy v0.0.1 // indirect
 	github.com/klauspost/compress v1.13.6 // indirect
 	github.com/montanaflynn/stats v0.0.0-20171201202039-1bf9dbcd8cbe // indirect

go.sum

+github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
+github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
+github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
+github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
-github.com/docker/docker v26.1.1+incompatible h1:oI+4kkAgIwwb54b9OC7Xc3hSgu1RlJA/Lln/DF72djQ=
-github.com/docker/docker v26.1.1+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
 github.com/golang/snappy v0.0.1 h1:Qgr9rKW7uDUkrbSmQeiDsGa8SjGyCOGtuasMWwvp2P4=
 github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/google/go-cmp v0.5.2 h1:X2ev0eStA3AbceY54o37/0PQ/UWqKEiiO2dKL5OPaFM=
@@ -10,6 +14,8 @@ github.com/klauspost/compress v1.13.6 h1:P76CopJELS0TiO2mebmnzgWaajssP/EszplttgQ
 github.com/klauspost/compress v1.13.6/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
 github.com/montanaflynn/stats v0.0.0-20171201202039-1bf9dbcd8cbe h1:iruDEfMl2E6fbMZ9s0scYfZQ84/6SPL6zC8ACM2oIL0=
 github.com/montanaflynn/stats v0.0.0-20171201202039-1bf9dbcd8cbe/go.mod h1:wL8QJuTMNUDYhXwkmfOly8iTdp5TEcJFWZD2D7SIkUc=
+github.com/redis/go-redis/v9 v9.5.1 h1:H1X4D3yHPaYrkL5X06Wh6xNVM/pX0Ft4RV0vMGvLBh8=
+github.com/redis/go-redis/v9 v9.5.1/go.mod h1:hdY0cQFCN4fnSYT6TkisLufl/4W5UIXyv0b/CLO2V2M=
 github.com/xdg-go/pbkdf2 v1.0.0 h1:Su7DPu48wXMwC3bs7MCNG+z4FhcyEuz5dlvchbq0B0c=
 github.com/xdg-go/pbkdf2 v1.0.0/go.mod h1:jrpuAogTd400dnrH08LKmI/xc1MbPOebTwRqcT5RDeI=
 github.com/xdg-go/scram v1.1.2 h1:FHX5I5B4i4hKRVRBCFRxq1iQRej7WO3hhBuJf+UUySY=

sortset/sortset.go

+package sortset
+
+import (
+	"context"
+	"fmt"
+	"github.com/redis/go-redis/v9"
+)
+
+type WorkerInfo struct {
+	WorkerId string
+	ModelId  string
+	WaitTime int
+}
+
+func getKey(modelId string) string {
+	return fmt.Sprintf("model:%s", modelId)
+}
+
+type SortSet struct {
+	client *redis.Client
+}
+
+func NewSortSet(client *redis.Client) *SortSet {
+	return &SortSet{
+		client: client,
+	}
+}
+
+func (s *SortSet) AddWorker(worker *WorkerInfo) error {
+	key := getKey(worker.ModelId)
+	_, err := s.client.ZAdd(context.Background(), key, redis.Z{Score: float64(worker.WaitTime), Member: worker.WorkerId}).Result()
+	return err
+}
+
+func (s *SortSet) FindWorker(modelid string, limit int) ([]*WorkerInfo, error) {
+	key := getKey(modelid)
+	workers, err := s.client.ZRevRangeWithScores(context.Background(), key, 0, int64(limit)).Result()
+	if err != nil {
+		return nil, err
+	}
+	var result []*WorkerInfo
+	for _, worker := range workers {
+		result = append(result, &WorkerInfo{
+			WorkerId: worker.Member.(string),
+			ModelId:  modelid,
+			WaitTime: int(worker.Score),
+		})
+	}
+	return result, nil
+}

sortset/sortset_test.go

+package sortset
+
+import (
+	"fmt"
+	"github.com/google/uuid"
+	"github.com/redis/go-redis/v9"
+	"math/rand"
+	"strconv"
+	"testing"
+)
+
+func getRandId(max int) string {
+	return strconv.Itoa(rand.Intn(max) + 1)
+}
+
+func initData(id string) []*WorkerInfo {
+	w := make([]*WorkerInfo, 100000)
+	for i := 0; i < 100000; i++ {
+		w[i] = &WorkerInfo{
+			WorkerId: uuid.NewString(),
+			ModelId:  id,
+			WaitTime: rand.Intn(100),
+		}
+	}
+	return w
+}
+func BenchmarkSortSet_AddWorker(b *testing.B) {
+	client := redis.NewClient(&redis.Options{
+		Addr: "0.0.0.0:6379",
+	})
+	sortSet := NewSortSet(client)
+	mid := "13"
+	w := initData(mid)
+	b.ResetTimer()
+	// start benchmark
+	for i := 0; i < b.N; i++ {
+		if err := sortSet.AddWorker(w[i%100000]); err != nil {
+			panic(fmt.Sprintf("add worker failed: %v", err))
+		}
+	}
+}
+
+func BenchmarkSortSet_FindWorker(b *testing.B) {
+	client := redis.NewClient(&redis.Options{
+		Addr: "0.0.0.0:6379",
+	})
+	sortSet := NewSortSet(client)
+	mid := "13"
+	b.ResetTimer()
+	// start benchmark
+	for i := 0; i < b.N; i++ {
+		if _, err := sortSet.FindWorker(mid, 10); err != nil {
+			panic(fmt.Sprintf("find worker failed: %v", err))
+		}
+	}
+}