Add more locks to execution contexts handling

Also refactor getting a context with a world
name.

Author: inancgumus

common/frame.go

@@ -345,13 +345,7 @@ func (f *Frame) hasContext(world string) bool {
 	f.executionContextMu.RLock()
 	defer f.executionContextMu.RUnlock()
 
-	switch world {
-	case mainWorld:
-		return f.mainExecutionContext != nil
-	case utilityWorld:
-		return f.utilityExecutionContext != nil
-	}
-	return false // Should never reach here!
+	return f.getExecCtx(world) != nil
 }
 
 func (f *Frame) hasLifecycleEventFired(event LifecycleEvent) bool {
@@ -517,20 +511,18 @@ func (f *Frame) waitForFunction(apiCtx context.Context, world string, predicateF
 		"fid:%s furl:%q world:%s pt:%s timeout:%s",
 		f.ID(), f.URL(), world, polling, timeout)
 
-	rt := k6common.GetRuntime(f.ctx)
 	f.waitForExecutionContext(world)
 
 	f.executionContextMu.RLock()
 	defer f.executionContextMu.RUnlock()
 
-	execCtx := f.mainExecutionContext
-	if world == utilityWorld {
-		execCtx = f.utilityExecutionContext
-	}
+	execCtx := f.getExecCtx(world)
 	injected, err := execCtx.getInjectedScript(apiCtx)
 	if err != nil {
 		return nil, err
 	}
+
+	rt := k6common.GetRuntime(f.ctx)
 	pageFn := rt.ToValue(`
 		(injected, predicate, polling, timeout, ...args) => {
 			return injected.waitForPredicateFunction(predicate, polling, timeout, ...args);
@@ -1241,7 +1233,12 @@ func (f *Frame) SetContent(html string, opts goja.Value) {
 		document.write(html);
 		document.close();
 	}`
-	f.waitForExecutionContext(utilityExecutionContext)
+
+	f.waitForExecutionContext(utilityWorld)
+
+	f.executionContextMu.RLock()
+	defer f.executionContextMu.RUnlock()
+
 	_, err := f.utilityExecutionContext.evaluate(f.ctx, true, true, rt.ToValue(js), rt.ToValue(html))
 	if err != nil {
 		k6common.Throw(rt, err)
@@ -1377,6 +1374,9 @@ func (f *Frame) WaitForFunction(pageFunc goja.Value, opts goja.Value, args ...go
 		k6common.Throw(rt, fmt.Errorf("failed parsing options: %w", err))
 	}
 
+	f.executionContextMu.RLock()
+	defer f.executionContextMu.RUnlock()
+
 	handle, err := f.waitForFunction(f.ctx, utilityWorld, pageFunc, parsedOpts.Polling, parsedOpts.Interval, parsedOpts.Timeout, args...)
 	if err != nil {
 		k6common.Throw(rt, err)
@@ -1445,6 +1445,32 @@ func (f *Frame) WaitForTimeout(timeout int64) {
 	}
 }
 
+func (f *Frame) adoptBackendNodeID(world string, id cdp.BackendNodeID) (*ElementHandle, error) {
+	f.executionContextMu.RLock()
+	defer f.executionContextMu.RUnlock()
+
+	ec := f.getExecCtx(world)
+	return ec.adoptBackendNodeID(id)
+}
+
+func (f *Frame) evaluate(world string, apiCtx context.Context, forceCallable bool, returnByValue bool, pageFunc goja.Value, args ...goja.Value) (res interface{}, err error) {
+	f.executionContextMu.RLock()
+	defer f.executionContextMu.RUnlock()
+
+	ec := f.getExecCtx(world)
+	return ec.evaluate(apiCtx, forceCallable, returnByValue, pageFunc, args...)
+}
+
+// getExecCtx returns an execution context using a given world name.
+// Unsafe to use concurrently.
+func (f *Frame) getExecCtx(world string) frameExecutionContext {
+	ec := f.mainExecutionContext
+	if world == utilityWorld {
+		ec = f.utilityExecutionContext
+	}
+	return ec
+}
+
 // frameExecutionContext represents a JS execution context that belongs to Frame.
 type frameExecutionContext interface {
 	// adoptBackendNodeID adopts specified backend node into this execution

common/frame_test.go

@@ -52,7 +52,7 @@ func TestFrameNilDocument(t *testing.T) {
 	// document() waits for the main execution context
 	ok := make(chan struct{}, 1)
 	go func() {
-		frame.setContext(mainExecutionContext, stub)
+		frame.setContext(mainWorld, stub)
 		ok <- struct{}{}
 	}()
 	select {

common/page.go

@@ -193,7 +193,7 @@ func (p *Page) evaluateOnNewDocument(source string) {
 	// TODO: implement
 }
 
-func (p *Page) getFrameElement(f *Frame) (*ElementHandle, error) {
+func (p *Page) getFrameElement(f *Frame) (handle *ElementHandle, _ error) {
 	if f == nil {
 		p.logger.Debugf("Page:getFrameElement", "sid:%v frame:nil", p.sessionID())
 	} else {
@@ -220,8 +220,7 @@ func (p *Page) getFrameElement(f *Frame) (*ElementHandle, error) {
 	if parent == nil {
 		return nil, errors.New("frame has been detached 2")
 	}
-	handle, err := parent.mainExecutionContext.adoptBackendNodeID(backendNodeId)
-	return handle, err
+	return parent.adoptBackendNodeID(mainWorld, backendNodeId)
 }
 
 func (p *Page) getOwnerFrame(apiCtx context.Context, h *ElementHandle) cdp.FrameID {

common/screenshotter.go

@@ -47,7 +47,7 @@ func newScreenshotter(ctx context.Context) *screenshotter {
 
 func (s *screenshotter) fullPageSize(p *Page) (*Size, error) {
 	rt := k6common.GetRuntime(s.ctx)
-	result, err := p.frameManager.mainFrame.mainExecutionContext.evaluate(s.ctx, true, true, rt.ToValue(`
+	result, err := p.frameManager.mainFrame.evaluate(mainWorld, s.ctx, true, true, rt.ToValue(`
         () => {
             if (!document.body || !document.documentElement) {
                 return null;