fix retry on markers (#5441)

* fix retry on markers

Signed-off-by: Alan Protasio <[email protected]>

* lint + changelog

Signed-off-by: Alan Protasio <[email protected]>

---------

Signed-off-by: Alan Protasio <[email protected]>

Author: alanprot

Diff for: pkg/storage/bucket/s3/bucket_client.go

@@ -59,6 +59,17 @@ func NewBucketReaderClient(cfg Config, name string, logger log.Logger) (objstore
 	}, nil
 }
 
+// NewBucketWithRetries wraps the original bucket into the BucketWithRetries
+func NewBucketWithRetries(bucket objstore.Bucket, operationRetries int, retryMinBackoff time.Duration, retryMaxBackoff time.Duration, logger log.Logger) (objstore.Bucket, error) {
+	return &BucketWithRetries{
+		logger:           logger,
+		bucket:           bucket,
+		operationRetries: operationRetries,
+		retryMinBackoff:  retryMinBackoff,
+		retryMaxBackoff:  retryMaxBackoff,
+	}, nil
+}
+
 func newS3Config(cfg Config) (s3.Config, error) {
 	sseCfg, err := cfg.SSE.BuildThanosConfig()
 	if err != nil {
@@ -166,7 +177,11 @@ func (b *BucketWithRetries) Upload(ctx context.Context, name string, r io.Reader
 	if !ok {
 		// Skip retry if incoming Reader is not seekable to avoid
 		// loading entire content into memory
-		return b.bucket.Upload(ctx, name, r)
+		err := b.bucket.Upload(ctx, name, r)
+		if err != nil {
+			level.Warn(b.logger).Log("msg", "skip upload retry as reader is not seekable", "file", name, "err", err)
+		}
+		return err
 	}
 	return b.retry(ctx, func() error {
 		if _, err := rs.Seek(0, io.SeekStart); err != nil {

Diff for: pkg/storage/bucket/s3/bucket_client_test.go

@@ -60,21 +60,46 @@ func TestBucketWithRetries_ShouldRetry(t *testing.T) {
 func TestBucketWithRetries_UploadSeekable(t *testing.T) {
 	t.Parallel()
 
-	m := mockBucket{
-		FailCount: 3,
-	}
-	b := BucketWithRetries{
-		logger:           log.NewNopLogger(),
-		bucket:           &m,
-		operationRetries: 5,
-		retryMinBackoff:  10 * time.Millisecond,
-		retryMaxBackoff:  time.Second,
+	cases := map[string]struct {
+		readerFactory func(i string) io.Reader
+		input         string
+		err           error
+	}{
+		"should retry when seekable": {
+			err:   nil,
+			input: "test input",
+			readerFactory: func(i string) io.Reader {
+				return bytes.NewReader([]byte(i))
+			},
+		},
+
+		"should not retry when seekable": {
+			err:   fmt.Errorf("failed upload: 2"),
+			input: "test input",
+			readerFactory: func(i string) io.Reader {
+				return bytes.NewBuffer([]byte(i))
+			},
+		},
 	}
 
-	input := []byte("test input")
-	err := b.Upload(context.Background(), "dummy", bytes.NewReader(input))
-	require.NoError(t, err)
-	require.Equal(t, input, m.uploadedContent)
+	for name, tc := range cases {
+		t.Run(name, func(*testing.T) {
+			m := mockBucket{
+				FailCount: 3,
+			}
+			b := BucketWithRetries{
+				logger:           log.NewNopLogger(),
+				bucket:           &m,
+				operationRetries: 5,
+				retryMinBackoff:  10 * time.Millisecond,
+				retryMaxBackoff:  time.Second,
+			}
+
+			err := b.Upload(context.Background(), "dummy", tc.readerFactory(tc.input))
+			require.Equal(t, tc.err, err)
+			require.Equal(t, tc.input, string(m.uploadedContent))
+		})
+	}
 }
 
 func TestBucketWithRetries_UploadNonSeekable(t *testing.T) {

Diff for: pkg/storage/tsdb/bucketindex/markers_bucket_client.go

@@ -38,12 +38,12 @@ func (b *globalMarkersBucket) Upload(ctx context.Context, name string, r io.Read
 	}
 
 	// Upload it to the global marker's location.
-	if err := b.parent.Upload(ctx, globalMarkPath, bytes.NewBuffer(body)); err != nil {
+	if err := b.parent.Upload(ctx, globalMarkPath, bytes.NewReader(body)); err != nil {
 		return err
 	}
 
 	// Upload it to the original location too.
-	return b.parent.Upload(ctx, name, bytes.NewBuffer(body))
+	return b.parent.Upload(ctx, name, bytes.NewReader(body))
 }
 
 // Delete implements objstore.Bucket.

Diff for: pkg/storage/tsdb/bucketindex/markers_bucket_client_test.go

@@ -3,9 +3,11 @@ package bucketindex
 import (
 	"bytes"
 	"context"
+	"errors"
 	"strings"
 	"testing"
 
+	"github.com/go-kit/log"
 	"github.com/oklog/ulid"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/testutil"
@@ -14,6 +16,8 @@ import (
 	"github.com/thanos-io/objstore"
 	"github.com/thanos-io/thanos/pkg/block/metadata"
 
+	"github.com/cortexproject/cortex/pkg/storage/bucket/s3"
+
 	"github.com/cortexproject/cortex/pkg/storage/bucket"
 	cortex_testutil "github.com/cortexproject/cortex/pkg/storage/tsdb/testutil"
 )
@@ -161,6 +165,48 @@ func TestGlobalMarkersBucket_isMark(t *testing.T) {
 	}
 }
 
+func TestBucketWithGlobalMarkers_ShouldRetryUpload(t *testing.T) {
+	ctx := context.Background()
+	block1 := ulid.MustNew(1, nil)
+
+	bkt, _ := cortex_testutil.PrepareFilesystemBucket(t)
+
+	// Fail the global markers and the non-global marker
+	prefixErrors := []string{"", "marker"}
+
+	tests := []struct {
+		mark       string
+		globalpath string
+	}{
+		{
+			mark:       metadata.DeletionMarkFilename,
+			globalpath: "markers/" + block1.String() + "-deletion-mark.json",
+		},
+		{
+			mark:       metadata.NoCompactMarkFilename,
+			globalpath: "markers/" + block1.String() + "-no-compact-mark.json",
+		},
+	}
+
+	for _, tc := range tests {
+		for _, p := range prefixErrors {
+			t.Run(tc.mark+"/"+p, func(t *testing.T) {
+				mBucket := &cortex_testutil.MockBucketFailure{
+					Bucket:         bkt,
+					UploadFailures: map[string]error{p: errors.New("test")},
+				}
+				bkt, _ = s3.NewBucketWithRetries(mBucket, 5, 0, 0, log.NewNopLogger())
+				bkt = BucketWithGlobalMarkers(bkt)
+				originalPath := block1.String() + "/" + tc.mark
+				err := bkt.Upload(ctx, originalPath, strings.NewReader("{}"))
+				require.Equal(t, errors.New("test"), err)
+				require.Equal(t, mBucket.UploadCalls.Load(), int32(5))
+			})
+		}
+
+	}
+}
+
 func TestBucketWithGlobalMarkers_ShouldWorkCorrectlyWithBucketMetrics(t *testing.T) {
 	reg := prometheus.NewPedanticRegistry()
 	ctx := context.Background()

Diff for: pkg/storage/tsdb/testutil/objstore.go

@@ -10,6 +10,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/require"
 	"github.com/thanos-io/objstore"
+	"go.uber.org/atomic"
 
 	"github.com/cortexproject/cortex/pkg/util"
 
@@ -37,6 +38,10 @@ type MockBucketFailure struct {
 
 	DeleteFailures []string
 	GetFailures    map[string]error
+	UploadFailures map[string]error
+
+	UploadCalls atomic.Int32
+	GetCalls    atomic.Int32
 }
 
 func (m *MockBucketFailure) Delete(ctx context.Context, name string) error {
@@ -47,6 +52,7 @@ func (m *MockBucketFailure) Delete(ctx context.Context, name string) error {
 }
 
 func (m *MockBucketFailure) Get(ctx context.Context, name string) (io.ReadCloser, error) {
+	m.GetCalls.Add(1)
 	for prefix, err := range m.GetFailures {
 		if strings.HasPrefix(name, prefix) {
 			return nil, err
@@ -59,6 +65,20 @@ func (m *MockBucketFailure) Get(ctx context.Context, name string) (io.ReadCloser
 	return m.Bucket.Get(ctx, name)
 }
 
+func (m *MockBucketFailure) Upload(ctx context.Context, name string, r io.Reader) error {
+	m.UploadCalls.Add(1)
+	for prefix, err := range m.UploadFailures {
+		if strings.HasPrefix(name, prefix) {
+			return err
+		}
+	}
+	if e, ok := m.GetFailures[name]; ok {
+		return e
+	}
+
+	return m.Bucket.Upload(ctx, name, r)
+}
+
 func (m *MockBucketFailure) WithExpectedErrs(expectedFunc objstore.IsOpFailureExpectedFunc) objstore.Bucket {
 	if ibkt, ok := m.Bucket.(objstore.InstrumentedBucket); ok {
 		return &MockBucketFailure{Bucket: ibkt.WithExpectedErrs(expectedFunc), DeleteFailures: m.DeleteFailures, GetFailures: m.GetFailures}