Implement percentage node replacement for load balancer

This is a hack to allow us to re-route a portion of traffic from an overloaded region to alternative nodes

balancer/balancer.go

@@ -126,4 +126,8 @@ type Config struct {
 	MistHost                 string
 	OwnRegion                string
 	OwnRegionTagAdjust       int
+
+	ReplaceHostMatch   string
+	ReplaceHostList    []string
+	ReplaceHostPercent int
 }

balancer/mist/mist_balancer.go

@@ -6,11 +6,13 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math/rand"
 	"net/http"
 	"net/url"
 	"os"
 	"os/exec"
 	"regexp"
+	"strings"
 	"sync"
 	"time"
 
@@ -335,6 +337,8 @@ func (b *MistBalancer) queryMistForClosestNode(ctx context.Context, playbackID,
 	return node, nil
 }
 
+var nodeHostRegex = regexp.MustCompile(`^.+?\.`) // matches the first part of the hostname before the first dot
+
 // return the best node available for a given stream. will return any node if nobody has the stream.
 func (b *MistBalancer) GetBestNode(ctx context.Context, redirectPrefixes []string, playbackID, lat, lon, fallbackPrefix string, isStudioReq bool) (string, string, error) {
 	var nodeAddr, fullPlaybackID, fallbackAddr string
@@ -366,6 +370,12 @@ func (b *MistBalancer) GetBestNode(ctx context.Context, redirectPrefixes []strin
 
 	// good path: we found the stream and a good node to play it back, yay!
 	if nodeAddr != "" {
+		if b.config.ReplaceHostMatch != "" && len(b.config.ReplaceHostList) > 0 && rand.Intn(100) < b.config.ReplaceHostPercent {
+			if strings.Contains(nodeHostRegex.FindString(nodeAddr), b.config.ReplaceHostMatch) {
+				nodeAddr = nodeHostRegex.ReplaceAllString(nodeAddr, b.config.ReplaceHostList[rand.Intn(len(b.config.ReplaceHostList))])
+			}
+		}
+
 		return nodeAddr, fullPlaybackID, nil
 	}
 

config/cli.go

@@ -91,6 +91,10 @@ type Cli struct {
 	SerfQueueSize                   int
 	SerfEventBuffer                 int
 	SerfMaxQueueDepth               int
+
+	LBReplaceHostMatch   string
+	LBReplaceHostPercent int
+	LBReplaceHostList    []string
 }
 
 // Return our own URL for callback trigger purposes

main.go

@@ -135,6 +135,9 @@ func main() {
 	fs.StringVar(&cli.SerfMembersEndpoint, "serf-members-endpoint", "", "Endpoint to get the current members in the cluster")
 	fs.StringVar(&cli.EventsEndpoint, "events-endpoint", "", "Endpoint to send proxied events from catalyst-api into catalyst")
 	fs.StringVar(&cli.CatalystApiURL, "catalyst-api-url", "", "Endpoint for externally deployed catalyst-api; if not set, use local catalyst-api")
+	fs.StringVar(&cli.LBReplaceHostMatch, "lb-replace-host-match", "", "What to match on the hostname for node replacement e.g. sto")
+	config.CommaSliceFlag(fs, &cli.LBReplaceHostList, "lb-replace-host-list", []string{}, "List of hostnames to replace with for node replacement")
+	fs.IntVar(&cli.LBReplaceHostPercent, "lb-replace-host-percent", 0, "Percentage of matching requests to replace host on")
 	pprofPort := fs.Int("pprof-port", 6061, "Pprof listen port")
 
 	fs.String("send-audio", "", "[DEPRECATED] ignored, will be removed")
@@ -205,6 +208,10 @@ func main() {
 		NodeName:                 cli.NodeName,
 		OwnRegion:                cli.OwnRegion,
 		OwnRegionTagAdjust:       cli.OwnRegionTagAdjust,
+
+		ReplaceHostMatch:   cli.LBReplaceHostMatch,
+		ReplaceHostPercent: cli.LBReplaceHostPercent,
+		ReplaceHostList:    cli.LBReplaceHostList,
 	}
 	broker = misttriggers.NewTriggerBroker()