cloud/tasks/persistence/s3.go - main - githubmirror/nbs

cloud/tasks/persistence/s3.go
431 строка10 KB
Sergei
issue-6142: [Disk Manager] Allow to specify S3 storage class for images (#6143)
09 июн 2026, 20:06
09 июн 2026, 20:0639cc8b0
package persistence

import (
	"bytes"
	"context"
	"encoding/json"
	"io/ioutil"
	"os"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/client"
	aws_credentials "github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	aws_s3 "github.com/aws/aws-sdk-go/service/s3"
	"github.com/ydb-platform/nbs/cloud/tasks/errors"
	"github.com/ydb-platform/nbs/cloud/tasks/logging"
	"github.com/ydb-platform/nbs/cloud/tasks/metrics"
	persistence_config "github.com/ydb-platform/nbs/cloud/tasks/persistence/config"
)

////////////////////////////////////////////////////////////////////////////////

func withComponentLoggingField(ctx context.Context) context.Context {
	return logging.WithComponent(ctx, logging.ComponentS3)
}

////////////////////////////////////////////////////////////////////////////////

type s3ClientRetryer struct {
	client.DefaultRetryer
	metrics *s3Metrics
}

func (r *s3ClientRetryer) RetryRules(req *request.Request) time.Duration {
	r.metrics.OnRetry(req)

	return r.DefaultRetryer.RetryRules(req)
}

////////////////////////////////////////////////////////////////////////////////

type S3Client struct {
	s3          *aws_s3.S3
	callTimeout time.Duration
	metrics     *s3Metrics
}

const (
	s3ErrCodeQuotaLimitExceeded = "QuotaLimitExceeded"
)

func NewS3Client(
	endpoint string,
	region string,
	credentials S3Credentials,
	callTimeout time.Duration,
	registry metrics.Registry,
	maxRetriableErrorCount uint64,
	availabilityMonitoring *AvailabilityMonitoring,
) (*S3Client, error) {

	s3Metrics := newS3Metrics(
		callTimeout,
		registry,
		availabilityMonitoring,
	)

	sessionConfig := &aws.Config{
		Credentials: aws_credentials.NewStaticCredentials(
			credentials.ID,
			credentials.Secret,
			"", // token - only required for temporary security credentials retrieved via STS, we don't need that
		),
		Endpoint:         &endpoint,
		Region:           &region,
		S3ForcePathStyle: aws.Bool(true), // otherwise, we get DNS DDOS errors in tests
		Retryer: &s3ClientRetryer{
			DefaultRetryer: client.DefaultRetryer{
				NumMaxRetries: int(maxRetriableErrorCount),
			},
			metrics: s3Metrics,
		},
	}

	session, err := session.NewSession(sessionConfig)
	if err != nil {
		return nil, errors.NewRetriableError(err)
	}

	return &S3Client{
		s3:          aws_s3.New(session),
		callTimeout: callTimeout,
		metrics:     s3Metrics,
	}, nil
}

func NewS3ClientFromConfig(
	config *persistence_config.S3Config,
	registry metrics.Registry,
	availabilityMonitoring *AvailabilityMonitoring,
) (*S3Client, error) {

	credentials, err := NewS3CredentialsFromFile(config.GetCredentialsFilePath())
	if err != nil {
		return nil, err
	}

	callTimeout, err := time.ParseDuration(config.GetCallTimeout())
	if err != nil {
		return nil, errors.NewNonRetriableErrorf(
			"failed to parse callTimeout: %w",
			err,
		)
	}

	return NewS3Client(
		config.GetEndpoint(),
		config.GetRegion(),
		credentials,
		callTimeout,
		registry,
		config.GetMaxRetriableErrorCount(),
		availabilityMonitoring,
	)
}

////////////////////////////////////////////////////////////////////////////////

func (c *S3Client) CreateBucket(
	ctx context.Context,
	bucket string,
) (err error) {

	ctx = withComponentLoggingField(ctx)
	logging.Info(ctx, "creating bucket %v in s3", bucket)

	ctx, cancel := context.WithTimeout(ctx, c.callTimeout)
	defer cancel()

	defer c.metrics.StatCall(ctx, "CreateBucket", bucket, "")(&err)

	_, err = c.s3.CreateBucketWithContext(ctx, &aws_s3.CreateBucketInput{
		Bucket: &bucket,
	})
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok {
			switch aerr.Code() {
			case aws_s3.ErrCodeBucketAlreadyOwnedByYou:
				// Bucket is already created
				return nil
			}
		}

		return errors.NewRetriableError(err)
	}

	return nil
}

////////////////////////////////////////////////////////////////////////////////

func (c *S3Client) BucketExists(
	ctx context.Context,
	bucket string,
) (exists bool, err error) {

	ctx = withComponentLoggingField(ctx)
	logging.Info(ctx, "checking if bucket %v exists in s3", bucket)

	ctx, cancel := context.WithTimeout(ctx, c.callTimeout)
	defer cancel()

	defer c.metrics.StatCall(ctx, "GetBucket", bucket, "")(&err)
	_, err = c.s3.HeadBucketWithContext(ctx, &aws_s3.HeadBucketInput{
		Bucket: &bucket,
	})
	if err != nil {
		if aerr, ok := err.(awserr.RequestFailure); ok {
			if aerr.StatusCode() == 404 {
				return false, nil
			}
		}

		return false, errors.NewRetriableError(err)
	}

	return true, nil
}

