common/pkg/restry-pool/resty.go

package restry_pool

import (
	"context"
	"github.com/corpix/uarand"
	"github.com/go-resty/resty/v2"
	"net"
	"net/http"
	"sync"
	"time"
)

var rp *restyPool

func init() {
	rp = NewRestyPool()
}

func GetRestyClient(conn net.Conn) (cli *resty.Client) {
	return rp.Get(conn)
}

func PutRestyClient(cli *resty.Client) {
	rp.Put(cli)
}

type restyPool struct {
	resty     sync.Pool
	transport sync.Pool
}

const (
	UserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36 Edg/123.0.0.0"
)

func beforeRequest(cli *resty.Client, req *resty.Request) error {
	cli.SetHeader("User-Agent", uarand.GetRandom())
	return nil
}
func NewRestyPool() *restyPool {
	return &restyPool{
		resty: sync.Pool{
			New: func() any {
				return resty.New().OnBeforeRequest(beforeRequest)
			},
		},
		transport: sync.Pool{
			New: func() any {
				return &http.Transport{
					MaxIdleConns:        100,
					IdleConnTimeout:     90 * time.Second,
					TLSHandshakeTimeout: 5 * time.Second,
				}
			},
		},
	}
}

func (r *restyPool) Get(conn net.Conn) (cli *resty.Client) {
	cli = r.resty.Get().(*resty.Client)
	transport := r.transport.Get().(*http.Transport)
	transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
		return conn, nil
	}
	return cli.SetTransport(transport)
}

func (r *restyPool) Put(cli *resty.Client) {
	transport, _ := cli.Transport()
	r.transport.Put(transport)
	r.resty.Put(cli)
}