package media import ( "bytes" "context" "errors" "fmt" "io" "sort" "sync" "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.uber.org/zap" ) // multipartUploader uploads a file to S3. // // TODO: extract to s3 package type multipartUploader struct { s3 S3Client logger *zap.SugaredLogger } type uploadResult struct { completedPart types.CompletedPart size int64 } const ( targetPartSizeBytes = 5 * 1024 * 1024 // 5MB readBufferSizeBytes = 32_768 // 32Kb ) func newMultipartUploader(s3Client S3Client, logger *zap.SugaredLogger) *multipartUploader { return &multipartUploader{s3: s3Client, logger: logger} } // Upload uploads to an S3 bucket in 5MB parts. It buffers data internally // until a part is ready to send over the network. Parts are sent as soon as // they exceed the minimum part size of 5MB. // // TODO: expire after configurable period. func (u *multipartUploader) Upload(ctx context.Context, r io.Reader, bucket, key, contentType string) (int64, error) { var uploaded bool input := s3.CreateMultipartUploadInput{ Bucket: aws.String(bucket), Key: aws.String(key), ContentType: aws.String(contentType), } output, err := u.s3.CreateMultipartUpload(ctx, &input) if err != nil { return 0, fmt.Errorf("error creating multipart upload: %v", err) } // abort the upload if possible, logging any errors, on exit. defer func() { if uploaded { return } input := s3.AbortMultipartUploadInput{ Bucket: aws.String(bucket), Key: aws.String(key), UploadId: output.UploadId, } // if the context was cancelled, just use the background context. ctxToUse := ctx if ctxToUse.Err() != nil { ctxToUse = context.Background() } _, abortErr := u.s3.AbortMultipartUpload(ctxToUse, &input) if abortErr != nil { u.logger.Errorf("uploader: error aborting upload: %v", abortErr) } else { u.logger.Infof("aborted upload, key = %s", key) } }() uploadResultChan := make(chan uploadResult) uploadErrorChan := make(chan error, 1) // uploadPart uploads an individual part. uploadPart := func(wg *sync.WaitGroup, buf []byte, partNum int32) { defer wg.Done() partLen := int64(len(buf)) u.logger.With("key", key, "partNum", partNum, "partLen", partLen).Debug("uploading part") input := s3.UploadPartInput{ Body: bytes.NewReader(buf), Bucket: aws.String(bucket), Key: aws.String(key), PartNumber: partNum, UploadId: output.UploadId, ContentLength: partLen, } output, uploadErr := u.s3.UploadPart(ctx, &input) if uploadErr != nil { // TODO: retry on failure uploadErrorChan <- uploadErr return } u.logger.With("key", key, "partNum", partNum, "partLen", partLen, "etag", *output.ETag).Debug("uploaded part") uploadResultChan <- uploadResult{ completedPart: types.CompletedPart{ETag: output.ETag, PartNumber: partNum}, size: partLen, } } wgDone := make(chan struct{}) var wg sync.WaitGroup wg.Add(1) // done when the reader goroutine returns go func() { wg.Wait() wgDone <- struct{}{} }() readChan := make(chan error, 1) go func() { defer wg.Done() var closing bool currPart := bytes.NewBuffer(make([]byte, 0, targetPartSizeBytes+readBufferSizeBytes)) partNum := int32(1) buf := make([]byte, readBufferSizeBytes) for { n, readErr := r.Read(buf) if readErr == io.EOF { closing = true } else if readErr != nil { readChan <- readErr return } _, _ = currPart.Write(buf[:n]) if closing || currPart.Len() >= targetPartSizeBytes { part := make([]byte, currPart.Len()) copy(part, currPart.Bytes()) currPart.Truncate(0) wg.Add(1) go uploadPart(&wg, part, partNum) partNum++ } if closing { return } } }() results := make([]uploadResult, 0, 64) outer: for { select { case readErr := <-readChan: if readErr != io.EOF { return 0, fmt.Errorf("reader error: %v", readErr) } case uploadResult := <-uploadResultChan: results = append(results, uploadResult) case uploadErr := <-uploadErrorChan: return 0, fmt.Errorf("error while uploading part: %v", uploadErr) case <-ctx.Done(): return 0, ctx.Err() case <-wgDone: break outer } } if len(results) == 0 { return 0, errors.New("no parts available to upload") } completedParts := make([]types.CompletedPart, 0, 64) var uploadedBytes int64 for _, result := range results { completedParts = append(completedParts, result.completedPart) uploadedBytes += result.size } // the parts may be out of order, especially with slow network conditions: sort.Slice(completedParts, func(i, j int) bool { return completedParts[i].PartNumber < completedParts[j].PartNumber }) completeInput := s3.CompleteMultipartUploadInput{ Bucket: aws.String(bucket), Key: aws.String(key), UploadId: output.UploadId, MultipartUpload: &types.CompletedMultipartUpload{Parts: completedParts}, } if _, err = u.s3.CompleteMultipartUpload(ctx, &completeInput); err != nil { return 0, fmt.Errorf("error completing upload: %v", err) } u.logger.With("key", key, "numParts", len(completedParts), "len", uploadedBytes).Debug("completed upload") uploaded = true return uploadedBytes, nil }