Add context done case

If the context is done before all the tasks are read from the queue,
the remaining tasks on the queue will not be read and so the done
channel will not be closed causing the goroutine to wait forever.

This is solved by adding a context.Done() check so that when the
context is closed we can carry on with shutting down of the iteration.

browser/page_mapping.go

@@ -452,9 +452,11 @@ func mapPageOn(vu moduleVU, p *common.Page) func(common.PageOnEventName, sobek.C
 			}
 		}
 
+		ctx := vu.Context()
+
 		// Run the the event handler in the task queue to
 		// ensure that the handler is executed on the event loop.
-		tq := vu.taskQueueRegistry.get(vu.Context(), p.TargetID())
+		tq := vu.taskQueueRegistry.get(ctx, p.TargetID())
 		eventHandler := func(event common.PageOnEvent) {
 			mapping := pageOnEvent.mapp(vu, event)
 
@@ -475,7 +477,10 @@ func mapPageOn(vu moduleVU, p *common.Page) func(common.PageOnEventName, sobek.C
 			})
 
 			if pageOnEvent.wait {
-				<-done
+				select {
+				case <-done:
+				case <-ctx.Done():
+				}
 			}
 		}