1 files changed, 35 insertions, 8 deletions

diff --git a/cmd/crawl/crawl.go b/cmd/crawl/crawl.go
index e31f63e..e7e8582 100644
--- a/cmd/crawl/crawl.go
+++ b/cmd/crawl/crawl.go
@@ -11,10 +11,13 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"os/signal"
+	"runtime/pprof"
 	"strconv"
 	"strings"
 	"sync"
 	"sync/atomic"
+	"syscall"
 	"time"
 
 	"git.autistici.org/ale/crawl"
@@ -30,6 +33,8 @@ var (
 	validSchemes         = flag.String("schemes", "http,https", "comma-separated list of allowed protocols")
 	alwaysIncludeRelated = flag.Bool("include-related", false, "always include related resources (css, images, etc)")
 	outputFile           = flag.String("output", "crawl.warc.gz", "output WARC file")
+
+	cpuprofile = flag.String("cpuprofile", "", "create cpu profile")
 )
 
 func extractLinks(c *crawl.Crawler, u string, depth int, resp *http.Response, err error) error {
@@ -147,14 +152,10 @@ func (c *crawlStats) Dump() {
 	fmt.Fprintf(os.Stderr, "stats: downloaded %d bytes (%.4g KB/s), status: %v\n", c.bytes, rate, c.states)
 }
 
-var (
-	stats *crawlStats
-
-	client *http.Client
-)
+var stats *crawlStats
 
 func fetch(urlstr string) (*http.Response, error) {
-	resp, err := client.Get(urlstr)
+	resp, err := crawl.DefaultClient.Get(urlstr)
 	if err == nil {
 		stats.Update(resp)
 	}
@@ -162,8 +163,6 @@ func fetch(urlstr string) (*http.Response, error) {
 }
 
 func init() {
-	client = &http.Client{}
-
 	stats = &crawlStats{
 		states: make(map[int]int),
 		start:  time.Now(),
@@ -191,6 +190,17 @@ func (b *byteCounter) Read(buf []byte) (int, error) {
 func main() {
 	flag.Parse()
 
+	if *cpuprofile != "" {
+		f, err := os.Create(*cpuprofile)
+		if err != nil {
+			log.Fatal("could not create CPU profile: ", err)
+		}
+		if err := pprof.StartCPUProfile(f); err != nil {
+			log.Fatal("could not start CPU profile: ", err)
+		}
+		defer pprof.StopCPUProfile()
+	}
+
 	outf, err := os.Create(*outputFile)
 	if err != nil {
 		log.Fatal(err)
@@ -216,9 +226,26 @@ func main() {
 	if err != nil {
 		log.Fatal(err)
 	}
+
+	// Set up signal handlers so we can terminate gently if possible.
+	var signaled atomic.Value
+	signaled.Store(false)
+	sigCh := make(chan os.Signal, 1)
+	go func() {
+		<-sigCh
+		log.Printf("exiting due to signal")
+		signaled.Store(true)
+		crawler.Stop()
+	}()
+	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
+
 	crawler.Run(*concurrency)
 
 	crawler.Close()
+
+	if signaled.Load().(bool) {
+		os.Exit(1)
+	}
 	if !*keepDb {
 		os.RemoveAll(*dbPath)
 	}