Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
7.6 kB · 188 lines
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189package s3
import ( "context" "fmt"
"github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/aws/aws-sdk-go-v2/service/s3/types" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" "go.opentelemetry.io/otel/trace" "golang.org/x/sync/errgroup" "stream.place/streamplace/pkg/log")
// maxCopyObjectSize is the largest object S3 will copy in a single// CopyObject call (5 GiB). Anything larger fails with EntityTooLarge and// must be copied part-by-part with a multipart upload driven by// UploadPartCopy.const maxCopyObjectSize = 5 * 1024 * 1024 * 1024
// copyPartSize is the byte range each UploadPartCopy transfers. The copy// happens server-side inside S3, so — unlike MultipartPartSize, which sizes// an in-process buffer — this never costs us memory; it's sized large to// keep the part (and request) count low while staying under S3's 5 GiB// per-part ceiling. At 1 GiB/part the 10000-part limit allows objects up to// ~10 TiB, far beyond any VOD.const copyPartSize = 1024 * 1024 * 1024
// copyConcurrency bounds how many UploadPartCopy requests run at once.// Mirrors multipartUploadConcurrency: server-side range copies serialize// badly otherwise, dragging a large VOD's finalize out for minutes.const copyConcurrency = 8
// copyAPI is the subset of *s3.Client that Copy uses. Pulled out so tests// can inject a fake; *s3.Client satisfies it.type copyAPI interface { HeadObject(context.Context, *s3.HeadObjectInput, ...func(*s3.Options)) (*s3.HeadObjectOutput, error) CopyObject(context.Context, *s3.CopyObjectInput, ...func(*s3.Options)) (*s3.CopyObjectOutput, error) CreateMultipartUpload(context.Context, *s3.CreateMultipartUploadInput, ...func(*s3.Options)) (*s3.CreateMultipartUploadOutput, error) UploadPartCopy(context.Context, *s3.UploadPartCopyInput, ...func(*s3.Options)) (*s3.UploadPartCopyOutput, error) CompleteMultipartUpload(context.Context, *s3.CompleteMultipartUploadInput, ...func(*s3.Options)) (*s3.CompleteMultipartUploadOutput, error) AbortMultipartUpload(context.Context, *s3.AbortMultipartUploadInput, ...func(*s3.Options)) (*s3.AbortMultipartUploadOutput, error)}
// Copy copies an object within bucket from srcKey to dstKey, preserving the// source's content type. Objects at or below maxCopyObjectSize use a single// CopyObject; larger objects exceed S3's single-copy limit, so they're// copied with a multipart upload whose parts are server-side UploadPartCopy// range copies. A HeadObject against the source missing key surfaces as a// NotFound error the caller can sniff for idempotency.func Copy(ctx context.Context, client *s3.Client, bucket, srcKey, dstKey string) error { return copyObject(ctx, client, bucket, srcKey, dstKey)}
func copyObject(ctx context.Context, client copyAPI, bucket, srcKey, dstKey string) error { ctx = log.WithLogValues(ctx, "func", "s3.Copy") ctx, span := s3Tracer.Start(ctx, "s3.Copy", trace.WithAttributes( attribute.String("bucket", bucket), attribute.String("src_key", srcKey), attribute.String("dst_key", dstKey), )) defer span.End()
head, err := client.HeadObject(ctx, &s3.HeadObjectInput{ Bucket: aws.String(bucket), Key: aws.String(srcKey), }) if err != nil { span.RecordError(err) return fmt.Errorf("head s3://%s/%s: %w", bucket, srcKey, err) } size := aws.ToInt64(head.ContentLength) contentType := aws.ToString(head.ContentType) span.SetAttributes(attribute.Int64("size_bytes", size))
// CopySource is "bucket/key"; keep it unescaped to match what the rest // of the codebase already sends (our keys are content hashes, UUIDs and // DIDs that the endpoint accepts raw). copySource := bucket + "/" + srcKey
if size <= maxCopyObjectSize { if _, err := client.CopyObject(ctx, &s3.CopyObjectInput{ Bucket: aws.String(bucket), Key: aws.String(dstKey), CopySource: aws.String(copySource), }); err != nil { span.RecordError(err) span.SetStatus(codes.Error, "copy_object") return fmt.Errorf("copy s3://%s/%s -> %s: %w", bucket, srcKey, dstKey, err) } return nil }
if err := multipartCopy(ctx, client, bucket, dstKey, copySource, contentType, size); err != nil { span.RecordError(err) span.SetStatus(codes.Error, "multipart_copy") return err } return nil}
// multipartCopy copies an over-5-GiB object by opening a multipart upload at// dstKey and filling it with server-side UploadPartCopy range copies of the// source. Parts run concurrently (bounded by copyConcurrency) but land at// their fixed part numbers, so the completed-parts list is already ordered.// Any failure aborts the upload so S3 doesn't retain orphaned parts.func multipartCopy(ctx context.Context, client copyAPI, bucket, dstKey, copySource, contentType string, size int64) error { create := &s3.CreateMultipartUploadInput{ Bucket: aws.String(bucket), Key: aws.String(dstKey), } if contentType != "" { create.ContentType = aws.String(contentType) } resp, err := client.CreateMultipartUpload(ctx, create) if err != nil { return fmt.Errorf("create multipart copy s3://%s/%s: %w", bucket, dstKey, err) } uploadID := aws.ToString(resp.UploadId)
// Pre-compute the (start,end) byte range of every part. Part numbers are // 1-based; CopySourceRange's end offset is inclusive. type byteRange struct{ start, end int64 } var ranges []byteRange for off := int64(0); off < size; off += copyPartSize { end := off + copyPartSize - 1 if end > size-1 { end = size - 1 } ranges = append(ranges, byteRange{start: off, end: end}) }
parts := make([]types.CompletedPart, len(ranges)) g, gctx := errgroup.WithContext(ctx) g.SetLimit(copyConcurrency) for i, r := range ranges { partNum := int32(i + 1) g.Go(func() error { res, err := client.UploadPartCopy(gctx, &s3.UploadPartCopyInput{ Bucket: aws.String(bucket), Key: aws.String(dstKey), UploadId: aws.String(uploadID), PartNumber: aws.Int32(partNum), CopySource: aws.String(copySource), CopySourceRange: aws.String(fmt.Sprintf("bytes=%d-%d", r.start, r.end)), }) if err != nil { return fmt.Errorf("upload part copy %d (bytes %d-%d) s3://%s/%s: %w", partNum, r.start, r.end, bucket, dstKey, err) } if res.CopyPartResult == nil { return fmt.Errorf("upload part copy %d s3://%s/%s: missing CopyPartResult", partNum, bucket, dstKey) } parts[i] = types.CompletedPart{ ETag: res.CopyPartResult.ETag, PartNumber: aws.Int32(partNum), } return nil }) } if err := g.Wait(); err != nil { // Best-effort cleanup so a failed copy doesn't leave parts S3 keeps // billing for. Use ctx (not the cancelled gctx) so the abort runs. _, _ = client.AbortMultipartUpload(ctx, &s3.AbortMultipartUploadInput{ Bucket: aws.String(bucket), Key: aws.String(dstKey), UploadId: aws.String(uploadID), }) return err }
if _, err := client.CompleteMultipartUpload(ctx, &s3.CompleteMultipartUploadInput{ Bucket: aws.String(bucket), Key: aws.String(dstKey), UploadId: aws.String(uploadID), MultipartUpload: &types.CompletedMultipartUpload{Parts: parts}, }); err != nil { _, _ = client.AbortMultipartUpload(ctx, &s3.AbortMultipartUploadInput{ Bucket: aws.String(bucket), Key: aws.String(dstKey), UploadId: aws.String(uploadID), }) return fmt.Errorf("complete multipart copy s3://%s/%s: %w", bucket, dstKey, err) } log.Log(ctx, "completed S3 multipart copy", "bucket", bucket, "key", dstKey, "parts", len(parts), "size", size) return nil}