diff --git a/cmd/cmd.go b/cmd/cmd.go index 1d7cfc8..47d9bb2 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -135,12 +135,22 @@ func Spray() { } go func() { - c := make(chan os.Signal, 2) - signal.Notify(c, os.Interrupt, syscall.SIGTERM) + exitChan := make(chan os.Signal, 2) + signal.Notify(exitChan, os.Interrupt, syscall.SIGTERM) + go func() { - <-c - logs.Log.Important("exit signal, save stat and exit") - canceler() + sigCount := 0 + for { + <-exitChan + sigCount++ + if sigCount == 1 { + logs.Log.Infof("Exit signal received, saving task and exiting...") + canceler() + } else if sigCount == 2 { + logs.Log.Infof("forcing exit...") + os.Exit(1) + } + } }() }() diff --git a/internal/option.go b/internal/option.go index 4ae66f6..5a33e3d 100644 --- a/internal/option.go +++ b/internal/option.go @@ -281,6 +281,7 @@ func (opt *Option) NewRunner() (*Runner, error) { logs.Log.SetColor(false) r.Color = false } + if !(opt.Quiet || opt.NoBar) { r.Progress = mpb.New(mpb.WithRefreshRate(100 * time.Millisecond)) logs.Log.SetOutput(r.Progress) diff --git a/internal/runner.go b/internal/runner.go index bd9d2b3..c412057 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -148,7 +148,7 @@ func (r *Runner) Prepare(ctx context.Context) error { }() if r.Count > 0 { - r.addBar(r.Count) + r.newBar(r.Count) } r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) { @@ -298,7 +298,7 @@ func (r *Runner) AddPool(task *Task) { r.Pools.Invoke(task) } -func (r *Runner) addBar(total int) { +func (r *Runner) newBar(total int) { if r.Progress == nil { return }