Properly encode values for identifier

.gitignore

@@ -22,4 +22,8 @@
 
 # End of https://www.gitignore.io/api/go
 
-.env
+.env
+
+# IntelliJ
+.idea/
+*.iml

balancer.go

@@ -0,0 +1,51 @@
+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, deficitFramePenalty, nullFramePenalty int
+
+	for i, node := range nodes {
+		playerPenalty = 0
+		cpuPenalty = 0
+		deficitFramePenalty = 0
+		nullFramePenalty = 0
+
+		if node.stats != nil {
+			playerPenalty = node.stats.ActivePlayers
+			cpuPenalty = int(math.Pow(1.05, 100*node.stats.Cpu.SystemLoad)*10 - 10)
+
+			if node.stats.Frames != nil && node.stats.Frames.Deficit != -1 {
+				deficitFramePenalty = int(math.Pow(1.03, 500*float64(node.stats.Frames.Deficit/3000))*600 - 600)
+				nullFramePenalty = int(math.Pow(1.03, 500*float64(node.stats.Frames.Nulled/3000))*300 - 300)
+				nullFramePenalty *= 2
+			}
+		}
+
+		penalties[i] = balancePenalty{node, playerPenalty + cpuPenalty + deficitFramePenalty + nullFramePenalty}
+	}
+
+	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
+}

decoder.go

@@ -6,6 +6,7 @@ import (
 	"encoding/binary"
 	"github.com/valyala/fastjson"
 	"io"
+	"time"
 )
 
 const trackInfoVersioned int32 = 1
@@ -109,7 +110,7 @@ func Decode(r io.Reader) (*TrackInfo, error) {
 		Author:     author,
 		URI:        url,
 		Stream:     stream == 1,
-		Length:     int(length),
+		Length:     time.Duration(length) * time.Millisecond,
 	}
 
 	return track, nil
@@ -128,7 +129,7 @@ func readString(r io.Reader) (string, error) {
 	return string(buf), nil
 }
 
