Custom Tags

websockets/params.go

@@ -7,21 +7,27 @@ import (
 
 	"github.com/dop251/goja"
 
+	"go.k6.io/k6/js/common"
 	httpModule "go.k6.io/k6/js/modules/k6/http"
 	"go.k6.io/k6/lib"
+	"go.k6.io/k6/metrics"
 )
 
 // wsParams represent the parameters bag for websocket
 type wsParams struct {
-	headers   http.Header
-	cookieJar *cookiejar.Jar
+	headers     http.Header
+	cookieJar   *cookiejar.Jar
+	tagsAndMeta *metrics.TagsAndMeta
 }
 
 // buildParams builds WebSocket params and configure some of them
 func buildParams(state *lib.State, rt *goja.Runtime, raw goja.Value) (*wsParams, error) {
+	tagsAndMeta := state.Tags.GetCurrentValues()
+
 	parsed := &wsParams{
-		headers:   make(http.Header),
-		cookieJar: state.CookieJar,
+		headers:     make(http.Header),
+		cookieJar:   state.CookieJar,
+		tagsAndMeta: &tagsAndMeta,
 	}
 
 	if raw == nil || goja.IsUndefined(raw) {
@@ -43,6 +49,10 @@ func buildParams(state *lib.State, rt *goja.Runtime, raw goja.Value) (*wsParams,
 			for _, key := range headersObj.Keys() {
 				parsed.headers.Set(key, headersObj.Get(key).String())
 			}
+		case "tags":
+			if err := common.ApplyCustomUserTags(rt, parsed.tagsAndMeta, params.Get(k)); err != nil {
+				return nil, fmt.Errorf("invalid WebSocket tags option: %w", err)
+			}
 		case "jar":
 			jarV := params.Get(k)
 			if goja.IsUndefined(jarV) || goja.IsNull(jarV) {

websockets/websockets.go

@@ -66,7 +66,7 @@ type webSocket struct {
 	vu             modules.VU
 	url            *url.URL
 	conn           *websocket.Conn
-	tagsAndMeta    metrics.TagsAndMeta
+	tagsAndMeta    *metrics.TagsAndMeta
 	tq             *taskqueue.TaskQueue
 	builtinMetrics *metrics.BuiltinMetrics
 	obj            *goja.Object // the object that is given to js to interact with the WebSocket
@@ -116,6 +116,7 @@ func (r *WebSocketsAPI) websocket(c goja.ConstructorCall) *goja.Object {
 		writeQueueCh:   make(chan message, 10),
 		eventListeners: newEventListeners(),
 		obj:            rt.NewObject(),
+		tagsAndMeta:    params.tagsAndMeta,
 	}
 
 	// Maybe have this after the goroutine below ?!?
@@ -240,29 +241,27 @@ func (w *webSocket) establishConnection(params *wsParams) {
 	conn, httpResponse, connErr := wsd.DialContext(ctx, w.url.String(), params.headers)
 	connectionEnd := time.Now()
 	connectionDuration := metrics.D(connectionEnd.Sub(start))
-	ctm := state.Tags.GetCurrentValues()
-	if state.Options.SystemTags.Has(metrics.TagIP) && conn.RemoteAddr() != nil {
+
+	systemTags := state.Options.SystemTags
+
+	if conn != nil && conn.RemoteAddr() != nil {
 		if ip, _, err := net.SplitHostPort(conn.RemoteAddr().String()); err == nil {
-			ctm.SetTag("ip", ip)
+			w.tagsAndMeta.SetSystemTagOrMetaIfEnabled(systemTags, metrics.TagIP, ip)
 		}
 	}
 
 	if httpResponse != nil {
 		defer func() {
 			_ = httpResponse.Body.Close()
 		}()
-		if state.Options.SystemTags.Has(metrics.TagStatus) {
-			ctm.SetTag("status", strconv.Itoa(httpResponse.StatusCode))
-		}
-		if state.Options.SystemTags.Has(metrics.TagSubproto) {
-			ctm.SetTag("subproto", httpResponse.Header.Get("Sec-WebSocket-Protocol"))
-		}
+
+		w.tagsAndMeta.SetSystemTagOrMetaIfEnabled(systemTags, metrics.TagStatus, strconv.Itoa(httpResponse.StatusCode))
+		subProtocol := httpResponse.Header.Get("Sec-WebSocket-Protocol")
+		w.tagsAndMeta.SetSystemTagOrMetaIfEnabled(systemTags, metrics.TagSubproto, subProtocol)
 	}
 	w.conn = conn
-	if state.Options.SystemTags.Has(metrics.TagURL) {
-		ctm.SetTag("url", w.url.String())
-	}
-	w.tagsAndMeta = ctm
+
+	w.tagsAndMeta.SetSystemTagOrMetaIfEnabled(systemTags, metrics.TagURL, w.url.String())
 
 	w.emitConnectionMetrics(ctx, start, connectionDuration)
 	if connErr != nil {

websockets/websockets_test.go

@@ -879,3 +879,90 @@ func TestCookiesDefaultJar(t *testing.T) {
 	assert.Equal(t, map[string]string{"someheader": "defaultjar"}, collected)
 	mu.Unlock()
 }
+
+func TestSystemTags(t *testing.T) {
+	t.Parallel()
+
+	testedSystemTags := []string{"status", "subproto", "url", "ip"}
+	for _, expectedTagStr := range testedSystemTags {
+		expectedTagStr := expectedTagStr
+		t.Run("only "+expectedTagStr, func(t *testing.T) {
+			t.Parallel()
+			expectedTag, err := metrics.SystemTagString(expectedTagStr)
+			require.NoError(t, err)
+
+			ts := newTestState(t)
+			sr := ts.tb.Replacer.Replace
+			ts.vu.StateField.Options.SystemTags = metrics.ToSystemTagSet([]string{expectedTagStr})
+
+			err = ts.ev.Start(func() error {
+				_, runErr := ts.rt.RunString(sr(`
+				var ws = new WebSocket("WSBIN_URL/ws-echo")
+				ws.onopen = () => {
+					ws.send("test")
+				}
+				ws.onmesage = (data) => {
+					if (!data=="test") {
+						throw new Error ("echo'd data doesn't match our message!");
+					}
+					ws.close()
+				}
+			`))
+				return runErr
+			})
+			require.NoError(t, err)
+
+			containers := metrics.GetBufferedSamples(ts.samples)
+			require.NotEmpty(t, containers)
+			for _, sampleContainer := range containers {
+				require.NotEmpty(t, sampleContainer.GetSamples())
+				for _, sample := range sampleContainer.GetSamples() {
+					var dataToCheck map[string]string
+					if metrics.NonIndexableSystemTags.Has(expectedTag) {
+						dataToCheck = sample.Metadata
+					} else {
+						dataToCheck = sample.Tags.Map()
+					}
+
+					require.NotEmpty(t, dataToCheck)
+					for emittedTag := range dataToCheck {
+						assert.Equal(t, expectedTagStr, emittedTag)
+					}
+				}
+			}
+		})
+	}
+}
+
+func TestCustomTags(t *testing.T) {
+	t.Parallel()
+
+	ts := newTestState(t)
+	sr := ts.tb.Replacer.Replace
+	err := ts.ev.Start(func() error {
+		_, err := ts.rt.RunString(sr(`
+    var ws = new WebSocket("WSBIN_URL/ws-echo", null, {tags: {lorem: "ipsum", version: 13}})
+    ws.onopen = () => {
+      ws.send("something")
+      ws.close()
+    }
+	`))
+		return err
+	})
+	require.NoError(t, err)
+	samples := metrics.GetBufferedSamples(ts.samples)
+	assertSessionMetricsEmitted(t, samples, "", sr("WSBIN_URL/ws-echo"), http.StatusSwitchingProtocols, "")
+
+	for _, sampleContainer := range samples {
+		require.NotEmpty(t, sampleContainer.GetSamples())
+		for _, sample := range sampleContainer.GetSamples() {
+			dataToCheck := sample.Tags.Map()
+
+			require.NotEmpty(t, dataToCheck)
+
+			assert.Equal(t, "ipsum", dataToCheck["lorem"])
+			assert.Equal(t, "13", dataToCheck["version"])
+			assert.NotEmpty(t, dataToCheck["url"])
+		}
+	}
+}