Transfer - Send chunks if they are bigger than GiB

jfrog · Sep 14, 2023 · 7a80247 · 7a80247
1 parent 3a430cf
commit 7a80247
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 3 deletions.
diff --git a/artifactory/commands/transferfiles/api/types.go b/artifactory/commands/transferfiles/api/types.go
@@ -23,6 +23,10 @@ const (
 	Phase1 int = 0
 	Phase2 int = 1
 	Phase3 int = 2
+
+	maxFilesInChunk = 16
+	// 1GiB
+	maxBytesInChunk = 1 << (30)
 )
 
 type TargetAuth struct {
@@ -101,3 +105,18 @@ func (uc *UploadChunk) AppendUploadCandidateIfNeeded(file FileRepresentation, bu
 	}
 	uc.UploadCandidates = append(uc.UploadCandidates, file)
 }
+
+// Return true if the chunk contains at least 16 files or at least 1GiB in total
+func (uc *UploadChunk) IsChunkFull() bool {
+	if len(uc.UploadCandidates) >= maxFilesInChunk {
+		return true
+	}
+	var totalSize int64 = 0
+	for _, uploadCandidate := range uc.UploadCandidates {
+		totalSize += uploadCandidate.Size
+		if totalSize > maxBytesInChunk {
+			return true
+		}
+	}
+	return false
+}
diff --git a/artifactory/commands/transferfiles/api/types_test.go b/artifactory/commands/transferfiles/api/types_test.go
@@ -0,0 +1,49 @@
+package api
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAppendUploadCandidateIfNeeded(t *testing.T) {
+	uploadChunk := &UploadChunk{}
+
+	// Regular file
+	uploadChunk.AppendUploadCandidateIfNeeded(FileRepresentation{Name: "regular-file"}, false)
+	assert.Len(t, uploadChunk.UploadCandidates, 1)
+
+	// Build info
+	uploadChunk.AppendUploadCandidateIfNeeded(FileRepresentation{Name: "build-info.json"}, true)
+	assert.Len(t, uploadChunk.UploadCandidates, 2)
+
+	// Directory in build info - should be skipped
+	uploadChunk.AppendUploadCandidateIfNeeded(FileRepresentation{}, true)
+	assert.Len(t, uploadChunk.UploadCandidates, 2)
+}
+
+var isChunkFullCases = []struct {
+	files  []FileRepresentation
+	isFull bool
+}{
+	{[]FileRepresentation{}, false},
+	{[]FileRepresentation{{Name: "fat-jim", Size: 10737418}}, false},
+	{[]FileRepresentation{{Name: "fat-vinny", Size: 1073741825}}, true},
+}
+
+func TestIsChunkFull(t *testing.T) {
+	for _, testCase := range isChunkFullCases {
+		t.Run("", func(t *testing.T) {
+			uploadChunk := &UploadChunk{UploadCandidates: testCase.files}
+			assert.Equal(t, testCase.isFull, uploadChunk.IsChunkFull())
+		})
+	}
+	t.Run("", func(t *testing.T) {
+		uploadChunk := &UploadChunk{}
+		for i := 0; i < 17; i++ {
+			uploadChunk.AppendUploadCandidateIfNeeded(FileRepresentation{Name: fmt.Sprintf("%d", i)}, false)
+		}
+		assert.True(t, uploadChunk.IsChunkFull())
+	})
+}
diff --git a/artifactory/commands/transferfiles/fulltransfer.go b/artifactory/commands/transferfiles/fulltransfer.go
@@ -239,7 +239,7 @@ func (m *fullTransferPhase) handleFoundFile(pcWrapper producerConsumerWrapper,
 		return
 	}
 	curUploadChunk.AppendUploadCandidateIfNeeded(file, m.buildInfoRepo)
-	if len(curUploadChunk.UploadCandidates) == uploadChunkSize {
+	if curUploadChunk.IsChunkFull() {
 		_, err = pcWrapper.chunkUploaderProducerConsumer.AddTaskWithError(uploadChunkWhenPossibleHandler(&m.phaseBase, *curUploadChunk, uploadChunkChan, errorsChannelMng), pcWrapper.errorsQueue.AddError)
 		if err != nil {
 			return

diff --git a/artifactory/commands/transferfiles/transfer.go b/artifactory/commands/transferfiles/transfer.go
@@ -28,7 +28,6 @@ import (
 )
 
 const (
-	uploadChunkSize = 16
 	// Size of the channel where the transfer's go routines write the transfer errors
 	fileWritersChannelSize       = 500000
 	retries                      = 600

diff --git a/artifactory/commands/transferfiles/utils.go b/artifactory/commands/transferfiles/utils.go
@@ -419,7 +419,7 @@ func uploadByChunks(files []api.FileRepresentation, uploadTokensChan chan Upload
 			continue
 		}
 		curUploadChunk.AppendUploadCandidateIfNeeded(file, base.buildInfoRepo)
-		if len(curUploadChunk.UploadCandidates) == uploadChunkSize {
+		if curUploadChunk.IsChunkFull() {
 			_, err = pcWrapper.chunkUploaderProducerConsumer.AddTaskWithError(uploadChunkWhenPossibleHandler(&base, curUploadChunk, uploadTokensChan, errorsChannelMng), pcWrapper.errorsQueue.AddError)
 			if err != nil {
 				return