tamabox/internal/storage/s3_retry.go

package storage

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/smithy-go"
)

const (
	s3AttemptTimeout      = 60 * time.Second
	s3MaxAttempts         = 3
	s3InitialRetryBackoff = 100 * time.Millisecond
)

// PutObject 上传对象。每次尝试独立限时 60 秒，最多执行 3 次；
// SDK 内部重试被关闭，确保重试次数、Body 回绕和错误分类只由本层控制。
func PutObject(ctx context.Context, client *s3.Client, input *s3.PutObjectInput) error {
	if client == nil {
		return errors.New("nil S3 client")
	}
	if input == nil {
		return errors.New("nil S3 PutObject input")
	}
	return doS3RequestWithRetry(ctx, input.Body, func(attemptCtx context.Context) error {
		_, err := client.PutObject(attemptCtx, input, func(options *s3.Options) {
			options.Retryer = aws.NopRetryer{}
		})
		return err
	})
}

// DeleteObject 删除对象。S3 Delete 本身幂等；404 也表示目标已经不存在，
// 因而不重试并返回成功。其他错误采用与 PutObject 相同的超时和重试规则。
func DeleteObject(ctx context.Context, client *s3.Client, bucket, key string) error {
	if client == nil {
		return errors.New("nil S3 client")
	}
	err := doS3RequestWithRetry(ctx, nil, func(attemptCtx context.Context) error {
		_, err := client.DeleteObject(attemptCtx, &s3.DeleteObjectInput{
			Bucket: aws.String(bucket),
			Key:    aws.String(key),
		}, func(options *s3.Options) {
			options.Retryer = aws.NopRetryer{}
		})
		return err
	})
	if err != nil {
		if statusCode, ok := s3ErrorHTTPStatus(err); ok && statusCode == 404 {
			return nil
		}
	}
	return err
}

func doS3RequestWithRetry(ctx context.Context, body io.Reader, operation func(context.Context) error) error {
	var (
		seeker       io.Seeker
		bodyStart    int64
		bodyStartErr error
	)
	if body != nil {
		if candidate, ok := body.(io.Seeker); ok {
			seeker = candidate
			// 记录调用时的位置，而不是武断跳到 0，兼容从切片区间上传的 Reader。
			bodyStart, bodyStartErr = seeker.Seek(0, io.SeekCurrent)
		}
	}

	for attempt := 0; attempt < s3MaxAttempts; attempt++ {
		if err := ctx.Err(); err != nil {
			return err
		}
		if attempt > 0 && seeker != nil {
			if bodyStartErr != nil {
				return fmt.Errorf("determine S3 request body start: %w", bodyStartErr)
			}
			if _, err := seeker.Seek(bodyStart, io.SeekStart); err != nil {
				return fmt.Errorf("rewind S3 request body: %w", err)
			}
		}

		attemptCtx, cancel := context.WithTimeout(ctx, s3AttemptTimeout)
		err := operation(attemptCtx)
		cancel()
		if err == nil {
			return nil
		}
		if ctx.Err() != nil {
			return ctx.Err()
		}
		if attempt == s3MaxAttempts-1 || !retryableS3Error(err) {
			return err
		}

		backoff := s3InitialRetryBackoff << attempt
		if err := waitForS3Retry(ctx, backoff); err != nil {
			return err
		}
	}
	return nil
}

func waitForS3Retry(ctx context.Context, delay time.Duration) error {
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}

type httpStatusCoder interface {
	HTTPStatusCode() int
}

func s3ErrorHTTPStatus(err error) (int, bool) {
	var statusErr httpStatusCoder
	if errors.As(err, &statusErr) {
		return statusErr.HTTPStatusCode(), true
	}
	return 0, false
}

func retryableS3Error(err error) bool {
	if err == nil {
		return false
	}
	if statusCode, ok := s3ErrorHTTPStatus(err); ok {
		return statusCode == 408 || statusCode == 429 || statusCode >= 500 && statusCode <= 599
	}
	if errors.Is(err, context.Canceled) {
		return false
	}
	// 父 context 已结束时会在调用方提前返回；这里的 DeadlineExceeded 是
	// 单次尝试或传输层超时，可以在父 context 仍有效时重试。
	if errors.Is(err, context.DeadlineExceeded) {
		return true
	}
	var networkErr net.Error
	if errors.As(err, &networkErr) {
		return true
	}
	if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
		return true
	}

	// 某些兼容 S3 的实现可能只返回错误码而没有可提取的 HTTP 状态。
	var apiErr smithy.APIError
	if errors.As(err, &apiErr) {
		switch apiErr.ErrorCode() {
		case "RequestTimeout", "RequestTimeoutException",
			"Throttling", "ThrottlingException", "TooManyRequestsException", "SlowDown",
			"InternalError", "InternalServerError", "ServiceUnavailable", "ServiceUnavailableException":
			return true
		}
	}
	return false
}