go-omni/httpc/client.go

// Package httpc provides a pooled HTTP client with an end-to-end deadline and
// conservative, bounded retries. Redirects are returned to the caller.
package httpc

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/logx"
)

var ErrClosed = errors.New("httpc: closed")
var ErrResponseTooLarge = errors.New("httpc: response exceeds limit")

type Config struct {
	Timeout          time.Duration     `json:"timeout"`
	MaxRetries       int               `json:"maxRetries"`
	Backoff          time.Duration     `json:"backoff"`
	MaxBackoff       time.Duration     `json:"maxBackoff"`
	MaxConnsPerHost  int               `json:"maxConnsPerHost"`
	MaxResponseBytes int64             `json:"maxResponseBytes"`
	DisableRequestID bool              `json:"disableRequestId"`
	DisableLogging   bool              `json:"disableLogging"`
	Transport        http.RoundTripper `json:"-"` // Borrowed. The caller owns its lifecycle.
}

type Client struct {
	config Config
	client *http.Client
	owned  *http.Transport
	closed atomic.Bool
}

func New(cfg Config) (*Client, error) {
	if cfg.Timeout < 0 || cfg.Backoff < 0 || cfg.MaxBackoff < 0 || cfg.MaxRetries < 0 || cfg.MaxRetries > 5 || cfg.MaxConnsPerHost < 0 || cfg.MaxResponseBytes < 0 {
		return nil, errors.New("httpc: invalid configuration")
	}
	if cfg.Timeout == 0 {
		cfg.Timeout = 10 * time.Second
	}
	if cfg.Backoff == 0 {
		cfg.Backoff = 100 * time.Millisecond
	}
	if cfg.MaxBackoff == 0 {
		cfg.MaxBackoff = 2 * time.Second
	}
	if cfg.Backoff > cfg.MaxBackoff {
		return nil, errors.New("httpc: backoff exceeds maxBackoff")
	}
	if cfg.MaxConnsPerHost == 0 {
		cfg.MaxConnsPerHost = 64
	}
	if cfg.MaxResponseBytes == 0 {
		cfg.MaxResponseBytes = 1 << 20
	}
	if cfg.MaxResponseBytes > 1<<30 || cfg.MaxConnsPerHost > 65536 {
		return nil, errors.New("httpc: response or connection limit exceeds supported capacity")
	}
	c := &Client{config: cfg}
	transport := cfg.Transport
	if transport == nil {
		// Own a private pool; never mutate the process-wide DefaultTransport.
		base, ok := http.DefaultTransport.(*http.Transport)
		if !ok {
			return nil, errors.New("httpc: provide Transport when DefaultTransport is replaced")
		}
		c.owned = base.Clone()
		c.owned.MaxConnsPerHost = cfg.MaxConnsPerHost
		c.owned.MaxIdleConnsPerHost = cfg.MaxConnsPerHost
		c.owned.MaxIdleConns = cfg.MaxConnsPerHost * 2
		transport = c.owned
	}
	c.client = &http.Client{Transport: transport, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
	return c, nil
}

// Do keeps the total deadline active until the response body is closed or
// consumed. Callers own response.Body. Only GET, HEAD and OPTIONS are retried,
// and only when their bodies can be replayed. MaxRetries=0 disables retries.
func (c *Client) Do(request *http.Request) (*http.Response, error) {
	var requestError error
	if c.closed.Load() {
		requestError = ErrClosed
	} else if request == nil || request.URL == nil {
		requestError = errors.New("httpc: nil request or URL")
	}
	if requestError != nil {
		// Match http.Client.Do's ownership rule even when rejection happens
		// before the underlying transport takes ownership of the body.
		if request != nil && request.Body != nil {
			_ = request.Body.Close()
		}
		return nil, requestError
	}
	ctx, cancel := context.WithTimeout(request.Context(), c.config.Timeout)
	base := request.Clone(ctx)
	if base.Header == nil {
		base.Header = make(http.Header)
	}
	if !c.config.DisableRequestID && base.Header.Get("X-Request-ID") == "" {
		if id := logx.RequestID(ctx); id != "" {
			base.Header.Set("X-Request-ID", id)
		}
	}
	replayable := (base.Method == "" || base.Method == "GET" || base.Method == "HEAD" || base.Method == "OPTIONS") && (base.Body == nil || base.Body == http.NoBody || base.GetBody != nil)
	for attempt := 0; ; attempt++ {
		req := base.Clone(ctx)
		if attempt > 0 && base.Body != nil && base.Body != http.NoBody {
			body, err := base.GetBody()
			if err != nil {
				cancel()
				return nil, err
			}
			req.Body = body
		}
		started := time.Now()
		response, err := c.client.Do(req)
		status := 0
		if response != nil {
			status = response.StatusCode
		}
		if !c.config.DisableLogging {
			logx.InfoContext(ctx, "http.client", "method", req.Method, "host", req.URL.Hostname(), "status", status, "attempt", attempt+1, "durationMs", time.Since(started).Seconds()*1000, "failed", err != nil)
		}
		retry := replayable && attempt < c.config.MaxRetries && ctx.Err() == nil && (err != nil || status == 429 || status == 502 || status == 503 || status == 504)
		delay := c.config.Backoff
		for i := 0; i < attempt && delay < c.config.MaxBackoff; i++ {
			if delay > c.config.MaxBackoff/2 {
				delay = c.config.MaxBackoff
			} else {
				delay *= 2
			}
		}
		if response != nil && response.Header.Get("Retry-After") != "" {
			if value, ok := retryAfter(response.Header.Get("Retry-After")); ok {
				delay = value
			}
		} else if delay > 1 {
			delay = delay/2 + time.Duration(rand.Int63n(int64(delay/2)+1))
		}
		if delay > c.config.MaxBackoff {
			retry = false
		}
		if !retry {
			if err != nil {
				cancel()
				return response, err
			}
			response.Body = &responseBody{ReadCloser: response.Body, cancel: cancel}
			return response, nil
		}
		if response != nil {
			if response.ContentLength >= 0 && response.ContentLength <= 32<<10 {
				_, _ = io.Copy(io.Discard, response.Body)
			}
			_ = response.Body.Close()
		}
		timer := time.NewTimer(delay)
		select {
		case <-ctx.Done():
			timer.Stop()
			cancel()
			return nil, ctx.Err()
		case <-timer.C:
		}
	}
}

type responseBody struct {
	io.ReadCloser
	cancel context.CancelFunc
}

func (b *responseBody) Read(data []byte) (int, error) {
	n, err := b.ReadCloser.Read(data)
	if err != nil {
		b.cancel()
	}
	return n, err
}
func (b *responseBody) Close() error { defer b.cancel(); return b.ReadCloser.Close() }

func retryAfter(value string) (time.Duration, bool) {
	if n, err := strconv.ParseInt(strings.TrimSpace(value), 10, 32); err == nil && n >= 0 {
		return time.Duration(n) * time.Second, true
	}
	if at, err := http.ParseTime(value); err == nil {
		return max(time.Until(at), 0), true
	}
	return 0, false
}

type StatusError struct{ StatusCode int }

func (e *StatusError) Error() string { return fmt.Sprintf("httpc: HTTP status %d", e.StatusCode) }

// JSON bounds response memory, closes the body, and decodes only 2xx responses.
// Error bodies, URLs, credentials and query strings are not included in logs.
func (c *Client) JSON(ctx context.Context, method, url string, input, output any) error {
	if ctx == nil {
		return errors.New("httpc: nil context")
	}
	var body io.Reader
	if input != nil {
		data, err := json.Marshal(input)
		if err != nil {
			return err
		}
		body = bytes.NewReader(data)
	}
	request, err := http.NewRequestWithContext(ctx, method, url, body)
	if err != nil {
		return err
	}
	if input != nil {
		request.Header.Set("Content-Type", "application/json")
	}
	response, err := c.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return &StatusError{StatusCode: response.StatusCode}
	}
	data, err := io.ReadAll(io.LimitReader(response.Body, c.config.MaxResponseBytes+1))
	if err != nil {
		return err
	}
	if int64(len(data)) > c.config.MaxResponseBytes {
		return ErrResponseTooLarge
	}
	if output == nil || len(bytes.TrimSpace(data)) == 0 {
		return nil
	}
	return json.Unmarshal(data, output)
}

func (c *Client) Close() error {
	c.closed.Store(true)
	if c.owned != nil {
		c.owned.CloseIdleConnections()
	}
	return nil
}