Add compare func for schedule replica

cfg.go

@@ -47,10 +47,11 @@ type Cfg struct {
 	// MaxRPCTimeout rpc max timeout
 	MaxRPCTimeout time.Duration
 
-	StorageNode bool
-	LeaseTTL    int64
-	Schedulers  []Scheduler
-	Handler     RoleChangeHandler
+	StorageNode             bool
+	LeaseTTL                int64
+	Schedulers              []Scheduler
+	Handler                 RoleChangeHandler
+	ResourceSortCompareFunc func(Resource, Resource) int
 
 	EnableScaleOnNewStore bool
 }

prophet_options.go

@@ -282,3 +282,10 @@ func WithScaleOnNewStore() Option {
 		opts.cfg.EnableScaleOnNewStore = true
 	}
 }
+
+// WithResourceSortCompareFunc set resource sort compare
+func WithResourceSortCompareFunc(value func(Resource, Resource) int) Option {
+	return func(opts *options) {
+		opts.cfg.ResourceSortCompareFunc = value
+	}
+}

runtime.go

@@ -1,6 +1,7 @@
 package prophet
 
 import (
+	"sort"
 	"sync"
 )
 
@@ -131,15 +132,15 @@ func (rc *Runtime) RandLeaderResource(id uint64, kind ResourceKind) *ResourceRun
 	rc.RLock()
 	defer rc.RUnlock()
 
-	return randResource(rc.leaders[id], kind)
+	return randResource(rc.leaders[id], kind, rc.p.cfg.ResourceSortCompareFunc)
 }
 
 // RandFollowerResource returns the random follower resource
 func (rc *Runtime) RandFollowerResource(id uint64, kind ResourceKind) *ResourceRuntime {
 	rc.RLock()
 	defer rc.RUnlock()
 
-	return randResource(rc.followers[id], kind)
+	return randResource(rc.followers[id], kind, rc.p.cfg.ResourceSortCompareFunc)
 }
 
 func (rc *Runtime) getContainerWithoutLock(id uint64) *ContainerRuntime {
@@ -160,8 +161,19 @@ func (rc *Runtime) getResourceWithoutLock(id uint64) *ResourceRuntime {
 	return resource.Clone()
 }
 
-func randResource(resources map[uint64]*ResourceRuntime, kind ResourceKind) *ResourceRuntime {
+func randResource(resources map[uint64]*ResourceRuntime, kind ResourceKind, compare func(Resource, Resource) int) *ResourceRuntime {
+	var values []*ResourceRuntime
 	for _, res := range resources {
+		values = append(values, res)
+	}
+
+	if compare != nil {
+		sort.Slice(values, func(i, j int) bool {
+			return compare(values[i].meta, values[j].meta) < 0
+		})
+	}
+
+	for _, res := range values {
 		if res.leaderPeer == nil {
 			log.Fatalf("rand resource %d without leader", res.meta.ID())
 		}

runtime_test.go

@@ -0,0 +1,44 @@
+package prophet
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRandResource(t *testing.T) {
+	m := make(map[uint64]*ResourceRuntime)
+
+	res := newTestResource()
+	res.ResID = 2
+	m[2] = &ResourceRuntime{
+		meta:       res,
+		leaderPeer: &Peer{},
+	}
+
+	res = newTestResource()
+	res.ResID = 10
+	m[10] = &ResourceRuntime{
+		meta:       res,
+		leaderPeer: &Peer{},
+	}
+
+	res = newTestResource()
+	res.ResID = 1
+	m[1] = &ResourceRuntime{
+		meta:       res,
+		leaderPeer: &Peer{},
+	}
+
+	value := randResource(m, ReplicaKind, func(res1, res2 Resource) int {
+		if res1.ID() == res2.ID() {
+			return 0
+		} else if res1.ID() > res2.ID() {
+			return 1
+		} else {
+			return -1
+		}
+	})
+	assert.NotNil(t, value, "TestRandResource failed")
+	assert.Equal(t, uint64(1), value.meta.ID(), "TestRandResource failed")
+}