common/pkg/restry-pool/resty.go

73 lines
1.5 KiB
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)
}