feat: optimize aws s3_v2 Write() code (29d3c47f) · Commits · GitLab.org / container-registry

registry/storage/driver/s3-aws/v2/s3.go

+60 −54

Original line number	Diff line number	Diff line
		@@ -1239,6 +1239,8 @@ type writer struct {
		closed bool
		committed bool
		canceled bool
		chunkSize int64
		multipartCopyMaxConcurrency int
		}

		func (d driver) newWriter(key, uploadID string, parts []s3.Part) storagedriver.FileWriter {
		@@ -1253,6 +1255,8 @@ func (d driver) newWriter(key, uploadID string, parts []s3.Part) storagedriver
		parts: parts,
		size: size,
		buffer: new(bytes.Buffer),
		chunkSize: d.ChunkSize,
		multipartCopyMaxConcurrency: d.MultipartCopyMaxConcurrency,
		}
		}

		@@ -1274,9 +1278,9 @@ func (w *writer) Write(p []byte) (int, error) {
		return 0, storagedriver.ErrAlreadyCanceled
		}

		// If the last written part is smaller than minChunkSize, we need to make a
		// new multipart upload
		if len(w.parts) > 0 && *w.parts[len(w.parts)-1].Size < common.MinChunkSize {
		// If the length of the last written part is different than chunkSize,
		// we need to make a new multipart upload to even things out.
		if len(w.parts) > 0 && *w.parts[len(w.parts)-1].Size != w.chunkSize {
		var completedUploadedParts completedParts
		for _, part := range w.parts {
		completedUploadedParts = append(completedUploadedParts, &s3.CompletedPart{
		@@ -1326,40 +1330,53 @@ func (w *writer) Write(p []byte) (int, error) {
		}
		w.uploadID = *resp.UploadId

		// If the entire written file is smaller than minChunkSize, we need to make
		// a new part from scratch
		if w.size < common.MinChunkSize {
		partCount := (w.size + w.chunkSize - 1) / w.chunkSize
		w.parts = make([]*s3.Part, 0, partCount)
		partsMutex := new(sync.Mutex)

		g, gctx := errgroup.WithContext(ctx)

		// Reduce the client/server exposure to long lived connections regardless of
		// how many requests per second are allowed.
		g.SetLimit(w.multipartCopyMaxConcurrency)

		for i := int64(0); i < partCount; i++ {
		g.Go(func() error {
		// Check if any other goroutine has failed
		select {
		case <-gctx.Done():
		return gctx.Err()
		default:
		}

		startByte := w.chunkSize * i
		endByte := startByte + w.chunkSize - 1
		if endByte > w.size-1 {
		// NOTE(prozlach): Special case when there is simply not
		// enough data for a full chunk. It handles both cases when
		// there is only one chunk and multiple chunks plus partial
		// a one. We just slurp in what we have and carry on with
		// the data passed to the Write() call.
		byteRange := fmt.Sprintf("bytes=%d-%d", startByte, endByte)
		resp, err := w.driver.S3.GetObjectWithContext(
		ctx,
		gctx,
		&s3.GetObjectInput{
		Bucket: aws.String(w.driver.Bucket),
		Key: aws.String(w.key),
		Range: aws.String(byteRange),
		})
		if err != nil {
		return 0, err
		return fmt.Errorf("fetching object from backend during re-upload: %w", err)
		}
		defer resp.Body.Close()
		w.parts = nil

		_, err = w.buffer.ReadFrom(resp.Body)
		if err != nil {
		return 0, err
		return fmt.Errorf("reading remaining bytes during data re-upload: %w", err)
		}
		} else {
		// nolint: revive // max-control-nesting
		if w.size > common.MaxChunkSize {
		// Calculate how many parts we need
		partCount := w.size/common.MaxChunkSize + 1
		// NOTE(prozlach): We need to prevent a situatin when the last
		// part is maller than the minimum part size, so we even things
		// out:
		partSize := w.size / partCount
		w.parts = make([]*s3.Part, partCount)

		startByte := int64(0)
		endByte := partSize

		// Copy each part individually
		for i := int64(0); i < partCount; i++ {
		return nil
		}

		// Part numbers are positive integers in the range 1 <= n <= 10000
		partNumber := i + 1

		@@ -1368,7 +1385,7 @@ func (w *writer) Write(p []byte) (int, error) {
		byteRange := fmt.Sprintf("bytes=%d-%d", startByte, endByte)

		copyPartResp, err := w.driver.S3.UploadPartCopyWithContext(
		ctx,
		gctx,
		&s3.UploadPartCopyInput{
		Bucket: aws.String(w.driver.Bucket),
		CopySource: aws.String(w.driver.Bucket + "/" + w.key),
		@@ -1378,7 +1395,17 @@ func (w *writer) Write(p []byte) (int, error) {
		UploadId: resp.UploadId,
		})
		if err != nil {
		return 0, fmt.Errorf("re-uploading chunk as multipart upload part: %w", err)
		return fmt.Errorf("re-uploading chunk as multipart upload part: %w", err)
		}

		// NOTE(prozlach): We can't pre-allocate parts slice and then
		// address it with `i` if we want to handle the case where
		// there is not enough data for a single chunk. At the expense
		// of this small extra loop and extra mutex we get the same
		// path for both cases.
		partsMutex.Lock()
		for int64(len(w.parts))-1 < i {
		w.parts = append(w.parts, (*s3.Part)(nil))
		}

		// Add this part to our list
		@@ -1389,35 +1416,14 @@ func (w *writer) Write(p []byte) (int, error) {
		// ending byte, hence `+ 1`.
		Size: aws.Int64(endByte - startByte + 1),
		}
		partsMutex.Unlock()

		startByte = endByte + 1
		endByte += startByte + partSize
		if endByte > w.size-1 {
		endByte = w.size - 1
		}
		}
		} else {
		// For objects smaller than maxPartSize, use the original code
		copyPartResp, err := w.driver.S3.UploadPartCopyWithContext(
		ctx,
		&s3.UploadPartCopyInput{
		Bucket: aws.String(w.driver.Bucket),
		CopySource: aws.String(w.driver.Bucket + "/" + w.key),
		Key: aws.String(w.key),
		PartNumber: aws.Int64(1),
		UploadId: resp.UploadId,
		return nil
		})
		if err != nil {
		return 0, fmt.Errorf("re-uploading object as first part of the new multipart upload: %w", err)
		}
		w.parts = []*s3.Part{
		{
		ETag: copyPartResp.CopyPartResult.ETag,
		PartNumber: aws.Int64(1),
		Size: aws.Int64(w.size),
		},
		}
		}

		if err := g.Wait(); err != nil {
		return 0, err
		}
		}