Release v0.1.0: now with starlark script support.

You may send SIGHUP signal to the process to make it reread the script.
It will use older loaded version if parsing of new file failed.
This commit is contained in:
2019-11-17 03:51:15 +08:00
parent 2740b7119e
commit fbc2301287
7 changed files with 232 additions and 35 deletions

154
mapping/elf/core.go Normal file
View File

@ -0,0 +1,154 @@
package elf
import (
"flag"
"fmt"
"log"
"os"
"os/signal"
"sync"
"sync/atomic"
"syscall"
"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{}{},
}
parsedFunc *starlark.StringDict
parseThread = &starlark.Thread{
Name: "parseThread",
}
reloadSignal = make(chan os.Signal)
pool *threadPool
counter uint64
)
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 ""
}
if r, ok := ret.(starlark.String); ok {
return string(r)
}
log.Println("Script returned unexpected result:", ret.String())
return ""
}