navigator/mapping/elf/core.go

package elf

import (
	"flag"
	"fmt"
	"log"
	"math/rand"
	"os"
	"os/signal"
	"sort"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"go.starlark.net/starlark"
)

type threadPool struct {
	pool  []*starlark.Thread
	mutex *sync.Mutex
}

var (
	scriptFile = flag.String("mapping-script", "map.starlark",
		"Starlark script used to do mapping logic")
	threadPoolSize = flag.Int("script-thread-pool", 128,
		"Thread pool size for starlark execution engine")
	builtinFunc  = make(starlark.StringDict)
	requiredFunc = map[string]struct{}{
		"getMapping": struct{}{},
		"getNodes":   struct{}{},
		"getSuffix":  struct{}{},
	}
	parsedFunc  *starlark.StringDict
	parseThread = &starlark.Thread{
		Name: "parseThread",
	}
	reloadSignal = make(chan os.Signal)
	pool         *threadPool
	counter      uint64
	rnd          = rand.New(rand.NewSource(time.Now().UnixNano()))
)

func newThreadPool(size int) *threadPool {
	ret := &threadPool{
		pool:  make([]*starlark.Thread, size),
		mutex: &sync.Mutex{},
	}
	for i := range ret.pool {
		counter += 1
		ret.pool[i] = &starlark.Thread{
			Name: fmt.Sprintf("ElfThread-%d", counter),
		}
	}
	return ret
}

func (pool *threadPool) Get() (ret *starlark.Thread) {
	pool.mutex.Lock()
	defer pool.mutex.Unlock()
	if len(pool.pool) == 0 {
		return &starlark.Thread{
			Name: fmt.Sprintf("ElfThread-%d", atomic.AddUint64(&counter, 1)),
		}
	}
	ret = pool.pool[len(pool.pool)-1]
	pool.pool = pool.pool[:len(pool.pool)-1]
	return
}

func (pool *threadPool) Put(thread *starlark.Thread) {
	pool.mutex.Lock()
	defer pool.mutex.Unlock()
	if len(pool.pool) == *threadPoolSize {
		return
	}
	pool.pool = append(pool.pool, thread)
}

func RegisterFunc(name string, impl func(*starlark.Thread, *starlark.Builtin,
	starlark.Tuple, []starlark.Tuple) (starlark.Value, error)) {
	if builtinFunc.Has(name) {
		panic(fmt.Errorf("Function %s has already been declared as: %s",
			name, builtinFunc[name].String()))
	}
	builtinFunc[name] = starlark.NewBuiltin(name, impl)
	log.Println("Registered function", name, "to starlark")
}

func RequireFunc(name string) {
	requiredFunc[name] = struct{}{}
}

func Initialize() {
	globals, err := starlark.ExecFile(
		parseThread, *scriptFile, nil, builtinFunc)
	if err != nil {
		panic(fmt.Errorf("Unable to parse starlark file: %s", err.Error()))
	}
	parsedFunc = &globals
	for name := range requiredFunc {
		if !parsedFunc.Has(name) {
			panic(fmt.Errorf("Required function %s not provided", name))
		}
		if _, ok := globals[name].(starlark.Callable); !ok {
			panic(fmt.Errorf(
				"Variable %s is not callable as required function", name))
		}

	}
	signal.Notify(reloadSignal, syscall.SIGHUP)
	go func() {
		for {
			<-reloadSignal
			reload()
		}
	}()
	pool = newThreadPool(*threadPoolSize)
}

func reload() {
	globals, err := starlark.ExecFile(
		parseThread, *scriptFile, nil, builtinFunc)
	if err != nil {
		log.Println("Reload failed: Unable to parse starlark file:",
			err.Error())
		return
	}
	for name := range requiredFunc {
		if !globals.Has(name) {
			log.Println(
				"Reload failed: Required function", name, "not provided")
			return
		}
		if _, ok := globals[name].(starlark.Callable); !ok {
			log.Println(
				"Variable", name, "is not callable as required function")
			return
		}
	}
	log.Println("Reload success, new rules applied :)")
	parsedFunc = &globals
}

func GetMapping(ip string) string {
	thread := pool.Get()
	ret, err := starlark.Call(thread, (*parsedFunc)["getMapping"],
		starlark.Tuple{starlark.String(ip)}, nil)
	if err != nil {
		log.Println("Starlark execute error:", err.Error())
		return ""
	}
	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 ""
}

func GetNodes() (res []string) {
	thread := pool.Get()
	ret, err := starlark.Call(thread, (*parsedFunc)["getNodes"], nil, nil)
	if err != nil {
		log.Println("Starlark execute error:", err.Error())
		return nil
	}
	if r, ok := ret.(*starlark.List); ok {
		res = make([]string, r.Len())
		for i := 0; i < r.Len(); i++ {
			v := r.Index(i)
			if s, ok := v.(starlark.String); ok {
				res[i] = string(s)
			} else {
				log.Println("Script returned unexpected result:",
					ret.String())
			}
		}
	} else {
		log.Println("Script returned unexpected result:", ret.String())
	}
	return
}

func GetSuffix(ip string) string {
	thread := pool.Get()
	ret, err := starlark.Call(thread, (*parsedFunc)["getSuffix"],
		starlark.Tuple{starlark.String(ip)}, nil)
	if err != nil {
		log.Println("Starlark execute error:", err.Error())
		return ""
	}
	if r, ok := ret.(starlark.String); ok {
		return string(r)
	}
	log.Println("Script returned unexpected result:", ret.String())
	return ""
}