gavalink/lavalink.go

package gavalink

import (
	"context"
	"errors"
	"golang.org/x/sync/errgroup"
	"log"
	"net/http"
	"os"
	"runtime"
	"sync"
	"time"
)

// Log sets the log.Logger gavalink will write to
var Log *log.Logger

func init() {
	Log = log.New(os.Stdout, "(gavalink) ", 0)
}

// Lavalink manages a connection to Lavalink Nodes
type Lavalink struct {
	shards int
	userID string

	nodes []*Node

	players   map[string]*Player
	playersMu sync.RWMutex

	// Event handlers
	handlersMu   sync.RWMutex
	handlers     map[string][]*eventHandlerInstance
	onceHandlers map[string][]*eventHandlerInstance

	capabilities map[string]interface{}

	BestNodeFunc func([]*Node) (*Node, error)
}

var (
	errNoNodes          = errors.New("No nodes present")
	errNodeNotFound     = errors.New("Couldn't find that node")
	errPlayerNotFound   = errors.New("Couldn't find a player for that guild")
	errVolumeOutOfRange = errors.New("Volume is out of range, must be within [0, 1000]")
	errInvalidVersion   = errors.New("This library requires Lavalink >= 3")
	errUnknownPayload   = errors.New("Lavalink sent an unknown payload")
)

// NewLavalink creates a new Lavalink manager
func NewLavalink(shards int, userID string) *Lavalink {
	return &Lavalink{
		shards:  shards,
		userID:  userID,
		players: make(map[string]*Player),

		BestNodeFunc: BestNodeByPenalties,
	}
}

// AddNodes adds a node to the Lavalink manager
// This function calls all of the node connect methods at once.
// TODO perhaps add a pool/max at a time limit?
func (l *Lavalink) AddNodes(nodeConfigs ...NodeConfig) error {
	client := &http.Client{
		Timeout: 60 * time.Second,
	}

	eg, ctx := errgroup.WithContext(context.Background())

	for _, c := range nodeConfigs {
		eg.Go(func() error {
			n := &Node{
				config:  c,
				manager: l,
				client:  client,
			}

			err := n.open(ctx)

			if err != nil {
				return err
			}

			l.nodes = append(l.nodes, n)

			return nil
		})
	}

	return eg.Wait()
}

// RemoveNode removes a node from the manager
func (l *Lavalink) removeNode(node *Node) error {
	idx := -1
	for i, n := range l.nodes {
		if n == node {
			idx = i
			break
		}
	}
	if idx == -1 {
		return errNodeNotFound
	}

	node.stop()

	danglingPlayers := make([]*Player, 0)

	l.playersMu.RLock()
	for _, player := range l.players {
		if player.node == node {
			player.node = nil

			n, err := l.BestNode()

			if err != nil || n == nil {
				danglingPlayers = append(danglingPlayers, player)
				continue
			}

			player.ChangeNode(n)
		}
	}
	l.playersMu.RUnlock()

	if len(danglingPlayers) > 0 {
		for _, player := range danglingPlayers {
			player.Destroy()
		}
	}

	// temp var for easier reading
	n := l.nodes
	z := len(n) - 1

	n[idx] = n[z] // swap idx with last
	n = n[:z]

	l.nodes = n
	return nil
}

// BestNode returns the Node with the lowest latency
func (l *Lavalink) BestNode() (*Node, error) {
	if len(l.nodes) < 1 {
		return nil, errNoNodes
	}

	return l.BestNodeFunc(l.nodes)
}

// GetPlayer gets a player for a guild
func (l *Lavalink) GetPlayer(guild string) (*Player, error) {
	l.playersMu.RLock()
	defer l.playersMu.RUnlock()

	p, ok := l.players[guild]

	if !ok {
		return nil, errPlayerNotFound
	}

	return p, nil
}

// Add capabilities mappings to the client, letting the server know what we support
func (l *Lavalink) AddCapability(key string, i interface{}) {
	if l.capabilities == nil {
		l.capabilities = make(map[string]interface{})
	}

	l.capabilities[key] = i
}

func gavalinkUserAgent() string {
	return "Gavalink (v1.0, " + runtime.Version() + ")"
}