Implement objectSizer for minio object and timingReadCloser

pkg/objstore/objstore.go

@@ -12,8 +12,9 @@ import (
 	"strings"
 	"time"
 
-	"github.com/go-kit/log"
-	"github.com/go-kit/log/level"
+	"github.com/go-kit/kit/log"
+	"github.com/go-kit/kit/log/level"
+	"github.com/minio/minio-go/v7"
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
@@ -145,6 +146,12 @@ func TryToGetSize(r io.Reader) (int64, error) {
 		return int64(f.Len()), nil
 	case *strings.Reader:
 		return f.Size(), nil
+	case *minio.Object:
+		objectInfo, err := f.Stat()
+		if err != nil {
+			return 0, errors.Wrap(err, "minio.Object.Stat()")
+		}
+		return objectInfo.Size, nil
 	case ObjectSizer:
 		return f.ObjectSize()
 	}
@@ -500,9 +507,16 @@ type timingReadCloser struct {
 	duration          *prometheus.HistogramVec
 	failed            *prometheus.CounterVec
 	isFailureExpected IsOpFailureExpectedFunc
+
+	objSize    int64
+	objSizeErr error
 }
 
 func newTimingReadCloser(rc io.ReadCloser, op string, dur *prometheus.HistogramVec, failed *prometheus.CounterVec, isFailureExpected IsOpFailureExpectedFunc) *timingReadCloser {
+	// Since TryToGetSize can only reliably return size before doing any read calls,
+	// we call during "construction" and remember the results.
+	objSize, objSizeErr := TryToGetSize(rc)
+
 	// Initialize the metrics with 0.
 	dur.WithLabelValues(op)
 	failed.WithLabelValues(op)
@@ -513,6 +527,8 @@ func newTimingReadCloser(rc io.ReadCloser, op string, dur *prometheus.HistogramV
 		duration:          dur,
 		failed:            failed,
 		isFailureExpected: isFailureExpected,
+		objSize:           objSize,
+		objSizeErr:        objSizeErr,
 	}
 }
 
@@ -539,3 +555,7 @@ func (rc *timingReadCloser) Read(b []byte) (n int, err error) {
 	}
 	return n, err
 }
+
+func (rc *timingReadCloser) ObjectSize() (int64, error) {
+	return rc.objSize, rc.objSizeErr
+}

pkg/objstore/objstore_test.go

@@ -6,8 +6,12 @@ package objstore
 import (
 	"bytes"
 	"io"
+	"os"
+	"path"
+	"strings"
 	"testing"
 
+	"github.com/pkg/errors"
 	promtest "github.com/prometheus/client_golang/prometheus/testutil"
 
 	"github.com/thanos-io/thanos/pkg/testutil"
@@ -86,3 +90,73 @@ func TestTracingReader(t *testing.T) {
 	testutil.Ok(t, err)
 	testutil.Equals(t, int64(11), size)
 }
+
+func TestTryToGetSize(t *testing.T) {
+	tmpFile, err := os.OpenFile(path.Join(os.TempDir(), "test_try_get_size"), os.O_CREATE|os.O_RDWR, os.ModePerm)
+	testutil.Ok(t, err)
+	defer tmpFile.Close()
+	_, err = tmpFile.WriteString("abc")
+	testutil.Ok(t, err)
+
+	buf := make([]byte, 5)
+	bbuff := bytes.NewBuffer(buf)
+
+	sreader := strings.NewReader("hello world")
+
+	buf2 := make([]byte, 100)
+	breader := bytes.NewReader(buf2)
+
+	r := bytes.NewBuffer(make([]byte, 20))
+	trc := newTimingReadCloser(NopCloserWithSize(r), "test", nil, nil, nil)
+	for _, tcase := range []struct {
+		name    string
+		reader  io.Reader
+		expSize int64
+		expErr  error
+	}{
+		{
+			name:    "local file",
+			reader:  tmpFile,
+			expSize: 3,
+		},
+		{
+			name:    "bytes buffer",
+			reader:  bbuff,
+			expSize: 5,
+		},
+		{
+			name:    "bytes reader",
+			reader:  breader,
+			expSize: 100,
+		},
+		{
+			name:    "string reader",
+			reader:  sreader,
+			expSize: 11,
+		},
+		{
+			name:    "timing read closer",
+			reader:  trc,
+			expSize: 20,
+		},
+		{
+			name:   "other io.Reader implementation",
+			reader: &testReader{},
+			expErr: errors.Errorf("unsupported type of io.Reader: %T", &testReader{}),
+		},
+	} {
+		t.Run(tcase.name, func(t *testing.T) {
+			size, err := TryToGetSize(tcase.reader)
+			if tcase.expErr == nil {
+				testutil.Ok(t, err)
+				testutil.Equals(t, tcase.expSize, size)
+			} else {
+				testutil.Equals(t, tcase.expErr.Error(), err.Error())
+			}
+		})
+	}
+}
+
+type testReader struct{}
+
+func (r *testReader) Read([]byte) (n int, err error) { return 0, nil }