feat(exporter): custom args

Signed-off-by: xiayu.lyt <xiayu.lyt@alibaba-inc.com>

pkg/exporter/cmd/server.go

@@ -233,8 +233,8 @@ type EventSinkConfig struct {
 }
 
 type ProbeConfig struct {
-	Name string      `yaml:"name" mapstructure:"name"`
-	Args interface{} `yaml:"args" mapstructure:"args"`
+	Name string                 `yaml:"name" mapstructure:"name"`
+	Args map[string]interface{} `yaml:"args" mapstructure:"args"`
 }
 
 type inspServer struct {

pkg/exporter/probe/event.go

@@ -2,37 +2,107 @@ package probe
 
 import (
 	"fmt"
+	"reflect"
 
 	"github.com/alibaba/kubeskoop/pkg/exporter/nettop"
 	"golang.org/x/exp/slog"
 )
 
 var (
-	availableEventProbe = make(map[string]EventProbeCreator)
+	availableEventProbe = make(map[string]*EventProbeCreator)
 )
 
-type EventProbeCreator func(sink chan<- *Event, args map[string]interface{}) (EventProbe, error)
+// type EventProbeCreator func(sink chan<- *Event, args map[string]interface{}) (EventProbe, error)
+type EventProbeCreator struct {
+	f reflect.Value
+	s *reflect.Type
+}
+
+func NewEventProbeCreator(creator interface{}) (*EventProbeCreator, error) {
+	t := reflect.TypeOf(creator)
+	if t.Kind() != reflect.Func {
+		return nil, fmt.Errorf("metric probe creator %#v is not a func", creator)
+	}
 
-func MustRegisterEventProbe(name string, creator EventProbeCreator) {
+	err := validateProbeCreatorReturnValue[EventProbe](reflect.TypeOf(creator))
+	if err != nil {
+		return nil, err
+	}
+
+	if t.NumIn() != 1 && t.NumIn() != 2 {
+		return nil, fmt.Errorf("input parameter count of creator should be either 2 or 3")
+	}
+
+	ct := t.In(0)
+	et := reflect.TypeOf((*Event)(nil))
+	if ct.Kind() != reflect.Chan || ct.ChanDir() != reflect.SendDir || ct.Elem() != et {
+		return nil, fmt.Errorf("first input parameter should be a send channel of *Event")
+	}
+
+	ret := &EventProbeCreator{
+		f: reflect.ValueOf(creator),
+	}
+
+	if t.NumIn() == 2 {
+		st := t.In(1)
+		if st.Kind() != reflect.Struct && st.Kind() != reflect.Map {
+			return nil, fmt.Errorf("input parameter should be struct, but %s", st.Kind())
+		}
+		if st.Kind() == reflect.Map && st.Key().Kind() != reflect.String {
+			return nil, fmt.Errorf("map key type of input parameter should be string")
+		}
+		ret.s = &st
+	}
+
+	return ret, nil
+}
+
+func (e *EventProbeCreator) Call(sink chan<- *Event, args map[string]interface{}) (EventProbe, error) {
+	in := []reflect.Value{
+		reflect.ValueOf(sink),
+	}
+	if e.s != nil {
+		s, err := createStructFromTypeWithArgs(*e.s, args)
+		if err != nil {
+			return nil, err
+		}
+		in = append(in, s)
+	}
+
+	result := e.f.Call(in)
+	// return parameter count and type has been checked in NewEventProbeCreator
+	ret := result[0].Interface().(EventProbe)
+	err := result[1].Interface()
+	if err == nil {
+		return ret, nil
+	}
+	return ret, err.(error)
+}
+
+func MustRegisterEventProbe(name string, creator interface{}) {
 	if _, ok := availableEventProbe[name]; ok {
 		panic(fmt.Errorf("duplicated event probe %s", name))
 	}
 
-	availableEventProbe[name] = creator
+	c, err := NewEventProbeCreator(creator)
+	if err != nil {
+		panic(fmt.Errorf("error register event probe %s: %s", name, err))
+	}
+
+	availableEventProbe[name] = c
 }
 
 func NewEventProbe(name string, simpleProbe SimpleProbe) EventProbe {
 	return NewProbe(name, simpleProbe)
 }
 
-func CreateEventProbe(name string, sink chan<- *Event, _ interface{}) (EventProbe, error) {
+func CreateEventProbe(name string, sink chan<- *Event, args map[string]interface{}) (EventProbe, error) {
 	creator, ok := availableEventProbe[name]
 	if !ok {
 		return nil, fmt.Errorf("undefined probe %s", name)
 	}
 
-	//TODO reflect creator's arguments
-	return creator(sink, nil)
+	return creator.Call(sink, args)
 }
 
 func ListEventProbes() []string {

pkg/exporter/probe/metrics.go

@@ -3,6 +3,7 @@ package probe
 import (
 	"errors"
 	"fmt"
+	"reflect"
 
 	log "github.com/sirupsen/logrus"
 
@@ -14,28 +15,89 @@ const LegacyMetricsNamespace = "inspector"
 const MetricsNamespace = "kubeskoop"
 
 var (
-	availableMetricsProbes = make(map[string]MetricsProbeCreator)
+	availableMetricsProbes = make(map[string]*MetricsProbeCreator)
 	ErrUndeclaredMetrics   = errors.New("undeclared metrics")
 )
 
-type MetricsProbeCreator func(args map[string]interface{}) (MetricsProbe, error)
+// type MetricsProbeCreator func(args map[string]interface{}) (MetricsProbe, error)
+type MetricsProbeCreator struct {
+	f reflect.Value
+	s *reflect.Type
+}
+
+func NewMetricProbeCreator(creator interface{}) (*MetricsProbeCreator, error) {
+	t := reflect.TypeOf(creator)
+	if t.Kind() != reflect.Func {
+		return nil, fmt.Errorf("metric probe creator %#v is not a func", creator)
+	}
+
+	err := validateProbeCreatorReturnValue[MetricsProbe](reflect.TypeOf(creator))
+	if err != nil {
+		return nil, err
+	}
+
+	if t.NumIn() > 1 {
+		return nil, fmt.Errorf("input parameter count of creator should be either 0 or 1")
+	}
+
+	ret := &MetricsProbeCreator{
+		f: reflect.ValueOf(creator),
+	}
+
+	if t.NumIn() == 1 {
+		st := t.In(0)
+		if st.Kind() != reflect.Struct && st.Kind() != reflect.Map {
+			return nil, fmt.Errorf("input parameter should be struct, but %s", st.Kind())
+		}
+		if st.Kind() == reflect.Map && st.Key().Kind() != reflect.String {
+			return nil, fmt.Errorf("map key type of input parameter should be string")
+		}
+		ret.s = &st
+	}
+
+	return ret, nil
+}
+
+func (m *MetricsProbeCreator) Call(args map[string]interface{}) (MetricsProbe, error) {
+	var in []reflect.Value
+	if m.s != nil {
+		s, err := createStructFromTypeWithArgs(*m.s, args)
+		if err != nil {
+			return nil, err
+		}
+		in = append(in, s)
+	}
 
-func MustRegisterMetricsProbe(name string, creator MetricsProbeCreator) {
+	result := m.f.Call(in)
+	// return parameter count and type has been checked in NewMetricProbeCreator
+	ret := result[0].Interface().(MetricsProbe)
+	err := result[1].Interface()
+	if err == nil {
+		return ret, nil
+	}
+	return ret, err.(error)
+}
+
+func MustRegisterMetricsProbe(name string, creator interface{}) {
 	if _, ok := availableMetricsProbes[name]; ok {
-		panic(fmt.Errorf("duplicated event probe %s", name))
+		panic(fmt.Errorf("duplicated metric probe %s", name))
+	}
+
+	c, err := NewMetricProbeCreator(creator)
+	if err != nil {
+		panic(fmt.Errorf("error register metric probe %s: %s", name, err))
 	}
 
-	availableMetricsProbes[name] = creator
+	availableMetricsProbes[name] = c
 }
 
-func CreateMetricsProbe(name string, _ interface{}) (MetricsProbe, error) {
+func CreateMetricsProbe(name string, args map[string]interface{}) (MetricsProbe, error) {
 	creator, ok := availableMetricsProbes[name]
 	if !ok {
 		return nil, fmt.Errorf("undefined probe %s", name)
 	}
 
-	//TODO reflect creator's arguments
-	return creator(nil)
+	return creator.Call(args)
 }
 
 func ListMetricsProbes() []string {

pkg/exporter/probe/procsnmp/procsnmp.go

@@ -86,7 +86,7 @@ func init() {
 	probe.MustRegisterMetricsProbe(IP, newSnmpProbeCreator(IP))
 }
 
-func newSnmpProbeCreator(probeName string) probe.MetricsProbeCreator {
+func newSnmpProbeCreator(probeName string) func(args map[string]interface{}) (probe.MetricsProbe, error) {
 	return func(args map[string]interface{}) (probe.MetricsProbe, error) {
 		p := &procSNMP{
 			name: probeName,

pkg/exporter/probe/util.go

@@ -0,0 +1,34 @@
+package probe
+
+import (
+	"fmt"
+	"github.com/mitchellh/mapstructure"
+	"reflect"
+)
+
+func validateProbeCreatorReturnValue[T interface{}](t reflect.Type) error {
+	if t.Kind() != reflect.Func {
+		return fmt.Errorf("%s is not Func type", t)
+	}
+
+	if t.NumOut() != 2 {
+		return fmt.Errorf("expect return value count 2, but actual %d", t.NumOut())
+	}
+
+	it := reflect.TypeOf((*T)(nil)).Elem()
+	if !t.Out(0).Implements(it) {
+		return fmt.Errorf("arg 0 should implement interface %s", it)
+	}
+
+	et := reflect.TypeOf((*error)(nil)).Elem()
+	if !t.Out(1).Implements(et) {
+		return fmt.Errorf("arg 1 should implement error")
+	}
+	return nil
+}
+
+func createStructFromTypeWithArgs(st reflect.Type, args map[string]interface{}) (reflect.Value, error) {
+	v := reflect.New(st)
+	err := mapstructure.Decode(args, v.Interface())
+	return v.Elem(), err
+}

pkg/exporter/probe/util_test.go

@@ -0,0 +1,108 @@
+package probe
+
+import (
+	"golang.org/x/exp/slices"
+	"reflect"
+	"testing"
+)
+
+type test interface {
+	TestA()
+	TestB()
+}
+
+type aa struct {
+}
+
+func (a *aa) TestA() {
+}
+
+func (a *aa) TestB() {
+
+}
+
+func funcA() (test, error) {
+	return nil, nil
+}
+
+func funcB() (*aa, error) {
+	return nil, nil
+}
+
+func funcC() test {
+	return nil
+}
+
+func funcD() (string, error) {
+	return "", nil
+}
+
+func funcE() (test, string) {
+	return nil, ""
+}
+
+func TestValidateProbeCreatorReturnValue(t *testing.T) {
+	testcases := []struct {
+		name string
+		f    interface{}
+		err  bool
+	}{
+		{"funcA", funcA, false},
+		{"funcB", funcB, false},
+		{"funcC", funcC, true},
+		{"funcD", funcD, true},
+		{"funcE", funcE, true},
+	}
+
+	for _, c := range testcases {
+		err := validateProbeCreatorReturnValue[test](reflect.TypeOf(c.f))
+		if c.err && err == nil {
+			t.Errorf("testcase %s except error, but nil", c.name)
+			t.Fail()
+		}
+
+		if !c.err && err != nil {
+			t.Errorf("testcase %s except nil error, but %s", c.name, err.Error())
+			t.Fail()
+		}
+	}
+}
+
+func TestCreateStructFromTypeWithArgs(t *testing.T) {
+	type s struct {
+		ArgA string
+		ArgB uint32
+		ArgC []string
+	}
+
+	st := reflect.TypeOf(s{})
+
+	m := map[string]interface{}{
+		"argA": "a",
+		"argB": 1,
+		"argC": []string{"a", "b"},
+		"argD": 1,
+	}
+	r, err := createStructFromTypeWithArgs(st, m)
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err.Error())
+	}
+	sr := r.Interface().(s)
+	if sr.ArgA != "a" || sr.ArgB != 1 || !slices.Equal(sr.ArgC, []string{"a", "b"}) {
+		t.Fatalf("expect %+v, actual %+v", m, sr)
+	}
+
+	m = map[string]interface{}{
+		"a": "a",
+		"b": "b",
+		"c": "c",
+	}
+	r, err = createStructFromTypeWithArgs(reflect.TypeOf(map[string]string{}), m)
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err.Error())
+	}
+	mr := r.Interface().(map[string]string)
+	if len(mr) != 3 || mr["a"] != "a" || mr["b"] != "b" || mr["c"] != "c" {
+		t.Fatalf("expect %+v, actual %+v", m, mr)
+	}
+}