diff --git a/balancer.go b/balancer.go new file mode 100644 index 0000000..7a1cc54 --- /dev/null +++ b/balancer.go @@ -0,0 +1,43 @@ +package gavalink + +import ( + "math" + "sort" +) + +type balancePenalty struct { + node *Node + penalty int +} + +func BestNodeByPenalties(nodes []*Node) (*Node, error) { + penalties := make([]*balancePenalty, len(nodes)) + + var playerPenalty, cpuPenalty int + + for i, node := range nodes { + if node.stats == nil { + penalties[i] = &balancePenalty{node, 1} + continue + } + + playerPenalty = node.stats.ActivePlayers + cpuPenalty = int(math.Pow(1.05, 100*node.stats.Cpu.SystemLoad)*10 - 10) + + penalties[i] = &balancePenalty{node, playerPenalty + cpuPenalty} + } + + sort.SliceStable(penalties, func(i, j int) bool { + return penalties[i].penalty < penalties[j].penalty + }) + + return penalties[0].node, nil +} + +func BestNodeByLoad(n []*Node) (*Node, error) { + sort.SliceStable(n, func(i, j int) bool { + return n[i].stats.Cpu.LavalinkLoad < n[j].stats.Cpu.LavalinkLoad + }) + + return n[0], nil +} diff --git a/lavalink.go b/lavalink.go index abd329b..4a4369b 100644 --- a/lavalink.go +++ b/lavalink.go @@ -5,7 +5,6 @@ import ( "log" "net/http" "os" - "sort" "time" ) @@ -44,7 +43,7 @@ func NewLavalink(shards string, userID string) *Lavalink { userID: userID, players: make(map[string]*Player), - BestNodeFunc: BestNodeByLoad, + BestNodeFunc: BestNodeByPenalties, } } @@ -134,11 +133,3 @@ func (lavalink *Lavalink) GetPlayer(guild string) (*Player, error) { return p, nil } - -func BestNodeByLoad(n []*Node) (*Node, error) { - sort.SliceStable(n, func(i, j int) bool { - return n[i].load < n[j].load - }) - - return n[0], nil -} diff --git a/node.go b/node.go index cc09283..9a1ea1e 100644 --- a/node.go +++ b/node.go @@ -29,12 +29,34 @@ type NodeConfig struct { // Node wraps a Lavalink Node type Node struct { config NodeConfig - load float32 + stats *RemoteStats manager *Lavalink wsConn *websocket.Conn client *http.Client } +type RemoteStats struct { + Op string `json:"op"` + Players int `json:"players"` + ActivePlayers int `json:"playingPlayers"` + Uptime int64 `json:"uptime"` + Memory *MemoryStats `json:"memory"` + Cpu *CpuStats `json:"cpu"` +} + +type MemoryStats struct { + Free uint64 `json:"free"` + Used uint64 `json:"used"` + Allocated uint64 `json:"allocated"` + Reserveable uint64 `json:"reserveable"` +} + +type CpuStats struct { + Cores int `json:"cores"` + SystemLoad float64 `json:"systemLoad"` + LavalinkLoad float64 `json:"lavalinkLoad"` +} + func (node *Node) open() error { header := http.Header{} @@ -108,16 +130,25 @@ func (node *Node) listen() { continue } - node.onEvent(v) + node.onEvent(v, msg) } } -func (node *Node) onEvent(v *fastjson.Value) error { +func (node *Node) onEvent(v *fastjson.Value, msg []byte) error { op := jsonStringValue(v, "op") switch op { + case opStats: + node.stats = &RemoteStats{} + + err := json.Unmarshal(msg, &node.stats) + + if err != nil { + return err + } case opPlayerUpdate: player, err := node.manager.GetPlayer(jsonStringValue(v, "guildId")) + if err != nil { return err } @@ -158,8 +189,6 @@ func (node *Node) onEvent(v *fastjson.Value) error { } return player.handler.OnVoiceProcessed(player, data, v.GetBool("hotword"), v.GetBool("override")) - case opStats: - node.load = float32(v.GetFloat64("cpu", "lavalinkLoad")) default: return errUnknownPayload }