Don't block the event loop on context end and related fixes

Previously setTimeout (and co) might block the event loop in cases where
timeout hasn't triggered but the context has been done.

While fixing that some other potential races were found. Most of them
due to trying to write less code - but it is actually important that if
we are using the event loop we do all the stuff on it.

Additionally fixing a panic in init context if a timer gets interrupted
- all due to using the logger from the lib.State.

The test added did tease out a bunch of this issues, and does reproduce
the original issue *on* occasion. Unfortunately a better reproducible
test seems to be very hard to write

timers/timers.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"github.com/mstoykov/k6-taskqueue-lib/taskqueue"
+	"github.com/sirupsen/logrus"
 
 	"github.com/dop251/goja"
 	"go.k6.io/k6/js/modules"
@@ -212,10 +213,21 @@ func (e *Timers) closeTaskQueue() {
 	// so that we do not execute it twice
 	e.taskQueueCh = nil
 
-	// wait for this to happen so we don't need to hit the event loop again
-	// instead this just closes the queue
-	ch <- struct{}{}
-	<-ch
+	select {
+	case ch <- struct{}{}:
+		// wait for this to happen so we don't need to hit the event loop again
+		// instead this just closes the queue
+		<-ch
+	case <-e.vu.Context().Done(): // still shortcircuit if the context is done as we might block otherwise
+	}
+}
+
+// logger is helper to get a logger either from the state or the initenv
+func (e *Timers) logger() logrus.FieldLogger {
+	if state := e.vu.State(); state != nil {
+		return state.Logger
+	}
+	return e.vu.InitEnv().Logger
 }
 
 func (e *Timers) setupTaskQueueCloserOnIterationEnd() {
@@ -229,23 +241,28 @@ func (e *Timers) setupTaskQueueCloserOnIterationEnd() {
 			// lets report timers won't be executed and clean the fields for the next execution
 			// we need to do this on the event loop as we don't want to have a race
 			q.Queue(func() error {
-				logger := e.vu.State().Logger
+				logger := e.logger()
 				for _, timer := range e.queue.queue {
-					logger.Warnf("%s %d was stopped because the VU iteration was interrupted", timer.name, timer.id)
+					logger.Warnf("%s %d was stopped because the VU iteration was interrupted",
+						timer.name, timer.id)
 				}
 
 				// TODO: use `clear` when we only support go 1.21 and above
 				e.timers = make(map[uint64]time.Time)
+				e.queue.stopTimer()
 				e.queue = new(timerQueue)
 				e.taskQueue = nil
 				return nil
 			})
+			q.Close()
 		case <-ch:
+			e.timers = make(map[uint64]time.Time)
+			e.queue.stopTimer()
+			e.queue = new(timerQueue)
 			e.taskQueue = nil
+			q.Close()
 			close(ch)
 		}
-		e.queue.stopTimer()
-		q.Close()
 	}()
 }
 

timers/timers_test.go

@@ -1,6 +1,7 @@
 package timers
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -138,3 +139,54 @@ func TestSetIntervalOrder(t *testing.T) {
 		log = log[:0]
 	}
 }
+
+func TestSetTimeoutContextCancel(t *testing.T) {
+	t.Parallel()
+	runtime := modulestest.NewRuntime(t)
+	err := runtime.SetupModuleSystem(map[string]any{"k6/x/timers": New()}, nil, nil)
+	require.NoError(t, err)
+
+	rt := runtime.VU.Runtime()
+	var log []string
+	interruptChannel := make(chan struct{})
+	require.NoError(t, rt.Set("print", func(s string) { log = append(log, s) }))
+	require.NoError(t, rt.Set("interrupt", func() {
+		select {
+		case interruptChannel <- struct{}{}:
+		default:
+		}
+	}))
+
+	_, err = rt.RunString(`globalThis.setTimeout = require("k6/x/timers").setTimeout;`)
+	require.NoError(t, err)
+
+	for i := 0; i < 2000; i++ {
+		ctx, cancel := context.WithCancel(context.Background())
+		runtime.CancelContext = cancel
+		runtime.VU.CtxField = ctx
+		runtime.VU.RuntimeField.ClearInterrupt()
+		const interruptMsg = "definitely an interrupt"
+		go func() {
+			<-interruptChannel
+			time.Sleep(time.Millisecond)
+			runtime.CancelContext()
+			runtime.VU.RuntimeField.Interrupt(interruptMsg)
+		}()
+		_, err = runtime.RunOnEventLoop(`
+			(async () => {
+				let poll = async (resolve, reject) => {
+					await (async () => 5);
+					setTimeout(poll, 1, resolve, reject);
+					interrupt();
+				}
+				setTimeout(async () => {
+					await new Promise(poll)
+				}, 0)
+			})()
+		`)
+		if err != nil {
+			require.ErrorContains(t, err, interruptMsg)
+		}
+		require.Empty(t, log)
+	}
+}