Repository: bufbuild/buf
Pull request #1390 (2022-09-08 10:53:25)

Add `LimitWriteBucket`
Adds a `LimitWriteBucket` that limits the total size of bytes that can be written to the bucket.


Patch:
diff --git a/private/pkg/storage/errors.go b/private/pkg/storage/errors.go
index fffe0b48c..566ad20bf 100644
--- a/private/pkg/storage/errors.go
+++ b/private/pkg/storage/errors.go
@@ -16,6 +16,7 @@ package storage
 
 import (
 	"errors"
+	"fmt"
 	"strings"
 
 	"github.com/bufbuild/buf/private/pkg/normalpath"
@@ -58,6 +59,15 @@ func IsExistsMultipleLocations(err error) bool {
 	return errors.As(err, &asErr)
 }
 
+// IsWriteLimitReached returns true if the error is of writes exceeding the limit of the bucket.
+func IsWriteLimitReached(err error) bool {
+	if err == nil {
+		return false
+	}
+	asErr := &errWriteLimitReached{}
+	return errors.As(err, &asErr)
+}
+
 // errorExistsMultipleLocations is the error returned if a path exists in multiple locations.
 type errorExistsMultipleLocations struct {
 	Path          string
@@ -68,3 +78,16 @@ type errorExistsMultipleLocations struct {
 func (e *errorExistsMultipleLocations) Error() string {
 	return e.Path + " exists in multiple locations: " + strings.Join(e.ExternalPaths, " ")
 }
+
+// errWriteLimitReached is the error returned if the write limit is reached.
+//
+// See [LimitWriteBucket].
+type errWriteLimitReached struct {
+	Limit       int64
+	ExceedingBy int64
+}
+
+// Error implements error.
+func (e *errWriteLimitReached) Error() string {
+	return fmt.Sprintf("write limit reached: limit: %d, exceeding by: %d", e.Limit, e.ExceedingBy)
+}
diff --git a/private/pkg/storage/limit.go b/private/pkg/storage/limit.go
new file mode 100644
index 000000000..0c2073140
--- /dev/null
+++ b/private/pkg/storage/limit.go
@@ -0,0 +1,92 @@
+// Copyright 2020-2022 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package storage
+
+import (
+	"context"
+
+	"go.uber.org/atomic"
+)
+
+// LimitWriteBucket returns a [WriteBucket] that writes to [writeBucket]
+// but stops with an error after [limit] bytes are written.
+//
+// The error can be checked using [IsWriteLimitReached].
+//
+// A negative [limit] is same as 0 limit.
+func LimitWriteBucket(writeBucket WriteBucket, limit int) WriteBucket {
+	if limit < 0 {
+		limit = 0
+	}
+	return newLimitedWriteBucket(writeBucket, int64(limit))
+}
+
+type limitedWriteBucket struct {
+	WriteBucket
+	currentSize *atomic.Int64
+	limit       int64
+}
+
+func newLimitedWriteBucket(bucket WriteBucket, limit int64) *limitedWriteBucket {
+	return &limitedWriteBucket{
+		WriteBucket: bucket,
+		currentSize: atomic.NewInt64(0),
+		limit:       limit,
+	}
+}
+
+func (w *limitedWriteBucket) Put(ctx context.Context, path string) (WriteObjectCloser, error) {
+	writeObjectCloser, err := w.WriteBucket.Put(ctx, path)
+	if err != nil {
+		return nil, err
+	}
+	return newLimitedWriteObjectCloser(writeObjectCloser, w.currentSize, w.limit), nil
+}
+
+type limitedWriteObjectCloser struct {
+	WriteObjectCloser
+
+	bucketSize *atomic.Int64
+	limit      int64
+}
+
+func newLimitedWriteObjectCloser(
+	writeObjectCloser WriteObjectCloser,
+	bucketSize *atomic.Int64,
+	limit int64,
+) *limitedWriteObjectCloser {
+	return &limitedWriteObjectCloser{
+		WriteObjectCloser: writeObjectCloser,
+		bucketSize:        bucketSize,
+		limit:             limit,
+	}
+}
+
+func (o *limitedWriteObjectCloser) Write(p []byte) (int, error) {
+	writeSize := int64(len(p))
+	newBucketSize := o.bucketSize.Add(writeSize)
+	if newBucketSize > o.limit {
+		o.bucketSize.Sub(writeSize)
+		return 0, &errWriteLimitReached{
+			Limit:       o.limit,
+			ExceedingBy: newBucketSize - o.limit,
+		}
+	}
+	writtenSize, err := o.WriteObjectCloser.Write(p)
+	if int64(writtenSize) < writeSize {
+		o.bucketSize.Sub(writeSize - int64(writtenSize))
+	}
+	return writtenSize, err
+}


Test patch:
diff --git a/private/pkg/storage/storagetesting/storagetesting.go b/private/pkg/storage/storagetesting/storagetesting.go
index e31579099..927ca3804 100644
--- a/private/pkg/storage/storagetesting/storagetesting.go
+++ b/private/pkg/storage/storagetesting/storagetesting.go
@@ -24,6 +24,8 @@ import (
 	"path/filepath"
 	"runtime"
 	"sort"
+	"strconv"
+	"sync"
 	"testing"
 
 	"github.com/bufbuild/buf/private/pkg/command"
@@ -35,6 +37,7 @@ import (
 	"github.com/bufbuild/buf/private/pkg/tmp"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
 )
 
 const (
@@ -1496,4 +1499,37 @@ func RunTestSuite(
 			"1.proto",
 		)
 	})
+	t.Run("limit-write-bucket", func(t *testing.T) {
+		t.Parallel()
+		writeBucket := newWriteBucket(t, defaultProvider)
+		readBucket := writeBucketToReadBucket(t, writeBucket)
+		const limit = 2048
+		limitedWriteBucket := storage.LimitWriteBucket(writeBucket, limit)
+		var (
+			wg           sync.WaitGroup
+			writtenBytes atomic.Int64
+			triedBytes   atomic.Int64
+		)
+		for i := 0; i < 10; i++ {
+			wg.Add(1)
+			go func(i int) {
+				defer wg.Done()
+				data := bytes.Repeat([]byte("b"), i*100)
+				path := strconv.Itoa(i)
+				triedBytes.Add(int64(len(data)))
+				err := storage.PutPath(context.Background(), limitedWriteBucket, path, data)
+				if err != nil {
+					assert.True(t, storage.IsWriteLimitReached(err))
+					return
+				}
+				readData, err := storage.ReadPath(context.Background(), readBucket, path)
+				assert.NoError(t, err)
+				assert.Equal(t, readData, data)
+				writtenBytes.Add(int64(len(data)))
+			}(i)
+		}
+		wg.Wait()
+		require.Greater(t, triedBytes.Load(), int64(limit))
+		assert.LessOrEqual(t, writtenBytes.Load(), int64(limit))
+	})
 }