func (c *S3Client) GetObject(
	ctx context.Context,
	bucket string,
	key string,
) (o S3Object, err error) {

	ctx = withComponentLoggingField(ctx)
	logging.Info(ctx, "getting object from s3, bucket %v, key %v", bucket, key)

	ctx, cancel := context.WithTimeout(ctx, c.callTimeout)
	defer cancel()

	defer c.metrics.StatCall(ctx, "GetObject", bucket, key)(&err)

	res, err := c.s3.GetObjectWithContext(ctx, &aws_s3.GetObjectInput{
		Bucket: &bucket,
		Key:    &key,
	})
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok {
			switch aerr.Code() {
			case aws_s3.ErrCodeNoSuchKey:
				return S3Object{}, errors.NewSilentNonRetriableErrorf(
					"s3 object not found: %v",
					key,
				)
			case aws_s3.ErrCodeNoSuchBucket:
				return S3Object{}, errors.NewNonRetriableError(err)
			}
		}

		return S3Object{}, errors.NewRetriableError(err)
	}

	objData, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return S3Object{}, errors.NewRetriableError(err)
	}

	storageClass := aws.StringValue(res.StorageClass)
	if len(storageClass) == 0 {
		storageClass = aws_s3.StorageClassStandard
	}

	return S3Object{
		Data:         objData,
		Metadata:     res.Metadata,
		StorageClass: storageClass,
	}, nil
}

////////////////////////////////////////////////////////////////////////////////

func (c *S3Client) PutObject(
	ctx context.Context,
	bucket string,
	key string,
	object S3Object,
) (err error) {

	ctx = withComponentLoggingField(ctx)
	logging.Info(ctx, "putting object to s3, bucket %v, key %v", bucket, key)

	ctx, cancel := context.WithTimeout(ctx, c.callTimeout)
	defer cancel()

	defer c.metrics.StatCall(ctx, "PutObject", bucket, key)(&err)

	err = c.putObjectWithStorageClass(ctx, bucket, key, object)
	if err == nil {
		return nil
	}

	if !isQuotaLimitExceeded(err) {
		return wrapError(err)
	}

	c.metrics.OnQuotaExceeded(ctx, "PutObject")

	if len(object.StorageClass) == 0 {
		return wrapError(err)
	}

	logging.Debug(
		ctx,
		"s3 put object with storage class %q failed with quota "+
			"error, retrying without storage class",
		object.StorageClass,
	)

	err = c.putObject(ctx, bucket, key, object)
	return wrapError(err)
}

////////////////////////////////////////////////////////////////////////////////

func (c *S3Client) putObject(
	ctx context.Context,
	bucket string,
	key string,
	object S3Object,
) error {

	input := &aws_s3.PutObjectInput{
		Bucket:          &bucket,
		Key:             &key,
		Body:            bytes.NewReader(object.Data),
		Metadata:        object.Metadata,
		ContentEncoding: aws.String("application/octet-stream"),
	}

	_, err := c.s3.PutObjectWithContext(ctx, input)
	return err
}

////////////////////////////////////////////////////////////////////////////////

func (c *S3Client) putObjectWithStorageClass(
	ctx context.Context,
	bucket string,
	key string,
	object S3Object,
) error {

	input := &aws_s3.PutObjectInput{
		Bucket:          &bucket,
		Key:             &key,
		Body:            bytes.NewReader(object.Data),
		Metadata:        object.Metadata,
		ContentEncoding: aws.String("application/octet-stream"),
	}
	if len(object.StorageClass) != 0 {
		input.StorageClass = aws.String(object.StorageClass)
	}

	_, err := c.s3.PutObjectWithContext(ctx, input)
	return err
}

////////////////////////////////////////////////////////////////////////////////

func wrapError(err error) error {
	if err == nil {
		return nil
	}

	if aerr, ok := err.(awserr.Error); ok {
		switch aerr.Code() {
		case aws_s3.ErrCodeNoSuchBucket:
			return errors.NewNonRetriableError(err)
		}
	}

	return errors.NewRetriableError(err)
}

////////////////////////////////////////////////////////////////////////////////

func isQuotaLimitExceeded(err error) bool {
	aerr, ok := err.(awserr.Error)
	return ok && aerr.Code() == s3ErrCodeQuotaLimitExceeded
}

////////////////////////////////////////////////////////////////////////////////

func (c *S3Client) DeleteObject(
	ctx context.Context,
	bucket string,
	key string,
) (err error) {

	ctx = withComponentLoggingField(ctx)
	logging.Info(
		ctx,
		"deleting object from s3, bucket %v, key %v",
		bucket,
		key,
	)

	ctx, cancel := context.WithTimeout(ctx, c.callTimeout)
	defer cancel()

	defer c.metrics.StatCall(ctx, "DeleteObject", bucket, key)(&err)

	_, err = c.s3.DeleteObjectWithContext(ctx, &aws_s3.DeleteObjectInput{
		Bucket: &bucket,
		Key:    &key,
	})
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok {
			switch aerr.Code() {
			case aws_s3.ErrCodeNoSuchBucket:
				return errors.NewNonRetriableError(err)
			}
		}

		return errors.NewRetriableError(err)
	}

	return nil
}

////////////////////////////////////////////////////////////////////////////////

type S3Object struct {
	Data         []byte
	Metadata     map[string]*string
	StorageClass string
}

////////////////////////////////////////////////////////////////////////////////

type S3Credentials struct {
	ID     string `json:"id,omitempty"`
	Secret string `json:"secret,omitempty"`
}

func NewS3Credentials(id, secret string) S3Credentials {
	return S3Credentials{
		ID:     id,
		Secret: secret,
	}
}

func NewS3CredentialsFromFile(filePath string) (S3Credentials, error) {
	file, err := os.ReadFile(filePath)
	if err != nil {
		return S3Credentials{}, err
	}

	credentials := S3Credentials{}

	err = json.Unmarshal(file, &credentials)
	if err != nil {
		return S3Credentials{}, err
	}

	return credentials, nil
}