diff --git a/pkg/media/media.go b/pkg/media/media.go index 7a608ea7..258f4bae 100644 --- a/pkg/media/media.go +++ b/pkg/media/media.go @@ -9,6 +9,7 @@ import ( "fmt" "io" "sync" + "sync/atomic" "github.com/google/uuid" "github.com/pion/interceptor" @@ -69,6 +70,20 @@ type MediaManager struct { // is distributed asynchronously (~1 GoP later). See transcode_stream.go. transcoders map[string]*streamTranscoder transcodersMu sync.Mutex + + // Monotonic ingest-session epoch. Each live ingest session (one + // SegmentAndSignElem) claims a fresh value, stamped onto its context, so the + // per-DID transcoder rebuilds when a streamer reconnects rather than feeding + // the restarted media timeline into the previous session's continuous encoder. + // See withIngestSession / feedStreamTranscoder. + ingestSessionSeq atomic.Uint64 +} + +// nextIngestSession claims a fresh monotonic ingest-session epoch for a new live +// session. Epochs are strictly increasing, so a newer session always wins over a +// transcoder built for an older one. +func (mm *MediaManager) nextIngestSession() uint64 { + return mm.ingestSessionSeq.Add(1) } type NewSegmentNotification struct { diff --git a/pkg/media/segmenter.go b/pkg/media/segmenter.go index 15709159..65fdabbc 100644 --- a/pkg/media/segmenter.go +++ b/pkg/media/segmenter.go @@ -180,6 +180,13 @@ func (mm *MediaManager) SegmentAndSignElem(ctx context.Context, ms MediaSigner) tracer := otel.Tracer("signer") streamer := ms.Streamer() + // Stamp a fresh ingest-session epoch on the context that flows down to every + // segment of this session (onSegment → ValidateMP4 → feedStreamTranscoder). A + // new live session (RTMP/WHIP (re)connect) restarts the media timeline; the + // per-DID continuous transcoder keys on this epoch and rebuilds rather than + // feeding the restarted timeline into the previous session's encoder. + ctx = withIngestSession(ctx, mm.nextIngestSession()) + // muxl path: stream the fMP4 through the per-segment signer. Each GoP // arrives as a bare canonical .m4s, which ValidateMP4 verifies, archives // (as .m4s), and distributes. muxl-sign stamps the signing time into the diff --git a/pkg/media/transcode_stream.go b/pkg/media/transcode_stream.go index ad5095cb..79ce2d6b 100644 --- a/pkg/media/transcode_stream.go +++ b/pkg/media/transcode_stream.go @@ -40,8 +40,12 @@ type transcodeJob struct { // tail). Not safe for concurrent Feed; feed from one goroutine (the stream's // validate path delivers segments in order). type streamTranscoder struct { - mm *MediaManager - target string // codec being ADDED: "opus" (source AAC) or "aac" (source Opus) + mm *MediaManager + target string // codec being ADDED: "opus" (source AAC) or "aac" (source Opus) + sessionID uint64 // ingest-session epoch this transcoder was built for; a newer + // session rebuilds it (registry-owned: set under transcodersMu before the + // transcoder is published to the map, read only by needsReset under the same + // lock, so it needs no separate synchronization). cert []byte keyPEM []byte onComplete func(token any, completed []byte) @@ -66,25 +70,53 @@ type streamTranscoder struct { // segment before it's flushed and torn down (the stream is presumed ended). const streamTranscoderIdle = 30 * time.Second +// ingestSessionKey carries the per-ingest-session epoch down the validate path. +// SegmentAndSignElem stamps a fresh epoch on every live session's context (a new +// RTMP/WHIP connection = a new session, with a restarted media timeline); +// feedStreamTranscoder reads it so a reconnect rebuilds the continuous transcoder +// instead of feeding the restarted timeline into the previous session's +// still-running encoder — a large backwards PTS discontinuity that makes the +// encoder stop emitting audio and wedges the stream. +type ingestSessionKey struct{} + +func withIngestSession(ctx context.Context, epoch uint64) context.Context { + return context.WithValue(ctx, ingestSessionKey{}, epoch) +} + +// ingestSessionFromContext returns the ingest-session epoch stamped on ctx, or 0 +// if none — e.g. a segment replicated from another node (which is already +// dual-codec, so it never builds a transcoder) or a direct unit-test feed. +func ingestSessionFromContext(ctx context.Context) uint64 { + epoch, _ := ctx.Value(ingestSessionKey{}).(uint64) + return epoch +} + // feedStreamTranscoder routes one source segment into the stream's continuous // transcoder, creating it on first use. The completed dual-codec segment is // distributed asynchronously (≈1 GoP later) via distributeSegment. func (mm *MediaManager) feedStreamTranscoder(ctx context.Context, vs *validatedSegment, src []byte, target string, cert, keyPEM []byte) error { did := vs.repoDID + sessionID := ingestSessionFromContext(ctx) mm.transcodersMu.Lock() t := mm.transcoders[did] - if t != nil && t.needsReset(target) { - // Source codec swapped (the needed target flipped, e.g. a streamer - // dropped RTMP/AAC and picked up WHIP/Opus) or the pipeline failed. The - // old pipeline expects the previous source codec, so feeding it the new - // one would stall it. Flush + tear it down (async, so we don't block - // ingest — its tail segments still complete) and rebuild for the new - // codec. One seam at the swap, clean after. + if t != nil && t.needsReset(target, sessionID) { + // The live encoder is wrong for the incoming segment: + // - a newer ingest session took over — the streamer reconnected (a rapid + // stop/start), which restarts the media timeline. Feeding that into the + // previous session's continuous encoder is a large backwards PTS jump + // that makes it stop emitting audio and wedges the stream; + // - the source codec swapped (the needed target flipped, e.g. a streamer + // dropped RTMP/AAC and picked up WHIP/Opus), so the pipeline expects the + // previous codec; or + // - the pipeline failed. + // Flush + tear it down (async, so we don't block ingest — its tail segments + // still complete) and rebuild. One seam at the boundary, clean after. old := t delete(mm.transcoders, did) t = nil - log.Log(ctx, "stream source codec swapped, resetting transcoder", - "streamer", did, "from_target", old.target, "to_target", target) + log.Log(ctx, "resetting stream transcoder", + "streamer", did, "from_target", old.target, "to_target", target, + "from_session", old.sessionID, "to_session", sessionID) go func() { if err := old.Close(); err != nil { log.Error(ctx, "stream transcoder reset close failed", "streamer", did, "error", err) @@ -101,9 +133,10 @@ func (mm *MediaManager) feedStreamTranscoder(ctx context.Context, vs *validatedS log.Error(streamCtx, "distribute completed segment failed", "streamer", v.repoDID, "error", err) } }) + t.sessionID = sessionID t.reaper = time.AfterFunc(streamTranscoderIdle, func() { mm.reapStreamTranscoder(did, t) }) mm.transcoders[did] = t - log.Log(ctx, "stream transcoder started", "streamer", did, "target", target) + log.Log(ctx, "stream transcoder started", "streamer", did, "target", target, "session", sessionID) } mm.transcodersMu.Unlock() @@ -111,11 +144,24 @@ func (mm *MediaManager) feedStreamTranscoder(ctx context.Context, vs *validatedS return t.Feed(src, vs) } -// needsReset reports whether an existing per-stream transcoder must be torn -// down and rebuilt: the source codec swapped (target flipped — e.g. RTMP/AAC → -// WHIP/Opus mid-stream) or its pipeline has failed. -func (t *streamTranscoder) needsReset(target string) bool { - return t.target != target || t.failed() +// needsReset reports whether an existing per-stream transcoder must be torn down +// and rebuilt rather than fed the incoming segment: +// - a newer ingest session took over (sessionID advanced — the streamer +// reconnected, restarting the media timeline; the continuous encoder is still +// at the old timeline, and a backwards PTS jump would make it stop emitting +// audio and wedge the stream), +// - the source codec swapped (target flipped — e.g. RTMP/AAC → WHIP/Opus +// mid-stream), or +// - its pipeline has failed. +// +// A stale straggler from an OLDER session (sessionID < t.sessionID) does NOT +// reset — the newer session keeps its encoder. In practice this can't arise: a +// session's segments are all fed (synchronously, before its ingest returns) +// before the next session starts, so feeds never interleave across sessions. +// Guarding on strict advance rather than inequality just makes that explicit and +// avoids any reset thrash if they ever did. +func (t *streamTranscoder) needsReset(target string, sessionID uint64) bool { + return t.target != target || sessionID > t.sessionID || t.failed() } // failed reports whether the transcoder's pipeline has errored out. @@ -125,6 +171,13 @@ func (t *streamTranscoder) failed() bool { return t.err != nil } +// isClosed reports whether the transcoder has been torn down (flushed + stopped). +func (t *streamTranscoder) isClosed() bool { + t.mu.Lock() + defer t.mu.Unlock() + return t.closed +} + // reapStreamTranscoder flushes and removes an idle stream's transcoder (the // flush emits its final buffered segment). Safe to call once per timer fire. func (mm *MediaManager) reapStreamTranscoder(did string, t *streamTranscoder) { diff --git a/pkg/media/transcode_stream_test.go b/pkg/media/transcode_stream_test.go index 2878ae2a..9a458c06 100644 --- a/pkg/media/transcode_stream_test.go +++ b/pkg/media/transcode_stream_test.go @@ -8,25 +8,114 @@ import ( "strconv" "sync" "testing" + "time" "github.com/stretchr/testify/require" + "stream.place/streamplace/pkg/aqtime" "stream.place/streamplace/pkg/config" "stream.place/streamplace/pkg/crypto/signers" + "stream.place/streamplace/pkg/livehls" "stream.place/streamplace/pkg/muxl" "stream.place/streamplace/test/remote" ) -// TestStreamTranscoderNeedsReset covers the source-codec-swap detection: a -// streamer dropping RTMP/AAC and picking up WHIP/Opus (or vice versa) flips the -// needed target, which must reset the continuous transcoder rather than feed -// the new codec into a pipeline built for the old one. +// TestStreamTranscoderNeedsReset covers the rebuild triggers for the per-DID +// continuous transcoder: a newer ingest session (a streamer reconnecting — a +// rapid stop/start that restarts the media timeline), a source-codec swap (RTMP/ +// AAC ↔ WHIP/Opus flips the needed target), or a failed pipeline. In each case +// the live encoder is wrong for the incoming segment and must be torn down +// rather than fed (feeding a restarted timeline into it wedges the stream). func TestStreamTranscoderNeedsReset(t *testing.T) { - tr := &streamTranscoder{target: "opus"} // adding Opus to an AAC source - require.False(t, tr.needsReset("opus"), "same source codec keeps the encoder running") - require.True(t, tr.needsReset("aac"), "source codec swapped (AAC→Opus) must reset") + tr := &streamTranscoder{target: "opus", sessionID: 1} // adding Opus to an AAC source - failed := &streamTranscoder{target: "aac", err: errors.New("pipeline died")} - require.True(t, failed.needsReset("aac"), "a failed pipeline rebuilds on the next segment") + // Same codec + same session: keep the continuous encoder running. + require.False(t, tr.needsReset("opus", 1), "same codec + same session keeps the encoder running") + + // Source codec swapped mid-session (AAC→Opus flips the needed target): reset. + require.True(t, tr.needsReset("aac", 1), "source codec swapped (AAC→Opus) must reset") + + // A newer ingest session took over (streamer reconnected → restarted media + // timeline): reset even with the same codec — the rapid-restart fix. + require.True(t, tr.needsReset("opus", 2), "a newer ingest session must rebuild the transcoder") + + // A stale straggler from an OLDER session must NOT reset the newer encoder. + require.False(t, tr.needsReset("opus", 0), "an older/stale session must not reset the current encoder") + + // A failed pipeline rebuilds on the next segment regardless. + failed := &streamTranscoder{target: "aac", sessionID: 5, err: errors.New("pipeline died")} + require.True(t, failed.needsReset("aac", 5), "a failed pipeline rebuilds on the next segment") +} + +// TestFeedStreamTranscoderRebuildsOnNewSession is the end-to-end regression for +// the rapid stop/start wedge: a streamer disconnects and reconnects within the +// transcoder's idle window, so the registry would otherwise feed the second +// session's restarted media timeline into the first session's still-running +// continuous encoder (a large backwards PTS jump → the encoder stops emitting +// audio → "emitted segment missing audio track" → segments dropped → the stream +// wedges). With the ingest-session epoch, the second session must get a FRESH +// transcoder and the first session's must be flushed + torn down. +// +// It drives the real registry path (feedStreamTranscoder, the same one +// ValidateMP4 uses): two sessions over the same DID + codec, with the same +// fixture re-fed for the second session — re-feeding restarts the source PTS at +// zero, exactly the discontinuity a reconnect produces. +func TestFeedStreamTranscoderRebuildsOnNewSession(t *testing.T) { + ctx := context.Background() + ms := newBareSegmentSigner(t) + segs := allSignedBareSegments(t, ctx, ms, getFixture("h264-opus-frag.mp4")) + require.GreaterOrEqual(t, len(segs), 2, "fixture should produce multiple segments") + + keyPEM, err := signers.MarshalES256KPrivateKeyPEM(ms.Signer) + require.NoError(t, err) + + // A real-enough MediaManager: a temp data dir so completed segments archive + // without touching the repo, and an (unused) live-window map. Completed + // segments are unpublished, so distributeSegment archives them but folds + // nothing into the live window and notifies no subscribers — no blocking. + mm := &MediaManager{ + cli: &config.CLI{BroadcasterHost: "test.example.com", DataDir: t.TempDir()}, + transcoders: map[string]*streamTranscoder{}, + liveWindows: map[string]*livehls.Writer{}, + } + + const did = "did:web:didweb.example" + base := time.Unix(1700000000, 0).UTC() + feedSession := func(epoch uint64, startIdx int) { + sctx := withIngestSession(ctx, epoch) + for i, seg := range segs { + // Distinct per-segment StartTime so archived filenames don't collide. + vs := &validatedSegment{ + repoDID: did, + meta: &SegmentMetadata{StartTime: aqtime.FromTime(base.Add(time.Duration(startIdx+i) * time.Second))}, + local: true, + } + require.NoError(t, mm.feedStreamTranscoder(sctx, vs, seg, "aac", ms.Cert, keyPEM), + "feed session %d segment %d", epoch, i) + } + } + + // Session 1. + s1 := mm.nextIngestSession() + feedSession(s1, 0) + t1 := mm.transcoders[did] + require.NotNil(t, t1, "session 1 built a transcoder") + require.Equal(t, s1, t1.sessionID) + + // Session 2: same DID + codec, fresh epoch (the reconnect). The first feed of + // this session must reset the registry to a brand-new transcoder. + s2 := mm.nextIngestSession() + feedSession(s2, len(segs)) + t2 := mm.transcoders[did] + require.NotNil(t, t2, "session 2 built a transcoder") + require.NotSame(t, t1, t2, + "a new ingest session must rebuild the transcoder, not reuse the previous session's continuous encoder") + require.Equal(t, s2, t2.sessionID) + + // The previous session's transcoder is flushed + torn down (async on reset). + require.Eventually(t, t1.isClosed, 20*time.Second, 20*time.Millisecond, + "the previous session's transcoder must be flushed + torn down on reconnect") + + require.NoError(t, t2.Close(), "the rebuilt transcoder drains cleanly") } // allSignedBareSegments signs the fragmented fixture per-segment and returns -- 2.51.2 From 65cd4a50fb3e58277a595032a7be06329a0d44e4 Mon Sep 17 00:00:00 2001 From: Eli Mallon Date: Mon, 1 Jun 2026 15:35:34 -0700 Subject: [PATCH 2/2] blob: use multipart server-side copy for >5 GiB Move A single CopyObject is capped at 5 GiB, so finalizing VODs larger than that failed. Add s3.Copy, which HEADs the source and either issues a single CopyObject (<=5 GiB) or drives a multipart UploadPartCopy for larger objects, and have S3Store.Move delegate to it. The multipart path runs server-side range copies concurrently (bounded), preserves content type, and aborts the upload on any part failure so S3 doesn't retain orphaned parts. Co-Authored-By: Claude Opus 4.8 --- pkg/blob/s3.go | 14 +-- pkg/blob/store.go | 2 +- pkg/s3/copy.go | 188 ++++++++++++++++++++++++++++++++++++++ pkg/s3/copy_test.go | 214 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 410 insertions(+), 8 deletions(-) create mode 100644 pkg/s3/copy.go create mode 100644 pkg/s3/copy_test.go diff --git a/pkg/blob/s3.go b/pkg/blob/s3.go index ab45bb19..22f0bd27 100644 --- a/pkg/blob/s3.go +++ b/pkg/blob/s3.go @@ -19,7 +19,9 @@ import ( // Writes go to a hidden staging prefix (.staging/) so that // in-progress multipart uploads can't collide with the final // content-addressed key. Complete renames staging -> the configured -// key via CopyObject + DeleteObject. +// key via a server-side Copy (single CopyObject, or a multipart +// UploadPartCopy for objects past S3's 5 GiB single-copy cap) + +// DeleteObject. type S3Store struct { client *awss3.Client bucket string @@ -84,12 +86,10 @@ func (s *S3Store) NewWriter(ctx context.Context, key, contentType string) (Write } func (s *S3Store) Move(ctx context.Context, srcKey, dstKey string) error { - _, err := s.client.CopyObject(ctx, &awss3.CopyObjectInput{ - Bucket: aws.String(s.bucket), - Key: aws.String(dstKey), - CopySource: aws.String(s.bucket + "/" + srcKey), - }) - if err != nil { + // Copy handles the >5 GiB VODs that a single CopyObject can't (it falls + // back to a multipart server-side copy), and HEADs the source first so a + // missing key surfaces as a NotFound we can treat idempotently below. + if err := s3pkg.Copy(ctx, s.client, s.bucket, srcKey, dstKey); err != nil { if isS3NotFound(err) { // Idempotency: maybe a previous Move already renamed // source -> dest. If dest exists, we're done. diff --git a/pkg/blob/store.go b/pkg/blob/store.go index c54e1df5..fb5a4dfe 100644 --- a/pkg/blob/store.go +++ b/pkg/blob/store.go @@ -51,7 +51,7 @@ type Store interface { // Move relocates the blob from srcKey to dstKey atomically (where // the underlying storage permits — POSIX rename on FileStore, - // CopyObject+DeleteObject on S3Store). If dstKey already exists, + // server-side Copy+DeleteObject on S3Store). If dstKey already exists, // it is overwritten. Returns nil if srcKey does not exist after a // successful Move (idempotency for retried renames). Move(ctx context.Context, srcKey, dstKey string) error diff --git a/pkg/s3/copy.go b/pkg/s3/copy.go new file mode 100644 index 00000000..b7f8e4db --- /dev/null +++ b/pkg/s3/copy.go @@ -0,0 +1,188 @@ +package 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 +} diff --git a/pkg/s3/copy_test.go b/pkg/s3/copy_test.go new file mode 100644 index 00000000..e98bc524 --- /dev/null +++ b/pkg/s3/copy_test.go @@ -0,0 +1,214 @@ +package s3 + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "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" + "github.com/stretchr/testify/require" +) + +// fakeCopyClient is an in-memory stand-in for the subset of *s3.Client that +// Copy uses. HeadObject reports a configurable size/content-type (so the +// >5 GiB multipart path can be exercised without allocating gigabytes), and +// the multipart ops record the ranges they were asked to copy. +type fakeCopyClient struct { + headSize int64 + headContentType string + headErr error + + copyDelay time.Duration + failPart int32 // if >0, UploadPartCopy for this part number returns an error + + srcBytes []byte // for the small-object CopyObject path + dst map[string][]byte // CopyObject writes here when non-nil + + mu sync.Mutex + copyObjectCalls int + createCalls int + createdType string + partRanges map[int32][2]int64 + completedParts []types.CompletedPart + aborted bool + inFlight int + maxInFlight int +} + +func (f *fakeCopyClient) HeadObject(context.Context, *s3.HeadObjectInput, ...func(*s3.Options)) (*s3.HeadObjectOutput, error) { + if f.headErr != nil { + return nil, f.headErr + } + out := &s3.HeadObjectOutput{ContentLength: aws.Int64(f.headSize)} + if f.headContentType != "" { + out.ContentType = aws.String(f.headContentType) + } + return out, nil +} + +func (f *fakeCopyClient) CopyObject(_ context.Context, in *s3.CopyObjectInput, _ ...func(*s3.Options)) (*s3.CopyObjectOutput, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.copyObjectCalls++ + if f.dst != nil { + f.dst[aws.ToString(in.Key)] = f.srcBytes + } + return &s3.CopyObjectOutput{}, nil +} + +func (f *fakeCopyClient) CreateMultipartUpload(_ context.Context, in *s3.CreateMultipartUploadInput, _ ...func(*s3.Options)) (*s3.CreateMultipartUploadOutput, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.createCalls++ + f.createdType = aws.ToString(in.ContentType) + return &s3.CreateMultipartUploadOutput{UploadId: aws.String("test-upload-id")}, nil +} + +func (f *fakeCopyClient) UploadPartCopy(_ context.Context, in *s3.UploadPartCopyInput, _ ...func(*s3.Options)) (*s3.UploadPartCopyOutput, error) { + f.mu.Lock() + f.inFlight++ + if f.inFlight > f.maxInFlight { + f.maxInFlight = f.inFlight + } + f.mu.Unlock() + + if f.copyDelay > 0 { + time.Sleep(f.copyDelay) + } + + num := aws.ToInt32(in.PartNumber) + var start, end int64 + if _, err := fmt.Sscanf(aws.ToString(in.CopySourceRange), "bytes=%d-%d", &start, &end); err != nil { + return nil, fmt.Errorf("bad CopySourceRange %q: %w", aws.ToString(in.CopySourceRange), err) + } + + f.mu.Lock() + defer f.mu.Unlock() + f.inFlight-- + if f.failPart > 0 && num == f.failPart { + return nil, fmt.Errorf("simulated failure on part %d", num) + } + f.partRanges[num] = [2]int64{start, end} + return &s3.UploadPartCopyOutput{ + CopyPartResult: &types.CopyPartResult{ETag: aws.String(fmt.Sprintf("etag-%d", num))}, + }, nil +} + +func (f *fakeCopyClient) CompleteMultipartUpload(_ context.Context, in *s3.CompleteMultipartUploadInput, _ ...func(*s3.Options)) (*s3.CompleteMultipartUploadOutput, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.completedParts = in.MultipartUpload.Parts + return &s3.CompleteMultipartUploadOutput{}, nil +} + +func (f *fakeCopyClient) AbortMultipartUpload(context.Context, *s3.AbortMultipartUploadInput, ...func(*s3.Options)) (*s3.AbortMultipartUploadOutput, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.aborted = true + return &s3.AbortMultipartUploadOutput{}, nil +} + +// TestCopySmallObjectUsesCopyObject verifies sub-5-GiB objects take the +// single-call CopyObject path and land byte-for-byte at the destination. +func TestCopySmallObjectUsesCopyObject(t *testing.T) { + src := []byte("hello world, a small object") + fake := &fakeCopyClient{ + headSize: int64(len(src)), + headContentType: "video/mp4", + srcBytes: src, + dst: map[string][]byte{}, + partRanges: map[int32][2]int64{}, + } + require.NoError(t, copyObject(context.Background(), fake, "bucket", "src", "dst")) + require.Equal(t, 1, fake.copyObjectCalls) + require.Equal(t, 0, fake.createCalls, "small object must not open a multipart upload") + require.Equal(t, src, fake.dst["dst"]) +} + +// TestCopyAtThresholdUsesCopyObject pins the boundary: an object exactly at +// maxCopyObjectSize still copies in one shot. +func TestCopyAtThresholdUsesCopyObject(t *testing.T) { + fake := &fakeCopyClient{headSize: maxCopyObjectSize, partRanges: map[int32][2]int64{}} + require.NoError(t, copyObject(context.Background(), fake, "bucket", "src", "dst")) + require.Equal(t, 1, fake.copyObjectCalls) + require.Equal(t, 0, fake.createCalls) +} + +// TestCopyLargeObjectUsesMultipart verifies an over-5-GiB object is copied +// with a multipart UploadPartCopy whose parts tile the source exactly, +// preserve content type, finish in ascending order, and run concurrently. +func TestCopyLargeObjectUsesMultipart(t *testing.T) { + const size = int64(12) * 1024 * 1024 * 1024 // 12 GiB -> spans many parts + fake := &fakeCopyClient{ + headSize: size, + headContentType: "video/mp4", + copyDelay: 5 * time.Millisecond, // force overlap so concurrency is observable + partRanges: map[int32][2]int64{}, + } + require.NoError(t, copyObject(context.Background(), fake, "bucket", "src", "dst")) + + require.Equal(t, 0, fake.copyObjectCalls, "large object must not use single CopyObject") + require.Equal(t, 1, fake.createCalls) + require.Equal(t, "video/mp4", fake.createdType, "content type must be preserved") + + wantParts := int((size + copyPartSize - 1) / copyPartSize) + require.Len(t, fake.partRanges, wantParts) + require.Len(t, fake.completedParts, wantParts) + + // Parts must tile [0,size) contiguously with no gaps or overlaps, and + // every part except the last is exactly copyPartSize. + var covered int64 + for n := int32(1); n <= int32(wantParts); n++ { + r, ok := fake.partRanges[n] + require.Truef(t, ok, "missing part %d", n) + require.Equalf(t, covered, r[0], "part %d start should continue from previous end", n) + require.GreaterOrEqual(t, r[1], r[0]) + if int(n) < wantParts { + require.Equalf(t, int64(copyPartSize), r[1]-r[0]+1, "non-final part %d should be a full part", n) + } + covered = r[1] + 1 + } + require.Equal(t, size, covered, "parts must cover the whole object") + + // CompleteMultipartUpload requires ascending part numbers. + for i := 1; i < len(fake.completedParts); i++ { + require.Less(t, + aws.ToInt32(fake.completedParts[i-1].PartNumber), + aws.ToInt32(fake.completedParts[i].PartNumber)) + } + require.Greater(t, fake.maxInFlight, 1, "expected part copies to run concurrently") + require.False(t, fake.aborted) +} + +// TestCopyLargeObjectAbortsOnPartError verifies a failed part copy aborts +// the multipart upload (so S3 doesn't retain orphaned parts) and never +// completes it. +func TestCopyLargeObjectAbortsOnPartError(t *testing.T) { + const size = int64(12) * 1024 * 1024 * 1024 + fake := &fakeCopyClient{ + headSize: size, + partRanges: map[int32][2]int64{}, + failPart: 3, + } + err := copyObject(context.Background(), fake, "bucket", "src", "dst") + require.Error(t, err) + require.Contains(t, err.Error(), "upload part copy 3") + require.True(t, fake.aborted, "a failed part must abort the upload") + require.Empty(t, fake.completedParts, "must not complete after a part failure") +} + +// TestCopyHeadErrorPropagates verifies a missing source surfaces the +// HeadObject error (which blob.Move sniffs as NotFound for idempotency) +// before any copy is attempted. +func TestCopyHeadErrorPropagates(t *testing.T) { + fake := &fakeCopyClient{headErr: fmt.Errorf("api error NotFound: object missing")} + err := copyObject(context.Background(), fake, "bucket", "missing", "dst") + require.Error(t, err) + require.Contains(t, err.Error(), "NotFound") + require.Equal(t, 0, fake.copyObjectCalls) + require.Equal(t, 0, fake.createCalls) +}