add acceptance test + refactor

Signed-off-by: Ashwanth Goli <[email protected]>
thanos-io · Oct 10, 2024 · 8a2e26e · 8a2e26e
1 parent 58428f2
commit 8a2e26e
Show file tree

Hide file tree

Showing 12 changed files with 94 additions and 16 deletions.
diff --git a/objstore.go b/objstore.go
@@ -832,10 +832,10 @@ func (t *timingReaderWriterTo) WriteTo(w io.Writer) (n int64, err error) {
 
 type ObjectSizerReadCloser struct {
 	io.ReadCloser
-	Size int64
+	Size func() (int64, error)
 }
 
 // ObjectSize implement ObjectSizer.
 func (o ObjectSizerReadCloser) ObjectSize() (int64, error) {
-	return o.Size, nil
+	return o.Size()
 }
diff --git a/providers/azure/azure.go b/providers/azure/azure.go
@@ -273,7 +273,13 @@ func (b *Bucket) getBlobReader(ctx context.Context, name string, httpRange blob.
 		return nil, errors.Wrapf(err, "cannot download blob, address: %s", blobClient.URL())
 	}
 	retryOpts := azblob.RetryReaderOptions{MaxRetries: int32(b.readerMaxRetries)}
-	return objstore.ObjectSizerReadCloser{ReadCloser: resp.NewRetryReader(ctx, &retryOpts), Size: *resp.ContentLength}, nil
+
+	return objstore.ObjectSizerReadCloser{
+		ReadCloser: resp.NewRetryReader(ctx, &retryOpts),
+		Size: func() (int64, error) {
+			return *resp.ContentLength, nil
+		},
+	}, nil
 }
 
 // Get returns a reader for the given object name.

diff --git a/providers/bos/bos.go b/providers/bos/bos.go
@@ -308,7 +308,12 @@ func (b *Bucket) getRange(_ context.Context, bucketName, objectKey string, off,
 		return nil, err
 	}
 
-	return objstore.ObjectSizerReadCloser{ReadCloser: obj.Body, Size: obj.ContentLength}, err
+	return objstore.ObjectSizerReadCloser{
+		ReadCloser: obj.Body,
+		Size: func() (int64, error) {
+			return obj.ContentLength, nil
+		},
+	}, err
 }
 
 func configFromEnv() Config {

diff --git a/providers/cos/cos.go b/providers/cos/cos.go
@@ -320,8 +320,12 @@ func (b *Bucket) getRange(ctx context.Context, name string, off, length int64) (
 		return nil, err
 	}
 	// Add size info into reader to pass it to Upload function.
-	r := objstore.ObjectSizerReadCloser{ReadCloser: resp.Body, Size: resp.ContentLength}
-	return r, nil
+	return objstore.ObjectSizerReadCloser{
+		ReadCloser: resp.Body,
+		Size: func() (int64, error) {
+			return resp.ContentLength, nil
+		},
+	}, nil
 }
 
 // Get returns a reader for the given object name.

diff --git a/providers/filesystem/filesystem.go b/providers/filesystem/filesystem.go
@@ -176,7 +176,9 @@ func (b *Bucket) GetRange(ctx context.Context, name string, off, length int64) (
 	if length == -1 {
 		return objstore.ObjectSizerReadCloser{
 			ReadCloser: f,
-			Size:       size,
+			Size: func() (int64, error) {
+				return size, nil
+			},
 		}, nil
 	}
 
@@ -185,7 +187,9 @@ func (b *Bucket) GetRange(ctx context.Context, name string, off, length int64) (
 			Reader: io.LimitReader(f, length),
 			f:      f,
 		},
-		Size: min(length, size),
+		Size: func() (int64, error) {
+			return min(length, size), nil
+		},
 	}, nil
 }
 

diff --git a/providers/gcs/gcs.go b/providers/gcs/gcs.go
@@ -231,7 +231,12 @@ func (b *Bucket) Get(ctx context.Context, name string) (io.ReadCloser, error) {
 		return r, err
 	}
 
-	return objstore.ObjectSizerReadCloser{ReadCloser: r, Size: r.Attrs.Size}, nil
+	return objstore.ObjectSizerReadCloser{
+		ReadCloser: r,
+		Size: func() (int64, error) {
+			return r.Attrs.Size, nil
+		},
+	}, nil
 }
 
 // GetRange returns a new range reader for the given object name and range.
@@ -241,7 +246,13 @@ func (b *Bucket) GetRange(ctx context.Context, name string, off, length int64) (
 		return r, err
 	}
 
-	return objstore.ObjectSizerReadCloser{ReadCloser: r, Size: r.Attrs.Size}, nil
+	sz := r.Remain()
+	return objstore.ObjectSizerReadCloser{
+		ReadCloser: r,
+		Size: func() (int64, error) {
+			return sz, nil
+		},
+	}, nil
 }
 
 // Attributes returns information about the specified object.

diff --git a/providers/obs/obs.go b/providers/obs/obs.go
@@ -299,7 +299,12 @@ func (b *Bucket) getRange(_ context.Context, name string, off, length int64) (io
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to get object")
 	}
-	return objstore.ObjectSizerReadCloser{ReadCloser: output.Body, Size: output.ContentLength}, nil
+	return objstore.ObjectSizerReadCloser{
+		ReadCloser: output.Body,
+		Size: func() (int64, error) {
+			return output.ContentLength, nil
+		},
+	}, nil
 }
 
 // Exists checks if the given object exists in the bucket.

