Complete refactor of the extension (#25)

* Refactor

Signed-off-by: Daniel González Lopes <danielgonzalezlopes@gmail.com>

* Add k6 version

Signed-off-by: Daniel González Lopes <danielgonzalezlopes@gmail.com>

* wip

* wip2

* wip

* crocospans->cloud

Signed-off-by: Daniel González Lopes <danielgonzalezlopes@gmail.com>

* Fix examples

Signed-off-by: Daniel González Lopes <danielgonzalezlopes@gmail.com>

* Refactor README

Signed-off-by: Daniel González Lopes <danielgonzalezlopes@gmail.com>

* Add new line

Signed-off-by: Daniel González Lopes <danielgonzalezlopes@gmail.com>

Signed-off-by: Daniel González Lopes <danielgonzalezlopes@gmail.com>
Co-authored-by: Nedyalko Andreev <n@andreev.sh>

Author: dgzlopes

.gitignore

@@ -2,4 +2,5 @@
 examples/k6
 k6
 k6-distributed-tracing
-vendor
+vendor
+gen

Makefile

@@ -0,0 +1,12 @@
+DOCKER_BUILD=docker build
+
+DOCKER_RUN=docker run
+
+.PHONY: build
+build:
+	xk6 build master --with github.com/grafana/xk6-distributed-tracing="${PWD}/../xk6-distributed-tracing"
+
+.PHONY: proto
+proto:
+	$(DOCKER_RUN) -v ${PWD}/crocospans:/defs namely/protoc-all -f *.proto -l go
+	cp -r ${PWD}/crocospans/gen/pb-go/*.pb.go ${PWD}/crocospans

README.md

@@ -1,32 +1,13 @@
 # xk6-distributed-tracing
 
-> ⚠️ **This is a proof of concept** ⚠️  
-> It won't be supported by the k6 team.  
-> It may also break in the future as xk6 evolves.
+This extension adds distributed tracing support to [k6](https://github.com/grafana/k6)! 
 
-</div>
+That means that if you're testing an instrumented system, you can use this extension to start the traces on k6. 
 
-This extension adds distributed tracing support to [k6](https://github.com/grafana/k6)! That means, that if you're testing a system that is instrumented, you can use this extension to start the traces on k6.
+Currently, it supports HTTP requests and the following propagation formats: `w3c`, `b3`, and `jaeger`.
 
 It is implemented using the [xk6](https://github.com/grafana/xk6) extension system.
 
-<div align="center">
-
-![Example trace](/media/trace.png)
-*Trace started on k6*
-
-</div>
-
-## Features
-
-The extension is built on top of [OpenTelemetry](https://opentelemetry.io/), and gives k6:
-- An instrumented HTTP client.
-- The hability to: 
-    - Export spans (your tracing data).
-        - Supported exporters: `jaeger`, `zipkin`, `otlp`, `noop`, `stdout`
-    - Propagate context.
-        - Supported protocols: `w3c`, `b3`, `jaeger`, `ot`
-
 ## Build
 
 To build a `k6` binary with this extension, first ensure you have the prerequisites:
@@ -56,7 +37,7 @@ import { sleep } from 'k6';
 
 export let options = {
   vus: 1,
-  duration: '10s',
+  iterations: 10,
 };
 
 export function setup() {
@@ -65,30 +46,23 @@ export function setup() {
 
 export default function() {
   const http = new Http({
-    exporter: "jaeger",
     propagator: "w3c",
-    endpoint: "http://localhost:14268/api/traces"
   });
   const r = http.get('https://test-api.k6.io');
   console.log(`trace_id=${r.trace_id}`);
   sleep(1);
 }
-
-export function teardown(){
-  // Cleanly shutdown and flush telemetry when k6 exits.
-  tracing.shutdown();
-}
 ```
 
 Result output:
 
 ```bash
 $ ./k6 run script.js
 
-          /\      |‾‾| /‾‾/   /‾‾/   
-     /\  /  \     |  |/  /   /  /    
-    /  \/    \    |     (   /   ‾‾\  
-   /          \   |  |\  \ |  (‾)  | 
+          /\      |‾‾| /‾‾/   /‾‾/
+     /\  /  \     |  |/  /   /  /
+    /  \/    \    |     (   /   ‾‾\
+   /          \   |  |\  \ |  (‾)  |
   / __________ \  |__| \__\ \_____/ .io
 
   execution: local
@@ -98,17 +72,17 @@ $ ./k6 run script.js
   scenarios: (100.00%) 1 scenario, 1 max VUs, 40s max duration (incl. graceful stop):
            * default: 1 looping VUs for 10s (gracefulStop: 30s)
 
-INFO[0000] Running xk6-distributed-tracing v0.0.2  source=console                
-INFO[0000] trace-id=743fff0b96778539acb7139e72ea1e33
-INFO[0001] trace-id=365f4637a52526db1de2d30a5568ca3a
-INFO[0002] trace-id=c49e1df945049c5c3c8b59acc84d7d3b
-INFO[0003] trace-id=53e1937d56aa172b46d2310e3380dfe9
-INFO[0004] trace-id=d61e8757d35c9ca1780b88977ac56d72
-INFO[0005] trace-id=358e794ed636d268a918dcd2f3f9db0a
-INFO[0006] trace-id=992a959e09ee84f3905a215bec8b53a0
-INFO[0007] trace-id=aee11c64de11744ab5b66d5dd8ed361b
-INFO[0008] trace-id=c4dc45d857e99ede2bb902666457239d
-INFO[0009] trace-id=7623d10293d9f03c15deb8055935664e
+INFO[0000] Running xk6-distributed-tracing v0.2.0  source=console
+INFO[0000] trace_id=743fff0b96778539acb7139e72ea1e33
+INFO[0001] trace_id=365f4637a52526db1de2d30a5568ca3a
+INFO[0002] trace_id=c49e1df945049c5c3c8b59acc84d7d3b
+INFO[0003] trace_id=53e1937d56aa172b46d2310e3380dfe9
+INFO[0004] trace_id=d61e8757d35c9ca1780b88977ac56d72
+INFO[0005] trace_id=358e794ed636d268a918dcd2f3f9db0a
+INFO[0006] trace_id=992a959e09ee84f3905a215bec8b53a0
+INFO[0007] trace_id=aee11c64de11744ab5b66d5dd8ed361b
+INFO[0008] trace_id=c4dc45d857e99ede2bb902666457239d
+INFO[0009] trace_id=7623d10293d9f03c15deb8055935664e
 
 running (10.1s), 0/1 VUs, 10 complete and 0 interrupted iterations
 default ✓ [======================================] 1 VUs  10s
@@ -119,13 +93,13 @@ default ✓ [======================================] 1 VUs  10s
      data_sent..................: 1.7 kB 165 B/s
      http_req_blocked...........: avg=223.43µs min=146.53µs med=217.39µs max=314.54µs p(90)=276.68µs p(95)=295.61µs
      http_req_connecting........: avg=137.18µs min=87.22µs  med=130.17µs max=196.38µs p(90)=184.38µs p(95)=190.38µs
-     http_req_duration..........: avg=6.58ms   min=5.07ms   med=6.45ms   max=7.91ms   p(90)=7.83ms   p(95)=7.87ms  
+     http_req_duration..........: avg=6.58ms   min=5.07ms   med=6.45ms   max=7.91ms   p(90)=7.83ms   p(95)=7.87ms
      http_req_receiving.........: avg=187.27µs min=94.29µs  med=171.7µs  max=295.67µs p(90)=293.28µs p(95)=294.48µs
      http_req_sending...........: avg=128.07µs min=94.64µs  med=121.77µs max=175.65µs p(90)=160.41µs p(95)=168.03µs
-     http_req_tls_handshaking...: avg=0s       min=0s       med=0s       max=0s       p(90)=0s       p(95)=0s      
-     http_req_waiting...........: avg=6.27ms   min=4.83ms   med=6.13ms   max=7.64ms   p(90)=7.56ms   p(95)=7.6ms   
+     http_req_tls_handshaking...: avg=0s       min=0s       med=0s       max=0s       p(90)=0s       p(95)=0s
+     http_req_waiting...........: avg=6.27ms   min=4.83ms   med=6.13ms   max=7.64ms   p(90)=7.56ms   p(95)=7.6ms
      http_reqs..................: 10     0.991797/s
-     iteration_duration.........: avg=916.48ms min=65.67µs  med=1s       max=1s       p(90)=1s       p(95)=1s      
+     iteration_duration.........: avg=916.48ms min=65.67µs  med=1s       max=1s       p(90)=1s       p(95)=1s
      iterations.................: 10     0.991797/s
      vus........................: 1      min=1 max=1
      vus_max....................: 1      min=1 max=1

client/client.go

@@ -2,38 +2,42 @@ package client
 
 import (
 	"context"
+	"fmt"
 	"net/http"
-	"net/http/httptrace"
+	"time"
 
 	"github.com/dop251/goja"
 	"go.k6.io/k6/js/modules"
 	k6HTTP "go.k6.io/k6/js/modules/k6/http"
-	"go.opentelemetry.io/contrib/instrumentation/net/http/httptrace/otelhttptrace"
-	"go.opentelemetry.io/otel"
-	"go.opentelemetry.io/otel/attribute"
-	"go.opentelemetry.io/otel/propagation"
-	"go.opentelemetry.io/otel/trace"
+	"go.k6.io/k6/metrics"
 )
 
+type Options struct {
+	Propagator string
+}
+
 type TracingClient struct {
 	vu          modules.VU
 	httpRequest HttpRequestFunc
+
+	options Options
 }
 
 type HTTPResponse struct {
-	*k6HTTP.Response
-	TraceID string
+	*k6HTTP.Response `js:"-"`
+	TraceID          string
 }
 
 type (
 	HttpRequestFunc func(method string, url goja.Value, args ...goja.Value) (*k6HTTP.Response, error)
 	HttpFunc        func(ctx context.Context, url goja.Value, args ...goja.Value) (*k6HTTP.Response, error)
 )
 
-func New(vu modules.VU, requestFunc HttpRequestFunc) *TracingClient {
+func New(vu modules.VU, requestFunc HttpRequestFunc, options Options) *TracingClient {
 	return &TracingClient{
 		httpRequest: requestFunc,
 		vu:          vu,
+		options:     options,
 	}
 }
 
@@ -72,44 +76,76 @@ func (c *TracingClient) Options(url goja.Value, args ...goja.Value) (*HTTPRespon
 	return c.WithTrace(requestToHttpFunc(http.MethodOptions, c.httpRequest), "HTTP OPTIONS", url, args...)
 }
 
-func (c *TracingClient) WithTrace(fn HttpFunc, spanName string, url goja.Value, args ...goja.Value) (*HTTPResponse, error) {
-	ctx, _, span := startTraceAndSpan(c.vu.Context(), spanName)
-	defer span.End()
-
-	id := span.SpanContext().TraceID().String()
-
-	ctx, val := getTraceHeadersArg(ctx)
-
-	args = append(args, val)
-	res, err := fn(ctx, url, args...)
-	span.SetAttributes(attribute.String("http.method", res.Request.Method), attribute.Int("http.status_code", res.Response.Status), attribute.String("http.url", res.Request.URL))
-	// TODO: extract the textmap from the response
-	return &HTTPResponse{Response: res, TraceID: id}, err
-}
-
-func startTraceAndSpan(ctx context.Context, name string) (context.Context, trace.Tracer, trace.Span) {
-	trace := otel.Tracer("xk6/http")
-	ctx, span := trace.Start(ctx, name)
-	return ctx, trace, span
+func isNilly(val goja.Value) bool {
+	return val == nil || goja.IsNull(val) || goja.IsUndefined(val)
 }
 
-func getTraceHeadersArg(ctx context.Context) (context.Context, goja.Value) {
-	vm := goja.New()
-
-	ctx = httptrace.WithClientTrace(ctx, otelhttptrace.NewClientTrace(ctx))
+func (c *TracingClient) WithTrace(fn HttpFunc, spanName string, url goja.Value, args ...goja.Value) (*HTTPResponse, error) {
+	state := c.vu.State()
+	if state == nil {
+		return nil, fmt.Errorf("HTTP requests can only be made in the VU context")
+	}
 
-	h := http.Header{}
+	traceID, _, err := EncodeTraceID(TraceID{
+		Prefix:            K6Prefix,
+		Code:              K6Code_Cloud,
+		UnixTimestampNano: uint64(time.Now().UnixNano()) / uint64(time.Millisecond),
+	})
+	if err != nil {
+		return nil, err
+	}
 
-	otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(h))
+	tracingHeaders, err := GenerateHeaderBasedOnPropagator(c.options.Propagator, traceID)
+	if err != nil {
+		return nil, err
+	}
 
-	headers := map[string][]string{}
-	for key, header := range h {
-		headers[key] = header
+	// This makes sure that the tracing header will always be added correctly to
+	// the HTTP request headers, whether they were explicitly specified by the
+	// user in the script or not.
+	//
+	// First we make sure to either get the existing request params, or create
+	// them from scratch if they were not specified:
+	rt := c.vu.Runtime()
+	var params *goja.Object
+	if len(args) < 2 {
+		params = rt.NewObject()
+		if len(args) == 0 {
+			args = []goja.Value{goja.Null(), params}
+		} else {
+			args = append(args, params)
+		}
+	} else {
+		jsParams := args[1]
+		if isNilly(jsParams) {
+			params = rt.NewObject()
+			args[1] = params
+		} else {
+			params = jsParams.ToObject(rt)
+		}
+	}
+	// Then we either augment the existing params.headers or create them:
+	var headers *goja.Object
+	if jsHeaders := params.Get("headers"); isNilly(jsHeaders) {
+		headers = rt.NewObject()
+		params.Set("headers", headers)
+	} else {
+		headers = jsHeaders.ToObject(rt)
+	}
+	for key, val := range tracingHeaders {
+		headers.Set(key, val)
 	}
 
-	val := vm.ToValue(map[string]map[string][]string{
-		"headers": headers,
+	// TODO: set span_id as well as some other metadata?
+	state.Tags.Modify(func(tagsAndMeta *metrics.TagsAndMeta) {
+		tagsAndMeta.SetMetadata("trace_id", traceID)
 	})
+	defer state.Tags.Modify(func(tagsAndMeta *metrics.TagsAndMeta) {
+		tagsAndMeta.DeleteMetadata("trace_id")
+	})
+
+	// This calls the actual request() function from k6/http with our augmented arguments
+	res, e := fn(c.vu.Context(), url, args...)
 
-	return ctx, val
+	return &HTTPResponse{Response: res, TraceID: traceID}, e
 }