Decouple the trace id struct generation from its encoding (#3019)

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

@@ -228,25 +228,18 @@ func (c *Client) instrumentedCall(call func(args ...goja.Value) error, args ...g
 }
 
 func (c *Client) generateTraceContext() (http.Header, string, error) {
-	traceID := TraceID{
-		Prefix:     k6Prefix,
-		Code:       k6CloudCode,
-		Time:       time.Now(),
-		randSource: c.randSource,
-	}
-
-	encodedTraceID, err := traceID.Encode()
+	traceID, err := newTraceID(k6Prefix, k6CloudCode, time.Now(), c.randSource)
 	if err != nil {
-		return http.Header{}, "", fmt.Errorf("failed to encode the generated trace ID; reason: %w", err)
+		return http.Header{}, "", fmt.Errorf("failed to generate trace ID; reason: %w", err)
 	}
 
 	// Produce a trace header in the format defined by the configured propagator.
-	traceContextHeader, err := c.propagator.Propagate(encodedTraceID)
+	traceContextHeader, err := c.propagator.Propagate(traceID)
 	if err != nil {
 		return http.Header{}, "", fmt.Errorf("failed to propagate trace ID; reason: %w", err)
 	}
 
-	return traceContextHeader, encodedTraceID, nil
+	return traceContextHeader, traceID, nil
 }
 
 // instrumentArguments: expects args to be in the format expected by the

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

@@ -20,84 +20,47 @@ const (
 
 	// metadataTraceIDKeyName is the key name of the traceID in the output metadata.
 	metadataTraceIDKeyName = "trace_id"
-)
-
-// TraceID represents a trace-id as defined by the [W3c specification], and
-// used by w3c, b3 and jaeger propagators. See Considerations for trace-id field [generation]
-// for more information.
-//
-// [W3c specification]: https://www.w3.org/TR/trace-context/#trace-id
-// [generation]: https://www.w3.org/TR/trace-context/#considerations-for-trace-id-field-generation
-type TraceID struct {
-	// Prefix is the first 2 bytes of the trace-id, and is used to identify the
-	// vendor of the trace-id.
-	Prefix int16
 
-	// Code is the third byte of the trace-id, and is used to identify the
-	// vendor's specific trace-id format.
-	Code int8
-
-	// Time is the time at which the trace-id was generated.
-	//
-	// The time component is used as a source of randomness, and to ensure
-	// uniqueness of the trace-id.
-	//
-	// When encoded, it should be in a format occupying the last 8 bytes of
-	// the trace-id, and should ideally be encoded as nanoseconds.
-	Time time.Time
-
-	// randSource holds the randomness source to use when encoding the
-	// trace-id. The `rand.Reader` should be your default pick. But
-	// you can replace it with a different source for testing purposes.
-	randSource io.Reader
-}
+	// traceIDEncodedSize is the size of the encoded traceID.
+	traceIDEncodedSize = 16
+)
 
-// Encode encodes the TraceID into a hex string.
+// newTraceID generates a new hexadecimal-encoded trace ID as defined by the [W3C specification].
 //
-// The trace id is first encoded as a 16 bytes sequence, as follows:
-// 1. Up to 2 bytes are encoded as the Prefix
-// 2. The third byte is the Code.
-// 3. Up to the following 8 bytes are UnixTimestampNano.
-// 4. The remaining bytes are filled with random bytes.
+// `prefix` is the first 2 bytes of the trace ID, and is used to identify the
+// vendor of the trace ID. `code` is the third byte of the trace ID, and is
+// used to identify the type of the trace ID. `t` is the time at which the trace
+// ID was generated. `randSource` is the source of randomness used to fill the rest
+// of bytes of the trace ID.
 //
