Metrics: Fix objstore_bucket_operation_transferred_bytes never reco…

CHANGELOG.md

@@ -17,6 +17,7 @@ We use *breaking :warning:* to mark changes that are not backward compatible (re
 
 ### Fixed
 - [#287](https://github.com/thanos-io/objstore/pull/287) S3: Fix goroutine leak in `Iter` and `IterWithAttributes` when iteration stops early.
+- [#283](https://github.com/thanos-io/objstore/pull/283) Metrics: Fix `objstore_bucket_operation_transferred_bytes` never recorded when the underlying reader is read via `io.ReaderAt` (e.g. concurrent multipart uploads).
 - [#196](https://github.com/thanos-io/objstore/pull/196) GCS: fix error check in Exists method when object does not exist.
 - [#153](https://github.com/thanos-io/objstore/pull/153) Metrics: Fix `objstore_bucket_operation_duration_seconds_*` for `get` and `get_range` operations.
 - [#141](https://github.com/thanos-io/objstore/pull/142) S3: Fix missing encryption configuration for `Bucket.Exists()` and `Bucket.Attributes()` calls.

objstore.go

@@ -903,7 +903,7 @@ func (b *metricBucket) Upload(ctx context.Context, name string, r io.Reader, opt
 	defer trc.Close()
 	err := b.bkt.Upload(ctx, name, trc, opts...)
 	if err != nil {
-		if !b.metrics.isOpFailureExpected(err) && ctx.Err() != context.Canceled {
+		if !b.metrics.isOpFailureExpected(err) && !errors.Is(ctx.Err(), context.Canceled) {
 			b.metrics.opsFailures.WithLabelValues(op).Inc()
 		}
 		return err
@@ -959,16 +959,67 @@ type timingReader struct {
 
 	alreadyGotErr bool
 
-	start             time.Time
-	op                string
-	readBytes         int64
+	start         time.Time
+	op            string
+	readBytes     int64
+	readAtTracker *readAtTracker
+
 	duration          *prometheus.HistogramVec
 	failed            *prometheus.CounterVec
 	isFailureExpected IsOpFailureExpectedFunc
 	fetchedBytes      *prometheus.CounterVec
 	transferredBytes  *prometheus.HistogramVec
 }
 
+type readAtTracker struct {
+	offsetResults     map[int64]readAtOffsetResult
+	mtx               sync.Mutex
+	isExpectedFailure IsOpFailureExpectedFunc
+}
+
+func newReadAtTracker(isExpectedFailure IsOpFailureExpectedFunc) *readAtTracker {
+	return &readAtTracker{
+		offsetResults:     make(map[int64]readAtOffsetResult),
+		isExpectedFailure: isExpectedFailure,
+	}
+}
+
+func (t *readAtTracker) updateMetrics(off int64, n int, err error) {
+	if t == nil {
+		return
+	}
+
+	t.mtx.Lock()
+	defer t.mtx.Unlock()
+	t.offsetResults[off] = readAtOffsetResult{n: n, err: err}
+}
+
+func (t *readAtTracker) close() (readBytes int64, anyErr bool, anyRealFailure bool) {
+	if t == nil {
+		return 0, false, false
+	}
+
+	t.mtx.Lock()
+	defer t.mtx.Unlock()
+
+	for _, res := range t.offsetResults {
+		readBytes += int64(res.n)
+		if err := res.err; err != nil && err != io.EOF {
+			anyErr = true
+			if !t.isExpectedFailure(err) && !errors.Is(err, context.Canceled) {
+				anyRealFailure = true
+			}
+		}
+	}
+
+	return readBytes, anyErr, anyRealFailure
+}
+
+type readAtOffsetResult struct {
+	n   int
+	err error
+}
+
 func newTimingReader(start time.Time, r io.Reader, closeReader bool, op string, dur *prometheus.HistogramVec, failed *prometheus.CounterVec, isFailureExpected IsOpFailureExpectedFunc, fetchedBytes *prometheus.CounterVec, transferredBytes *prometheus.HistogramVec) io.ReadCloser {
 	// Initialize the metrics with 0.
 	dur.WithLabelValues(op)
@@ -993,6 +1044,8 @@ func newTimingReader(start time.Time, r io.Reader, closeReader bool, op string,
 	_, isSeeker := r.(io.Seeker)
 	_, isReaderAt := r.(io.ReaderAt)
 	if isSeeker && isReaderAt {
+		trc.readAtTracker = newReadAtTracker(isFailureExpected)
+
 		// The assumption is that in most cases when io.ReaderAt() is implemented then
 		// io.Seeker is implemented too (e.g. os.File).
 		return &timingReaderSeekerReaderAt{timingReaderSeeker: timingReaderSeeker{timingReader: trc}}
@@ -1024,15 +1077,26 @@ func (r *timingReader) Close() error {
 		}
 	}
 
-	// Track duration and transferred bytes only if no error occurred.
-	if !r.alreadyGotErr {
-		r.duration.WithLabelValues(r.op).Observe(time.Since(r.start).Seconds())
-		r.transferredBytes.WithLabelValues(r.op).Observe(float64(r.readBytes))
+	if r.alreadyGotErr {
+		return closeErr
+	}
 
-		// Trick to tracking metrics multiple times in case Close() gets called again.
+	// Track duration and transferred bytes only if no error occurred.
+	readAtBytes, readAtAnyErr, readAtRealFailure := r.readAtTracker.close()
+	if readAtAnyErr {
+		if readAtRealFailure {
+			r.failed.WithLabelValues(r.op).Inc()
+		}
 		r.alreadyGotErr = true
+		return closeErr
 	}
 
+	r.duration.WithLabelValues(r.op).Observe(time.Since(r.start).Seconds())
+	r.transferredBytes.WithLabelValues(r.op).Observe(float64(r.readBytes + readAtBytes))
+
+	// Trick to tracking metrics multiple times in case Close() gets called again.
+	r.alreadyGotErr = true
+
 	return closeErr
 }
 
@@ -1070,7 +1134,12 @@ type timingReaderSeekerReaderAt struct {
 }
 
 func (rsc *timingReaderSeekerReaderAt) ReadAt(p []byte, off int64) (int, error) {
-	return (rsc.Reader).(io.ReaderAt).ReadAt(p, off)
+	n, err := (rsc.Reader).(io.ReaderAt).ReadAt(p, off)
+	rsc.readAtTracker.updateMetrics(off, n, err)
+	if rsc.fetchedBytes != nil {
+		rsc.fetchedBytes.WithLabelValues(rsc.op).Add(float64(n))
+	}
+	return n, err
 }
 
 type timingReaderWriterTo struct {

objstore_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	promtest "github.com/prometheus/client_golang/prometheus/testutil"
+	dto "github.com/prometheus/client_model/go"
 	"go.uber.org/atomic"
 )
 
@@ -491,6 +492,168 @@ func (r dummyReader) Read(_ []byte) (int, error) {
 	return 0, r.err
 }
 
+// flakyReaderAt fails the first ReadAt() call at a given offset and succeeds afterward,
+// simulating a transient error that gets retried by a caller (e.g. an SDK's multipart uploader).
+type flakyReaderAt struct {
+	*bytes.Reader
+	failedOnce bool
+}
+
+func (f *flakyReaderAt) ReadAt(p []byte, off int64) (int, error) {
+	if !f.failedOnce {
+		f.failedOnce = true
+		return 0, errors.New("transient error")
+	}
+	return f.Reader.ReadAt(p, off)
+}
+
+// alwaysFailReaderAt fails every ReadAt() call, simulating a permanent, never-retried error.
+type alwaysFailReaderAt struct {
+	*bytes.Reader
+	err error
+}
+
+func (f *alwaysFailReaderAt) ReadAt(_ []byte, _ int64) (int, error) {
+	return 0, f.err
+}
+
+func histogramSumAndCount(t testing.TB, obs prometheus.Observer) (sum float64, count uint64) {
+	var dm dto.Metric
+	testutil.Ok(t, obs.(prometheus.Histogram).Write(&dm))
+	return dm.GetHistogram().GetSampleSum(), dm.GetHistogram().GetSampleCount()
+}
+
+func TestTimingReader_ReadAtUnexpectedErrorRecordsFailure(t *testing.T) {
+	data := []byte("hello world")
+	fr := &alwaysFailReaderAt{Reader: bytes.NewReader(data), err: errors.New("permanent disk error")}
+
+	m := WrapWithMetrics(NewInMemBucket(), nil, "")
+	tr := newTimingReader(time.Now(), fr, true, OpUpload, m.metrics.opsDuration, m.metrics.opsFailures, func(err error) bool {
+		return false
+	}, nil, m.metrics.opsTransferredBytes)
+
+	ra, ok := tr.(io.ReaderAt)
+	testutil.Assert(t, ok)
+
+	_, err := ra.ReadAt(make([]byte, len(data)), 0)
+	testutil.NotOk(t, err)
+
+	testutil.Ok(t, tr.Close())
+
+	// A permanent ReadAt error that is never retried successfully must be counted as a failure,
+	// and must not be reported as a successful transfer.
+	testutil.Equals(t, float64(1), promtest.ToFloat64(m.metrics.opsFailures.WithLabelValues(OpUpload)))
+	_, count := histogramSumAndCount(t, m.metrics.opsTransferredBytes.WithLabelValues(OpUpload))
+	testutil.Equals(t, uint64(0), count)
+}
+
+func TestTimingReader_ReadAtExpectedErrorNotCounted(t *testing.T) {
+	readerErr := errors.New("expected sentinel error")
+
+	for name, tc := range map[string]struct {
+		err               error
+		isFailureExpected IsOpFailureExpectedFunc
+	}{
+		"error explicitly marked as expected": {
+			err:               readerErr,
+			isFailureExpected: func(err error) bool { return errors.Is(err, readerErr) },
+		},
+		"io.EOF": {
+			err:               io.EOF,
+			isFailureExpected: func(err error) bool { return false },
+		},
+		"context.Canceled": {
+			err:               context.Canceled,
+			isFailureExpected: func(err error) bool { return false },
+		},
+	} {
+		t.Run(name, func(t *testing.T) {
+			data := []byte("hello world")
+			fr := &alwaysFailReaderAt{Reader: bytes.NewReader(data), err: tc.err}
+
+			m := WrapWithMetrics(NewInMemBucket(), nil, "")
+			tr := newTimingReader(time.Now(), fr, true, OpUpload, m.metrics.opsDuration, m.metrics.opsFailures, tc.isFailureExpected, nil, m.metrics.opsTransferredBytes)
+
+			ra, ok := tr.(io.ReaderAt)
+			testutil.Assert(t, ok)
+
+			_, err := ra.ReadAt(make([]byte, len(data)), 0)
+			testutil.Equals(t, tc.err, err)
+
+			testutil.Ok(t, tr.Close())
+
+			// None of these errors must be counted as a failure, even though the ReadAt() call
+			// never succeeded.
+			testutil.Equals(t, float64(0), promtest.ToFloat64(m.metrics.opsFailures.WithLabelValues(OpUpload)))
+		})
+	}
+}
+
+func TestTimingReader_ReadAtRetryAtSameOffsetDoesNotDoubleCount(t *testing.T) {
+	data := []byte("hello world")
+	fr := &flakyReaderAt{Reader: bytes.NewReader(data)}
+
+	m := WrapWithMetrics(NewInMemBucket(), nil, "")
+	tr := newTimingReader(time.Now(), fr, true, OpUpload, m.metrics.opsDuration, m.metrics.opsFailures, func(err error) bool {
+		return false
+	}, nil, m.metrics.opsTransferredBytes)
+
+	ra, ok := tr.(io.ReaderAt)
+	testutil.Assert(t, ok)
+
+	buf := make([]byte, len(data))
+	_, err := ra.ReadAt(buf, 0)
+	testutil.NotOk(t, err)
+
+	n, err := ra.ReadAt(buf, 0)
+	testutil.Ok(t, err)
+	testutil.Equals(t, len(data), n)
+
+	testutil.Ok(t, tr.Close())
+
+	// A transient ReadAt error that was retried and ultimately succeeded at the same offset must
+	// not be counted as a failure, nor double-count the bytes read across both attempts.
+	testutil.Equals(t, float64(0), promtest.ToFloat64(m.metrics.opsFailures.WithLabelValues(OpUpload)))
+	sum, count := histogramSumAndCount(t, m.metrics.opsTransferredBytes.WithLabelValues(OpUpload))
+	testutil.Equals(t, float64(len(data)), sum)
+	testutil.Equals(t, uint64(1), count)
+}
+
+func TestTimingReader_ReadAtHappyPathAcrossMultipleOffsets(t *testing.T) {
+	data := []byte("hello world!") // 12 bytes
+	r := bytes.NewReader(data)
+
+	m := WrapWithMetrics(NewInMemBucket(), nil, "")
+	tr := newTimingReader(time.Now(), r, true, OpUpload, m.metrics.opsDuration, m.metrics.opsFailures, func(err error) bool {
+		return false
+	}, m.metrics.opsFetchedBytes, m.metrics.opsTransferredBytes)
+
+	ra, ok := tr.(io.ReaderAt)
+	testutil.Assert(t, ok)
+
+	// Read every chunk of the object at a distinct offset, all succeeding on the first
+	// attempt, simulating a multipart upload's parts.
+	const chunkSize = 4
+	for off := 0; off < len(data); off += chunkSize {
+		buf := make([]byte, chunkSize)
+		n, err := ra.ReadAt(buf, int64(off))
+		testutil.Ok(t, err)
+		testutil.Equals(t, chunkSize, n)
+	}
+
+	testutil.Ok(t, tr.Close())
+
+	testutil.Equals(t, float64(0), promtest.ToFloat64(m.metrics.opsFailures.WithLabelValues(OpUpload)))
+
+	// transferredBytes must sum across all offsets, not just reflect one of them.
+	sum, count := histogramSumAndCount(t, m.metrics.opsTransferredBytes.WithLabelValues(OpUpload))
+	testutil.Equals(t, float64(len(data)), sum)
+	testutil.Equals(t, uint64(1), count)
+
+	// fetchedBytes must be incremented for ReadAt() calls too, same as for Read().
+	testutil.Equals(t, float64(len(data)), promtest.ToFloat64(m.metrics.opsFetchedBytes.WithLabelValues(OpUpload)))
+}
+
 func TestTimingReader_ShouldCorrectlyWrapFile(t *testing.T) {
 	// Create a test file.
 	testFilepath := filepath.Join(t.TempDir(), "test")