doc_ai_backed/pkg/httpclient/client.go

package httpclient

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"math"
	"net"
	"net/http"
	"time"

	"gitea.com/bitwsd/document_ai/pkg/log"
)

// RetryConfig 重试配置
type RetryConfig struct {
	MaxRetries      int           // 最大重试次数
	InitialInterval time.Duration // 初始重试间隔
	MaxInterval     time.Duration // 最大重试间隔
	SkipTLSVerify   bool          // 是否跳过TLS验证
}

// DefaultRetryConfig 默认重试配置
var DefaultRetryConfig = RetryConfig{
	MaxRetries:      2,
	InitialInterval: 100 * time.Millisecond,
	MaxInterval:     5 * time.Second,
	SkipTLSVerify:   true,
}

// Client HTTP客户端封装
type Client struct {
	client *http.Client
	config RetryConfig
}

// NewClient 创建新的HTTP客户端
func NewClient(config *RetryConfig) *Client {
	cfg := DefaultRetryConfig
	if config != nil {
		cfg = *config
	}

	tr := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: cfg.SkipTLSVerify,
		},
	}

	return &Client{
		client: &http.Client{
			Transport: tr,
		},
		config: cfg,
	}
}

// RequestWithRetry 执行带重试的HTTP请求
func (c *Client) RequestWithRetry(ctx context.Context, method, url string, body io.Reader, headers map[string]string) (*http.Response, error) {
	var lastErr error

	for attempt := 0; attempt < c.config.MaxRetries; attempt++ {
		if attempt > 0 {
			backoff := c.calculateBackoff(attempt)
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(backoff):
			}

			// 如果body是可以Seek的，则重置到开始位置
			if seeker, ok := body.(io.Seeker); ok {
				_, err := seeker.Seek(0, io.SeekStart)
				if err != nil {
					return nil, fmt.Errorf("failed to reset request body: %w", err)
				}
			}

			log.Info(ctx, "func", "RequestWithRetry", "msg", "正在重试请求",
				"attempt", attempt+1, "max_retries", c.config.MaxRetries, "backoff", backoff)
		}

		// 检查 context 是否已经取消
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}

		req, err := http.NewRequestWithContext(ctx, method, url, body)
		if err != nil {
			lastErr = fmt.Errorf("create request failed: %w", err)
			continue
		}

		for k, v := range headers {
			req.Header.Set(k, v)
		}

		resp, err := c.client.Do(req)
		if err != nil {
			lastErr = fmt.Errorf("request failed: %w", err)

			// 如果是超时错误，直接返回不再重试
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				return nil, fmt.Errorf("request timeout: %w", err)
			}

			log.Error(ctx, "func", "RequestWithRetry", "msg", "请求失败",
				"error", err, "attempt", attempt+1)
			continue
		}

		if resp.StatusCode != http.StatusOK {
			body, _ := io.ReadAll(resp.Body)
			fmt.Println(string(body))
			resp.Body.Close()
			lastErr = fmt.Errorf("unexpected status code: %d", resp.StatusCode)
			log.Error(ctx, "func", "RequestWithRetry", "msg", "请求返回非200状态码",
				"status", resp.StatusCode, "attempt", attempt+1)
			continue
		}

		return resp, nil
	}

	return nil, fmt.Errorf("max retries reached: %w", lastErr)
}

// calculateBackoff 计算退避时间
func (c *Client) calculateBackoff(attempt int) time.Duration {
	backoff := float64(c.config.InitialInterval) * math.Pow(2, float64(attempt))
	if backoff > float64(c.config.MaxInterval) {
		backoff = float64(c.config.MaxInterval)
	}
	return time.Duration(backoff)
}