Go并发爬虫小案例

package main

import (
    "fmt"
    "io/ioutil"
    "net/http"
    "regexp"
    "strconv"
    "strings"
    "sync"
    "time"
)

// 并发爬思路：
// 初始化数据管道
// 爬虫写入：所有协程向管道中添加图片链接
// 下载协程：从管道里读取链接并下载
// 任务统计协程：检查所有任务是否都完成，完成则关闭数据管道

var (
    // 存放图片链接的数据管道
    chanImageUrls chan string
    waitGroup     sync.WaitGroup
    // 用于监控协程
    chanTask chan string
    // 正则匹配图片链接规则
    reImg = `/uploads/[^"]+?(\.((jpg)|(png)|(jpeg)|(gif)|(bmp)))`
    // 爬取网站根地址【彼岸图网】
    baseUrl = `https://pic.netbian.com/`
    // 爬取页数
    downloadPage = 5
    // 下载路径
    downloadDir = "D:\\GoProject\\src\\test_app\\pic\\"
)

func HandleError(err error, why string) {
    // 异常处理
    if err != nil {
        fmt.Println(why, err)
    }
}

func DownloadFile(url string, filename string) (ok bool) {
    // 下载图片，传入的图片叫什么
    resp, err := http.Get(url)
    HandleError(err, "http.get.url")
    defer resp.Body.Close()
    bytes, err := ioutil.ReadAll(resp.Body)
    HandleError(err, "resp.body")
    filename = downloadDir + filename
    // 写入数据文件
    err = ioutil.WriteFile(filename, bytes, 0666)
    if err != nil {
        return false
    } else {
        return true
    }
}

func main() {
    // 初始化管道
    chanImageUrls = make(chan string, 1000000)
    chanTask = make(chan string, downloadPage)
    page := ""
    // 爬虫协程
    for i := 1; i < downloadPage+1; i++ {
        waitGroup.Add(1)
        if i < 2 {
            page = ""
        } else {
            page = "_" + strconv.Itoa(i)
        }

        go getImgUrls(baseUrl + "index" + page + ".html")

    }
    // 下载协程：从管道中读取链接并下载
    for i := 0; i < downloadPage; i++ {
        waitGroup.Add(1)
        go DownloadImg()
    }
    // 任务统计协程，统计所有任务是否都完成，完成则关闭管道
    waitGroup.Add(1)
    go CheckOK()

    // 等待所有协程结束
    waitGroup.Wait()
}

func DownloadImg() {
    // 下载图片
    for url := range chanImageUrls {
        filename := GetFilenameFromUrl(url)
        ok := DownloadFile(url, filename)
        if ok {
            fmt.Printf("%s 下载成功\n", filename)
        } else {
            fmt.Printf("%s 下载失败\n", filename)
        }
    }
    waitGroup.Done()
}

func GetFilenameFromUrl(url string) (filename string) {
    // 获取文件名
    // 返回最后一个/的位置
    lastIndex := strings.LastIndex(url, "/")
    // 切出来
    filename = url[lastIndex+1:]
    // 时间戳解决重名
    timePrefix := strconv.Itoa(int(time.Now().UnixNano()))
    filename = timePrefix + "_" + filename
    return
}

func CheckOK() {
    // 任务统计协程，当所有协程完成下载，则关闭数据通道
    var count int
    for {
        url := <-chanTask
        fmt.Printf("%s 完成了爬取任务\n", url)
        count++
        if count == downloadPage {
            close(chanImageUrls)
            break
        }
    }
    waitGroup.Done()
}

func getImgUrls(url string) {
    // 爬图片链接到管道
    urls := getImgs(url)
    // 遍历切片里所有链接，存入数据管道
    for _, pic_url := range urls {
        chanImageUrls <- pic_url
    }
    // 标识当前协程完成
    // 每完成一个任务，写一条数据
    // 用于监控协程知道已经完成了几个任务
    chanTask <- url

    waitGroup.Done()
}

func getImgs(url string) (urls []string) {
    // 获取当前页图片链接
    pageStr := GetPageStr(url)
    re := regexp.MustCompile(reImg)
    results := re.FindAllStringSubmatch(pageStr, -1)
    fmt.Printf("共找到%d条结果\n", len(results))
    for _, result := range results {
        urls = append(urls, baseUrl+result[0])
    }
    return
}

func GetPageStr(url string) (pageStr string) {
    // 抽取根据url获取内容
    resp, err := http.Get(url)
    HandleError(err, "http.Get url")
    defer resp.Body.Close()
    // 2.读取页面内容
    pageBytes, err := ioutil.ReadAll(resp.Body)
    HandleError(err, "ioutil.ReadAll")
    // 字节转字符串
    pageStr = string(pageBytes)
    return pageStr
}