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 (
	"context"
	"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 Client struct {
	config Config
	client *http.Client
	owned  *http.Transport
	closed atomic.Bool
}

// New creates a reusable client. Close it after callers and response streams drain.
func New(cfg Config) (*Client, error) {
	normalized, err := normalizeConfig(cfg)
	if err != nil {
		return nil, err
	}
	c := &Client{config: normalized}
	transport := cfg.Transport
	if transport == nil {
		c.owned, err = ownedTransport(normalized)
		if err != nil {
			return nil, err
		}
		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
		if attempt > 0 {
			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)
		var delay time.Duration
		if retry {
			delay, retry = c.retryDelay(attempt, response)
		}
		if !retry {
			if err != nil {
				cancel()
				return response, err
			}
			response.Body = &responseBody{ReadCloser: response.Body, cancel: cancel}
			return response, nil
		}
		if response != nil {
			drainSmallBody(response, 32<<10)
			_ = 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) }

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

// retryDelay is deliberately outside the success path: no jitter/random work
// for calls that cannot or do not need to retry.
func (c *Client) retryDelay(attempt int, response *http.Response) (time.Duration, bool) {
	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 {
		if value, ok := retryAfter(response.Header.Get("Retry-After")); ok {
			return value, value <= c.config.MaxBackoff
		}
	}
	if delay > 1 {
		delay = delay/2 + time.Duration(rand.Int63n(int64(delay/2)+1))
	}
	return delay, delay <= c.config.MaxBackoff
}

// Drain only a known small error body. Unknown/large bodies are closed promptly;
// never allocate or consume an unbounded error document just to retain a pool slot.
func drainSmallBody(response *http.Response, limit int64) {
	limit = min(limit, 32<<10)
	if response.Body != nil && response.ContentLength >= 0 && response.ContentLength <= limit {
		_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, limit+1))
	}
}