Repository: pingcap/tidb
File: pkg/objstore/s3store/s3.go (Go)

Code before:
		input.ACL = types.ObjectCannedACL(options.Acl)
	}
	if options.Sse != "" {
		input.ServerSideEncryption = types.ServerSideEncryption(options.Sse)
	}
	if options.SseKmsKeyId != "" {
		input.SSEKMSKeyId = aws.String(options.SseKmsKeyId)
	}
	if options.StorageClass != "" {
		input.StorageClass = types.StorageClass(options.StorageClass)
	}
	return input
}

// WriteFile writes data to a file to storage.
func (rs *S3Storage) WriteFile(ctx context.Context, file string, data []byte) error {
	input := buildPutObjectInput(rs.options, file, data)
	// we don't need to calculate contentMD5 if s3 object lock enabled.
	// since aws-go-sdk already did it in #computeBodyHashes
	// https://github.com/aws/aws-sdk-go/blob/bcb2cf3fc2263c8c28b3119b07d2dbb44d7c93a0/service/s3/body_hash.go#L30
	_, err := rs.svc.PutObject(ctx, input)
	if err != nil {
		return errors.Trace(err)
	}
	// Use the proper waiter pattern in AWS SDK v2
	waiter := s3.NewObjectExistsWaiter(rs.svc)
	hinput := &s3.HeadObjectInput{
		Bucket: aws.String(rs.options.Bucket),
		Key:    aws.String(rs.options.Prefix + file),
	}
	err = waiter.Wait(ctx, hinput, 30*time.Second)
	rs.accessRec.RecWrite(len(data))
	return errors.Trace(err)
}

// ReadFile implements Storage.ReadFile.
func (rs *S3Storage) ReadFile(ctx context.Context, file string) ([]byte, error) {
	backoff := 10 * time.Millisecond
	remainRetry := 5
	contRetry := func() bool {
		if remainRetry <= 0 {
			return false
		}
		time.Sleep(backoff)
		remainRetry -= 1
		return true
	}

	// The errors cannot be handled by the SDK because they happens during reading the HTTP response body.
	// We cannot use `utils.WithRetry[V2]` here because cyclinic deps.
	for {

Review comment (suggestion):
The RecWrite call should only be executed after a successful PutObject. If PutObject returns an error (line 796), the function returns early, so RecWrite is never called on error. However, if the error handling changes in the future, this could be problematic. Consider moving RecWrite before the return statement on line 799 to make the intent clearer, or add a comment explaining that RecWrite only happens on success.

Code after:
		input.ACL = types.ObjectCannedACL(options.Acl)
	}
	if options.Sse != "" {
		input.ServerSideEncryption = types.ServerSideEncryption(options.Sse)
	}
	if options.SseKmsKeyId != "" {
		input.SSEKMSKeyId = aws.String(options.SseKmsKeyId)
	}
	if options.StorageClass != "" {
		input.StorageClass = types.StorageClass(options.StorageClass)
	}
	return input
}

// WriteFile writes data to a file to storage.
func (rs *S3Storage) WriteFile(ctx context.Context, file string, data []byte) error {
	input := buildPutObjectInput(rs.options, file, data)
	// we don't need to calculate contentMD5 if s3 object lock enabled.
	// since aws-go-sdk already did it in #computeBodyHashes
	// https://github.com/aws/aws-sdk-go/blob/bcb2cf3fc2263c8c28b3119b07d2dbb44d7c93a0/service/s3/body_hash.go#L30
	_, err := rs.svc.PutObject(ctx, input)
	if err != nil {
		return errors.Trace(err)
	}
	rs.accessRec.RecWrite(len(data))
	return nil
}

// ReadFile implements Storage.ReadFile.
func (rs *S3Storage) ReadFile(ctx context.Context, file string) ([]byte, error) {
	backoff := 10 * time.Millisecond
	remainRetry := 5
	contRetry := func() bool {
		if remainRetry <= 0 {
			return false
		}
		time.Sleep(backoff)
		remainRetry -= 1
		return true
	}

	// The errors cannot be handled by the SDK because they happens during reading the HTTP response body.
	// We cannot use `utils.WithRetry[V2]` here because cyclinic deps.
	for {
		data, err := rs.doReadFile(ctx, file)
		if err != nil {
			log.Warn("ReadFile: failed to read file.",
				zap.String("file", file), logutil.ShortError(err), zap.Int("remained", remainRetry))
			if !isHTTP2ConnAborted(err) {
				return nil, err
			}