admin/spray
1
0
Fork 0
mirror of https://github.com/chainreactors/spray.git synced 2025-09-15 19:50:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

优化代码退出的逻辑, 防止结果无法正确输出, 或管道无法正确关闭

Browse Source
This commit is contained in:
M09Ic 2022-11-17 16:58:14 +08:00
parent 1895496e0b
commit 2399c4ff5c
2 changed files with 12 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
internal/pool.go
View File

@ -52,7 +52,6 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
return pool.buildPathRequest(s) return pool.buildPathRequest(s)
} }
pool.check = func() { pool.check = func() {
pool.wg.Add(1)
_ = pool.pool.Invoke(newUnit(pkg.RandPath(), CheckSource)) _ = pool.pool.Invoke(newUnit(pkg.RandPath(), CheckSource))
if pool.failedCount > pool.BreakThreshold { if pool.failedCount > pool.BreakThreshold {
@ -67,7 +66,6 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
} }
pool.check = func() { pool.check = func() {
pool.wg.Add(1)
_ = pool.pool.Invoke(newUnit(pkg.RandHost(), CheckSource)) _ = pool.pool.Invoke(newUnit(pkg.RandHost(), CheckSource))
if pool.failedCount > pool.BreakThreshold { if pool.failedCount > pool.BreakThreshold {
@ -83,7 +81,6 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
req, err := pool.genReq(unit.path) req, err := pool.genReq(unit.path)
if err != nil { if err != nil {
logs.Log.Error(err.Error()) logs.Log.Error(err.Error())
return
} }
var bl *pkg.Baseline var bl *pkg.Baseline
@ -148,7 +145,6 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
pool.bar.Done() pool.bar.Done()
} }
pool.wg.Done()
}) })
pool.pool = p pool.pool = p
@ -160,6 +156,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
} }
} else { } else {
pool.BaseCompare(bl) pool.BaseCompare(bl)
pool.wg.Done()
} }
} }
@ -260,6 +257,7 @@ Loop:
break Loop break Loop
} }
} }
//p.wg.Add(100)
p.wg.Wait() p.wg.Wait()
p.Close() p.Close()
} }
@ -363,7 +361,7 @@ func (p *Pool) recover() {
} }
func (p *Pool) Close() { func (p *Pool) Close() {
for !p.analyzeDone { for p.analyzeDone {
time.Sleep(time.Duration(100) * time.Millisecond) time.Sleep(time.Duration(100) * time.Millisecond)
} }
close(p.tempCh) close(p.tempCh)

11
internal/runner.go
View File

@ -159,7 +159,11 @@ func (r *Runner) Outputting() {
for { for {
select { select {
case bl := <-r.OutputCh: case bl, ok := <-r.OutputCh:
if !ok {
return
}
if bl.IsValid { if bl.IsValid {
outFunc(bl) outFunc(bl)
} else { } else {
@ -172,7 +176,10 @@ func (r *Runner) Outputting() {
go func() { go func() {
for { for {
select { select {
case bl := <-r.FuzzyCh: case bl, ok := <-r.FuzzyCh:
if !ok {
return
}
if r.Fuzzy { if r.Fuzzy {
logs.Log.Console("[baseline.fuzzy] " + bl.String() + "\n") logs.Log.Console("[baseline.fuzzy] " + bl.String() + "\n")
} }