goProject/.svn/pristine/48/48c5b123aaa4f5248637d444f7ab2c0fef970cc8.svn-base

// 适用于大量http(s)请求，连接复用
package webUtil

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptrace"
	"net/url"
	"time"
)

type Client struct {
	traceCtx context.Context
	client   *http.Client
}

type TransportOPT struct {
	//超时时间
	Timeout time.Duration

	//代理字符串，如"http://127.0.0.1:1080"
	ProxyString string

	//最大保持连接数
	MaxIdleConns int

	//每个主机的最大保持连接数
	MaxIdleConnsPerHost int

	//单主机最大连接数
	MaxConnsPerHost int

	//保持连接的超时时间
	IdleConnTimeout time.Duration

	//禁止保持连接
	DisableKeepAlives bool
}

// GET请求
func (c *Client) Get(urlStr string, header map[string]string) (result []byte, err error) {
	req, err := http.NewRequestWithContext(c.traceCtx, http.MethodGet, urlStr, nil)
	if err != nil {
		//fmt.Println(err)
		return
	}

	// 处理头部（包括默认头部，以及传入的头部集合）
	if header == nil {
		req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	} else {
		for k, v := range header {
			req.Header.Add(k, v)
		}
	}

	res, err := c.client.Do(req)
	if err != nil {
		//fmt.Println(err)
		return
	}
	defer res.Body.Close()

	result, err = ioutil.ReadAll(res.Body)

	return
}

// POST请求
func (c *Client) PostWithMap(urlStr string, data map[string]string, header map[string]string) (result []byte, err error) {
	// 组装POST数据
	postValues := url.Values{}
	for key, value := range data {
		postValues.Set(key, value)
	}
	postDataStr := postValues.Encode()
	byteData := []byte(postDataStr)

	var req *http.Request
	req, err = http.NewRequestWithContext(c.traceCtx, http.MethodPost, urlStr, bytes.NewReader(byteData))
	if err != nil {
		//fmt.Println(err)
		return
	}

	// 处理头部（包括默认头部，以及传入的头部集合）
	if header == nil {
		req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	} else {
		for k, v := range header {
			req.Header.Add(k, v)
		}
	}

	var res *http.Response
	res, err = c.client.Do(req)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer res.Body.Close()

	result, err = ioutil.ReadAll(res.Body)

	return
}

// POST请求
func (c *Client) PostWithByte(urlStr string, data []byte, header map[string]string) (result []byte, err error) {
	var req *http.Request
	req, err = http.NewRequestWithContext(c.traceCtx, http.MethodPost, urlStr, bytes.NewReader(data))
	if err != nil {
		//fmt.Println(err)
		return
	}

	// 处理头部（包括默认头部，以及传入的头部集合）
	if header == nil {
		req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	} else {
		for k, v := range header {
			req.Header.Add(k, v)
		}
	}

	var res *http.Response
	res, err = c.client.Do(req)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer res.Body.Close()

	result, err = ioutil.ReadAll(res.Body)

	return
}

// 新建Client对象
// transportOPT参数说明
//
//	Timeout - 连接绝对超时时间
//	ProxyString - 代理字符串
//	MaxIdleConns - 最大空闲连接数
//	MaxIdleConnsPerHost - 每个目标主机的最大空闲连接数
//	MaxConnsPerHost - 每个目标主机的最大连接数
//	IdleConnTimeout - 空闲连接超时时间
//	DisableKeepAlives - 禁用连接保持(要使用连接复用，此值不传入或传入false；否则连接可能不会复用)
func NewClient(transportOPT *TransportOPT) (c *Client) {
	c = &Client{}

	//代理
	getProxy := func() func(*http.Request) (*url.URL, error) {
		if (transportOPT != nil) && len(transportOPT.ProxyString) > 0 {
			uri, _ := url.Parse(transportOPT.ProxyString)
			return http.ProxyURL(uri)
		}

		return nil
	}

	// 默认参数
	timeout := 30 * time.Second
	maxIdleConns := 60000
	maxIdleConnsPerHost := 30000
	maxConnsPerHost := 30000
	idleConnTimeout := time.Minute * 1
	disableKeepAlives := false

	if transportOPT != nil {
		// 根据传入参数修改默认参数
		if transportOPT.Timeout > 0 {
			timeout = transportOPT.Timeout
		}
		if transportOPT.MaxIdleConns > 0 {
			maxIdleConns = transportOPT.MaxIdleConns
		}
		if transportOPT.MaxIdleConnsPerHost > 0 {
			maxIdleConnsPerHost = transportOPT.MaxIdleConnsPerHost
		}
		if transportOPT.MaxConnsPerHost > 0 {
			maxConnsPerHost = transportOPT.MaxConnsPerHost
		}
		if transportOPT.IdleConnTimeout > 0 {
			idleConnTimeout = transportOPT.IdleConnTimeout
		}

		disableKeepAlives = transportOPT.DisableKeepAlives
	}

	clientTrace := &httptrace.ClientTrace{
		//GotConn: func(gci httptrace.GotConnInfo) {
		//	fmt.Printf("conn was reused: %t\n", gci.Reused)
		//},
	}

	c.traceCtx = httptrace.WithClientTrace(context.Background(), clientTrace)
	c.client = &http.Client{
		Timeout: timeout,

		Transport: &http.Transport{
			Proxy:               getProxy(),
			MaxIdleConns:        maxIdleConns,
			MaxIdleConnsPerHost: maxIdleConnsPerHost,
			MaxConnsPerHost:     maxConnsPerHost,
			IdleConnTimeout:     idleConnTimeout,
			DisableKeepAlives:   disableKeepAlives,
			// ForceAttemptHTTP2:   true,
		},
	}

	return
}