crawler.go

package main

import (
	"errors"
	"log"
	"os"
	"sync"
)

var fmt = log.New(os.Stdout, "", 0)

type Fetcher interface {
	Fetch(url string) (body string, urls []string, err error)
}

type Crawler struct {
	Fetcher
	sync.Mutex
	sync.WaitGroup
	urls map[string]struct{}
}

func NewCrawler(fetcher Fetcher) *Crawler {
	var mx sync.Mutex
	var wg sync.WaitGroup
	urls := map[string]struct{}{}
	return &Crawler{fetcher, mx, wg, urls}
}

func (c *Crawler) CrawlAndWait(url string, depth int) {
	c.Add(1)
	c.Crawl(url, depth)
	c.Wait()
}

func (c *Crawler) Crawl(url string, depth int) {
	defer c.Done()

	if depth <= 0 {
		return
	}

	body, urls, err := c.Fetch(url)

	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("found: %s %q\n", url, body)

	c.Lock()
	c.urls[url] = struct{}{}
	for _, u := range urls {
		if _, ok := c.urls[u]; !ok {
			c.urls[u] = struct{}{}
			c.Add(1)
			go c.Crawl(u, depth-1)
		}
	}
	c.Unlock()
	return
}

func main() {
	c := NewCrawler(fetcher)
	c.CrawlAndWait("http://golang.org/", 4)
}

// fakeFetcher is Fetcher that returns canned results.
type fakeFetcher map[string]*fakeResult

type fakeResult struct {
	body string
	urls []string
}

func (f fakeFetcher) Fetch(url string) (string, []string, error) {
	if res, ok := f[url]; ok {
		return res.body, res.urls, nil
	}
	return "", nil, errors.New("not found: " + url)
}

// fetcher is a populated fakeFetcher.
var fetcher = fakeFetcher{
	"http://golang.org/": &fakeResult{
		"The Go Programming Language",
		[]string{
			"http://golang.org/pkg/",
			"http://golang.org/cmd/",
		},
	},
	"http://golang.org/pkg/": &fakeResult{
		"Packages",
		[]string{
			"http://golang.org/",
			"http://golang.org/cmd/",
			"http://golang.org/pkg/fmt/",
			"http://golang.org/pkg/os/",
		},
	},
	"http://golang.org/pkg/fmt/": &fakeResult{
		"Package fmt",
		[]string{
			"http://golang.org/",
			"http://golang.org/pkg/",
		},
	},
	"http://golang.org/pkg/os/": &fakeResult{
		"Package os",
		[]string{
			"http://golang.org/",
			"http://golang.org/pkg/",
		},
	},
}

answer to http://tour.golang.org/#73