diff --git a/pkg/s3/concat.go b/pkg/s3/concat.go index d4f8b4220..dcab300d0 100644 --- a/pkg/s3/concat.go +++ b/pkg/s3/concat.go @@ -3,7 +3,6 @@ package s3 import ( "bytes" "context" - "errors" "fmt" "io" @@ -13,15 +12,21 @@ import ( "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" ) -// ErrConcatPartTooSmall is returned by ConcatWithHeader when a non-final source -// object is smaller than S3's 5 MB minimum part size and therefore can't be a -// server-side UploadPartCopy part. The caller is expected to fall back to a -// download-and-rewrite assembly. In production (10-minute cutover objects) this -// never triggers; it only guards against pathologically short streams. -var ErrConcatPartTooSmall = errors.New("s3 concat: non-final source object below 5MB minimum part size") +// concatPartSize is the uniform length of every non-final part the concat +// emits. R2 — unlike AWS/minio — rejects CompleteMultipartUpload with "All +// non-trailing parts must have the same length" unless parts are uniform, so +// the destination byte stream is sliced into fixed windows rather than one +// part per source object. 16MB (matching MultipartPartSize) keeps the +// boundary-window downloads small while allowing a 10000-part object to +// reach ~156GB. +const concatPartSize = MultipartPartSize + +// s3MaxParts is the multipart part-count ceiling shared by AWS and R2. +const s3MaxParts = 10000 // concatAPI is the subset of *s3.Client that ConcatWithHeader uses. Pulled out // so tests can inject a fake; *s3.Client satisfies it. @@ -37,15 +42,12 @@ type concatAPI interface { // // The layout exists to give live-to-VOD finalize a content blob shaped exactly // like an uploaded VOD ([init][segments…]) without re-uploading the (possibly -// many-GB) stream: the synthesized init header rides in part 1, and the bulk of -// the bytes are copied with UploadPartCopy and never pass through this process. -// -// Part 1 (uploaded) is header plus just enough leading source bytes to clear -// S3's 5 MB minimum, so it costs ~5 MB of memory + one upload. Every remaining -// source object becomes one server-side UploadPartCopy part. S3 requires every -// part except the last to be ≥5 MB; a non-final source object below that bound -// can't be copied as its own part, so ConcatWithHeader returns -// ErrConcatPartTooSmall and the caller falls back to a full rewrite. +// many-GB) stream: the destination byte stream is sliced into concatPartSize +// windows (uniform non-trailing parts, which R2 requires); a window that falls +// entirely inside one source object becomes a server-side UploadPartCopy, and +// only the windows containing the header or an object boundary are downloaded +// and re-uploaded. That bounds the through-process traffic to roughly one +// window per source object regardless of stream length. func ConcatWithHeader(ctx context.Context, client *s3.Client, bucket string, header []byte, srcKeys []string, dstKey, contentType string) error { return concatWithHeader(ctx, client, bucket, header, srcKeys, dstKey, contentType) } @@ -78,6 +80,23 @@ func concatWithHeader(ctx context.Context, client concatAPI, bucket string, head sizes[i] = aws.ToInt64(head.ContentLength) } + // Destination layout: header at [0, headerLen), then each source object at + // offsets[i]. Every part is exactly concatPartSize except the trailing one. + headerLen := int64(len(header)) + offsets := make([]int64, len(srcKeys)) + total := headerLen + for i := range srcKeys { + offsets[i] = total + total += sizes[i] + } + if total == 0 { + return fmt.Errorf("s3 concat: nothing to assemble (empty header and sources)") + } + numParts := (total + concatPartSize - 1) / concatPartSize + if numParts > s3MaxParts { + return fmt.Errorf("s3 concat: %d bytes needs %d parts, exceeding the %d-part limit", total, numParts, s3MaxParts) + } + create := &s3.CreateMultipartUploadInput{ Bucket: aws.String(bucket), Key: aws.String(dstKey), @@ -101,103 +120,70 @@ func concatWithHeader(ctx context.Context, client concatAPI, bucket string, head }) } - var parts []types.CompletedPart - var partNum int32 - - // --- Part 1: header + leading source bytes until >= minPartSize. --- - // We pull only enough source bytes to clear the 5 MB floor (bounding memory - // to ~minPartSize), tracking which object we stopped in so the copy phase - // resumes from exactly there. - buf := bytes.NewBuffer(make([]byte, 0, 2*minPartSize+len(header))) - buf.Write(header) - idx := 0 // index of the source object the copy phase resumes at - var consumedInIdx int64 // bytes of srcKeys[idx] already pulled into part 1 - for idx < len(srcKeys) && int64(buf.Len()) < minPartSize { - remaining := int64(minPartSize) - int64(buf.Len()) - avail := sizes[idx] - consumedInIdx - take := remaining - if take > avail { - take = avail - } - // Look-ahead: a partial take that leaves this (non-final) object with a - // sub-5MB remainder would make that remainder an illegal small middle - // copy part. Absorb the whole object into part 1 instead — costing a - // little extra upload but keeping every copy part legal. - if take < avail && (avail-take) < int64(minPartSize) && idx != len(srcKeys)-1 { - take = avail - } - if take > 0 { - body, err := getRange(ctx, client, bucket, srcKeys[idx], consumedInIdx, consumedInIdx+take-1) - if err != nil { - abort() - span.RecordError(err) - return err + // Parts land at fixed indices, so the completed list is already ordered. + // Downloaded (boundary) windows hold up to concatPartSize bytes each, but + // there's at most ~one per source object; copyConcurrency bounds how many + // are in memory at once. + parts := make([]types.CompletedPart, numParts) + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(copyConcurrency) + for p := int64(0); p < numParts; p++ { + partNum := int32(p + 1) + wStart := p * concatPartSize + wEnd := min(wStart+concatPartSize, total) // exclusive + g.Go(func() error { + // A window entirely inside one source object copies server-side. + for i := range srcKeys { + if offsets[i] <= wStart && wEnd <= offsets[i]+sizes[i] { + 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(bucket + "/" + srcKeys[i]), + CopySourceRange: aws.String(fmt.Sprintf("bytes=%d-%d", wStart-offsets[i], wEnd-1-offsets[i])), + }) + if err != nil { + return fmt.Errorf("upload part copy %d s3://%s/%s: %w", partNum, bucket, srcKeys[i], err) + } + if res.CopyPartResult == nil { + return fmt.Errorf("upload part copy %d s3://%s/%s: missing CopyPartResult", partNum, bucket, srcKeys[i]) + } + parts[p] = types.CompletedPart{ + ETag: res.CopyPartResult.ETag, + PartNumber: aws.Int32(partNum), + } + return nil + } } - buf.Write(body) - consumedInIdx += take - } - if consumedInIdx >= sizes[idx] { - idx++ - consumedInIdx = 0 - } + // Mixed window (contains the header and/or spans object boundaries): + // assemble it in memory and upload it as a regular part. + buf := make([]byte, 0, wEnd-wStart) + if wStart < headerLen { + buf = append(buf, header[wStart:min(wEnd, headerLen)]...) + } + for i := range srcKeys { + s := max(wStart, offsets[i]) + e := min(wEnd, offsets[i]+sizes[i]) + if s >= e { + continue + } + body, err := getRange(gctx, client, bucket, srcKeys[i], s-offsets[i], e-1-offsets[i]) + if err != nil { + return err + } + buf = append(buf, body...) + } + return uploadPartAt(gctx, client, bucket, dstKey, uploadID, partNum, buf, &parts[p]) + }) } - - partNum++ - if err := uploadPart(ctx, client, bucket, dstKey, uploadID, partNum, buf.Bytes(), &parts); err != nil { + if err := g.Wait(); err != nil { abort() span.RecordError(err) - span.SetStatus(codes.Error, "upload_part_1") + span.SetStatus(codes.Error, "upload_parts") return err } - // --- Remaining objects: one server-side UploadPartCopy part each. --- - for i := idx; i < len(srcKeys); i++ { - start := int64(0) - if i == idx { - start = consumedInIdx // resume mid-object if part 1 stopped here - } - if start >= sizes[i] { - continue // wholly absorbed into part 1 - } - partSize := sizes[i] - start - isLast := i == len(srcKeys)-1 - if !isLast && partSize < minPartSize { - abort() - err := fmt.Errorf("%w (object %s contributes %d bytes)", ErrConcatPartTooSmall, srcKeys[i], partSize) - span.RecordError(err) - return err - } - if partSize > maxCopyObjectSize { - abort() - err := fmt.Errorf("s3 concat: object %s range %d bytes exceeds single-part copy limit", srcKeys[i], partSize) - span.RecordError(err) - return err - } - partNum++ - res, err := client.UploadPartCopy(ctx, &s3.UploadPartCopyInput{ - Bucket: aws.String(bucket), - Key: aws.String(dstKey), - UploadId: aws.String(uploadID), - PartNumber: aws.Int32(partNum), - CopySource: aws.String(bucket + "/" + srcKeys[i]), - CopySourceRange: aws.String(fmt.Sprintf("bytes=%d-%d", start, sizes[i]-1)), - }) - if err != nil { - abort() - span.RecordError(err) - span.SetStatus(codes.Error, "upload_part_copy") - return fmt.Errorf("upload part copy %d s3://%s/%s: %w", partNum, bucket, srcKeys[i], err) - } - if res.CopyPartResult == nil { - abort() - return fmt.Errorf("upload part copy %d s3://%s/%s: missing CopyPartResult", partNum, bucket, srcKeys[i]) - } - parts = append(parts, types.CompletedPart{ - ETag: res.CopyPartResult.ETag, - PartNumber: aws.Int32(partNum), - }) - } - if _, err := client.CompleteMultipartUpload(ctx, &s3.CompleteMultipartUploadInput{ Bucket: aws.String(bucket), Key: aws.String(dstKey), @@ -231,7 +217,7 @@ func getRange(ctx context.Context, client concatAPI, bucket, key string, start, return body, nil } -func uploadPart(ctx context.Context, client concatAPI, bucket, dstKey, uploadID string, partNum int32, body []byte, parts *[]types.CompletedPart) error { +func uploadPartAt(ctx context.Context, client concatAPI, bucket, dstKey, uploadID string, partNum int32, body []byte, out *types.CompletedPart) error { res, err := client.UploadPart(ctx, &s3.UploadPartInput{ Bucket: aws.String(bucket), Key: aws.String(dstKey), @@ -242,9 +228,9 @@ func uploadPart(ctx context.Context, client concatAPI, bucket, dstKey, uploadID if err != nil { return fmt.Errorf("upload part %d s3://%s/%s: %w", partNum, bucket, dstKey, err) } - *parts = append(*parts, types.CompletedPart{ + *out = types.CompletedPart{ ETag: res.ETag, PartNumber: aws.Int32(partNum), - }) + } return nil } diff --git a/pkg/s3/concat_test.go b/pkg/s3/concat_test.go index d02af9e28..fd983f796 100644 --- a/pkg/s3/concat_test.go +++ b/pkg/s3/concat_test.go @@ -9,6 +9,7 @@ import ( "sort" "strconv" "strings" + "sync" "testing" "github.com/aws/aws-sdk-go-v2/aws" @@ -18,10 +19,12 @@ import ( // fakeConcatS3 is an in-memory stand-in for the subset of S3 ConcatWithHeader // uses. CompleteMultipartUpload reconstructs the destination object from the -// recorded parts in part-number order and enforces S3's rule that every part -// except the last is ≥5MB, so a ConcatWithHeader bug that emits a small middle -// part fails loudly here rather than only against real S3. +// recorded parts in part-number order and enforces the strictest real-backend +// rules: every part except the last is ≥5MB (AWS), and all non-trailing parts +// have the same length (R2), so a ConcatWithHeader bug that emits uneven parts +// fails loudly here rather than only against real R2. type fakeConcatS3 struct { + mu sync.Mutex objects map[string][]byte parts map[int32][]byte // partNumber -> bytes, for the single in-flight upload assembled map[string][]byte @@ -60,7 +63,9 @@ func (f *fakeConcatS3) CreateMultipartUpload(_ context.Context, _ *awss3.CreateM func (f *fakeConcatS3) UploadPart(_ context.Context, in *awss3.UploadPartInput, _ ...func(*awss3.Options)) (*awss3.UploadPartOutput, error) { body, _ := io.ReadAll(in.Body) + f.mu.Lock() f.parts[aws.ToInt32(in.PartNumber)] = body + f.mu.Unlock() return &awss3.UploadPartOutput{ETag: aws.String(fmt.Sprintf("etag-%d", aws.ToInt32(in.PartNumber)))}, nil } @@ -73,13 +78,17 @@ func (f *fakeConcatS3) UploadPartCopy(_ context.Context, in *awss3.UploadPartCop return nil, fmt.Errorf("copy from missing object %s", key) } start, end := parseRange(aws.ToString(in.CopySourceRange), len(b)) + f.mu.Lock() f.parts[aws.ToInt32(in.PartNumber)] = append([]byte(nil), b[start:end+1]...) + f.mu.Unlock() return &awss3.UploadPartCopyOutput{ CopyPartResult: &types.CopyPartResult{ETag: aws.String(fmt.Sprintf("etag-%d", aws.ToInt32(in.PartNumber)))}, }, nil } func (f *fakeConcatS3) CompleteMultipartUpload(_ context.Context, in *awss3.CompleteMultipartUploadInput, _ ...func(*awss3.Options)) (*awss3.CompleteMultipartUploadOutput, error) { + f.mu.Lock() + defer f.mu.Unlock() nums := make([]int32, 0, len(in.MultipartUpload.Parts)) for _, p := range in.MultipartUpload.Parts { nums = append(nums, aws.ToInt32(p.PartNumber)) @@ -92,6 +101,9 @@ func (f *fakeConcatS3) CompleteMultipartUpload(_ context.Context, in *awss3.Comp if !isLast && len(body) < minPartSize { return nil, fmt.Errorf("EntityTooSmall: part %d is %d bytes (<5MB) and not last", n, len(body)) } + if !isLast && len(body) != len(f.parts[nums[0]]) { + return nil, fmt.Errorf("InvalidPart: all non-trailing parts must have the same length (part %d is %d bytes, part %d is %d bytes)", n, len(body), nums[0], len(f.parts[nums[0]])) + } out = append(out, body...) } f.assembled[aws.ToString(in.Key)] = out @@ -99,7 +111,9 @@ func (f *fakeConcatS3) CompleteMultipartUpload(_ context.Context, in *awss3.Comp } func (f *fakeConcatS3) AbortMultipartUpload(_ context.Context, _ *awss3.AbortMultipartUploadInput, _ ...func(*awss3.Options)) (*awss3.AbortMultipartUploadOutput, error) { + f.mu.Lock() f.aborted = true + f.mu.Unlock() return &awss3.AbortMultipartUploadOutput{}, nil } @@ -137,19 +151,21 @@ func TestConcatWithHeader(t *testing.T) { name string objects map[string][]byte order []string - wantErr error }{ { - name: "large objects partial first part", + // Objects bigger than concatPartSize: interior windows are pure + // server-side copies, boundary windows are assembled in memory. + name: "objects spanning multiple copy windows", objects: map[string][]byte{ - "a": filled('a', 10*mb), - "b": filled('b', 8*mb), + "a": filled('a', 40*mb), + "b": filled('b', 33*mb), "c": filled('c', 7*mb), }, order: []string{"a", "b", "c"}, }, { - name: "medium objects absorb whole first object", + // Objects smaller than one window: every window is mixed. + name: "objects smaller than a window", objects: map[string][]byte{ "a": filled('a', 6*mb), "b": filled('b', 6*mb), @@ -173,14 +189,15 @@ func TestConcatWithHeader(t *testing.T) { order: []string{"a", "b"}, }, { - name: "small middle object falls back", + // A sub-5MB middle object used to be unrepresentable as its own + // copy part; with uniform windows it just rides in a mixed window. + name: "small middle object", objects: map[string][]byte{ "a": filled('a', 6*mb), "b": filled('b', 2*mb), "c": filled('c', 6*mb), }, - order: []string{"a", "b", "c"}, - wantErr: ErrConcatPartTooSmall, + order: []string{"a", "b", "c"}, }, } @@ -188,15 +205,6 @@ func TestConcatWithHeader(t *testing.T) { t.Run(tc.name, func(t *testing.T) { fake := newFakeConcatS3(tc.objects) err := concatWithHeader(context.Background(), fake, "bucket", header, tc.order, "dst", "video/mp4") - if tc.wantErr != nil { - if !errors.Is(err, tc.wantErr) { - t.Fatalf("want error %v, got %v", tc.wantErr, err) - } - if !fake.aborted { - t.Fatalf("expected multipart upload to be aborted on too-small part") - } - return - } if err != nil { t.Fatalf("concatWithHeader: %v", err) } diff --git a/pkg/s3/live_test.go b/pkg/s3/live_test.go new file mode 100644 index 000000000..acaf71f5e --- /dev/null +++ b/pkg/s3/live_test.go @@ -0,0 +1,173 @@ +package s3 + +import ( + "bytes" + "context" + "io" + "net/url" + "os" + "strings" + "sync" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + awss3 "github.com/aws/aws-sdk-go-v2/service/s3" + "github.com/stretchr/testify/require" +) + +// liveConfig builds a Config from the S3_* env vars, tolerating the bucket +// being embedded as the endpoint URL's path (the format of Eli's ~/r2.env / +// ~/rustfs.env files). Returns false if the env vars aren't set. +func liveConfig(t *testing.T) (Config, bool) { + endpoint := os.Getenv("S3_ENDPOINT") + if endpoint == "" { + return Config{}, false + } + cfg := Config{ + Endpoint: endpoint, + Bucket: os.Getenv("S3_BUCKET"), + AccessKeyID: os.Getenv("S3_ACCESS_KEY_ID"), + SecretAccessKey: os.Getenv("S3_SECRET_ACCESS_KEY"), + Region: os.Getenv("S3_REGION"), + } + u, err := url.Parse(endpoint) + require.NoError(t, err) + if p := strings.Trim(u.Path, "/"); p != "" { + if cfg.Bucket == "" { + cfg.Bucket = p + } + u.Path = "" + cfg.Endpoint = u.String() + } + if cfg.Region == "" { + cfg.Region = "auto" + } + return cfg, true +} + +// keyRecorder captures the object keys the uploader creates so the test can +// fetch them back. +type keyRecorder struct { + mu sync.Mutex + keys []string +} + +func (r *keyRecorder) RecordStart(ctx context.Context, userDID, bucket, key, livestreamURI string, started time.Time) (string, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.keys = append(r.keys, key) + return key, nil +} + +func (r *keyRecorder) RecordComplete(ctx context.Context, id string, parts int32, size int64) error { + return nil +} + +// TestLiveS3Uploader exercises the real multipart upload path against a real +// S3-compatible endpoint. Skipped unless S3_ENDPOINT / S3_ACCESS_KEY_ID / +// S3_SECRET_ACCESS_KEY are set, e.g.: +// +// set -a; . ~/r2.env; go test -count=1 -run TestLiveS3Uploader -v ./pkg/s3 +// +// Segment sizes are deliberately irregular so the flushed parts have unequal +// sizes, like real live segments do — R2 (unlike AWS/minio) rejects multipart +// completes whose non-final parts aren't all the same size. +func TestLiveS3Uploader(t *testing.T) { + cfg, ok := liveConfig(t) + if !ok { + t.Skip("set S3_ENDPOINT etc. to run the live uploader test") + } + ctx := context.Background() + rec := &keyRecorder{} + up := NewS3Uploader(cfg, "did:test:live", "live-test/", time.Hour, rec) + + // Irregular segments: parts flush at >=5MB boundaries with different sizes. + var want []byte + for _, mb := range []int{2, 2, 2, 3, 3, 1, 2} { // parts: 6MB, 6MB(3+3), then 3MB tail + seg := bytes.Repeat([]byte("streamplace-live-s3-test"), mb*1024*1024/24) + seg = append(seg, bytes.Repeat([]byte{0xAB}, 137*mb)...) // knock sizes off round numbers + require.NoError(t, up.AddSegment(ctx, seg)) + want = append(want, seg...) + } + require.NoError(t, up.Close(ctx)) + + require.Len(t, rec.keys, 1) + key := rec.keys[0] + t.Logf("uploaded key %s", key) + + client := NewClient(cfg) + out, err := client.GetObject(ctx, &awss3.GetObjectInput{ + Bucket: aws.String(cfg.Bucket), + Key: aws.String(key), + }) + require.NoError(t, err) + got, err := io.ReadAll(out.Body) + require.NoError(t, err) + require.NoError(t, out.Body.Close()) + require.Equal(t, len(want), len(got)) + require.True(t, bytes.Equal(want, got)) + + _, err = client.DeleteObject(ctx, &awss3.DeleteObjectInput{ + Bucket: aws.String(cfg.Bucket), + Key: aws.String(key), + }) + require.NoError(t, err) +} + +// TestLiveConcatWithHeader exercises the live-to-VOD finalize assembly +// (header + UploadPartCopy concat) against a real endpoint. Same env gating +// as TestLiveS3Uploader. +func TestLiveConcatWithHeader(t *testing.T) { + cfg, ok := liveConfig(t) + if !ok { + t.Skip("set S3_ENDPOINT etc. to run the live concat test") + } + ctx := context.Background() + client := NewClient(cfg) + + put := func(key string, size int) []byte { + body := bytes.Repeat([]byte("streamplace-live-concat!"), size/24) + _, err := client.PutObject(ctx, &awss3.PutObjectInput{ + Bucket: aws.String(cfg.Bucket), + Key: aws.String(key), + Body: bytes.NewReader(body), + }) + require.NoError(t, err) + return body + } + keys := []string{"live-test/concat-src-0", "live-test/concat-src-1", "live-test/concat-src-2"} + cleanup := append([]string(nil), keys...) + defer func() { + for _, k := range cleanup { + _, _ = client.DeleteObject(ctx, &awss3.DeleteObjectInput{ + Bucket: aws.String(cfg.Bucket), + Key: aws.String(k), + }) + } + }() + + header := bytes.Repeat([]byte{0x42}, 1024) + want := append([]byte(nil), header...) + // Deliberately unequal source sizes, like real 10-minute cutover objects. + // The first two exceed concatPartSize so interior windows exercise real + // server-side UploadPartCopy; the boundaries exercise the mixed windows. + for i, size := range []int{40 * 1024 * 1024, 33 * 1024 * 1024, 7 * 1024 * 1024} { + want = append(want, put(keys[i], size)...) + } + + dstKey := "live-test/concat-dst.m4s" + cleanup = append(cleanup, dstKey) + require.NoError(t, ConcatWithHeader(ctx, client, cfg.Bucket, header, keys, dstKey, "video/mp4")) + + out, err := client.GetObject(ctx, &awss3.GetObjectInput{ + Bucket: aws.String(cfg.Bucket), + Key: aws.String(dstKey), + }) + require.NoError(t, err) + got, err := io.ReadAll(out.Body) + require.NoError(t, err) + require.NoError(t, out.Body.Close()) + require.Equal(t, len(want), len(got)) + require.True(t, bytes.Equal(want, got)) +} diff --git a/pkg/s3/s3.go b/pkg/s3/s3.go index 1eb340f8e..18641624d 100644 --- a/pkg/s3/s3.go +++ b/pkg/s3/s3.go @@ -91,6 +91,15 @@ func (u *S3Uploader) getLivestreamURI() string { // S3 requires each part except the last to be at least 5MB. const minPartSize = 5 * 1024 * 1024 +// liveUploadPartSize is the exact size of every non-final part the live +// uploader flushes. R2 — unlike AWS/minio — rejects CompleteMultipartUpload +// with "All non-trailing parts must have the same length" unless parts are +// uniform, so the buffer is sliced at fixed boundaries instead of flushing +// whatever segments accumulated past the 5MB minimum. The S3 minimum keeps +// flushes prompt for a live stream; at 5MB a 10000-part object still spans +// ~48GB, far beyond one 10-minute cutover object. +const liveUploadPartSize = minPartSize + type activeUpload struct { key string uploadID string @@ -254,8 +263,12 @@ func (u *S3Uploader) uploadLoop(ctx context.Context) { return nil } - flushBuffer := func() error { - if current == nil || len(current.buf) == 0 { + // flushPart uploads exactly the first n buffered bytes as the next part, + // keeping the remainder buffered. Callers pass liveUploadPartSize for every + // part except the object's final flush (completeUpload passes whatever is + // left) — R2 requires all non-trailing parts to have the same length. + flushPart := func(n int) error { + if current == nil || n == 0 { return nil } current.partNum++ @@ -266,18 +279,18 @@ func (u *S3Uploader) uploadLoop(ctx context.Context) { Key: aws.String(current.key), UploadId: aws.String(current.uploadID), PartNumber: aws.Int32(partNum), - Body: bytes.NewReader(current.buf), + Body: bytes.NewReader(current.buf[:n]), }) if err != nil { return fmt.Errorf("uploading part %d: %w", partNum, err) } - log.Debug(ctx, "uploaded S3 part", "key", current.key, "part", partNum, "size", len(current.buf)) + log.Debug(ctx, "uploaded S3 part", "key", current.key, "part", partNum, "size", n) current.parts = append(current.parts, types.CompletedPart{ ETag: resp.ETag, PartNumber: aws.Int32(partNum), }) - current.totalSize += int64(len(current.buf)) - current.buf = current.buf[:0] + current.totalSize += int64(n) + current.buf = append(current.buf[:0], current.buf[n:]...) return nil } @@ -285,7 +298,7 @@ func (u *S3Uploader) uploadLoop(ctx context.Context) { if current == nil { return nil } - if err := flushBuffer(); err != nil { + if err := flushPart(len(current.buf)); err != nil { return fmt.Errorf("error flushing buffer: %w", err) } if len(current.parts) == 0 { @@ -345,9 +358,10 @@ func (u *S3Uploader) uploadLoop(ctx context.Context) { // Append segment data to buffer current.buf = append(current.buf, seg...) - // Flush if buffer is large enough for a part - if len(current.buf) >= minPartSize { - if err := flushBuffer(); err != nil { + // Flush full-size parts; a sub-part remainder stays buffered until the + // next segment or the object's completing flush. + for len(current.buf) >= liveUploadPartSize { + if err := flushPart(liveUploadPartSize); err != nil { return err } } diff --git a/pkg/s3/uploader_test.go b/pkg/s3/uploader_test.go index 523fd3dab..e2531a9d2 100644 --- a/pkg/s3/uploader_test.go +++ b/pkg/s3/uploader_test.go @@ -3,6 +3,7 @@ package s3 import ( "context" "fmt" + "io" "sync" "testing" "time" @@ -15,8 +16,9 @@ import ( // fakeUploadAPI is an in-memory stand-in for the multipart subset of *s3.Client // the upload loop drives. It hands back dummy upload IDs / etags and never errs. type fakeUploadAPI struct { - mu sync.Mutex - creates int + mu sync.Mutex + creates int + partSizes []int } func (f *fakeUploadAPI) CreateMultipartUpload(_ context.Context, _ *awss3.CreateMultipartUploadInput, _ ...func(*awss3.Options)) (*awss3.CreateMultipartUploadOutput, error) { @@ -27,7 +29,14 @@ func (f *fakeUploadAPI) CreateMultipartUpload(_ context.Context, _ *awss3.Create return &awss3.CreateMultipartUploadOutput{UploadId: aws.String(fmt.Sprintf("up-%d", n))}, nil } -func (f *fakeUploadAPI) UploadPart(_ context.Context, _ *awss3.UploadPartInput, _ ...func(*awss3.Options)) (*awss3.UploadPartOutput, error) { +func (f *fakeUploadAPI) UploadPart(_ context.Context, in *awss3.UploadPartInput, _ ...func(*awss3.Options)) (*awss3.UploadPartOutput, error) { + body, err := io.ReadAll(in.Body) + if err != nil { + return nil, err + } + f.mu.Lock() + f.partSizes = append(f.partSizes, len(body)) + f.mu.Unlock() return &awss3.UploadPartOutput{ETag: aws.String("etag")}, nil } @@ -145,6 +154,41 @@ func TestS3UploaderCutoverCompletesObject(t *testing.T) { require.NotEqual(t, keys[0], keys[1], "post-cutover object must have a distinct key") } +// TestS3UploaderUniformParts proves the uploader slices its buffer into +// exactly liveUploadPartSize parts regardless of segment sizes, with only the +// object's final part smaller. R2 rejects multipart completes whose +// non-trailing parts differ in length, so "flush whatever accumulated past +// 5MB" (the old behavior) breaks against R2 with real, variable-size +// segments. +func TestS3UploaderUniformParts(t *testing.T) { + fc := &fakeUploadAPI{} + rec := &fakeRecorder{} + u := newS3Uploader(fc, "bucket", "did:plc:test", "did:plc:test/", time.Hour, rec) + + ctx := context.Background() + total := 0 + for _, mb := range []int{2, 3, 4, 3} { // 12MB in irregular chunks + require.NoError(t, u.AddSegment(ctx, make([]byte, mb*1024*1024))) + total += mb * 1024 * 1024 + } + require.NoError(t, u.Close(ctx)) + + fc.mu.Lock() + sizes := append([]int(nil), fc.partSizes...) + fc.mu.Unlock() + require.NotEmpty(t, sizes) + got := 0 + for i, s := range sizes { + got += s + if i < len(sizes)-1 { + require.Equal(t, liveUploadPartSize, s, "non-trailing part %d must be exactly liveUploadPartSize", i+1) + } else { + require.LessOrEqual(t, s, liveUploadPartSize) + } + } + require.Equal(t, total, got, "flushed parts must cover every byte exactly once") +} + // TestS3UploaderCloseIdempotent exercises the lifecycle fix that re-enabled // live S3 upload: Close must be safe to call repeatedly and concurrently (it // was a plain close(segCh) before, which panicked on the second call), and diff --git a/pkg/vod/finalize_livestream.go b/pkg/vod/finalize_livestream.go index ae2f133cf..d61b8bef7 100644 --- a/pkg/vod/finalize_livestream.go +++ b/pkg/vod/finalize_livestream.go @@ -168,19 +168,12 @@ func captureInitHeader(ctx context.Context, store blob.Store, key string) ([]byt } // assembleContentBlob writes header ++ objects to contentKey. For an S3 store -// this is a near-entirely server-side concat (UploadPartCopy); a non-final -// object below S3's 5 MB part floor (only short dev streams) falls back to a -// full rewrite, as does any non-S3 store. +// this is a near-entirely server-side concat (uniform UploadPartCopy windows; +// only header/boundary windows pass through this process); any non-S3 store +// falls back to a full rewrite. func assembleContentBlob(ctx context.Context, store blob.Store, header []byte, keys []string, contentKey string) error { if s3store, ok := store.(*blob.S3Store); ok { - err := s3pkg.ConcatWithHeader(ctx, s3store.Client(), s3store.Bucket(), header, keys, contentKey, "video/mp4") - if err == nil { - return nil - } - if !errors.Is(err, s3pkg.ErrConcatPartTooSmall) { - return err - } - log.Warn(ctx, "s3 concat fell back to rewrite (object below 5MB part floor)", "key", contentKey) + return s3pkg.ConcatWithHeader(ctx, s3store.Client(), s3store.Bucket(), header, keys, contentKey, "video/mp4") } return assembleViaWriter(ctx, store, header, keys, contentKey) }