-func jsonStringValue(v *fastjson.Value, keys ... string) string {
+func jsonStringValue(v *fastjson.Value, keys ...string) string {
 	value := v.Get(keys...)
 
 	if value == nil {

lavalink.go

@@ -3,8 +3,9 @@ package gavalink
 import (
 	"errors"
 	"log"
+	"net/http"
 	"os"
-	"sort"
+	"time"
 )
 
 // Log sets the log.Logger gavalink will write to
@@ -19,8 +20,12 @@ type Lavalink struct {
 	shards string
 	userID string
 
-	nodes   []Node
+	nodes   []*Node
 	players map[string]*Player
+
+	capabilities map[string]interface{}
+
+	BestNodeFunc func([]*Node) (*Node, error)
 }
 
 var (
@@ -36,21 +41,27 @@ var (
 // NewLavalink creates a new Lavalink manager
 func NewLavalink(shards string, userID string) *Lavalink {
 	return &Lavalink{
-		shards: shards,
-		userID: userID,
-		/*		nodes:   make([]Node, 1),*/
+		shards:  shards,
+		userID:  userID,
 		players: make(map[string]*Player),
+
+		BestNodeFunc: BestNodeByPenalties,
 	}
 }
 
 // AddNodes adds a node to the Lavalink manager
 func (lavalink *Lavalink) AddNodes(nodeConfigs ...NodeConfig) error {
-	nodes := make([]Node, len(nodeConfigs))
+	nodes := make([]*Node, len(nodeConfigs))
+
+	client := &http.Client{
+		Timeout: 60 * time.Second,
+	}
 
 	for i, c := range nodeConfigs {
-		n := Node{
+		n := &Node{
 			config:  c,
 			manager: lavalink,
+			client:  client,
 		}
 
 		err := n.open()
@@ -71,7 +82,7 @@ func (lavalink *Lavalink) AddNodes(nodeConfigs ...NodeConfig) error {
 func (lavalink *Lavalink) removeNode(node *Node) error {
 	idx := -1
 	for i, n := range lavalink.nodes {
-		if n == *node {
+		if n == node {
 			idx = i
 			break
 		}
@@ -82,6 +93,18 @@ func (lavalink *Lavalink) removeNode(node *Node) error {
 
 	node.stop()
 
+	for _, player := range lavalink.players {
+		if player.node == node {
+			n, err := lavalink.BestNode()
+
+			if err != nil {
+				continue
+			}
+
+			player.ChangeNode(n)
+		}
+	}
+
 	// temp var for easier reading
 	n := lavalink.nodes
 	z := len(n) - 1
@@ -99,11 +122,7 @@ func (lavalink *Lavalink) BestNode() (*Node, error) {
 		return nil, errNoNodes
 	}
 
-	sort.SliceStable(lavalink.nodes, func(i, j int) bool {
-		return lavalink.nodes[i].load < lavalink.nodes[j].load
-	})
-
-	return &lavalink.nodes[0], nil
+	return lavalink.BestNodeFunc(lavalink.nodes)
 }
 
 // GetPlayer gets a player for a guild
@@ -116,3 +135,12 @@ func (lavalink *Lavalink) GetPlayer(guild string) (*Player, error) {
 
 	return p, nil
 }
+
+// Add capabilities mappings to the client, letting the server know what we support
+func (lavalink *Lavalink) AddCapability(key string, i interface{}) {
+	if lavalink.capabilities == nil {
+		lavalink.capabilities = make(map[string]interface{})
+	}
+
+	lavalink.capabilities[key] = i
+}

model.go

@@ -1,9 +1,11 @@
 package gavalink
 
 import (
+	"encoding/json"
 	"io"
 	"net/http"
 	"os"
+	"time"
 )
 
 const (
@@ -27,7 +29,7 @@ type Tracks struct {
 	// NoMatches, or LoadFailed
 	Type         string        `json:"loadType"`
 	PlaylistInfo *PlaylistInfo `json:"playlistInfo"`
-	Tracks       []Track       `json:"tracks"`
+	Tracks       []*Track      `json:"tracks"`
 }
 
 // PlaylistInfo contains information about a loaded playlist
@@ -42,20 +44,46 @@ type PlaylistInfo struct {
 // Track contains information about a loaded track
 type Track struct {
 	// Data contains the base64 encoded Lavaplayer track
-	Data string    `json:"track"`
-	Info TrackInfo `json:"info"`
+	Data string     `json:"track"`
+	Info *TrackInfo `json:"info"`
 }
 
 // TrackInfo contains more data about a loaded track
 type TrackInfo struct {
-	Identifier string `json:"identifier"`
-	Title      string `json:"title"`
-	Author     string `json:"author"`
-	URI        string `json:"uri"`
-	Seekable   bool   `json:"isSeekable"`
-	Stream     bool   `json:"isStream"`
-	Length     int    `json:"length"`
-	Position   int    `json:"position"`
+	Identifier string        `json:"identifier"`
+	Title      string        `json:"title"`
+	Author     string        `json:"author"`
+	URI        string        `json:"uri"`
+	Seekable   bool          `json:"isSeekable"`
+	Stream     bool          `json:"isStream"`
+	Length     time.Duration `json:"length"`
+	Position   int           `json:"position"`
+}
+
+func (t *TrackInfo) MarshalJSON() ([]byte, error) {
+	type Alias TrackInfo
+	return json.Marshal(&struct {
+		Length int64 `json:"length"`
+		*Alias
+	}{
+		Length: int64(t.Length / time.Millisecond),
+		Alias:  (*Alias)(t),
+	})
+}
+
+func (t *TrackInfo) UnmarshalJSON(data []byte) error {
+	type Alias TrackInfo
+	aux := &struct {
+		Length int64 `json:"length"`
+		*Alias
+	}{
+		Alias: (*Alias)(t),
+	}
+	if err := json.Unmarshal(data, &aux); err != nil {
+		return err
+	}
+	t.Length = time.Duration(aux.Length) * time.Millisecond
+	return nil
 }
 
 const (
@@ -90,7 +118,9 @@ type VoiceServerUpdate struct {
 type VoiceProcessingData struct {
 	io.ReadCloser
 
-	Client *http.Client
+	node *Node
+
+	UserID string
 	URL    string
 	File   string
 
@@ -98,7 +128,11 @@ type VoiceProcessingData struct {
 }
 
 func (v *VoiceProcessingData) open() error {
-	res, err := v.Client.Get(v.URL)
+	req, err := http.NewRequest(http.MethodGet, v.URL, nil)
+
+	req.Header.Set("Authorization", v.node.config.Password)
+
+	res, err := v.node.client.Do(req)
 
 	if err != nil {
 		return err

model_test.go

@@ -0,0 +1,33 @@
+package gavalink
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+)
+
+func TestTrackInfo_JSON(t *testing.T) {
+	i := &TrackInfo{
+		Length: 10 * time.Second,
+	}
+
+	b, err := json.Marshal(i)
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Log(string(b))
+
+	deserialize := &TrackInfo{}
+
+	if err = json.Unmarshal(b, &deserialize); err != nil {
+		t.Fatal(err)
+	}
+
+	t.Log("Deserialized length:", deserialize.Length)
+
+	if deserialize.Length != time.Second*10 {
+		t.Fatal("Expected deserialized time to be 10 seconds!")
+	}
+}

node.go

@@ -5,7 +5,9 @@ import (
 	"fmt"
 	"github.com/valyala/fastjson"
 	"net/http"
+	"net/url"
 	"strconv"
+	"strings"
 
 	"github.com/gorilla/websocket"
 )
@@ -29,12 +31,41 @@ 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"`
+	Frames        *FrameStats  `json:"frameStats"`
+}
+
+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"`
+}
+
+type FrameStats struct {
+	Sent    int `json:"sent"`
+	Nulled  int `json:"nulled"`
+	Deficit int `json:"deficit"`
+}
+
 func (node *Node) open() error {
 	header := http.Header{}
 
@@ -42,6 +73,22 @@ func (node *Node) open() error {
 	header.Set("Num-Shards", node.manager.shards)
 	header.Set("User-Id", node.manager.userID)
 
+	if node.manager.capabilities != nil {
+		v := make([]string, 0)
+
+		for k, vals := range node.manager.capabilities {
+			b, err := json.Marshal(vals)
+
+			if err != nil {
+				continue
+			}
+
+			v = append(v, k+"="+string(b))
+		}
+
+		header.Set("Capabilities", strings.Join(v, ";"))
+	}
+
 	ws, resp, err := websocket.DefaultDialer.Dial(node.config.WebSocket, header)
 
 	if err != nil {
@@ -108,16 +155,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
 		}
@@ -153,13 +209,13 @@ func (node *Node) onEvent(v *fastjson.Value) error {
 		track := jsonStringValue(v, "track")
 
 		data := &VoiceProcessingData{
-			URL:  fmt.Sprintf("%s/audio/%s", node.config.REST, track),
-			File: track,
+			node:   node,
+			UserID: jsonStringValue(v, "userId"),
+			URL:    fmt.Sprintf("%s/audio/%s", node.config.REST, track),
+			File:   track,
 		}
 
 		return player.handler.OnVoiceProcessed(player, data, v.GetBool("hotword"), v.GetBool("override"))
-	case opStats:
-		node.load = float32(v.GetFloat64("cpu", "lavalinkLoad"))
 	default:
 		return errUnknownPayload
 	}
@@ -183,11 +239,13 @@ func (node *Node) CreatePlayer(guildID string, sessionID string, event VoiceServ
 	}
 
 	player := &Player{
-		guildID: guildID,
-		manager: node.manager,
-		node:    node,
-		handler: handler,
-		vol:     100,
+		guildID:               guildID,
+		sessionID:             sessionID,
+		manager:               node.manager,
+		node:                  node,
+		handler:               handler,
+		vol:                   100,
+		lastVoiceServerUpdate: event,
 	}
 
 	node.manager.players[guildID] = player
@@ -204,9 +262,10 @@ func (node *Node) CreatePlayer(guildID string, sessionID string, event VoiceServ
 //
 // See the Lavaplayer Source Code for all valid options.
 func (node *Node) LoadTracks(query string) (*Tracks, error) {
-	url := fmt.Sprintf("%s/loadtracks?identifier=%s", node.config.REST, query)
+	v := url.Values{}
+	v.Set("identifier", query)
 
-	req, err := http.NewRequest(http.MethodGet, url, nil)
+	req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("%s/loadtracks?%s", node.config.REST, v.Encode()), nil)
 
 	if err != nil {
 		return nil, err

player.go

@@ -6,15 +6,17 @@ import (
 
 // Player is a Lavalink player
 type Player struct {
-	guildID  string
-	time     int
-	position int
-	paused   bool
-	vol      int
-	track    string
-	manager  *Lavalink
-	node     *Node
-	handler  EventHandler
+	guildID               string
+	sessionID             string
+	time                  int
+	position              int
+	paused                bool
+	vol                   int
+	track                 string
+	manager               *Lavalink
+	node                  *Node
+	handler               EventHandler
+	lastVoiceServerUpdate VoiceServerUpdate
 }
 
 // GuildID returns this player's Guild ID
@@ -166,6 +168,8 @@ func (player *Player) UserLeave(userId string) error {
 // To move a player to a new Node, first player.Destroy() it, and then
 // create a new player on the new node.
 func (player *Player) Forward(sessionID string, event VoiceServerUpdate) error {
+	player.sessionID = sessionID
+
 	msg := voiceUpdateMessage{
 		Op:        opVoiceUpdate,
 		GuildID:   player.guildID,
@@ -173,9 +177,19 @@ func (player *Player) Forward(sessionID string, event VoiceServerUpdate) error {
 		Event:     &event,
 	}
 
+	player.lastVoiceServerUpdate = event
+
 	return player.node.wsConn.WriteJSON(msg)
 }
 
+func (player *Player) ChangeNode(node *Node) error {
+	player.node = node
+
+	player.Forward(player.sessionID, player.lastVoiceServerUpdate)
+
+	return player.PlayAt(player.track, player.position, 0)
+}
+
 // Destroy will destroy this player
 func (player *Player) Destroy() error {
 	msg := basicMessage{