diff options
| -rw-r--r-- | classic.go | 21 | ||||
| -rw-r--r-- | run.go | 25 |
2 files changed, 26 insertions, 20 deletions
@@ -6,13 +6,12 @@ import ( "fmt" "log" "net/url" - "os" "sync" "golang.org/x/net/html" ) -func classic(startURL url.URL, maxConcurrency, maxURLs, maxDepth int) { +func classic(startURL url.URL, maxConcurrency, maxURLs, maxDepth int) map[url.URL]int { worklist := make(chan []packet) var numPendingSends int @@ -75,23 +74,7 @@ loop: cancel() wg.Wait() - // FIXME: eventually, when all crawlers terminate properly, we - // can move this out to main: that is, all crawlers will - // return the same seen map that will be processed by this - // code. - fmt.Println("Generating sitemap...") - sitemap, err := toSitemap(seen, maxDepth, maxURLs) - if err != nil { - log.Fatal(err) - } - - xmlFilename := fmt.Sprintf("xml/%s.xml", startURL.Host) - - if err := os.WriteFile(xmlFilename, []byte(sitemap), 0666); err != nil { - log.Fatal(err) - } - - fmt.Printf("Wrote sitemap to %s\n", xmlFilename) + return seen } func getBatch(u url.URL) []url.URL { @@ -2,8 +2,12 @@ package main import ( "context" + "errors" "fmt" + "io/fs" + "log" "net/url" + "os" "runtime/pprof" "strings" "time" @@ -35,10 +39,29 @@ func run(ctx context.Context, cmd *cli.Command) error { return err } + var seen map[url.URL]int getLeakProfile(func() { - classic(u, cmd.Int("concurrency"), cmd.Int("maxurls"), cmd.Int("depth")) + seen = classic(u, cmd.Int("concurrency"), cmd.Int("maxurls"), cmd.Int("depth")) }) + // Generate a sitemap. + fmt.Println("Generating sitemap...") + sitemap, err := toSitemap(seen, cmd.Int("depth"), cmd.Int("maxurls")) + if err != nil { + log.Fatal(err) + } + + if err := os.Mkdir("xml", 0750); err != nil && !errors.Is(err, fs.ErrExist) { + log.Fatal(err) + } + + xmlFilename := fmt.Sprintf("xml/%s.xml", u.Host) + if err := os.WriteFile(xmlFilename, []byte(sitemap), 0666); err != nil { + log.Fatal(err) + } + + fmt.Printf("Wrote sitemap to %s\n", xmlFilename) + return nil } |