diff --git a/main.go b/main.go index 36a0b2d..5e99a86 100644 --- a/main.go +++ b/main.go @@ -7,7 +7,6 @@ import ( "os" "os/signal" "syscall" - "time" "github.com/aykhans/dodo/config" "github.com/aykhans/dodo/requests" @@ -50,10 +49,6 @@ func main() { ctx, cancel := context.WithCancel(context.Background()) go listenForTermination(func() { cancel() }) - if requestConf.Duration > 0 { - time.AfterFunc(requestConf.Duration, func() { cancel() }) - } - responses, err := requests.Run(ctx, requestConf) if err != nil { if err == types.ErrInterrupt { diff --git a/requests/run.go b/requests/run.go index a076671..d3937dd 100644 --- a/requests/run.go +++ b/requests/run.go @@ -20,6 +20,9 @@ import ( // - ctx: The context for managing request lifecycle and cancellation. // - requestConfig: The configuration for the request, including timeout, proxies, and other settings. func Run(ctx context.Context, requestConfig *config.RequestConfig) (Responses, error) { + ctx, cancel := context.WithCancel(ctx) + defer cancel() + clients := getClients( ctx, requestConfig.Timeout, @@ -31,6 +34,10 @@ func Run(ctx context.Context, requestConfig *config.RequestConfig) (Responses, e return nil, types.ErrInterrupt } + if requestConfig.Duration > 0 { + time.AfterFunc(requestConfig.Duration, func() { cancel() }) + } + responses := releaseDodos(ctx, requestConfig, clients) if ctx.Err() != nil && len(responses) == 0 { return nil, types.ErrInterrupt