go goroutine池的使用

简介

• 使用goroutine池管理，防止线程过多导致cpu、内存爆满
• 使用网络库，批量访问网络资源
• 实现线程阻塞，等待goroutine 全部执行完毕

demo

package main

import (
	"fmt"
	"github.com/panjf2000/ants/v2"
	"golang.org/x/net/publicsuffix"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/cookiejar"
	"strconv"
	"sync"
	"time"
)

/*
协程池管理
*/

type Session struct {
	client *http.Client
}

// 批量设置 header
func (s *Session) SetHeaders(req *http.Request, headers map[string]string) {
	// 自定义Header
	// req.Header.Add("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9")
	// req.Header.Add("Connection", "keep-alive")
	// req.Header.Add("User-Agent", "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/85.0.4183.121 Safari/537.36")
	// req.Header.Add("Host", "quotes.toscrape.com")
	// req.Header.Add("Cookie", "xxxxxx")

	// Add将键、值对添加到标头。将与之关联的任何现有键值追加到该键值。键不区分大小写；它由CanonicalHeaderKey规范化。
	// req.Header.Add()
	// Set将与key关联的头条目设置为单个元素值。它将替换与键关联的任何现有值。键不区分大小写；它由规范化文本Proto.CanonicalTimeHeaderKey. 若要使用非规范键，请直接指定给映射。
	// req.Header.Set()
	// 值返回与给定键关联的所有值。它不区分大小写；文本Proto.CanonicalTimeHeaderKey用于规范化提供的密钥。要使用非规范键，请直接访问映射。返回的切片不是副本。
	// req.Header.Values()

	if headers != nil {
		for key, value := range headers {
			req.Header.Add(key, value)
		}
	}
}

// 默认headers头
func (s *Session) DefaultHeaders(req *http.Request) {
	header := map[string]string{
		"Accept":     "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9",
		"Connection": "keep-alive",
		"User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/85.0.4183.121 Safari/537.36",
	}

	// 设置 header
	s.SetHeaders(req, header)
}

// 创建HTTP客户端以及相关设置
func (s *Session) New() *http.Client {
	// 为了控制代理，TLS配置，保活，压缩和其他设置，请创建一个传输：
	tr := &http.Transport{
		MaxIdleConns:       10,
		IdleConnTimeout:    30 * time.Second,
		DisableCompression: true,
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second, // 限制建立TCP连接的时间
			KeepAlive: 30 * time.Second,
		}).DialContext,
		TLSHandshakeTimeout:   10 * time.Second, // 限制 TLS握手的时间
		ResponseHeaderTimeout: 30 * time.Second, // 限制读取response header的时间
		ExpectContinueTimeout: 10 * time.Second, // 限制client在发送包含 Expect: 100-continue的header到收到继续发送body的response之间的时间等待
	}

	// Transport中实现了连接池的功能，可以将连接保存下来以便下次访问此域名，其中也对连接的数量做出了一定的限制。
	// DisableKeepAlives这个字段可以用来关闭长连接，默认值为false，如果有特殊的需求，需要使用短连接，可以设置此字段为true：
	tr.DisableKeepAlives = false

	// 如果在生成Client时，没有给这个字段赋值，使其为nil的话，那么之后Client发起的请求将只会带上Request对象中指定的Cookie，请求响应中由服务器返回的Cookie也不会被保存。所以如果需要自动管理Cookie的话，我们还需要生成并设定一个CookieJar对象：
	// 这里的publicsuffix.List是一个域的公共后缀列表，是一个可选的选项，设置为nil代表不启用。但是不启用的情况下会使Cookie变得不安全：意味着foo.com的HTTP服务器可以为bar.com设置cookie。所以一般来说最好启用。
	// import "golang.org/x/net/publicsuffix"
	options := cookiejar.Options{
		PublicSuffixList: publicsuffix.List,
	}
	jar, err := cookiejar.New(&options)
	// jar, err := cookiejar.New(nil)
	if err != nil {
		log.Fatal(err)
	}

	// 要控制HTTP客户端标题，重定向策略和其他设置，请创建一个客户端：
	s.client = &http.Client{
		// Transport: tr,
		Timeout: 15 * time.Second,
		Jar:     jar,
	}

	return nil
}

func (s *Session) Get(url string, headers map[string]string) (*http.Response, error) {
	// 创建请求
	request, err := http.NewRequest("GET", url, nil)
	if err != nil {
		fmt.Println("http get error: ", err)
		return nil, err
	}

	// 设置 header
	s.DefaultHeaders(request)
	if headers != nil {
		s.SetHeaders(request, headers)
	}

	// 处理返回结果
	response, err := s.client.Do(request)
	if err != nil {
		fmt.Println("http get error: ", err)
		return nil, err
	}

	return response, nil
}

// 返回数据
// 返回字符串
func (s *Session) GetDataStr(response *http.Response) (string, error) {

	// 函数结束后关闭相关链接
	defer func() { _ = response.Body.Close() }()

	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		log.Fatal("read error", err)
		return "", err
	}

	// fmt.Println(response.Body)
	// fmt.Println(response.Header)
	// fmt.Println(response.ContentLength)
	// fmt.Println(response.Status)
	// fmt.Println(response.StatusCode)
	// fmt.Println(response.Cookies())

	return string(body), nil
}

// 读取数据
func reqData(s *Session, url string) {
	response, err := s.Get(url, nil)
	if err != nil {
		return
	}

	// 获取数据
	buf, err := s.GetDataStr(response)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println(buf)
}

func loopData() {
	// 阻塞主线程，防止直接停止
	var wg sync.WaitGroup

	// goroutine 池
	// size： 最大并发数
	pool, _ := ants.NewPool(10000)
	// 释放关闭此池。
	defer pool.Release()

	// 要访问的url
	s := &Session{}
	s.New()

	for i := 0; i < 100000; i++ {
		wg.Add(1)
		url := "http://127.0.0.1:8080/app04/async4?num=" + strconv.Itoa(i)

		// 任务提交
		// 提交任务通过调用 ants.Submit(func())方法：
		err := pool.Submit(func() {
			defer wg.Done()
			reqData(s, url)
		})
		if err != nil {
			log.Fatal(err)
		}
	}
	wg.Wait()

	// Running返回当前正在运行的goroutine的数量。
	n := pool.Running()
	fmt.Printf("running goroutines: %d\n", n)
	fmt.Printf("finish all tasks.\n")
}

func main() {
	startTime := time.Now()
	fmt.Println("start time: ",startTime.Format("2006-01-02 15:04"))
	loopData()
	endTime := time.Now()
	fmt.Println("end time: ",endTime.Format("2006-01-02 15:04"))
	fmt.Println(endTime.Sub(startTime).Seconds())
}