feat: add cache, outbound HTTP, request protection, tasks and module scaffolding

2026-09-23 13:38:25 +08:00
parent 467c890363
commit fe66221a9c
34 changed files with 3100 additions and 1 deletions
--- a/httpc/client.go
+++ b/httpc/client.go
@@ -0,0 +1,247 @@
+// 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) {
+	if c.closed.Load() {
+		return nil, ErrClosed
+	}
+	if request == nil || request.URL == nil {
+		return nil, errors.New("httpc: nil request or URL")
+	}
+	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
+}
--- a/httpc/client_test.go
+++ b/httpc/client_test.go
@@ -0,0 +1,131 @@
+package httpc
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"git.yuntuo.hk/k2safe/go-omni/logx"
+)
+
+func TestRetryCorrelationPoolAndWriteSafety(t *testing.T) {
+	var calls, connections atomic.Int32
+	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		n := calls.Add(1)
+		if r.Header.Get("X-Request-ID") != "request-1" {
+			t.Error("missing correlation")
+		}
+		if n == 1 || r.Method == "POST" {
+			w.WriteHeader(503)
+			io.WriteString(w, "busy")
+			return
+		}
+		io.WriteString(w, `{"ok":true}`)
+	}))
+	server.Config.ConnState = func(_ net.Conn, s http.ConnState) {
+		if s == http.StateNew {
+			connections.Add(1)
+		}
+	}
+	server.Start()
+	defer server.Close()
+	client, err := New(Config{MaxRetries: 2, Backoff: time.Millisecond, DisableLogging: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer client.Close()
+	ctx := logx.WithRequestID(context.Background(), "request-1")
+	var output struct {
+		OK bool `json:"ok"`
+	}
+	if err := client.JSON(ctx, "GET", server.URL, nil, &output); err != nil || !output.OK {
+		t.Fatal(err)
+	}
+	if calls.Load() != 2 || connections.Load() != 1 {
+		t.Fatalf("calls=%d connections=%d", calls.Load(), connections.Load())
+	}
+	if err := client.JSON(ctx, "POST", server.URL, map[string]string{"name": "value"}, nil); err == nil {
+		t.Fatal("expected 503")
+	}
+	if calls.Load() != 3 {
+		t.Fatal("POST was retried")
+	}
+	client.Close()
+	if err := client.JSON(ctx, "GET", server.URL, nil, nil); !errors.Is(err, ErrClosed) {
+		t.Fatal(err)
+	}
+}
+
+func TestTotalDeadlineIncludesBackoffAndBody(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/retry" {
+			w.Header().Set("Retry-After", "1")
+			w.WriteHeader(503)
+			return
+		}
+		w.WriteHeader(200)
+		w.(http.Flusher).Flush()
+		<-r.Context().Done()
+	}))
+	defer server.Close()
+	client, _ := New(Config{Timeout: 40 * time.Millisecond, MaxRetries: 3, DisableLogging: true})
+	defer client.Close()
+	started := time.Now()
+	err := client.JSON(context.Background(), "GET", server.URL+"/retry", nil, nil)
+	if !errors.Is(err, context.DeadlineExceeded) || time.Since(started) > time.Second {
+		t.Fatal(err, time.Since(started))
+	}
+	req, _ := http.NewRequest("GET", server.URL+"/body", nil)
+	response, err := client.Do(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer response.Body.Close()
+	if _, err := io.ReadAll(response.Body); err == nil {
+		t.Fatal("body deadline was canceled too early")
+	}
+}
+
+func TestLimitsRedirectAndNonReplayableBody(t *testing.T) {
+	var calls atomic.Int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		calls.Add(1)
+		switch r.URL.Path {
+		case "/redirect":
+			http.Redirect(w, r, "/large", 302)
+		case "/large":
+			io.WriteString(w, strings.Repeat("x", 17))
+		default:
+			w.WriteHeader(503)
+		}
+	}))
+	defer server.Close()
+	client, _ := New(Config{MaxRetries: 2, MaxResponseBytes: 16, DisableLogging: true})
+	defer client.Close()
+	if err := client.JSON(context.Background(), "GET", server.URL+"/large", nil, nil); !errors.Is(err, ErrResponseTooLarge) {
+		t.Fatal(err)
+	}
+	var status *StatusError
+	if err := client.JSON(context.Background(), "GET", server.URL+"/redirect", nil, nil); !errors.As(err, &status) || status.StatusCode != 302 {
+		t.Fatal(err)
+	}
+	if calls.Load() != 2 {
+		t.Fatal("followed redirect")
+	}
+	req, _ := http.NewRequest("GET", server.URL, io.NopCloser(strings.NewReader("body")))
+	res, err := client.Do(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	res.Body.Close()
+	if calls.Load() != 3 {
+		t.Fatal("non-replayable body was retried")
+	}
+}