From a84a54670bd15e90f31bc18a92fec4e3786bd13f Mon Sep 17 00:00:00 2001
From: ka3de <danijs12@hotmail.com>
Date: Fri, 10 Nov 2023 16:25:17 +0100
Subject: [PATCH] Propagate traces metadata to tracer

Because the traces registry and its underlying tracer are only
initialized on the first iterStart event for each VU, the browser
registry must keep a reference to the traces metadata until this event
is fired and the traces registry and tracer can be initialized.
---
 browser/module.go        | 12 +++++++++---
 browser/registry.go      | 18 +++++++++++-------
 browser/registry_test.go |  6 +++---
 3 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/browser/module.go b/browser/module.go
index 2648f60d5..ca20cf521 100644
--- a/browser/module.go
+++ b/browser/module.go
@@ -67,9 +67,15 @@ func (m *RootModule) NewModuleInstance(vu k6modules.VU) k6modules.Instance {
 	return &ModuleInstance{
 		mod: &JSModule{
 			Browser: mapBrowserToGoja(moduleVU{
-				VU:                vu,
-				pidRegistry:       m.PidRegistry,
-				browserRegistry:   newBrowserRegistry(context.Background(), vu, m.remoteRegistry, m.PidRegistry),
+				VU:          vu,
+				pidRegistry: m.PidRegistry,
+				browserRegistry: newBrowserRegistry(
+					context.Background(),
+					vu,
+					m.remoteRegistry,
+					m.PidRegistry,
+					m.tracesMetadata,
+				),
 				taskQueueRegistry: newTaskQueueRegistry(vu),
 			}),
 			Devices:         common.GetDevices(),
diff --git a/browser/registry.go b/browser/registry.go
index a6c3dff3c..f110c934d 100644
--- a/browser/registry.go
+++ b/browser/registry.go
@@ -173,8 +173,9 @@ func (r *remoteRegistry) isRemoteBrowser() (string, bool) {
 type browserRegistry struct {
 	vu k6modules.VU
 
-	tr     *tracesRegistry
-	trInit sync.Once
+	tr             *tracesRegistry
+	trInit         sync.Once
+	tracesMetadata map[string]string
 
 	mu sync.RWMutex
 	m  map[int64]*common.Browser
@@ -187,7 +188,7 @@ type browserRegistry struct {
 type browserBuildFunc func(ctx context.Context) (*common.Browser, error)
 
 func newBrowserRegistry(
-	ctx context.Context, vu k6modules.VU, remote *remoteRegistry, pids *pidRegistry,
+	ctx context.Context, vu k6modules.VU, remote *remoteRegistry, pids *pidRegistry, tracesMetadata map[string]string,
 ) *browserRegistry {
 	bt := chromium.NewBrowserType(vu)
 	builder := func(ctx context.Context) (*common.Browser, error) {
@@ -215,9 +216,10 @@ func newBrowserRegistry(
 	}
 
 	r := &browserRegistry{
-		vu:      vu,
-		m:       make(map[int64]*common.Browser),
-		buildFn: builder,
+		vu:             vu,
+		tracesMetadata: tracesMetadata,
+		m:              make(map[int64]*common.Browser),
+		buildFn:        builder,
 	}
 
 	exitSubID, exitCh := vu.Events().Global.Subscribe(
@@ -370,7 +372,9 @@ func (r *browserRegistry) initTracesRegistry() {
 	// Use a sync.Once so the traces registry is only initialized once
 	// per VU, as that is the scope for both browser and traces registry.
 	r.trInit.Do(func() {
-		r.tr = newTracesRegistry(browsertrace.NewTracer(r.vu.State().TracerProvider))
+		r.tr = newTracesRegistry(
+			browsertrace.NewTracer(r.vu.State().TracerProvider, r.tracesMetadata),
+		)
 	})
 }
 
diff --git a/browser/registry_test.go b/browser/registry_test.go
index 9806ccc9c..7356fd100 100644
--- a/browser/registry_test.go
+++ b/browser/registry_test.go
@@ -206,7 +206,7 @@ func TestBrowserRegistry(t *testing.T) {
 		var (
 			ctx             = context.Background()
 			vu              = k6test.NewVU(t)
-			browserRegistry = newBrowserRegistry(ctx, vu, remoteRegistry, &pidRegistry{})
+			browserRegistry = newBrowserRegistry(ctx, vu, remoteRegistry, &pidRegistry{}, nil)
 		)
 
 		vu.ActivateVU()
@@ -248,7 +248,7 @@ func TestBrowserRegistry(t *testing.T) {
 		var (
 			ctx             = context.Background()
 			vu              = k6test.NewVU(t)
-			browserRegistry = newBrowserRegistry(ctx, vu, remoteRegistry, &pidRegistry{})
+			browserRegistry = newBrowserRegistry(ctx, vu, remoteRegistry, &pidRegistry{}, nil)
 		)
 
 		vu.ActivateVU()
@@ -283,7 +283,7 @@ func TestBrowserRegistry(t *testing.T) {
 		var (
 			ctx             = context.Background()
 			vu              = k6test.NewVU(t)
-			browserRegistry = newBrowserRegistry(ctx, vu, remoteRegistry, &pidRegistry{})
+			browserRegistry = newBrowserRegistry(ctx, vu, remoteRegistry, &pidRegistry{}, nil)
 		)
 
 		vu.ActivateVU()