Fix: let lib/traceroute wait all goroutine finish

This commit is contained in:
2022-05-05 20:50:28 +08:00
parent e4af31660d
commit 0d2e6e1e17
3 changed files with 36 additions and 14 deletions

View File

@ -154,9 +154,12 @@ type TracerouteResult struct {
Hops []TracerouteHop
}
func notify(hop TracerouteHop, channels []chan TracerouteHop) {
func notify(hop TracerouteHop, channels []chan TracerouteHop, DoneChannels chan bool) {
// fmt.Print(hop)
for _, c := range channels {
c <- hop
<-DoneChannels
// fmt.Print("Done")
}
}
@ -166,7 +169,7 @@ func closeNotify(channels []chan TracerouteHop) {
}
}
func Traceroute(dest string, options *TracerouteOptions, c ...chan TracerouteHop) (result TracerouteResult, err error) {
func Traceroute(dest string, options *TracerouteOptions, d chan bool, c ...chan TracerouteHop) (result TracerouteResult, err error) {
result.Hops = []TracerouteHop{}
destAddrBytes, destIPAddr, err := destAddr(dest)
result.DestinationAddress = destAddrBytes
@ -237,7 +240,7 @@ func Traceroute(dest string, options *TracerouteOptions, c ...chan TracerouteHop
if err != nil {
if err, ok := err.(net.Error); ok && err.Timeout() {
// means timeout here
notify(TracerouteHop{Success: false, TTL: ttl}, c)
notify(TracerouteHop{Success: false, TTL: ttl}, c, d)
retry += 1
if retry > options.Retries() {
ttl += 1
@ -271,7 +274,7 @@ func Traceroute(dest string, options *TracerouteOptions, c ...chan TracerouteHop
}
if rm.Type == ipv4.ICMPTypeEchoReply || rm.Type == ipv4.ICMPTypeTimeExceeded {
notify(hop, c)
notify(hop, c, d)
result.Hops = append(result.Hops, hop)
}
@ -282,7 +285,5 @@ func Traceroute(dest string, options *TracerouteOptions, c ...chan TracerouteHop
closeNotify(c)
return result, nil
}
time.Sleep(time.Millisecond * 500)
}
}