Navigator now supports unweighted and weighted load balancing.

mapping/elf/core.go

@@ -4,11 +4,14 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"math/rand"
 	"os"
 	"os/signal"
+	"sort"
 	"sync"
 	"sync/atomic"
 	"syscall"
+	"time"
 
 	"go.starlark.net/starlark"
 )
@@ -36,6 +39,7 @@ var (
 	reloadSignal = make(chan os.Signal)
 	pool         *threadPool
 	counter      uint64
+	rnd          = rand.New(rand.NewSource(time.Now().UnixNano()))
 )
 
 func newThreadPool(size int) *threadPool {
@@ -147,9 +151,56 @@ func GetMapping(ip string) string {
 		log.Println("Starlark execute error:", err.Error())
 		return ""
 	}
-	if r, ok := ret.(starlark.String); ok {
+	switch r := ret.(type) {
+	case starlark.String:
 		return string(r)
+	case *starlark.List:
+		val := r.Index(rnd.Intn(r.Len()))
+		if s, ok := val.(starlark.String); ok {
+			return string(s)
+		}
+		goto scriptError
+	case starlark.Tuple:
+		val := r.Index(rnd.Intn(r.Len()))
+		if s, ok := val.(starlark.String); ok {
+			return string(s)
+		}
+		goto scriptError
+	case *starlark.Dict:
+		values := r.Items()
+		options := make([]string, r.Len())
+		tot := int64(0)
+		for i := range values {
+			if v, ok := values[i].Index(0).(starlark.String); ok {
+				options[i] = string(v)
+			} else {
+				goto scriptError
+			}
+			if v, ok := values[i].Index(1).(starlark.Int); ok {
+				if weight, ok := v.Int64(); ok {
+					tot += weight
+				} else {
+					goto scriptError
+				}
+			} else {
+				goto scriptError
+			}
+		}
+		sort.Strings(options)
+		pos := rnd.Int63n(tot)
+		for _, option := range options {
+			w, _, _ := r.Get(starlark.String(option))
+			v, _ := w.(starlark.Int).Int64()
+			pos -= v
+			if pos < 0 {
+				return option
+			}
+		}
+		log.Println("Unexpected code execution,",
+			"check random logic for *starlark.Dict")
+		return options[len(options)-1]
 	}
+scriptError:
 	log.Println("Script returned unexpected result:", ret.String())
 	return ""
 }