Go 学习笔记（9）：多并发爬虫下载图片

package main

import (
    "encoding/json"
    "fmt"
    "io"
    "net/http"
    "os"
    "sync"
)

var (
    imagePath = "../spider/images/"
    pwg       sync.WaitGroup // 用来对生产者计数
    cwg       sync.WaitGroup // 用来对消费者计数
)

// 图片信息结构体
type imageUrl struct {
    url    string
    name   string
    result string
}

// 定义一个接口返回的结构体，用来将json转化成结构体进行数据解析
type pageResult []struct {
    TypeCn  string `json:"typeCn"`
    Details []struct {
        Path   string `json:"path"`
        Width  int    `json:"width"`
        HashID string `json:"hashId"`
        Height int    `json:"height"`
        Desc   string `json:"desc"`
    } `json:"details"`
    ID    int    `json:"id"`
    Title string `json:"title"`
}

// 生产者，每个处理一个地址，并将解析到的图片信息存入通道
func producer(ch chan *imageUrl, url string, cname string) {
    defer pwg.Done()
    pageResultList, err := getPageImageUrls(url)
    if err != nil {
        fmt.Printf("%s:get url error:%s", cname, err)
        return
    }
    fmt.Printf("%s:get page %s urls ok.\n", cname, url)

    for _, result := range pageResultList {
        for _, detail := range result.Details {
            tmp := &imageUrl{
                url:  "https://image.dbbqb.com/" + detail.Path,
                name: detail.Desc,
            }
            ch <- tmp
        }
    }
}

// 消费者，每个处理一张图片信息，下载图片
func consumer(ch chan *imageUrl, cname string) {
    defer cwg.Done()
    for item := range ch {
        err := downloadImage(item.url, item.name)
        if err != nil {
            fmt.Printf("%s download image %s error:%s\n", cname, item.url, err)
        }
        fmt.Printf("%s download image %s, %s ok\n", cname, item.url, item.name)
    }
}

// 从一个页面地址中提取所有图片地址
func getPageImageUrls(pageUrl string) (pageResult, error) {
    resp, err := http.Get(pageUrl)
    if err != nil {
        return nil, err
    }
    defer resp.Body.Close()

    content, err := io.ReadAll(resp.Body)
    if err != nil {
        return nil, err
    }

    var ret pageResult

    err = json.Unmarshal(content, &ret)
    if err != nil {
        return nil, err
    }

    return ret, nil
}

func downloadImage(url, fileName string) error {
    // 发起 GET 请求获取图片数据
    resp, err := http.Get(url)
    if err != nil {
        return err
    }
    defer resp.Body.Close()

    // 创建保存图片的文件
    savePath := imagePath + fileName
    file, err := os.Create(savePath)
    if err != nil {
        return err
    }
    defer file.Close()

    // 将响应体的数据写入文件
    _, err = io.Copy(file, resp.Body)
    return err
}

func main() {
    // 创建一个通道，存储图片下载地址，容量给够
    itemCh := make(chan *imageUrl, 2000)

    // 创建多个生产者，每个地址生成一个，用来解析链接得到图片地址
    for start, i := 8, 0; i < 10; i++ {
        url := fmt.Sprintf("https://www.dbbqb.com/api/group?size=10&start=%d", start)
        start += 10
        go producer(itemCh, url, fmt.Sprintf("p-%d", i))
        pwg.Add(1)
    }

    // 创建多个消费者，数量自定义，用来下载图片
    num := 20
    for i := 0; i < num; i++ {
        go consumer(itemCh, fmt.Sprintf("c-%d", i))
        cwg.Add(1)
    }

    // 必须让生产者和消费都启动后才能等待生产者结束
    pwg.Wait()
    // 必须等生产者全部结束后关闭通道
    close(itemCh)
    cwg.Wait()
    fmt.Println("main done")
}

type pageResult []struct {
    TypeCn  string `json:"typeCn"`
    Details []struct {
        Path   string `json:"path"`
        Width  int    `json:"width"`
        HashID string `json:"hashId"`
        Height int    `json:"height"`
        Desc   string `json:"desc"`
    } `json:"details"`
    ID    int    `json:"id"`
    Title string `json:"title"`
}