diff options
| -rw-r--r-- | jwg.go | 88 | ||||
| -rw-r--r-- | main.go | 23 |
2 files changed, 101 insertions, 10 deletions
@@ -0,0 +1,88 @@ +package main + +import ( + "fmt" + "os" + "sort" + "strings" + "sync" + "time" +) + +type JobWaitGroup struct { + lock sync.RWMutex + jobs map[string]bool + wg sync.WaitGroup +} + +func (jwg *JobWaitGroup) Do(name string, fn func()) { + jwg.lock.Lock() + defer jwg.lock.Unlock() + + if jwg.jobs == nil { + jwg.jobs = make(map[string]bool) + } + if _, dup := jwg.jobs[name]; dup { + panic(fmt.Sprintf("job %q already exists", name)) + } + jwg.jobs[name] = true + //fmt.Fprintf(os.Stderr, "add %p %v\n", &jwg.wg, jwg.wg) + jwg.wg.Add(1) + go func() { + defer func() { + jwg.lock.Lock() + defer jwg.lock.Unlock() + jwg.jobs[name] = false + jwg.wg.Done() + //fmt.Fprintf(os.Stderr, "del %p %v\n", &jwg.wg, jwg.wg) + }() + fn() + }() +} + +func (jwg *JobWaitGroup) Wait() { + jwg.wg.Wait() +} + +func (jwg *JobWaitGroup) Status() (int, []string) { + jwg.lock.RLock() + + n := len(jwg.jobs) + var jobs []string + for job, active := range jwg.jobs { + if active { + jobs = append(jobs, job) + } + } + defer jwg.lock.RUnlock() + sort.Strings(jobs) + + return n, jobs +} + +func (jwg *JobWaitGroup) Watch(d time.Duration) { + ticker := time.NewTicker(d) + done := make(chan struct{}) + go func() { + jwg.Wait() + ticker.Stop() + done <- struct{}{} + }() + for { + select { + case <-ticker.C: + n, jobs := jwg.Status() + if len(jobs) == 0 { + panic("no active jobs, but wg still waiting") + } + line := fmt.Sprintf("%d/%d : %v", len(jobs), n, strings.Join(jobs, ", ")) + if len(line) > 70 { + line = line[:67] + "..." + } + fmt.Fprintf(os.Stderr, "\r%-70s", line) + case <-done: + fmt.Fprintf(os.Stderr, "\r%-70s\n", "done") + return + } + } +} @@ -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) } |