From e0de1e98fcfc093053941f1e50324119b16752c5 Mon Sep 17 00:00:00 2001 From: Eli Mallon Date: Mon, 1 Jun 2026 15:05:12 -0700 Subject: [PATCH 1/2] media: rebuild the stream transcoder on a streamer reconnect MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-DID continuous transcoder (a17b97f1) is only torn down on a 30s idle timeout, a codec swap, or a pipeline failure. A streamer who stops and restarts within that window — same DID, same codec — reuses the previous session's still-running encoder. The new RTMP/WHIP connection restarts the media timeline (PTS back near zero), the long-lived encoder sees a large backwards jump and stops emitting audio, so segments come out "missing audio track" and are dropped (consumed but never distributed) and the stream wedges until the encoder happens to re-sync (or never does). Fix: a monotonic ingest-session epoch. SegmentAndSignElem — the chokepoint all live ingest paths funnel through — stamps a fresh epoch onto the context that flows down to every segment of the session (→ ValidateMP4 → feedStreamTranscoder). needsReset now also fires when the incoming epoch is newer than the transcoder's, so a reconnect flushes + tears down the old transcoder (async, its tail still completes) and builds a fresh one for the restarted timeline — the same one-seam-at-the-boundary machinery as the codec-swap reset. Guarded on strict advance, so a stale straggler from an older session can't thrash it. TestFeedStreamTranscoderRebuildsOnNewSession drives the real registry path with two sessions over one DID+codec (re-feeding the fixture restarts the source PTS, the reconnect discontinuity) and asserts session 2 gets a fresh transcoder while session 1's is torn down; needsReset coverage extended for the session dimension. Co-Authored-By: Claude Opus 4.8 --- pkg/media/media.go | 15 ++++ pkg/media/segmenter.go | 7 ++ pkg/media/transcode_stream.go | 87 ++++++++++++++++++----- pkg/media/transcode_stream_test.go | 107 ++++++++++++++++++++++++++--- 4 files changed, 190 insertions(+), 26 deletions(-) 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) +} -- 2.51.2