golang练手小项目系列(2)-并发爬虫

实现

import (

"bytes"

"flag"

"fmt"

"golang.org/x/net/html"

"io"

"log"

"net/http"

"net/url"

"os"

"path/filepath"

"strings"

"sync"

"time"

)

type URLInfo struct {

url string

depth int

}

var base *url.URL

func forEachNode(n *html.Node, pre, post func(n *html.Node)){

if pre != nil{

pre(n)

}

for c := n.FirstChild; c != nil; c = c.NextSibling{

forEachNode(c, pre, post)

}

if post != nil{

post(n)

}

}

func linkNodes(n *html.Node) []*html.Node {

var links []*html.Node

visitNode := func(n *html.Node) {

if n.Type == html.ElementNode && n.Data == "a" {

links = append(links, n)

}

}

forEachNode(n, visitNode, nil)

return links

}

func linkURLs(linkNodes []*html.Node, base *url.URL) []string {

var urls []string

for _, n := range linkNodes {

for _, a := range n.Attr {

if a.Key != "href" {

continue

}

link, err := base.Parse(a.Val)

// ignore bad and non-local URLs

if err != nil {

log.Printf("skipping %q: %s", a.Val, err)

continue

}

if link.Host != base.Host {

//log.Printf("skipping %q: non-local host", a.Val)

continue

}

if strings.HasPrefix(link.String(), "javascript"){

continue

}

urls = append(urls, link.String())

}

}

return urls

}

func rewriteLocalLinks(linkNodes []*html.Node, base *url.URL) {

for _, n := range linkNodes {

for i, a := range n.Attr {

if a.Key != "href" {

continue

}

link, err := base.Parse(a.Val)

if err != nil || link.Host != base.Host {

continue // ignore bad and non-local URLs

}

link.Scheme = ""

link.Host = ""

link.User = nil

a.Val = link.String()

n.Attr[i] = a

}

}

}

func Extract(url string)(urls []string, err error){

timeout := time.Duration(10 * time.Second)

client := http.Client{

Timeout: timeout,

}

resp, err := client.Get(url)

if err != nil{

fmt.Println(err)

return nil, err

}

if resp.StatusCode != http.StatusOK{

resp.Body.Close()

return nil, fmt.Errorf("getting %s:%s", url, resp.StatusCode)

}

if err != nil{

return nil, fmt.Errorf("parsing %s as HTML: %v", url, err)

}

u, err := base.Parse(url)

if err != nil {

return nil, err

}

if base.Host != u.Host {

log.Printf("not saving %s: non-local", url)

return nil, nil

}

var body io.Reader

contentType := resp.Header["Content-Type"]

if strings.Contains(strings.Join(contentType, ","), "text/html") {

doc, err := html.Parse(resp.Body)

resp.Body.Close()

if err != nil {

return nil, fmt.Errorf("parsing %s as HTML: %v", u, err)

}

nodes := linkNodes(doc)

urls = linkURLs(nodes, u)

rewriteLocalLinks(nodes, u)

b := &bytes.Buffer{}

err = html.Render(b, doc)

if err != nil {

log.Printf("render %s: %s", u, err)

}

body = b

}

err = save(resp, body)

return urls, err

}

func crawl(url string) []string{

list, err := Extract(url)

if err != nil{

log.Print(err)

}

return list

}

func save(resp *http.Response, body io.Reader) error {

u := resp.Request.URL

filename := filepath.Join(u.Host, u.Path)

if filepath.Ext(u.Path) == "" {

filename = filepath.Join(u.Host, u.Path, "index.html")

}

err := os.MkdirAll(filepath.Dir(filename), 0777)

if err != nil {

return err

}

fmt.Println("filename:", filename)

file, err := os.Create(filename)

if err != nil {

return err

}

if body != nil {

_, err = io.Copy(file, body)

} else {

_, err = io.Copy(file, resp.Body)

}

if err != nil {

log.Print("save: ", err)

}

err = file.Close()

if err != nil {

log.Print("save: ", err)

}

return nil

}

func parallellyCrawl(initialLinks string, concurrency, depth int){

worklist := make(chan []URLInfo, 1)

unseenLinks := make(chan URLInfo, 1)

//值为1时表示进入unseenLinks队列，值为2时表示crawl完成

seen := make(map[string] int)

seenLock := sync.Mutex{}

var urlInfos []URLInfo

for _, url := range strings.Split(initialLinks, " "){

urlInfos = append(urlInfos, URLInfo{url, 1})

}

go func() {worklist <- urlInfos}()

go func() {

for{

time.Sleep(1 * time.Second)

seenFlag := true

seenLock.Lock()

for k := range seen{

if seen[k] == 1{

seenFlag = false

}

}

seenLock.Unlock()

if seenFlag && len(worklist) == 0{

close(unseenLinks)

close(worklist)

break

}

}

}()

for i := 0; i < concurrency; i++{

go func() {

for link := range unseenLinks{

foundLinks := crawl(link.url)

var urlInfos []URLInfo

for _, u := range foundLinks{

urlInfos = append(urlInfos, URLInfo{u, link.depth + 1})

}

go func(finishedUrl string) {

worklist <- urlInfos

seenLock.Lock()

seen[finishedUrl] = 2

seenLock.Unlock()

}(link.url)

}

}()

}

for list := range worklist{

for _, link := range list {

if link.depth > depth{

continue

}

seenLock.Lock()

_, ok := seen[link.url]

seenLock.Unlock()

if !ok{

seenLock.Lock()

seen[link.url] = 1

seenLock.Unlock()

unseenLinks <- link

}

}

}

fmt.Printf("共访问了%d个页面", len(seen))

}

func main() {

var maxDepth int

var concurrency int

var initialLink string

flag.IntVar(&maxDepth, "d", 3, "max crawl depth")

flag.IntVar(&concurrency, "c", 20, "number of crawl goroutines")

flag.StringVar(&initialLink, "u", "", "initial link")

flag.Parse()

u, err := url.Parse(initialLink)

if err != nil {

fmt.Fprintf(os.Stderr, "invalid url: %s\n", err)

os.Exit(1)

}

base = u

parallellyCrawl(initialLink, concurrency, maxDepth)

}