diff --git a/providers/oci/oci.go b/providers/oci/oci.go
@@ -134,7 +134,12 @@ func (b *Bucket) Get(ctx context.Context, name string) (io.ReadCloser, error) {
 	if err != nil {
 		return nil, err
 	}
-	return objstore.ObjectSizerReadCloser{ReadCloser: response.Content, Size: *response.ContentLength}, nil
+	return objstore.ObjectSizerReadCloser{
+		ReadCloser: response.Content,
+		Size: func() (int64, error) {
+			return *response.ContentLength, nil
+		},
+	}, nil
 }
 
 // GetRange returns a new range reader for the given object name and range.
@@ -164,7 +169,11 @@ func (b *Bucket) GetRange(ctx context.Context, name string, offset, length int64
 	if err != nil {
 		return nil, err
 	}
-	return objstore.ObjectSizerReadCloser{ReadCloser: response.Content, Size: *response.ContentLength}, nil
+	return objstore.ObjectSizerReadCloser{ReadCloser: response.Content,
+		Size: func() (int64, error) {
+			return *response.ContentLength, nil
+		},
+	}, nil
 }
 
 // Upload the contents of the reader as an object into the bucket.

diff --git a/providers/oss/oss.go b/providers/oss/oss.go
@@ -350,7 +350,12 @@ func (b *Bucket) getRange(_ context.Context, name string, off, length int64) (io
 
 	size, err := clientutil.ParseContentLength(resp.Response.Headers)
 	if err == nil {
-		return objstore.ObjectSizerReadCloser{ReadCloser: resp.Response, Size: size}, nil
+		return objstore.ObjectSizerReadCloser{
+			ReadCloser: resp.Response,
+			Size: func() (int64, error) {
+				return size, nil
+			},
+		}, nil
 	}
 
 	return resp.Response, nil

diff --git a/providers/s3/s3.go b/providers/s3/s3.go
@@ -449,7 +449,12 @@ func (b *Bucket) getRange(ctx context.Context, name string, off, length int64) (
 		return o, err
 	}
 
-	return objstore.ObjectSizerReadCloser{ReadCloser: o, Size: stat.Size}, nil
+	return objstore.ObjectSizerReadCloser{
+		ReadCloser: o,
+		Size: func() (int64, error) {
+			return stat.Size, nil
+		},
+	}, nil
 }
 
 // Get returns a reader for the given object name.

diff --git a/providers/swift/swift.go b/providers/swift/swift.go
@@ -262,7 +262,11 @@ func (c *Container) get(name string, headers swift.Headers, checkHash bool) (io.
 	if err != nil {
 		return nil, errors.Wrap(err, "open object")
 	}
-	return file, err
+
+	return objstore.ObjectSizerReadCloser{
+		ReadCloser: file,
+		Size:       file.Length,
+	}, nil
 }
 
 // Get returns a reader for the given object name.

diff --git a/testing.go b/testing.go
@@ -106,6 +106,11 @@ func AcceptanceTest(t *testing.T, bkt Bucket) {
 	rc1, err := bkt.Get(ctx, "id1/obj_1.some")
 	testutil.Ok(t, err)
 	defer func() { testutil.Ok(t, rc1.Close()) }()
+
+	sz, err := TryToGetSize(rc1)
+	testutil.Ok(t, err)
+	testutil.Equals(t, int64(11), sz, "expected size to be equal to 11")
+
 	content, err := io.ReadAll(rc1)
 	testutil.Ok(t, err)
 	testutil.Equals(t, "@test-data@", string(content))
@@ -118,6 +123,11 @@ func AcceptanceTest(t *testing.T, bkt Bucket) {
 	rc2, err := bkt.GetRange(ctx, "id1/obj_1.some", 1, 3)
 	testutil.Ok(t, err)
 	defer func() { testutil.Ok(t, rc2.Close()) }()
+
+	sz, err = TryToGetSize(rc2)
+	testutil.Ok(t, err)
+	testutil.Equals(t, int64(3), sz, "expected size to be equal to 3")
+
 	content, err = io.ReadAll(rc2)
 	testutil.Ok(t, err)
 	testutil.Equals(t, "tes", string(content))
@@ -126,6 +136,11 @@ func AcceptanceTest(t *testing.T, bkt Bucket) {
 	rcUnspecifiedLen, err := bkt.GetRange(ctx, "id1/obj_1.some", 1, -1)
 	testutil.Ok(t, err)
 	defer func() { testutil.Ok(t, rcUnspecifiedLen.Close()) }()
+
+	sz, err = TryToGetSize(rcUnspecifiedLen)
+	testutil.Ok(t, err)
+	testutil.Equals(t, int64(10), sz, "expected size to be equal to 10")
+
 	content, err = io.ReadAll(rcUnspecifiedLen)
 	testutil.Ok(t, err)
 	testutil.Equals(t, "test-data@", string(content))
@@ -141,6 +156,11 @@ func AcceptanceTest(t *testing.T, bkt Bucket) {
 	rcLength, err := bkt.GetRange(ctx, "id1/obj_1.some", 3, 9999)
 	testutil.Ok(t, err)
 	defer func() { testutil.Ok(t, rcLength.Close()) }()
+
+	sz, err = TryToGetSize(rcLength)
+	testutil.Ok(t, err)
+	testutil.Equals(t, int64(8), sz, "expected size to be equal to 8")
+
 	content, err = io.ReadAll(rcLength)
 	testutil.Ok(t, err)
 	testutil.Equals(t, "st-data@", string(content))