refactor: move sitemap-creation to 'run' and fix missing xml dir

2 files changed, 26 insertions, 20 deletions

diff --git a/classic.go b/classic.go
index 4a0bcc5..1e264a2 100644
--- a/classic.go
+++ b/classic.go
@@ -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 {
diff --git a/run.go b/run.go
index 5f56572..a1c64f7 100644
--- a/run.go
+++ b/run.go
@@ -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
 }