A Tour of Go の Exercise: Web Crawler を書いてみた

A Tour of Go

面白法人カヤックでGo書いてるギリ新卒のkoyoです！
次の新卒が入ってくる時期になってきましたね。新卒でGo書く人は A Tour of Go やる時期なのかなーと思います。
Goを覚えるにはとてもいい題材なのですが、Exerciseが難しい笑（なので新卒の皆さんは全部解けなくてもいいと思います。私も去年は解けなかった）
特に最後の Exercise: Web Crawler が難しいですよね。1年仕事で書いてると、理解も深まるもので、苦戦しましたが良さそうなコードを書くことができたので、記事を書いてみました！

コード

package main

import (
	"fmt"
	"sync"
)

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

type Crawler struct {
	cache *sync.Map
}

func NewCrawler() *Crawler {
	return &Crawler{
		cache: &sync.Map{},
	}
}

func (c *Crawler) Crawl(url string, depth int, fetcher Fetcher) {
	wg := &sync.WaitGroup{}
	c.crawl(url, depth, wg)
	wg.Wait()

	return
}

func (c *Crawler) crawl(url string, depth int, wg *sync.WaitGroup) {
	if depth <= 0 {
		return
	}

	if _, ok := c.cache.Load(url); ok {
		return
	}
	c.cache.Store(url, struct{}{})

	body, urls, err := fetcher.Fetch(url)
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("found: %s %q\n", url, body)
	wg.Add(len(urls))
	for _, u := range urls {
		go func(u string) {
			c.crawl(u, depth-1, wg)
			wg.Done()
		}(u)
	}
}

func main() {
	NewCrawler().Crawl("https://golang.org/", 4, fetcher)
}

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)
}

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/",
		},
	},
}

実行結果

found: https://golang.org/ "The Go Programming Language"
not found: https://golang.org/cmd/
found: https://golang.org/pkg/ "Packages"
found: https://golang.org/pkg/os/ "Package os"
found: https://golang.org/pkg/fmt/ "Package fmt"

Program exited.

簡単な解説

sync.WaitGroup を使って Crawler を定義する方法でやってみました。
channelを使うやり方や、クロージャで済ませる方法もあるかなと思いますね。ただクロージャだと可読性と拡張性が低くなるので、struct定義しちゃう方がいいかなーと思う派ですね。

実装のポイントは sync.Map{} を使うことです。 sync.Mutex と map[string]bool を組み合わせるやり方もあると思うんですが、これを使うと一つだけで済むのでラクですね。
簡単なキャッシュを実装したい場合に使えるのでオススメです！

厳密に考えるとするとエラー返しているので sync.WaitGroup ではなく errgroup.Group 使う方がいいかもですが、実装例の一つとして参考になればと思います。

おまけ: クロージャで解いてみた

package main

import (
	"fmt"
	"sync"
)

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

func Crawl(url string, depth int, fetcher Fetcher) {
	var (
		cache = &sync.Map{}
		wg    = &sync.WaitGroup{}
		crawl func(url string, depth int)
	)

	crawl = func(url string, depth int) {
		if depth <= 0 {
			return
		}

		if _, ok := cache.Load(url); ok {
			return
		}
		cache.Store(url, struct{}{})

		body, urls, err := fetcher.Fetch(url)
		if err != nil {
			fmt.Println(err)
			return
		}

		fmt.Printf("found: %s %q\n", url, body)
		wg.Add(len(urls))
		for _, u := range urls {
			go func(u string) {
				crawl(u, depth-1)
				wg.Done()
			}(u)
		}
	}

	crawl(url, depth)
	wg.Wait()
	return
}

func main() {
	Crawl("https://golang.org/", 4, fetcher)
}

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)
}

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/",
		},
	},
}