Add tests ensuring the tracing module sets and cleans up metadata pro…

…perly (#2894)

* Make TraceID randomness source selectable

In order to be able to test how the client updates VUs metadata with
a trace_id, we need to be able to predictably control the generated
trace ids.

Because the traceID encodes itself, we set a private an `io.Reader`
randomness source attribute on it that will be used during encoding.
That way we can for instance pass a reader that always produces the
same value in the context of tests.

* Add tests asserting tracing module sets and cleans up trace_id metadata

* Make use of the fixed random source in tests

js/modules/k6/experimental/tracing/client.go

@@ -2,6 +2,7 @@ package tracing
 
 import (
 	"fmt"
+	"math/rand"
 	"net/http"
 	"time"
 
@@ -36,13 +37,19 @@ type Client struct {
 	// uses it under the hood to emit the requests it
 	// instruments.
 	asyncRequestFunc HTTPAsyncRequestFunc
+
+	// randSource holds the client's random source, used
+	// to generate random values for the trace ID.
+	randSource *rand.Rand
 }
 
 type (
-	// HTTPRequestFunc is a type alias representing the prototype of k6's http module's request function
+	// HTTPRequestFunc is a type alias representing the prototype of
+	// k6's http module's request function
 	HTTPRequestFunc func(method string, url goja.Value, args ...goja.Value) (*httpmodule.Response, error)
 
-	// HTTPAsyncRequestFunc is a type alias representing the prototype of k6's http module's asyncRequest function
+	// HTTPAsyncRequestFunc is a type alias representing the prototype of
+	// k6's http module's asyncRequest function
 	HTTPAsyncRequestFunc func(method string, url goja.Value, args ...goja.Value) (*goja.Promise, error)
 )
 
@@ -64,6 +71,7 @@ func NewClient(vu modules.VU, opts options) (*Client, error) {
 		return nil,
 			fmt.Errorf("failed initializing tracing client, unable to require http.request method; reason: %w", err)
 	}
+
 	// Export the http module's syncRequest function goja.Callable as a Go function
 	var asyncRequestFunc HTTPAsyncRequestFunc
 	if err := rt.ExportTo(httpModuleObject.Get("asyncRequest"), &asyncRequestFunc); err != nil {
@@ -72,7 +80,13 @@ func NewClient(vu modules.VU, opts options) (*Client, error) {
 				err)
 	}
 
-	client := &Client{vu: vu, requestFunc: requestFunc, asyncRequestFunc: asyncRequestFunc}
+	client := &Client{
+		vu:               vu,
+		requestFunc:      requestFunc,
+		asyncRequestFunc: asyncRequestFunc,
+		randSource:       rand.New(rand.NewSource(time.Now().UnixNano())), //nolint:gosec
+	}
+
 	if err := client.Configure(opts); err != nil {
 		return nil,
 			fmt.Errorf("failed initializing tracing client, invalid configuration; reason: %w", err)
@@ -106,31 +120,11 @@ func (c *Client) Configure(opts options) error {
 	return nil
 }
 
-func (c *Client) generateTraceContext() (http.Header, string, error) {
-	traceID := TraceID{
-		Prefix: k6Prefix,
-		Code:   k6CloudCode,
-		Time:   time.Now(),
-	}
-
-	encodedTraceID, err := traceID.Encode()
-	if err != nil {
-		return http.Header{}, "", fmt.Errorf("failed to encode the generated trace ID; reason: %w", err)
-	}
-
-	// Produce a trace header in the format defined by the configured propagator.
-	traceContextHeader, err := c.propagator.Propagate(encodedTraceID)
-	if err != nil {
-		return http.Header{}, "", fmt.Errorf("failed to propagate trace ID; reason: %w", err)
-	}
-
-	return traceContextHeader, encodedTraceID, nil
-}
-
 // Request instruments the http module's request function with tracing headers,
 // and ensures the trace_id is emitted as part of the output's data points metadata.
 func (c *Client) Request(method string, url goja.Value, args ...goja.Value) (*httpmodule.Response, error) {
 	var result *httpmodule.Response
+
 	var err error
 	err = c.instrumentedCall(func(args ...goja.Value) error {
 		result, err = c.requestFunc(method, url, args...)
@@ -159,37 +153,6 @@ func (c *Client) AsyncRequest(method string, url goja.Value, args ...goja.Value)
 	return result, nil
 }
 
-func (c *Client) instrumentedCall(call func(args ...goja.Value) error, args ...goja.Value) error {
-	if len(args) == 0 {
-		args = []goja.Value{goja.Null()}
-	}
-
-	traceContextHeader, encodedTraceID, err := c.generateTraceContext()
-	if err != nil {
-		return err
-	}
-	// update the `params` argument with the trace context header
-	// so that it can be used by the http module's request function.
-	args, err = c.instrumentArguments(traceContextHeader, args...)
-	if err != nil {
-		return fmt.Errorf("failed to instrument request arguments; reason: %w", err)
-	}
-
-	// Add the trace ID to the VU's state, so that it can be
-	// used in the metrics emitted by the HTTP module.
-	c.vu.State().Tags.Modify(func(t *metrics.TagsAndMeta) {
-		t.SetMetadata(metadataTraceIDKeyName, encodedTraceID)
-	})
-	// Remove the trace ID from the VU's state, so that it doesn't leak into other requests.
-	defer func() {
-		c.vu.State().Tags.Modify(func(t *metrics.TagsAndMeta) {
-			t.DeleteMetadata(metadataTraceIDKeyName)
-		})
-	}()
-
-	return call(args...)
-}
-
 // Del instruments the http module's delete method.
 func (c *Client) Del(url goja.Value, args ...goja.Value) (*httpmodule.Response, error) {
 	return c.Request(http.MethodDelete, url, args...)
@@ -231,6 +194,61 @@ func (c *Client) Put(url goja.Value, args ...goja.Value) (*httpmodule.Response,
 	return c.Request(http.MethodPut, url, args...)
 }
 
+func (c *Client) instrumentedCall(call func(args ...goja.Value) error, args ...goja.Value) error {
+	if len(args) == 0 {
+		args = []goja.Value{goja.Null()}
+	}
+
+	traceContextHeader, encodedTraceID, err := c.generateTraceContext()
+	if err != nil {
+		return err
+	}
+
+	// update the `params` argument with the trace context header
+	// so that it can be used by the http module's request function.
+	args, err = c.instrumentArguments(traceContextHeader, args...)
+	if err != nil {
+		return fmt.Errorf("failed to instrument request arguments; reason: %w", err)
+	}
+
+	// Add the trace ID to the VU's state, so that it can be
+	// used in the metrics emitted by the HTTP module.
+	c.vu.State().Tags.Modify(func(t *metrics.TagsAndMeta) {
+		t.SetMetadata(metadataTraceIDKeyName, encodedTraceID)
+	})
+
+	// Remove the trace ID from the VU's state, so that it doesn't leak into other requests.
+	defer func() {
+		c.vu.State().Tags.Modify(func(t *metrics.TagsAndMeta) {
+			t.DeleteMetadata(metadataTraceIDKeyName)
+		})
+	}()
+
+	return call(args...)
+}
+
+func (c *Client) generateTraceContext() (http.Header, string, error) {
+	traceID := TraceID{
+		Prefix:     k6Prefix,
+		Code:       k6CloudCode,
+		Time:       time.Now(),
+		randSource: c.randSource,
+	}
+
+	encodedTraceID, err := traceID.Encode()
+	if err != nil {
+		return http.Header{}, "", fmt.Errorf("failed to encode the generated trace ID; reason: %w", err)
+	}
+
+	// Produce a trace header in the format defined by the configured propagator.
+	traceContextHeader, err := c.propagator.Propagate(encodedTraceID)
+	if err != nil {
+		return http.Header{}, "", fmt.Errorf("failed to propagate trace ID; reason: %w", err)
+	}
+
+	return traceContextHeader, encodedTraceID, nil
+}
+
 // instrumentArguments: expects args to be in the format expected by the
 // request method (body, params)
 func (c *Client) instrumentArguments(traceContext http.Header, args ...goja.Value) ([]goja.Value, error) {

js/modules/k6/experimental/tracing/client_test.go

@@ -1,13 +1,17 @@
 package tracing
 
 import (
+	"math/rand"
 	"net/http"
 	"testing"
 
 	"github.com/dop251/goja"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.k6.io/k6/js/modulestest"
+	"go.k6.io/k6/lib"
+	"go.k6.io/k6/lib/testutils/httpmultibin"
+	"go.k6.io/k6/metrics"
 )
 
 // traceParentHeaderName is the normalized trace header name.
@@ -18,6 +22,28 @@ const traceparentHeaderName string = "Traceparent"
 // testTraceID is a valid trace ID used in tests.
 const testTraceID string = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"
 
+// testMetadataTraceIDRandomness is the randomness part of the test trace ID encoded
+// as hexadecimal.
+//
+// It is used to test the randomness of the trace ID. As we use a fixed time
+// as the test setup's randomness source, we can can assert that the randomness
+// part of the trace ID is the same.
+//
+// Although the randomness part doesn't have a fixed size. We can assume that
+// it will be 4 bytes, as the Go time.Time is encoded as nanoseconds, and
+// the trace ID is encoded as 8 bytes.
+const testMetadataTraceIDRandomness string = "0194fdc2"
+
+// testTracePrefix is the prefix of the test trace ID encoded as hexadecimal.
+// It is equivalent to the first 2 bytes of the trace ID, which is always
+// set to `k6Prefix` in the context of tests.
+const testTracePrefix string = "dc07"
+
+// testTraceCode is the code of the test trace ID encoded as hexadecimal.
+// It is equivalent to the third byte of the trace ID, which is always set to `k6CloudCode`
+// in the context of tests.
+const testTraceCode string = "18"
+
 func TestClientInstrumentArguments(t *testing.T) {
 	t.Parallel()
 
@@ -144,6 +170,111 @@ func TestClientInstrumentArguments(t *testing.T) {
 	})
 }
 
+func TestClientInstrumentedCall(t *testing.T) {
+	t.Parallel()
+
+	testCase := newTestCase(t)
+	testCase.testSetup.MoveToVUContext(&lib.State{
+		Tags: lib.NewVUStateTags(&metrics.TagSet{}),
+	})
+	testCase.client.propagator = NewW3CPropagator(NewAlwaysOnSampler())
+
+	callFn := func(args ...goja.Value) error {
+		gotMetadataTraceID, gotTraceIDKey := testCase.client.vu.State().Tags.GetCurrentValues().Metadata["trace_id"]
+		assert.True(t, gotTraceIDKey)
+		assert.NotEmpty(t, gotMetadataTraceID)
+		assert.Equal(t, testTracePrefix, gotMetadataTraceID[:len(testTracePrefix)])
+		assert.Equal(t, testTraceCode, gotMetadataTraceID[len(testTracePrefix):len(testTracePrefix)+len(testTraceCode)])
+		assert.Equal(t, testMetadataTraceIDRandomness, gotMetadataTraceID[len(gotMetadataTraceID)-len(testMetadataTraceIDRandomness):])
+
+		return nil
+	}
+
+	// Assert there is no trace_id key in vu metadata before using intrumentedCall
+	_, hasTraceIDKey := testCase.client.vu.State().Tags.GetCurrentValues().Metadata["trace_id"]
+	assert.False(t, hasTraceIDKey)
+
+	// The callFn will assert that the trace_id key is present in vu metadata
+	// before returning
+	_ = testCase.client.instrumentedCall(callFn)
+
+	// Assert there is no trace_id key in vu metadata after using intrumentedCall
+	_, hasTraceIDKey = testCase.client.vu.State().Tags.GetCurrentValues().Metadata["trace_id"]
+	assert.False(t, hasTraceIDKey)
+}
+
+// This test ensures that the trace_id is added to the vu metadata when
+// and instrumented request is called; and that we can find it in the
+// produced samples.
+//
+// It also ensures that the trace_id is removed from the vu metadata
+// after the request is done.
+func TestCallingInstrumentedRequestEmitsTraceIdMetadata(t *testing.T) {
+	t.Parallel()
+
+	testCase := newTestSetup(t)
+	rt := testCase.TestRuntime.VU.Runtime()
+
+	// Making sure the instrumentHTTP is called in the init context
+	// before the test.
+	_, err := rt.RunString(`
+		let http = require('k6/http')
+		instrumentHTTP({propagator: 'w3c'})
+	`)
+	require.NoError(t, err)
+
+	// Move to VU context. Setup in way that the produced samples are
+	// written to a channel we own.
+	samples := make(chan metrics.SampleContainer, 1000)
+	httpBin := httpmultibin.NewHTTPMultiBin(t)
+	testCase.TestRuntime.MoveToVUContext(&lib.State{
+		BuiltinMetrics: metrics.RegisterBuiltinMetrics(testCase.TestRuntime.VU.InitEnvField.Registry),
+		Tags:           lib.NewVUStateTags(testCase.TestRuntime.VU.InitEnvField.Registry.RootTagSet()),
+		Transport:      httpBin.HTTPTransport,
+		BufferPool:     lib.NewBufferPool(),
+		Samples:        samples,
+		Options:        lib.Options{SystemTags: &metrics.DefaultSystemTagSet},
+	})
+
+	// Inject a function in the JS runtime to assert the trace_id key
+	// is present in the vu metadata.
+	err = rt.Set("assert_has_trace_id_metadata", func(expected bool, expectedTraceID string) {
+		gotTraceID, hasTraceID := testCase.TestRuntime.VU.State().Tags.GetCurrentValues().Metadata["trace_id"]
+		require.Equal(t, expected, hasTraceID)
+
+		if expectedTraceID != "" {
+			assert.Equal(t, testTracePrefix, gotTraceID[:len(testTracePrefix)])
+			assert.Equal(t, testTraceCode, gotTraceID[len(testTracePrefix):len(testTracePrefix)+len(testTraceCode)])
+			assert.Equal(t, testMetadataTraceIDRandomness, gotTraceID[len(gotTraceID)-len(testMetadataTraceIDRandomness):])
+		}
+	})
+	require.NoError(t, err)
+
+	// Assert there is no trace_id key in vu metadata before calling an instrumented
+	// function, and that it's cleaned up after the call.
+	t.Cleanup(testCase.TestRuntime.EventLoop.WaitOnRegistered)
+	err = testCase.TestRuntime.EventLoop.Start(func() error {
+		_, err = rt.RunString(httpBin.Replacer.Replace(`
+        	assert_has_trace_id_metadata(false)
+        	http.request("GET", "HTTPBIN_URL")
+        	assert_has_trace_id_metadata(false)
+        `))
+
+		return err
+	})
+	require.NoError(t, err)
+	close(samples)
+
+	var sampleRead bool
+	for sampleContainer := range samples {
+		for _, sample := range sampleContainer.GetSamples() {
+			require.NotEmpty(t, sample.Metadata["trace_id"])
+			sampleRead = true
+		}
+	}
+	require.True(t, sampleRead)
+}
+
 type tracingClientTestCase struct {
 	t                  *testing.T
 	testSetup          *modulestest.Runtime
@@ -153,7 +284,9 @@ type tracingClientTestCase struct {
 
 func newTestCase(t *testing.T) *tracingClientTestCase {
 	testSetup := modulestest.NewRuntime(t)
-	client := Client{vu: testSetup.VU}
+	// Here we provide the client with a fixed seed to ensure that the
+	// generated trace IDs random part is deterministic.
+	client := Client{vu: testSetup.VU, randSource: rand.New(rand.NewSource(0))} //nolint:gosec
 	traceContextHeader := http.Header{}
 	traceContextHeader.Add(traceparentHeaderName, testTraceID)
 

js/modules/k6/experimental/tracing/module_test.go

@@ -64,10 +64,9 @@ func newTestSetup(t *testing.T) testSetup {
 	rt := ts.VU.Runtime()
 	require.NoError(t, rt.Set("instrumentHTTP", m.Exports().Named["instrumentHTTP"]))
 
+	export := http.New().NewModuleInstance(ts.VU).Exports().Default
 	require.NoError(t, rt.Set("require", func(module string) *goja.Object {
 		require.Equal(t, "k6/http", module)
-		export := http.New().NewModuleInstance(ts.VU).Exports().Default
-
 		return rt.ToValue(export).ToObject(rt)
 	}))