From 1bb3b36df64051c505ca252ae17e51b24c4dbac5 Mon Sep 17 00:00:00 2001
From: Tyler <tystuyfzand@gmail.com>
Date: Sun, 13 Oct 2019 04:39:18 -0400
Subject: [PATCH] Implement proper (partial) penalty support

---
 balancer.go | 43 +++++++++++++++++++++++++++++++++++++++++++
 lavalink.go | 11 +----------
 node.go     | 39 ++++++++++++++++++++++++++++++++++-----
 3 files changed, 78 insertions(+), 15 deletions(-)
 create mode 100644 balancer.go

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
 	}