A Tour of Go: Exercise: Web Crawler

tour_of_go_exercise_web_crawler.go

// In this exercise you'll use Go's concurrency features to parallelize a web crawler.
//
// Modify the Crawl function to fetch URLs in parallel without fetching the same URL twice.
//
// Hint: you can keep a cache of the URLs that have been fetched on a map, but maps alone are not safe for concurrent use!

package main

import (
	"fmt"
	"sync"
)

type FetchMapItem struct {
	state int // 0 - not started, 1 - loading, 2 - done
	body  string
	urls  []string
	err   error
	ready sync.Mutex
}

type FetchMap struct {
	m map[string]*FetchMapItem
	l sync.Mutex
}

var fetchMap = FetchMap{make(map[string]*FetchMapItem), sync.Mutex{}}

func (fm *FetchMap) IsPresent(url string) bool {
	fm.l.Lock()
	defer fm.l.Unlock()
	item, ok := fm.m[url]
	return ok && item.state > 0
}

func (fm *FetchMap) AllCompleted() bool {
	fm.l.Lock()
	defer fm.l.Unlock()
	for _, item := range fm.m {
		if item.state == 1 {
			return false
		}
	}
	return true
}

func (fm *FetchMap) Get(url string) *FetchMapItem {
	fm.l.Lock()
	defer fm.l.Unlock()
	item, ok := fm.m[url]
	if !ok {
		fm.m[url] = &FetchMapItem{0, "", nil, nil, sync.Mutex{}}
		item = fm.m[url]
	}

	return item
}

type Fetcher interface {
	// Fetch returns the body of URL and
	// a slice of URLs found on that page.
	Fetch(url string) (body string, urls []string, err error)
}

type CrawlResult struct {
	cached bool
	url    string
	body   string
	err    error
}

// constb: This could've been a WaitGroup, but those aren't covered by the Tour
var crawlOps int

// Crawl uses fetcher to recursively crawl
// pages starting with url, to a maximum of depth.
func Crawl(url string, depth int, fetcher Fetcher, out chan CrawlResult) {
	if depth <= 0 {
		return
	}

	crawlOps++

	f := fetchMap.Get(url)
	if f.state == 0 {
		fetchMap.l.Lock()
		f.state = 1
		fetchMap.l.Unlock()
		go func() {
			f.ready.Lock()
			defer f.ready.Unlock()
			body, urls, err := fetcher.Fetch(url)

			f.state, f.body, f.urls, f.err = 2, body, urls, err

			out <- CrawlResult{false, url, body, err}

			if err == nil {
				for _, u := range urls {
					Crawl(u, depth-1, fetcher, out)
				}
			}
			crawlOps--
			if crawlOps == 0 {
				close(out)
			}
		}()
	} else {
		go func() {
			f.ready.Lock()
			defer f.ready.Unlock()
			out <- CrawlResult{true, url, f.body, f.err}
			crawlOps--
			if crawlOps == 0 {
				close(out)
			}
		}()
	}
}

func main() {
	ch := make(chan CrawlResult)
	go Crawl("https://golang.org/", 4, fetcher, ch)
	for r := range ch {
		if r.cached {
			fmt.Print("[cached] ")
		}
		if r.err != nil {
			fmt.Println(r.err)
		} else {
			fmt.Printf("found: %s %q\n", r.url, r.body)
		}
	}
}

// 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, fmt.Errorf("not found: %s", url)
}

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