Data Segment API

js/initcontext.go

@@ -46,6 +46,7 @@ import (
 	"go.k6.io/k6/js/modules/k6/html"
 	"go.k6.io/k6/js/modules/k6/http"
 	"go.k6.io/k6/js/modules/k6/metrics"
+	"go.k6.io/k6/js/modules/k6/segment"
 	"go.k6.io/k6/js/modules/k6/ws"
 	"go.k6.io/k6/lib"
 	"go.k6.io/k6/loader"
@@ -323,6 +324,7 @@ func getInternalJSModules() map[string]interface{} {
 		"k6/crypto":      crypto.New(),
 		"k6/crypto/x509": x509.New(),
 		"k6/data":        data.New(),
+		"k6/segment":     segment.New(),
 		"k6/encoding":    encoding.New(),
 		"k6/execution":   execution.New(),
 		"k6/net/grpc":    grpc.New(),

js/modules/k6/segment/module.go

@@ -0,0 +1,110 @@
+/*
+ *
+ * k6 - a next-generation load testing tool
+ * Copyright (C) 2021 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 segment exports a JS API for accessing execution segment info.
+package segment
+
+import (
+	"fmt"
+
+	"github.com/dop251/goja"
+	"go.k6.io/k6/js/common"
+	"go.k6.io/k6/js/modules"
+	"go.k6.io/k6/lib"
+)
+
+type (
+	// RootModule is the global module instance that will create module
+	// instances for each VU.
+	RootModule struct {
+		indexes sharedSegmentedIndexes
+	}
+
+	// ModuleInstance represents an instance of the segment module.
+	ModuleInstance struct {
+		vu            modules.VU
+		sharedIndexes *sharedSegmentedIndexes
+	}
+)
+
+var (
+	_ modules.Module   = &RootModule{}
+	_ modules.Instance = &ModuleInstance{}
+)
+
+// New returns a pointer to a new RootModule instance.
+func New() *RootModule {
+	return &RootModule{
+		indexes: sharedSegmentedIndexes{
+			data: make(map[string]*SegmentedIndex),
+		},
+	}
+}
+
+// NewModuleInstance implements the modules.Module interface to return
+// a new instance for each VU.
+func (rm *RootModule) NewModuleInstance(vu modules.VU) modules.Instance {
+	mi := &ModuleInstance{
+		vu:            vu,
+		sharedIndexes: &rm.indexes,
+	}
+	return mi
+}
+
+// Exports returns the exports of the segment module.
+func (mi *ModuleInstance) Exports() modules.Exports {
+	return modules.Exports{
+		Named: map[string]interface{}{
+			"SegmentedIndex":       mi.SegmentedIndex,
+			"SharedSegmentedIndex": mi.SharedSegmentedIndex,
+		},
+	}
+}
+
+// SegmentedIndex is a JS constructor for a SegmentedIndex.
+func (mi *ModuleInstance) SegmentedIndex(call goja.ConstructorCall) *goja.Object {
+	state := mi.vu.State()
+	rt := mi.vu.Runtime()
+	if state == nil {
+		common.Throw(rt, fmt.Errorf("getting instance information in the init context is not supported"))
+	}
+
+	// TODO: maybe replace with lib.GetExecutionState()?
+	tuple, err := lib.NewExecutionTuple(state.Options.ExecutionSegment, state.Options.ExecutionSegmentSequence)
+	if err != nil {
+		common.Throw(rt, err)
+	}
+
+	return rt.ToValue(NewSegmentedIndex(tuple)).ToObject(rt)
+}
+
+// SharedSegmentedIndex is a JS constructor for a SharedSegmentedIndex.
+func (mi *ModuleInstance) SharedSegmentedIndex(call goja.ConstructorCall) *goja.Object {
+	rt := mi.vu.Runtime()
+	name := call.Argument(0).String()
+	if len(name) == 0 {
+		common.Throw(rt, fmt.Errorf("empty name provided to SharedArray's constructor"))
+	}
+	si, err := mi.sharedIndexes.SegmentedIndex(mi.vu.State(), name)
+	if err != nil {
+		common.Throw(rt, err)
+	}
+	return rt.ToValue(si).ToObject(rt)
+}

js/modules/k6/segment/module_test.go

@@ -0,0 +1,120 @@
+package segment_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/dop251/goja"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.k6.io/k6/js/common"
+	"go.k6.io/k6/js/modules/k6/segment"
+	"go.k6.io/k6/js/modulestest"
+	"go.k6.io/k6/lib"
+	"go.k6.io/k6/lib/testutils"
+	"go.k6.io/k6/stats"
+)
+
+func TestModuleInstanceSegmentedIndex(t *testing.T) {
+	t.Parallel()
+
+	es, err := lib.NewExecutionSegmentFromString("0:1")
+	require.NoError(t, err)
+
+	ess, err := lib.NewExecutionSegmentSequenceFromString("0,1")
+	require.NoError(t, err)
+
+	state := &lib.State{
+		Options: lib.Options{
+			ExecutionSegment:         es,
+			ExecutionSegmentSequence: &ess,
+			SystemTags:               stats.NewSystemTagSet(stats.TagVU),
+		},
+		Tags:   lib.NewTagMap(nil),
+		Logger: testutils.NewLogger(t),
+	}
+
+	rt := goja.New()
+	rt.SetFieldNameMapper(common.FieldNameMapper{})
+
+	ctx := common.WithRuntime(context.Background(), rt)
+	ctx = lib.WithState(ctx, state)
+	m, ok := segment.New().NewModuleInstance(
+		&modulestest.VU{
+			RuntimeField: rt,
+			InitEnvField: &common.InitEnvironment{},
+			CtxField:     ctx,
+			StateField:   state,
+		},
+	).(*segment.ModuleInstance)
+	require.True(t, ok)
+	require.NoError(t, rt.Set("segment", m.Exports().Named))
+
+	si, err := rt.RunString(`
+		var index = new segment.SegmentedIndex();
+		var v = index.next()
+		if (v === undefined) {
+			throw('v is undefined')
+		}
+		if (v.scaled !== 1) {
+			throw('got unexpected value for scaled')
+		}
+		if (v.unscaled !== 1) {
+			throw('got unexpected value for unscaled')
+		}
+	`)
+	require.NoError(t, err)
+	assert.NotNil(t, si)
+}
+
+func TestModuleInstanceSharedSegmentedIndex(t *testing.T) {
+	t.Parallel()
+
+	es, err := lib.NewExecutionSegmentFromString("0:1")
+	require.NoError(t, err)
+
+	ess, err := lib.NewExecutionSegmentSequenceFromString("0,1")
+	require.NoError(t, err)
+
+	state := &lib.State{
+		Options: lib.Options{
+			ExecutionSegment:         es,
+			ExecutionSegmentSequence: &ess,
+			SystemTags:               stats.NewSystemTagSet(stats.TagVU),
+		},
+		Tags:   lib.NewTagMap(nil),
+		Logger: testutils.NewLogger(t),
+	}
+
+	rt := goja.New()
+	rt.SetFieldNameMapper(common.FieldNameMapper{})
+
+	ctx := common.WithRuntime(context.Background(), rt)
+	ctx = lib.WithState(ctx, state)
+	m, ok := segment.New().NewModuleInstance(
+		&modulestest.VU{
+			RuntimeField: rt,
+			InitEnvField: &common.InitEnvironment{},
+			CtxField:     ctx,
+			StateField:   state,
+		},
+	).(*segment.ModuleInstance)
+	require.True(t, ok)
+	require.NoError(t, rt.Set("segment", m.Exports().Named))
+
+	si, err := rt.RunString(`
+		var index = new segment.SharedSegmentedIndex('myarr');
+		var v = index.next()
+		if (v === undefined) {
+			throw('v is undefined')
+		}
+		if (v.scaled !== 1) {
+			throw('got unexpected value for scaled')
+		}
+		if (v.unscaled !== 1) {
+			throw('got unexpected value for unscaled')
+		}
+	`)
+	require.NoError(t, err)
+	assert.NotNil(t, si)
+}

js/modules/k6/segment/segment.go

@@ -0,0 +1,107 @@
+/*
+ *
+ * k6 - a next-generation load testing tool
+ * Copyright (C) 2021 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 segment
+
+import (
+	"sync"
+
+	"go.k6.io/k6/lib"
+)
+
+// SegmentedIndexResult wraps a computed segment.
+type SegmentedIndexResult struct {
+	Scaled, Unscaled int64
+}
+
+// SegmentedIndex wraps a lib.SegmentedIndex making it a concurrent-safe iterator.
+type SegmentedIndex struct {
+	index *lib.SegmentedIndex
+	rwm   sync.RWMutex
+}
+
+// NewSegmentedIndex returns a pointer to a new SegmentedIndex instance,
+// given a lib.ExecutionTuple.
+func NewSegmentedIndex(et *lib.ExecutionTuple) *SegmentedIndex {
+	return &SegmentedIndex{
+		index: lib.NewSegmentedIndex(et),
+	}
+}
+
+// Next goes to the next segment's point in the iterator.
+func (s *SegmentedIndex) Next() SegmentedIndexResult {
+	s.rwm.Lock()
+	defer s.rwm.Unlock()
+
+	scaled, unscaled := s.index.Next()
+	return SegmentedIndexResult{
+		Scaled:   scaled,
+		Unscaled: unscaled,
+	}
+}
+
+// Prev goes to the previous segment's point in the iterator.
+func (s *SegmentedIndex) Prev() SegmentedIndexResult {
+	s.rwm.Lock()
+	defer s.rwm.Unlock()
+
+	scaled, unscaled := s.index.Prev()
+	return SegmentedIndexResult{
+		Scaled:   scaled,
+		Unscaled: unscaled,
+	}
+}
+
+// GoTo TODO: document
+func (s *SegmentedIndex) GoTo(value int64) SegmentedIndexResult {
+	s.rwm.Lock()
+	defer s.rwm.Unlock()
+
+	scaled, unscaled := s.index.GoTo(value)
+	return SegmentedIndexResult{Scaled: scaled, Unscaled: unscaled}
+}
+
+type sharedSegmentedIndexes struct {
+	data map[string]*SegmentedIndex
+	mu   sync.RWMutex
+}
+
+func (s *sharedSegmentedIndexes) SegmentedIndex(state *lib.State, name string) (*SegmentedIndex, error) {
+	s.mu.RLock()
+	array, ok := s.data[name]
+	s.mu.RUnlock()
+	if ok {
+		return array, nil
+	}
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	array, ok = s.data[name]
+	if !ok {
+		tuple, err := lib.NewExecutionTuple(state.Options.ExecutionSegment, state.Options.ExecutionSegmentSequence)
+		if err != nil {
+			return nil, err
+		}
+		array = NewSegmentedIndex(tuple)
+		s.data[name] = array
+	}
+	return array, nil
+}

samples/segmented-index.js

@@ -0,0 +1,34 @@
+import http from 'k6/http';
+import { check, fail } from 'k6';
+import { SharedArray } from 'k6/data';
+import {
+	SegmentedIndex,
+	SharedSegmentedIndex,
+} from 'k6/segment';
+
+let data = new SharedArray('myarr', generateArray);
+
+const params = {
+		headers: { 'Content-type': 'application/json' },
+};
+
+export default function () {
+	let iterator = new SharedSegmentedIndex('myarr'); // maybe data.name ?
+	let index = iterator.next()
+
+	const reqBody = JSON.stringify(data[index.unscaled])
+
+	var res = http.post('https://httpbin.test.k6.io/anything', reqBody, params);
+	check(res, { 'status 200': (r) => r.status === 200 });
+
+	console.log(`Something: ${res.json().json.something}`)
+}
+
+function generateArray() {
+	let n = 10;
+	const arr = new Array(n);
+	for (let i = 0; i < n; i++) {
+		arr[i] = { something: 'something else ' + i, password: '12314561' };
+	}
+	return arr;
+}