From 09a768494cd04a7efbf741169433b42d67593afd Mon Sep 17 00:00:00 2001
From: Mihail Stoykov <M.Stoikov@gmail.com>
Date: Wed, 16 Oct 2024 16:22:49 +0300
Subject: [PATCH] refactor: use context.AfterFunc around VU management

This cuts down on the amount of goroutines needed to be up for each VU.
---
 js/runner.go                           | 17 +++++++----------
 lib/testutils/minirunner/minirunner.go |  6 ++----
 2 files changed, 9 insertions(+), 14 deletions(-)

diff --git a/js/runner.go b/js/runner.go
index 2a21eab6ef6..976f09e214f 100644
--- a/js/runner.go
+++ b/js/runner.go
@@ -368,10 +368,9 @@ func (r *Runner) HandleSummary(ctx context.Context, summary *lib.Summary) (map[s
 		return nil, err
 	}
 
-	go func() {
-		<-summaryCtx.Done()
+	_ = context.AfterFunc(summaryCtx, func() {
 		vu.Runtime.Interrupt(context.Canceled)
-	}()
+	})
 	vu.moduleVUImpl.ctx = summaryCtx
 
 	callbackResult := sobek.Undefined()
@@ -540,10 +539,9 @@ func (r *Runner) runPart(
 		return sobek.Undefined(), nil
 	}
 
-	go func() {
-		<-ctx.Done()
+	_ = context.AfterFunc(ctx, func() {
 		vu.Runtime.Interrupt(context.Canceled)
-	}()
+	})
 	vu.moduleVUImpl.ctx = ctx
 
 	groupPath, err := lib.NewGroupPath(lib.RootGroupPath, name)
@@ -706,9 +704,8 @@ func (u *VU) Activate(params *lib.VUActivationParams) lib.ActiveVU {
 		return avu.scIterGlobal
 	}
 
-	go func() {
-		// Wait for the run context to be over
-		<-ctx.Done()
+	// Wait for the run context to be over
+	context.AfterFunc(ctx, func() {
 		// Interrupt the JS runtime
 		u.Runtime.Interrupt(context.Canceled)
 		// Wait for the VU to stop running, if it was, and prevent it from
@@ -718,7 +715,7 @@ func (u *VU) Activate(params *lib.VUActivationParams) lib.ActiveVU {
 		if params.DeactivateCallback != nil {
 			params.DeactivateCallback(u)
 		}
-	}()
+	})
 
 	return avu
 }
diff --git a/lib/testutils/minirunner/minirunner.go b/lib/testutils/minirunner/minirunner.go
index d5fc66fe050..3f7b48d3cc8 100644
--- a/lib/testutils/minirunner/minirunner.go
+++ b/lib/testutils/minirunner/minirunner.go
@@ -172,9 +172,7 @@ func (vu *VU) Activate(params *lib.VUActivationParams) lib.ActiveVU {
 		return avu.scIterGlobal
 	}
 
-	go func() {
-		<-ctx.Done()
-
+	context.AfterFunc(ctx, func() {
 		// Wait for the VU to stop running, if it was, and prevent it from
 		// running again for this activation
 		avu.busy <- struct{}{}
@@ -182,7 +180,7 @@ func (vu *VU) Activate(params *lib.VUActivationParams) lib.ActiveVU {
 		if params.DeactivateCallback != nil {
 			params.DeactivateCallback(vu)
 		}
-	}()
+	})
 
 	return avu
 }