diff --git a/file_checker.go b/file_checker.go index 584a12d..c7dd62f 100644 --- a/file_checker.go +++ b/file_checker.go @@ -4,6 +4,7 @@ import ( "bytes" "io/ioutil" "net/url" + "sync" "time" "golang.org/x/net/html" @@ -44,6 +45,29 @@ func (c fileChecker) Check(f string) ([]urlResult, error) { return rs, nil } +func (c fileChecker) CheckMany(fs []string, rc chan fileResult) { + wg := sync.WaitGroup{} + + for _, f := range fs { + wg.Add(1) + + go func(f string) { + rs, err := c.Check(f) + + if err != nil { + rc <- fileResult{filename: f, err: err} + } else { + rc <- fileResult{filename: f, urlResults: rs} + } + + wg.Done() + }(f) + } + + wg.Wait() + close(rc) +} + func extractURLs(n *html.Node) []string { us := make(map[string]bool) ns := make([]*html.Node, 0, 1024) diff --git a/main.go b/main.go index 98b2ea1..125f078 100644 --- a/main.go +++ b/main.go @@ -23,23 +23,11 @@ func main() { rc := make(chan fileResult, len(fs)) c := newFileChecker(5 * time.Second) - for _, f := range fs { - go func(f string) { - rs, err := c.Check(f) - - if err != nil { - rc <- fileResult{filename: f, err: err} - } - - rc <- fileResult{filename: f, urlResults: rs} - }(f) - } + go c.CheckMany(fs, rc) ok := true - for i := 0; i < len(fs); i++ { - r := <-rc - + for r := range rc { if !r.Ok() { ok = false printToStderr(r.String(args.verbose))