Skip to content

Commit

Permalink
Unify extension registry
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
Ivan Mirić authored and imiric committed Dec 20, 2022
1 parent 7dc717b commit 8e543c0
Show file tree
Hide file tree
Showing 6 changed files with 187 additions and 67 deletions.
21 changes: 13 additions & 8 deletions cmd/outputs.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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,
Expand All @@ -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" {
Expand Down
3 changes: 3 additions & 0 deletions ext/doc.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
// Package ext contains the extension registry and all generic functionality for
// k6 extensions.
package ext
160 changes: 160 additions & 0 deletions ext/ext.go
Original file line number Diff line number Diff line change
@@ -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)
}
8 changes: 4 additions & 4 deletions js/jsmodules.go
Original file line number Diff line number Diff line change
@@ -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"
Expand Down Expand Up @@ -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
Expand Down
29 changes: 2 additions & 27 deletions js/modules/modules.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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
Expand Down
33 changes: 5 additions & 28 deletions output/extensions.go
Original file line number Diff line number Diff line change
@@ -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)
}

0 comments on commit 8e543c0

Please sign in to comment.