refactoring to NodeSelector Interface

livekit · Jul 8, 2022 · f71eb87 · f71eb87
1 parent cf88bde
commit f71eb87
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 4 deletions.
diff --git a/pkg/routing/selector/any.go b/pkg/routing/selector/any.go
@@ -1,6 +1,7 @@
 package selector
 
 import (
+	"github.com/livekit/livekit-server/pkg/config"
 	"github.com/livekit/protocol/livekit"
 )
 
@@ -17,3 +18,11 @@ func (s *AnySelector) SelectNode(nodes []*livekit.Node) (*livekit.Node, error) {
 
 	return SelectSortedNode(nodes, s.SortBy)
 }
+
+func (s *AnySelector) filterNodes(nodes []*livekit.Node) ([]*livekit.Node, error) {
+	return nodes, nil
+}
+
+func (s *AnySelector) SetConfig(conf *config.Config) error {
+	return nil
+}
diff --git a/pkg/routing/selector/cpuload.go b/pkg/routing/selector/cpuload.go
@@ -1,6 +1,7 @@
 package selector
 
 import (
+	"github.com/livekit/livekit-server/pkg/config"
 	"github.com/livekit/protocol/livekit"
 )
 
@@ -38,3 +39,7 @@ func (s *CPULoadSelector) SelectNode(nodes []*livekit.Node) (*livekit.Node, erro
 
 	return SelectSortedNode(nodes, s.SortBy)
 }
+
+func (s *CPULoadSelector) SetConfig(conf *config.Config) error {
+	return nil
+}
diff --git a/pkg/routing/selector/interfaces.go b/pkg/routing/selector/interfaces.go
@@ -14,6 +14,8 @@ var ErrUnsupportedSelector = errors.New("unsupported node selector")
 // NodeSelector selects an appropriate node to run the current session
 type NodeSelector interface {
 	SelectNode(nodes []*livekit.Node) (*livekit.Node, error)
+	filterNodes(nodes []*livekit.Node) ([]*livekit.Node, error)
+	SetConfig(conf *config.Config) error
 }
 
 func CreateNodeSelector(conf *config.Config) (NodeSelector, error) {
@@ -27,19 +29,22 @@ func CreateNodeSelector(conf *config.Config) (NodeSelector, error) {
 	case "cpuload":
 		return &CPULoadSelector{
 			CPULoadLimit: conf.NodeSelector.CPULoadLimit,
-			SortBy: conf.NodeSelector.SortBy,
+			SortBy:       conf.NodeSelector.SortBy,
 		}, nil
 	case "sysload":
 		return &SystemLoadSelector{
 			SysloadLimit: conf.NodeSelector.SysloadLimit,
-			SortBy: conf.NodeSelector.SortBy,
+			SortBy:       conf.NodeSelector.SortBy,
 		}, nil
 	case "regionaware":
 		s, err := NewRegionAwareSelector(conf.Region, conf.NodeSelector.Regions, conf.NodeSelector.SortBy)
 		if err != nil {
 			return nil, err
 		}
-		s.SysloadLimit = conf.NodeSelector.SysloadLimit
+		err = s.SetConfig(conf)
+		if err != nil {
+			return nil, err
+		}
 		return s, nil
 	case "random":
 		logger.Warnw("random node selector is deprecated, please switch to \"any\" or another selector", nil)

diff --git a/pkg/routing/selector/regionaware.go b/pkg/routing/selector/regionaware.go
@@ -24,7 +24,7 @@ func NewRegionAwareSelector(currentRegion string, regions []config.RegionConfig,
 	// build internal map of distances
 	s := &RegionAwareSelector{
 		CurrentRegion:   currentRegion,
-		regionDistances: make(map[string]float64),
+		regionDistances: make(map[string]float64, len(regions)),
 		regions:         regions,
 		SortBy:          sortBy,
 	}

diff --git a/pkg/routing/selector/sysload.go b/pkg/routing/selector/sysload.go
@@ -1,6 +1,7 @@
 package selector
 
 import (
+	"github.com/livekit/livekit-server/pkg/config"
 	"github.com/livekit/protocol/livekit"
 )
 
@@ -37,3 +38,8 @@ func (s *SystemLoadSelector) SelectNode(nodes []*livekit.Node) (*livekit.Node, e
 
 	return SelectSortedNode(nodes, s.SortBy)
 }
+
+func (s *SystemLoadSelector) SetConfig(conf *config.Config) error {
+	s.SysloadLimit = conf.NodeSelector.SysloadLimit
+	return nil
+}