Add support for zero chunk size values with tests

Author: jacobfederer

gcsfs/file.go

@@ -90,9 +90,10 @@ func NewGcsFileFromOldFH(
 	return res
 }
 
-// SetUploadChunkSizeByte The default chunk size for uploading files to GCS is 16MB.
-// This buffer size may use memory excessively. You can adjust it to a value that is a multiple of 256KiB.
-func (o *GcsFile) SetUploadChunkSizeByte(size int) {
+// SetUploadChunkSizeByte sets the chunk size for uploads.
+// To disable chunking, set this to 0.
+// For custom sizes, GCS requires multiples of 256KiB (262144 bytes).
+func (o *GcsFile) SetUploadChunkSizeByte(size *int) {
 	o.resource.uploadChunkSizeByte = size
 }
 

gcsfs/file_resource.go

@@ -46,7 +46,7 @@ type gcsFileResource struct {
 	name     string
 	fileMode os.FileMode
 
-	uploadChunkSizeByte int
+	uploadChunkSizeByte *int
 
 	currentGcsSize int64
 	offset         int64
@@ -188,13 +188,8 @@ func (o *gcsFileResource) WriteAt(b []byte, off int64) (n int, err error) {
 	// It will however require a download and upload of the original file but it
 	// can't be avoided if we should support seek-write-operations on GCS.
 
-	chunkSize := o.fs.UploadChunkSizeByte
-	if o.uploadChunkSizeByte != 0 {
-		chunkSize = o.uploadChunkSizeByte
-	}
-
-	if chunkSize != 0 {
-		w.SetChunkSize(chunkSize)
+	if size, ok := o.getEffectiveChunkSize(); ok {
+		w.SetChunkSize(size)
 	}
 
 	objAttrs, err := o.obj.Attrs(o.ctx)
@@ -235,6 +230,16 @@ func (o *gcsFileResource) WriteAt(b []byte, off int64) (n int, err error) {
 	return written, err
 }
 
+func (o *gcsFileResource) getEffectiveChunkSize() (int, bool) {
+	if o.uploadChunkSizeByte != nil {
+		return *o.uploadChunkSizeByte, true
+	}
+	if o.fs.UploadChunkSizeByte != nil {
+		return *o.fs.UploadChunkSizeByte, true
+	}
+	return 0, false
+}
+
 func min(x, y int) int {
 	if x < y {
 		return x

gcsfs/fs.go

@@ -43,7 +43,7 @@ type Fs struct {
 	rawGcsObjects map[string]*GcsFile
 
 	autoRemoveEmptyFolders bool // trigger for creating "virtual folders" (not required by GCSs)
-	UploadChunkSizeByte    int
+	UploadChunkSizeByte    *int
 }
 
 func NewGcsFs(ctx context.Context, client stiface.Client) *Fs {
@@ -63,7 +63,7 @@ func NewGcsFsWithSeparator(ctx context.Context, client stiface.Client, folderSep
 
 // SetUploadChunkSizeByte The default chunk size for uploading files to GCS is 16MB.
 // This buffer size may use memory excessively. You can adjust it to a value that is a multiple of 256KiB.
-func (fs *Fs) SetUploadChunkSizeByte(uploadChunkSize int) {
+func (fs *Fs) SetUploadChunkSizeByte(uploadChunkSize *int) {
 	fs.UploadChunkSizeByte = uploadChunkSize
 }
 

gcsfs/gcs.go

@@ -85,7 +85,10 @@ func (fs *GcsFs) Name() string {
 	return fs.source.Name()
 }
 
-func (fs *GcsFs) SetUploadChunkSizeByte(size int) {
+// SetUploadChunkSizeByte sets the chunk size for uploads.
+// To disable chunking, set this to 0.
+// For custom sizes, GCS requires multiples of 256KiB (262144 bytes).
+func (fs *GcsFs) SetUploadChunkSizeByte(size *int) {
 	fs.source.SetUploadChunkSizeByte(size)
 }
 

gcsfs/gcs_mocks.go

@@ -18,9 +18,9 @@ import (
 	"strings"
 
 	"cloud.google.com/go/storage"
+	"github.com/spf13/afero"
 	"google.golang.org/api/iterator"
 
-	"github.com/spf13/afero"
 	"github.com/spf13/afero/gcsfs/internal/stiface"
 )
 
@@ -31,15 +31,28 @@ func normSeparators(s string) string {
 
 type clientMock struct {
 	stiface.Client
-	fs afero.Fs
+	fs      afero.Fs
+	buckets map[string]*bucketMock
 }
 
 func newClientMock() *clientMock {
-	return &clientMock{fs: afero.NewMemMapFs()}
+	return &clientMock{
+		fs:      afero.NewMemMapFs(),
+		buckets: make(map[string]*bucketMock),
+	}
 }
 
 func (m *clientMock) Bucket(name string) stiface.BucketHandle {
-	return &bucketMock{bucketName: name, fs: m.fs}
+	if b, ok := m.buckets[name]; ok {
+		return b
+	}
+	b := &bucketMock{
+		bucketName: name,
+		fs:         m.fs,
+		objects:    make(map[string]*objectMock),
+	}
+	m.buckets[name] = b
+	return b
 }
 
 type bucketMock struct {
@@ -48,14 +61,24 @@ type bucketMock struct {
 	bucketName string
 
 	fs afero.Fs
+
+	objects map[string]*objectMock
 }
 
 func (m *bucketMock) Attrs(context.Context) (*storage.BucketAttrs, error) {
 	return &storage.BucketAttrs{}, nil
 }
 
 func (m *bucketMock) Object(name string) stiface.ObjectHandle {
-	return &objectMock{name: name, fs: m.fs}
+	if m.objects == nil {
+		m.objects = make(map[string]*objectMock)
+	}
+	if obj, ok := m.objects[name]; ok {
+		return obj
+	}
+	obj := &objectMock{name: name, fs: m.fs, lastChunkSize: -1}
+	m.objects[name] = obj
+	return obj
 }
 
 func (m *bucketMock) Objects(_ context.Context, q *storage.Query) (it stiface.ObjectIterator) {
@@ -67,10 +90,12 @@ type objectMock struct {
 
 	name string
 	fs   afero.Fs
+
+	lastChunkSize int // Track this for testing
 }
 
 func (o *objectMock) NewWriter(_ context.Context) stiface.Writer {
-	return &writerMock{name: o.name, fs: o.fs}
+	return &writerMock{name: o.name, fs: o.fs, parent: o}
 }
 
 func (o *objectMock) NewRangeReader(
@@ -146,12 +171,19 @@ func (o *objectMock) Attrs(_ context.Context) (*storage.ObjectAttrs, error) {
 type writerMock struct {
 	stiface.Writer
 
-	name string
-	fs   afero.Fs
+	name   string
+	fs     afero.Fs
+	parent *objectMock
 
 	file afero.File
 }
 
+func (w *writerMock) SetChunkSize(s int) {
+	if w.parent != nil {
+		w.parent.lastChunkSize = s
+	}
+}
+
 func (w *writerMock) Write(p []byte) (n int, err error) {
 	if w.name == "" {
 		return 0, ErrEmptyObjectName

gcsfs/gcs_test.go

@@ -861,3 +861,57 @@ func TestGcsRemoveAll(t *testing.T) {
 		}
 	})
 }
+
+func TestUploadChunkSize(t *testing.T) {
+	tests := []struct {
+		testName                string
+		chunkSizeAssignedVal    *int
+		expectedResultChunkSize int
+	}{
+		{
+			testName: "1MB setting",
+
+			chunkSizeAssignedVal: func() *int { val := 1 * 1024 * 1024; return &val }(),
+
+			expectedResultChunkSize: 1 * 1024 * 1024,
+		},
+		{
+			testName: "0 (disabled chunking)",
+
+			chunkSizeAssignedVal: func() *int { val := 0; return &val }(),
+
+			expectedResultChunkSize: 0,
+		},
+		{
+			testName: "nil (GCS default)",
+
+			chunkSizeAssignedVal:    nil,
+			expectedResultChunkSize: -1, // Sentinel value for "never called"
+		},
+	}
+
+	for i, tt := range tests {
+		t.Run(tt.testName, func(t *testing.T) {
+			// Create a fresh isolated mock for this sub-test
+			ctx := context.Background()
+			mockClient := newClientMock()
+			fs := &GcsFs{NewGcsFs(ctx, mockClient)}
+
+			fs.SetUploadChunkSizeByte(tt.chunkSizeAssignedVal)
+
+			fileName := fmt.Sprintf("test-chunk-%d.txt", i)
+			fullObjectName := filepath.Join(bucketName, fileName)
+
+			f, _ := fs.Create(fullObjectName)
+
+			// Trigger writer
+			f.Write([]byte("data"))
+
+			obj := mockClient.Bucket(bucketName).Object(fileName).(*objectMock)
+
+			if obj.lastChunkSize != tt.expectedResultChunkSize {
+				t.Errorf("Expected %d, got %d", tt.expectedResultChunkSize, obj.lastChunkSize)
+			}
+		})
+	}
+}