From 8e543c051e1e88a8ae0a6aeff51b22d00182fe8b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ivan=20Miri=C4=87?= Date: Tue, 6 Dec 2022 16:57:28 +0100 Subject: [PATCH] Unify extension registry Co-authored-by: HarrisChu <1726587+HarrisChu@users.noreply.github.com> This adds some structure and extracts common functionality for registering and retrieving extension information into a standalone package. Partly based on the work and feedback in #2754. --- cmd/outputs.go | 21 +++--- ext/doc.go | 3 + ext/ext.go | 160 ++++++++++++++++++++++++++++++++++++++++++ js/jsmodules.go | 8 +-- js/modules/modules.go | 29 +------- output/extensions.go | 33 ++------- 6 files changed, 187 insertions(+), 67 deletions(-) create mode 100644 ext/doc.go create mode 100644 ext/ext.go diff --git a/cmd/outputs.go b/cmd/outputs.go index 889a17a557c..3a83260c84d 100644 --- a/cmd/outputs.go +++ b/cmd/outputs.go @@ -6,6 +6,7 @@ import ( "sort" "strings" + "go.k6.io/k6/ext" "go.k6.io/k6/lib" "go.k6.io/k6/output" "go.k6.io/k6/output/cloud" @@ -18,9 +19,9 @@ import ( ) // TODO: move this to an output sub-module after we get rid of the old collectors? -func getAllOutputConstructors() (map[string]func(output.Params) (output.Output, error), error) { +func getAllOutputConstructors() (map[string]output.Constructor, error) { // Start with the built-in outputs - result := map[string]func(output.Params) (output.Output, error){ + result := map[string]output.Constructor{ "json": json.New, "cloud": cloud.New, "influxdb": influxdb.New, @@ -39,18 +40,22 @@ func getAllOutputConstructors() (map[string]func(output.Params) (output.Output, }, } - exts := output.GetExtensions() - for k, v := range exts { - if _, ok := result[k]; ok { - return nil, fmt.Errorf("invalid output extension %s, built-in output with the same type already exists", k) + exts := ext.Get(ext.OutputExtension) + for _, e := range exts { + if _, ok := result[e.Name]; ok { + return nil, fmt.Errorf("invalid output extension %s, built-in output with the same type already exists", e.Name) } - result[k] = v + m, ok := e.Module.(output.Constructor) + if !ok { + return nil, fmt.Errorf("unexpected output extension type %T", e.Module) + } + result[e.Name] = m } return result, nil } -func getPossibleIDList(constrs map[string]func(output.Params) (output.Output, error)) string { +func getPossibleIDList(constrs map[string]output.Constructor) string { res := make([]string, 0, len(constrs)) for k := range constrs { if k == "kafka" || k == "datadog" { diff --git a/ext/doc.go b/ext/doc.go new file mode 100644 index 00000000000..ef018b412a8 --- /dev/null +++ b/ext/doc.go @@ -0,0 +1,3 @@ +// Package ext contains the extension registry and all generic functionality for +// k6 extensions. +package ext diff --git a/ext/ext.go b/ext/ext.go new file mode 100644 index 00000000000..4ce9eec6095 --- /dev/null +++ b/ext/ext.go @@ -0,0 +1,160 @@ +package ext + +import ( + "fmt" + "reflect" + "runtime" + "runtime/debug" + "sort" + "strings" + "sync" +) + +// TODO: Make an ExtensionRegistry? +// +//nolint:gochecknoglobals +var ( + mx sync.RWMutex + extensions = make(map[ExtensionType]map[string]*Extension) +) + +// ExtensionType is the type of all supported k6 extensions. +type ExtensionType uint8 + +// All supported k6 extension types. +const ( + JSExtension ExtensionType = iota + 1 + OutputExtension +) + +func (e ExtensionType) String() string { + var s string + switch e { + case JSExtension: + s = "js" + case OutputExtension: + s = "output" + } + return s +} + +// Extension is a generic container for any k6 extension. +type Extension struct { + Name, Path, Version string + Type ExtensionType + Module interface{} +} + +func (e Extension) String() string { + return fmt.Sprintf("%s %s, %s [%s]", e.Path, e.Version, e.Name, e.Type) +} + +// Register a new extension with the given name and type. This function will +// panic if an unsupported extension type is provided, or if an extension of the +// same type and name is already registered. +func Register(name string, typ ExtensionType, mod interface{}) { + mx.Lock() + defer mx.Unlock() + + exts, ok := extensions[typ] + if !ok { + panic(fmt.Sprintf("unsupported extension type: %T", typ)) + } + + if _, ok := exts[name]; ok { + panic(fmt.Sprintf("extension already registered: %s", name)) + } + + path, version := extractModuleInfo(mod) + + exts[name] = &Extension{ + Name: name, + Type: typ, + Module: mod, + Path: path, + Version: version, + } +} + +// Get returns all extensions of the specified type. +func Get(typ ExtensionType) map[string]*Extension { + mx.RLock() + defer mx.RUnlock() + + exts, ok := extensions[typ] + if !ok { + panic(fmt.Sprintf("unsupported extension type: %T", typ)) + } + + result := make(map[string]*Extension, len(exts)) + + for name, ext := range exts { + result[name] = ext + } + + return result +} + +// GetAll returns all extensions, sorted by their import path and name. +func GetAll() []*Extension { + mx.RLock() + defer mx.RUnlock() + + js, out := extensions[JSExtension], extensions[OutputExtension] + result := make([]*Extension, 0, len(js)+len(out)) + + for _, e := range js { + result = append(result, e) + } + for _, e := range out { + result = append(result, e) + } + + sort.Slice(result, func(i, j int) bool { + if result[i].Path == result[j].Path { + return result[i].Name < result[j].Name + } + return result[i].Path < result[j].Path + }) + + return result +} + +// extractModuleInfo attempts to return the package path and version of the Go +// module that created the given value. +func extractModuleInfo(mod interface{}) (path, version string) { + t := reflect.TypeOf(mod) + + switch t.Kind() { + case reflect.Ptr: + if t.Elem() != nil { + path = t.Elem().PkgPath() + } + case reflect.Func: + path = runtime.FuncForPC(reflect.ValueOf(mod).Pointer()).Name() + default: + return + } + + buildInfo, ok := debug.ReadBuildInfo() + if !ok { + return + } + + for _, dep := range buildInfo.Deps { + depPath := strings.TrimSpace(dep.Path) + if strings.HasPrefix(path, depPath) { + if dep.Replace != nil { + return depPath, dep.Replace.Version + } + return depPath, dep.Version + } + } + + return +} + +func init() { + extensions[JSExtension] = make(map[string]*Extension) + extensions[OutputExtension] = make(map[string]*Extension) +} diff --git a/js/jsmodules.go b/js/jsmodules.go index d45d025740e..f58a48edbf0 100644 --- a/js/jsmodules.go +++ b/js/jsmodules.go @@ -1,7 +1,7 @@ package js import ( - "go.k6.io/k6/js/modules" + "go.k6.io/k6/ext" "go.k6.io/k6/js/modules/k6" "go.k6.io/k6/js/modules/k6/crypto" "go.k6.io/k6/js/modules/k6/crypto/x509" @@ -40,11 +40,11 @@ func getInternalJSModules() map[string]interface{} { func getJSModules() map[string]interface{} { result := getInternalJSModules() - external := modules.GetJSModules() + external := ext.Get(ext.JSExtension) // external is always prefixed with `k6/x` - for k, v := range external { - result[k] = v + for _, e := range external { + result[e.Name] = e.Module } return result diff --git a/js/modules/modules.go b/js/modules/modules.go index 4192096c37a..391ee5c01e1 100644 --- a/js/modules/modules.go +++ b/js/modules/modules.go @@ -4,21 +4,15 @@ import ( "context" "fmt" "strings" - "sync" "github.com/dop251/goja" + "go.k6.io/k6/ext" "go.k6.io/k6/js/common" "go.k6.io/k6/lib" ) const extPrefix string = "k6/x/" -//nolint:gochecknoglobals -var ( - modules = make(map[string]interface{}) - mx sync.RWMutex -) - // Register the given mod as an external JavaScript module that can be imported // by name. The name must be unique across all registered modules and must be // prefixed with "k6/x/", otherwise this function will panic. @@ -27,13 +21,7 @@ func Register(name string, mod interface{}) { panic(fmt.Errorf("external module names must be prefixed with '%s', tried to register: %s", extPrefix, name)) } - mx.Lock() - defer mx.Unlock() - - if _, ok := modules[name]; ok { - panic(fmt.Sprintf("module already registered: %s", name)) - } - modules[name] = mod + ext.Register(name, ext.JSExtension, mod) } // Module is the interface js modules should implement in order to get access to the VU @@ -43,19 +31,6 @@ type Module interface { NewModuleInstance(VU) Instance } -// GetJSModules returns a map of all registered js modules -func GetJSModules() map[string]interface{} { - mx.Lock() - defer mx.Unlock() - result := make(map[string]interface{}, len(modules)) - - for name, module := range modules { - result[name] = module - } - - return result -} - // Instance is what a module needs to return type Instance interface { Exports() Exports diff --git a/output/extensions.go b/output/extensions.go index cd6922a98a0..da37e0900f6 100644 --- a/output/extensions.go +++ b/output/extensions.go @@ -1,35 +1,12 @@ package output -import ( - "fmt" - "sync" -) +import "go.k6.io/k6/ext" -//nolint:gochecknoglobals -var ( - extensions = make(map[string]func(Params) (Output, error)) - mx sync.RWMutex -) - -// GetExtensions returns all registered extensions. -func GetExtensions() map[string]func(Params) (Output, error) { - mx.RLock() - defer mx.RUnlock() - res := make(map[string]func(Params) (Output, error), len(extensions)) - for k, v := range extensions { - res[k] = v - } - return res -} +// Constructor returns an instance of an output extension module. +type Constructor func(Params) (Output, error) // RegisterExtension registers the given output extension constructor. This // function panics if a module with the same name is already registered. -func RegisterExtension(name string, mod func(Params) (Output, error)) { - mx.Lock() - defer mx.Unlock() - - if _, ok := extensions[name]; ok { - panic(fmt.Sprintf("output extension already registered: %s", name)) - } - extensions[name] = mod +func RegisterExtension(name string, c Constructor) { + ext.Register(name, ext.OutputExtension, c) }