diff options
Diffstat (limited to 'main.go')
| -rw-r--r-- | main.go | 23 |
1 files changed, 13 insertions, 10 deletions
@@ -8,6 +8,7 @@ import ( "path/filepath" "strings" "sync" + "time" ) func fmtAddress(node, port string) string { @@ -23,7 +24,7 @@ func isIPv6(node string) bool { return ip != nil && ip.To4() == nil } -var wg sync.WaitGroup +var jwg JobWaitGroup func Emit(pt Point) { if pt == nil { @@ -33,23 +34,25 @@ func Emit(pt Point) { } func DoHostfile(fname string) { - defer wg.Done() hostname := filepath.Base(fname) + cfg, err := readConfigFile(fname) if err != nil { Emit(NewPoint("public", map[string]string{"host": hostname}, map[string]interface{}{"error": err.Error()})) return } + + addresses := make(map[string]struct{}) for _, address := range getAddresses(cfg) { - addressStr := fmtAddress(address.Node, address.Port) - wg.Add(2) - go Emit(DoAddress(hostname, "tcp4", addressStr)) - go Emit(DoAddress(hostname, "tcp6", addressStr)) + addresses[fmtAddress(address.Node, address.Port)] = struct{}{} + } + for address := range addresses { + jwg.Do(hostname+"/"+address+"/tcp4", func() { Emit(DoAddress(hostname, "tcp4", address)) }) + jwg.Do(hostname+"/"+address+"/tcp6", func() { Emit(DoAddress(hostname, "tcp6", address)) }) } } func DoAddress(host, network, address string) Point { - defer wg.Done() tags := map[string]string{ "host": host, "network": network, @@ -100,6 +103,7 @@ func Hello(addr *net.TCPAddr) (name, version string, err error) { return "", "", err } defer conn.Close() + conn.SetDeadline(time.Now().Add(10 * time.Second)) conn.CloseWrite() all, _ := ioutil.ReadAll(conn) line := strings.TrimRight(string(all), "\n") @@ -112,8 +116,7 @@ func Hello(addr *net.TCPAddr) (name, version string, err error) { func main() { for _, fname := range os.Args[1:] { - wg.Add(1) - go DoHostfile(fname) + DoHostfile(fname) } - wg.Wait() + jwg.Watch(time.Second) } |