fix ctrl+c cannot exit

This commit is contained in:
M09Ic 2024-07-24 04:21:43 +08:00
parent 44e88e0aa7
commit d5286eace5
3 changed files with 18 additions and 7 deletions

View File

@ -135,12 +135,22 @@ func Spray() {
} }
go func() { go func() {
c := make(chan os.Signal, 2) exitChan := make(chan os.Signal, 2)
signal.Notify(c, os.Interrupt, syscall.SIGTERM) signal.Notify(exitChan, os.Interrupt, syscall.SIGTERM)
go func() { go func() {
<-c sigCount := 0
logs.Log.Important("exit signal, save stat and exit") for {
canceler() <-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)
}
}
}() }()
}() }()

View File

@ -281,6 +281,7 @@ func (opt *Option) NewRunner() (*Runner, error) {
logs.Log.SetColor(false) logs.Log.SetColor(false)
r.Color = false r.Color = false
} }
if !(opt.Quiet || opt.NoBar) { if !(opt.Quiet || opt.NoBar) {
r.Progress = mpb.New(mpb.WithRefreshRate(100 * time.Millisecond)) r.Progress = mpb.New(mpb.WithRefreshRate(100 * time.Millisecond))
logs.Log.SetOutput(r.Progress) logs.Log.SetOutput(r.Progress)

View File

@ -148,7 +148,7 @@ func (r *Runner) Prepare(ctx context.Context) error {
}() }()
if r.Count > 0 { if r.Count > 0 {
r.addBar(r.Count) r.newBar(r.Count)
} }
r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) { r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) {
@ -298,7 +298,7 @@ func (r *Runner) AddPool(task *Task) {
r.Pools.Invoke(task) r.Pools.Invoke(task)
} }
func (r *Runner) addBar(total int) { func (r *Runner) newBar(total int) {
if r.Progress == nil { if r.Progress == nil {
return return
} }