Plugin support

cmd/common.go

@@ -104,6 +104,14 @@ func getNullString(flags *pflag.FlagSet, key string) null.String {
 	return null.NewString(v, flags.Changed(key))
 }
 
+func getStringSlice(flags *pflag.FlagSet, key string) []string {
+	v, err := flags.GetStringSlice(key)
+	if err != nil {
+		panic(err)
+	}
+	return v
+}
+
 func exactArgsWithMsg(n int, msg string) cobra.PositionalArgs {
 	return func(cmd *cobra.Command, args []string) error {
 		if len(args) != n {

cmd/options.go

@@ -82,6 +82,7 @@ func optionFlagSet() *pflag.FlagSet {
 	flags.StringSlice("tag", nil, "add a `tag` to be applied to all samples, as `[name]=[value]`")
 	flags.String("console-output", "", "redirects the console logging to the provided output file")
 	flags.Bool("discard-response-bodies", false, "Read but don't process or save HTTP response bodies")
+	flags.StringSlice("plugin", []string{}, "load a plugin at `path`")
 	return flags
 }
 
@@ -110,6 +111,8 @@ func getOptions(flags *pflag.FlagSet) (lib.Options, error) {
 		TeardownTimeout: types.NullDuration{Duration: types.Duration(10 * time.Second), Valid: false},
 
 		MetricSamplesBufferSize: null.NewInt(1000, false),
+
+		Plugins: getStringSlice(flags, "plugin"),
 	}
 
 	// Using Changed() because GetStringSlice() doesn't differentiate between empty and no value

cmd/run.go

@@ -39,16 +39,19 @@ import (
 	"github.com/spf13/afero"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
+	"golang.org/x/sync/errgroup"
 	null "gopkg.in/guregu/null.v3"
 
 	"github.com/loadimpact/k6/api"
 	"github.com/loadimpact/k6/core"
 	"github.com/loadimpact/k6/core/local"
 	"github.com/loadimpact/k6/js"
+	"github.com/loadimpact/k6/js/modules"
 	"github.com/loadimpact/k6/lib"
 	"github.com/loadimpact/k6/lib/consts"
 	"github.com/loadimpact/k6/lib/types"
 	"github.com/loadimpact/k6/loader"
+	"github.com/loadimpact/k6/plugin"
 	"github.com/loadimpact/k6/ui"
 )
 
@@ -107,6 +110,37 @@ a commandline interface for interacting with it.`,
 			Left:  func() string { return "    init" },
 		}
 
+		// Load plugins
+		cliConf, err := getConfig(cmd.Flags())
+		if err != nil {
+			return err
+		}
+
+		g, _ := errgroup.WithContext(context.Background())
+		if err != nil {
+			return err
+		}
+
+		plugins := []plugin.JavaScriptPlugin{}
+		pluginNames := []string{}
+		for _, pluginPath := range cliConf.Plugins {
+			jsPlugin, pluginErr := lib.LoadJavaScriptPlugin(pluginPath)
+			if pluginErr != nil {
+				return pluginErr
+			}
+
+			// Do the part that actually takes time in a runner group.
+			g.Go(jsPlugin.Setup)
+
+			modules.RegisterPluginModules(jsPlugin.GetModules())
+			plugins = append(plugins, jsPlugin)
+			pluginNames = append(pluginNames, jsPlugin.Name())
+		}
+
+		if err = g.Wait(); err != nil {
+			return err
+		}
+
 		// Create the Runner.
 		fprintf(stdout, "%s runner\r", initBar.String())
 		pwd, err := os.Getwd()
@@ -132,10 +166,6 @@ a commandline interface for interacting with it.`,
 
 		fprintf(stdout, "%s options\r", initBar.String())
 
-		cliConf, err := getConfig(cmd.Flags())
-		if err != nil {
-			return err
-		}
 		conf, err := getConsolidatedConfig(afero.NewOsFs(), cliConf, r)
 		if err != nil {
 			return err
@@ -249,6 +279,7 @@ a commandline interface for interacting with it.`,
 			}
 
 			fprintf(stdout, "  execution: %s\n", ui.ValueColor.Sprint("local"))
+			fprintf(stdout, "    plugins: %s\n", ui.ValueColor.Sprint(strings.Join(pluginNames, ", ")))
 			fprintf(stdout, "     output: %s%s\n", ui.ValueColor.Sprint(out), ui.ExtraColor.Sprint(link))
 			fprintf(stdout, "     script: %s\n", ui.ValueColor.Sprint(filename))
 			fprintf(stdout, "\n")
@@ -473,6 +504,11 @@ a commandline interface for interacting with it.`,
 			}
 		}
 
+		// Teardown plugins
+		for _, jsPlugin := range plugins {
+			_ = jsPlugin.Teardown()
+		}
+
 		if conf.Linger.Bool {
 			logrus.Info("Linger set; waiting for Ctrl+C...")
 			<-sigC

js/initcontext.go

@@ -112,6 +112,13 @@ func (i *InitContext) Require(arg string) goja.Value {
 			common.Throw(i.runtime, err)
 		}
 		return v
+	case strings.HasPrefix(arg, "k6-plugin/"):
+		// Try in the plugin store
+		v, err := i.requirePluginModule(arg)
+		if err != nil {
+			common.Throw(i.runtime, err)
+		}
+		return v
 	default:
 		// Fall back to loading from the filesystem.
 		v, err := i.requireFile(arg)
@@ -122,6 +129,14 @@ func (i *InitContext) Require(arg string) goja.Value {
 	}
 }
 
+func (i *InitContext) requirePluginModule(name string) (goja.Value, error) {
+	mod, ok := modules.PluginIndex[name]
+	if !ok {
+		return nil, errors.Errorf("unknown plugin module: %s", name)
+	}
+	return i.runtime.ToValue(common.Bind(i.runtime, mod, i.ctxPtr)), nil
+}
+
 func (i *InitContext) requireModule(name string) (goja.Value, error) {
 	mod, ok := modules.Index[name]
 	if !ok {

js/initcontext_test.go

@@ -27,6 +27,8 @@ import (
 	"net"
 	"net/http"
 	"net/http/httptest"
+	"os"
+	"os/exec"
 	"path/filepath"
 	"testing"
 	"time"
@@ -39,11 +41,22 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/loadimpact/k6/js/common"
+	"github.com/loadimpact/k6/js/modules"
 	"github.com/loadimpact/k6/lib"
 	"github.com/loadimpact/k6/lib/netext"
 	"github.com/loadimpact/k6/stats"
+	"github.com/loadimpact/k6/testhelpers/israce"
 )
 
+func getTemporaryFile(name string) (string, func()) {
+	tf, _ := ioutil.TempFile("", name)
+	cleanup := func() {
+		_ = os.Remove(tf.Name())
+	}
+
+	return tf.Name(), cleanup
+}
+
 func TestInitContextRequire(t *testing.T) {
 	t.Run("Modules", func(t *testing.T) {
 		t.Run("Nonexistent", func(t *testing.T) {
@@ -233,6 +246,65 @@ func TestInitContextRequire(t *testing.T) {
 			assert.NoError(t, err)
 		})
 	})
+
+	t.Run("Plugins", func(t *testing.T) {
+		if isWindows {
+			t.Skip("skipping test due to lack of plugin support in windows")
+			return
+		}
+
+		t.Run("leftpad", func(t *testing.T) {
+			path, cleanup := getTemporaryFile("leftpad.so")
+			defer cleanup()
+
+			args := []string{"build", "-buildmode=plugin", "-o", path}
+			if israce.Enabled {
+				args = append(args, "-race")
+			}
+			args = append(args, "github.com/loadimpact/k6/samples/leftpad")
+			cmd := exec.Command("go", args...)
+			if err := cmd.Run(); !assert.NoError(t, err, "compiling sample plugin error") {
+				return
+			}
+
+			plugin, err := lib.LoadJavaScriptPlugin(path)
+			if !assert.NoError(t, err) {
+				return
+			}
+			assert.Equal(t, "Leftpad", plugin.Name())
+
+			modules.RegisterPluginModules(plugin.GetModules())
+
+			b, err := getSimpleBundle("/script.js", `
+				import { leftpad } from "k6-plugin/leftpad";
+				export default function() {
+				  return leftpad('test', 10, '=');
+				}
+			`)
+			if !assert.NoError(t, err, "bundle error") {
+				return
+			}
+
+			bi, err := b.Instantiate()
+			if !assert.NoError(t, err, "instance error") {
+				return
+			}
+
+			val, err := bi.Default(goja.Undefined())
+			if !assert.NoError(t, err) {
+				return
+			}
+			assert.Equal(t, "======test", val.Export())
+		})
+		t.Run("Nonexistent", func(t *testing.T) {
+			_, err := getSimpleBundle("/script.js", `
+				import { leftpad } from "k6-plugin/nonexistent";
+				export default function() {}
+			`)
+			require.Error(t, err)
+			assert.Contains(t, err.Error(), "unknown plugin module: k6-plugin/nonexistent")
+		})
+	})
 }
 
 func createAndReadFile(t *testing.T, file string, content []byte, expectedLength int, binary bool) (*BundleInstance, error) {

js/modules/index.go

@@ -21,6 +21,8 @@
 package modules
 
 import (
+	"fmt"
+
 	"github.com/loadimpact/k6/js/modules/k6"
 	"github.com/loadimpact/k6/js/modules/k6/crypto"
 	"github.com/loadimpact/k6/js/modules/k6/crypto/x509"
@@ -42,3 +44,15 @@ var Index = map[string]interface{}{
 	"k6/html":        html.New(),
 	"k6/ws":          ws.New(),
 }
+
+// PluginIndex holds an index of plugin module implementations.
+var PluginIndex = map[string]interface{}{}
+
+// RegisterPluginModules takes care of registering a map of paths that a plugin exposes so they can be
+// loaded from the JavaScript VM.
+func RegisterPluginModules(modules map[string]interface{}) {
+	for path, impl := range modules {
+		importPath := fmt.Sprintf("k6-plugin/%s", path)
+		PluginIndex[importPath] = impl
+	}
+}

lib/options.go

@@ -285,6 +285,9 @@ type Options struct {
 
 	// Redirect console logging to a file
 	ConsoleOutput null.String `json:"-" envconfig:"K6_CONSOLE_OUTPUT"`
+
+	// Plugins to load
+	Plugins []string `json:"plugins" envconfig:"K6_PLUGINS"`
 }
 
 // Returns the result of overwriting any fields with any that are set on the argument.

lib/options_test.go

@@ -476,6 +476,11 @@ func TestOptionsEnv(t *testing.T) {
 			"true":  null.BoolFrom(true),
 			"false": null.BoolFrom(false),
 		},
+		{"Plugins", "K6_PLUGINS"}: {
+			"":              []string{},
+			"foo.so":        []string{"foo.so"},
+			"foo.so,bar.so": []string{"foo.so", "bar.so"},
+		},
 		// Thresholds
 		// External
 	}

lib/plugin.go

@@ -0,0 +1,64 @@
+/*
+ *
+ * k6 - a next-generation load testing tool
+ * Copyright (C) 2016 Load Impact
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ */
+
+package lib
+
+import (
+	"fmt"
+	goplugin "plugin"
+	"runtime"
+
+	"github.com/loadimpact/k6/plugin"
+	"github.com/pkg/errors"
+)
+
+// LoadJavaScriptPlugin tries to load a dynamic library that should conform to
+// the `plug.JavaScriptPlugin` interface.
+func LoadJavaScriptPlugin(path string) (plugin.JavaScriptPlugin, error) {
+	p, err := loadPlugin(path)
+	if err != nil {
+		return nil, err
+	}
+
+	jsSym, err := p.Lookup("JavaScriptPlugin")
+	if err != nil {
+		return nil, err
+	}
+
+	var jsPlugin plugin.JavaScriptPlugin
+	jsPlugin, ok := jsSym.(plugin.JavaScriptPlugin)
+	if !ok {
+		return nil, fmt.Errorf("could not cast plugin to the right type")
+	}
+
+	return jsPlugin, nil
+}
+
+func loadPlugin(path string) (*goplugin.Plugin, error) {
+	if runtime.GOOS == "windows" {
+		return nil, errors.New("plugins are not supported in Windows")
+	}
+
+	p, err := goplugin.Open(path)
+	if err != nil {
+		err = fmt.Errorf("error while loading plugin: %w", err)
+	}
+	return p, err
+}