fix(script): bind script HTTP requests to the run context

cmd/cli/main.go

@@ -75,7 +75,7 @@ func main() {
 		}),
 	)
 
-	srn.Start(ctx, stopCtrl)
+	srn.Start(stopCtrl)
 
 	switch *combinedConfig.Output {
 	case config.ConfigOutputTypeNone:

internal/sarin/runner.go

@@ -199,6 +199,10 @@ type sarin struct {
 	logError       bool
 	logFile        string
 
+	// jobsCtx ends the run, either when the duration expires or on Ctrl+C.
+	jobsCtx    context.Context //nolint:containedctx
+	jobsCancel context.CancelFunc
+
 	hostClients       []*fasthttp.HostClient
 	scriptHTTPClients []*scriptHTTPClient
 	responses         *SarinResponseData
@@ -272,15 +276,20 @@ func NewSarin(
 
 	scriptChain := script.NewChain(luaSources, jsSources)
 
+	jobsCtx, jobsCancel := context.WithCancel(ctx)
+
 	var scriptHTTPClients []*scriptHTTPClient
 	if !scriptChain.IsEmpty() {
-		scriptHTTPClients, err = newScriptHTTPClients(ctx, proxyURLs, workers)
+		scriptHTTPClients, err = newScriptHTTPClients(jobsCtx, proxyURLs, workers)
 		if err != nil {
+			jobsCancel()
 			return nil, err
 		}
 	}
 
 	srn := &sarin{
+		jobsCtx:           jobsCtx,
+		jobsCancel:        jobsCancel,
 		workers:           workers,
 		requestURL:        requestURL,
 		methods:           methods,
@@ -316,9 +325,7 @@ func (s sarin) GetResponses() *SarinResponseData {
 	return s.responses
 }
 
-func (s sarin) Start(ctx context.Context, stopCtrl *StopController) {
+func (s sarin) Start(stopCtrl *StopController) {
-	jobsCtx, jobsCancel := context.WithCancel(ctx)
-
 	var workersWG sync.WaitGroup
 	jobsCh := make(chan struct{}, max(s.workers, 1))
 
@@ -382,15 +389,14 @@ func (s sarin) Start(ctx context.Context, stopCtrl *StopController) {
 	s.startWorkers(&workersWG, jobsCh, &counter, sendLog, sendRespLog)
 
 	if runTUI {
-		//nolint:contextcheck // streamCtx must remain active until all workers complete to ensure all collected data is streamed
 		go s.streamProgress(streamCtx, stopCtrl, streamCh, totalRequests, &counter, tuiLogChannel, showProgressBar)
 	}
 
 	// Setup duration-based cancellation
-	s.setupDurationTimeout(ctx, jobsCancel)
+	s.setupDurationTimeout()
 	// Distribute jobs to workers.
 	// This blocks until all jobs are sent or the context is canceled.
-	s.sendJobs(jobsCtx, jobsCh)
+	s.sendJobs(jobsCh)
 
 	// Close the jobs channel so workers stop after completing their current job
 	close(jobsCh)
@@ -455,31 +461,31 @@ func (s sarin) startWorkers(wg *sync.WaitGroup, jobs <-chan struct{}, counter *a
 	}
 }
 
-func (s sarin) setupDurationTimeout(ctx context.Context, cancel context.CancelFunc) {
+func (s sarin) setupDurationTimeout() {
 	if s.totalDuration != nil {
 		go func() {
 			timer := time.NewTimer(*s.totalDuration)
 			defer timer.Stop()
 			select {
 			case <-timer.C:
-				cancel()
+				s.jobsCancel()
-			case <-ctx.Done():
+			case <-s.jobsCtx.Done():
 				// Context cancelled, cleanup
 			}
 		}()
 	}
 }
 
-func (s sarin) sendJobs(ctx context.Context, jobs chan<- struct{}) {
+func (s sarin) sendJobs(jobs chan<- struct{}) {
 	if s.totalRequests != nil && *s.totalRequests > 0 {
 		for range *s.totalRequests {
-			if ctx.Err() != nil {
+			if s.jobsCtx.Err() != nil {
 				break
 			}
 			jobs <- struct{}{}
 		}
 	} else {
-		for ctx.Err() == nil {
+		for s.jobsCtx.Err() == nil {
 			jobs <- struct{}{}
 		}
 	}

internal/sarin/scripthttp.go

@@ -17,6 +17,7 @@ const scriptHTTPDefaultTimeout = 30 * time.Second
 
 // scriptHTTPClient sends scripts' http.* requests. It is safe for concurrent use.
 type scriptHTTPClient struct {
+	ctx            context.Context //nolint:containedctx
 	client         *fasthttp.Client
 	insecureClient *fasthttp.Client
 	defaultTimeout time.Duration
@@ -29,7 +30,7 @@ var _ script.HTTPDoer = (*scriptHTTPClient)(nil)
 //   - types.ProxyDialError
 func newScriptHTTPClients(ctx context.Context, proxies []url.URL, maxConns uint) ([]*scriptHTTPClient, error) {
 	if len(proxies) == 0 {
-		return []*scriptHTTPClient{newScriptHTTPClient(fasthttp.DialDualStackTimeout, scriptHTTPDefaultTimeout, maxConns)}, nil
+		return []*scriptHTTPClient{newScriptHTTPClient(ctx, fasthttp.DialDualStackTimeout, scriptHTTPDefaultTimeout, maxConns)}, nil
 	}
 
 	clients := make([]*scriptHTTPClient, 0, len(proxies))
@@ -38,12 +39,12 @@ func newScriptHTTPClients(ctx context.Context, proxies []url.URL, maxConns uint)
 		if err != nil {
 			return nil, types.NewProxyDialError(proxy.String(), err)
 		}
-		clients = append(clients, newScriptHTTPClient(dial, scriptHTTPDefaultTimeout, maxConns))
+		clients = append(clients, newScriptHTTPClient(ctx, dial, scriptHTTPDefaultTimeout, maxConns))
 	}
 	return clients, nil
 }
 
-func newScriptHTTPClient(dial fasthttp.DialFuncWithTimeout, defaultTimeout time.Duration, maxConns uint) *scriptHTTPClient {
+func newScriptHTTPClient(ctx context.Context, dial fasthttp.DialFuncWithTimeout, defaultTimeout time.Duration, maxConns uint) *scriptHTTPClient {
 	newClient := func(insecure bool) *fasthttp.Client {
 		// Read and write timeouts stay unset because they would cap the per-call timeout,
 		// so dialWithDeadline is what bounds the TLS handshake.
@@ -60,6 +61,7 @@ func newScriptHTTPClient(dial fasthttp.DialFuncWithTimeout, defaultTimeout time.
 	}
 
 	return &scriptHTTPClient{
+		ctx:            ctx,
 		client:         newClient(false),
 		insecureClient: newClient(true),
 		defaultTimeout: defaultTimeout,
@@ -91,6 +93,10 @@ func dialWithDeadline(dial fasthttp.DialFuncWithTimeout) fasthttp.DialFuncWithTi
 // It can return the following errors:
 //   - types.ScriptHTTPRequestError
 func (c *scriptHTTPClient) Do(r *script.HTTPRequest) (*script.HTTPResponse, error) {
+	if err := c.ctx.Err(); err != nil {
+		return nil, types.NewScriptHTTPRequestError(r.Method, r.URL, err)
+	}
+
 	parsedURL, err := url.Parse(r.URL)
 	if err != nil {
 		return nil, types.NewScriptHTTPRequestError(r.Method, r.URL, err)