-// The resulting 16 bytes sequence is then encoded as a hex string.
-func (t TraceID) Encode() (string, error) {
-	if !t.isValid() {
-		return "", fmt.Errorf("failed to encode traceID: %v", t)
+// [W3C specification]: https://www.w3.org/TR/trace-context/#trace-id
+func newTraceID(prefix int16, code int8, t time.Time, randSource io.Reader) (string, error) {
+	if prefix != k6Prefix {
+		return "", fmt.Errorf("invalid prefix 0o%o, expected 0o%o", prefix, k6Prefix)
+	}
+
+	if (code != k6CloudCode) && (code != k6LocalCode) {
+		return "", fmt.Errorf("invalid code 0o%d, accepted values are 0o%d and 0o%d", code, k6CloudCode, k6LocalCode)
 	}
 
-	// TraceID is specified to be 16 bytes long.
-	buf := make([]byte, 16)
+	// Encode The trace ID into a binary buffer.
+	buf := make([]byte, traceIDEncodedSize)
+	n := binary.PutVarint(buf, int64(prefix))
+	n += binary.PutVarint(buf[n:], int64(code))
+	n += binary.PutVarint(buf[n:], t.UnixNano())
 
-	// The `PutVarint` and `PutUvarint` functions encode the given value into
-	// the provided buffer, and return the number of bytes written. Thus, it
-	// allows us to keep track of the number of bytes written, as we go, and
-	// to pack the values to use as less space as possible.
-	n := binary.PutVarint(buf, int64(t.Prefix))
-	n += binary.PutVarint(buf[n:], int64(t.Code))
-	n += binary.PutVarint(buf[n:], t.Time.UnixNano())
+	// Calculate the number of random bytes needed.
+	randomBytesSize := traceIDEncodedSize - n
 
-	// The rest of the space in the 16 bytes buffer, equivalent to the number
-	// of available bytes left after writing the prefix, code and timestamp (index n)
-	// is filled with random bytes.
-	randomness := make([]byte, 16-n)
-	err := binary.Read(t.randSource, binary.BigEndian, randomness)
+	// Generate the random bytes.
+	randomness := make([]byte, randomBytesSize)
+	err := binary.Read(randSource, binary.BigEndian, randomness)
 	if err != nil {
-		return "", fmt.Errorf("failed to read random bytes from os; reason: %w", err)
+		return "", fmt.Errorf("failed to generate random bytes from os; reason: %w", err)
 	}
 
+	// Combine the values and random bytes to form the encoded trace ID buffer.
 	buf = append(buf[:n], randomness...)
-	hx := hex.EncodeToString(buf)
-
-	return hx, nil
-}
-
-func (t TraceID) isValid() bool {
-	var (
-		isk6Prefix = t.Prefix == k6Prefix
-		isk6Cloud  = t.Code == k6CloudCode
-		isk6Local  = t.Code == k6LocalCode
-	)
 
-	return isk6Prefix && (isk6Cloud || isk6Local)
+	return hex.EncodeToString(buf), nil
 }

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

@@ -1,75 +1,98 @@
 package tracing
 
 import (
+	"bytes"
+	"io"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
-func TestTraceID_isValid(t *testing.T) {
+func TestNewTraceID(t *testing.T) {
 	t.Parallel()
 
-	type fields struct {
-		Prefix int16
-		Code   int8
-		Time   time.Time
-	}
+	testTime := time.Date(2022, time.January, 1, 0, 0, 0, 0, time.UTC)
+	testRandSourceFn := func() io.Reader { return bytes.NewReader([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}) }
+
+	// Precomputed hexadecimal representation of the binary values
+	// of the traceID components.
+	wantPrefixHexString := "dc07"
+	wantCloudCodeHexString := "18"
+	wantLocalCodeHexString := "42"
+	wantTimeHexString := "8080f8e1949cfec52d"
+	wantRandHexString := "01020304"
+
 	testCases := []struct {
-		name   string
-		fields fields
-		want   bool
+		name       string
+		prefix     int16
+		code       int8
+		t          time.Time
+		randSource io.Reader
+		wantErr    bool
 	}{
 		{
-			name: "traceID with k6 cloud code is valid",
-			fields: fields{
-				Prefix: k6Prefix,
-				Code:   k6CloudCode,
-				Time:   time.Unix(123456789, 0),
-			},
-			want: true,
+			name:       "valid traceID with cloud code should succeed",
+			prefix:     k6Prefix,
+			code:       k6CloudCode,
+			t:          testTime,
+			randSource: testRandSourceFn(),
+			wantErr:    false,
 		},
 		{
-			name: "traceID with k6 local code is valid",
-			fields: fields{
-				Prefix: k6Prefix,
-				Code:   k6LocalCode,
-				Time:   time.Unix(123456789, 0),
-			},
-			want: true,
+			name:       "valid traceID with local code should succeed",
+			prefix:     k6Prefix,
+			code:       k6LocalCode,
+			t:          testTime,
+			randSource: testRandSourceFn(),
+			wantErr:    false,
 		},
 		{
-			name: "traceID with prefix != k6Prefix is invalid",
-			fields: fields{
-				Prefix: 0,
-				Code:   k6CloudCode,
-				Time:   time.Unix(123456789, 0),
-			},
-			want: false,
+			name:       "traceID with invalid prefix should fail",
+			prefix:     0o123,
+			code:       k6CloudCode,
+			t:          testTime,
+			randSource: testRandSourceFn(),
+			wantErr:    true,
 		},
 		{
-			name: "traceID code with code != k6CloudCode and code != k6LocalCode is invalid",
-			fields: fields{
-				Prefix: k6Prefix,
-				Code:   0,
-				Time:   time.Unix(123456789, 0),
-			},
-			want: false,
+			name:       "traceID with invalid code should fail",
+			prefix:     k6Prefix,
+			code:       0o123,
+			t:          testTime,
+			randSource: testRandSourceFn(),
+			wantErr:    true,
 		},
 	}
+
 	for _, tc := range testCases {
 		tc := tc
 
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
-			tr := &TraceID{
-				Prefix: tc.fields.Prefix,
-				Code:   tc.fields.Code,
-				Time:   tc.fields.Time,
+			gotTraceID, gotErr := newTraceID(tc.prefix, tc.code, tc.t, tc.randSource)
+
+			if tc.wantErr {
+				require.Error(t, gotErr)
+				return
 			}
 
-			if got := tr.isValid(); got != tc.want {
-				t.Errorf("TraceID.isValid() = %v, want %v", got, tc.want)
+			prefixEndOffset := len(wantPrefixHexString)
+			assert.Equal(t, wantPrefixHexString, gotTraceID[:prefixEndOffset])
+
+			codeEndOffset := prefixEndOffset + len(wantCloudCodeHexString)
+			if tc.code == k6CloudCode {
+				assert.Equal(t, wantCloudCodeHexString, gotTraceID[prefixEndOffset:codeEndOffset])
+			} else {
+				assert.Equal(t, wantLocalCodeHexString, gotTraceID[prefixEndOffset:codeEndOffset])
 			}
+
+			timeEndOffset := codeEndOffset + len(wantTimeHexString)
+			assert.Equal(t, wantTimeHexString, gotTraceID[codeEndOffset:timeEndOffset])
+
+			assert.Equal(t, wantRandHexString, gotTraceID[timeEndOffset:])
 		})
 	}
 }