PBM-1498 hmac support for gcs fix (#1128)

Author: veceraj

pbm/storage/gcs/download.go

@@ -2,10 +2,8 @@ package gcs
 
 import (
 	"container/heap"
-	"context"
 	"io"
 	"net/http"
-	"time"
 
 	"google.golang.org/api/googleapi"
 
@@ -66,10 +64,7 @@ func (g *GCS) newPartReader(fname string, fsize int64, chunkSize int) *storage.P
 		Buf:       make([]byte, 32*1024),
 		L:         g.log,
 		GetChunk: func(fname string, arena *storage.Arena, cli interface{}, start, end int64) (io.ReadCloser, error) {
-			ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
-			defer cancel()
-
-			return cli.(gcsClient).getPartialObject(ctx, fname, start, end-start+1)
+			return cli.(gcsClient).getPartialObject(fname, arena, start, end-start+1)
 		},
 		GetSess: func() (interface{}, error) {
 			return g.client, nil // re-use the already-initialized client

pbm/storage/gcs/gcs.go

@@ -1,7 +1,6 @@
 package gcs
 
 import (
-	"context"
 	"io"
 	"path"
 	"reflect"
@@ -67,7 +66,7 @@ type gcsClient interface {
 	list(prefix, suffix string) ([]storage.FileInfo, error)
 	delete(name string) error
 	copy(src, dst string) error
-	getPartialObject(ctx context.Context, name string, start, length int64) (io.ReadCloser, error)
+	getPartialObject(name string, buf *storage.Arena, start, length int64) (io.ReadCloser, error)
 }
 
 type GCS struct {

pbm/storage/gcs/google_client.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"path"
 	"strings"
+	"time"
 
 	storagegcs "cloud.google.com/go/storage"
 	"github.com/googleapis/gax-go/v2"
@@ -216,7 +217,10 @@ func (g googleClient) copy(src, dst string) error {
 	return err
 }
 
-func (g googleClient) getPartialObject(ctx context.Context, name string, start, length int64) (io.ReadCloser, error) {
+func (g googleClient) getPartialObject(name string, buf *storage.Arena, start, length int64) (io.ReadCloser, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*60)
+	defer cancel()
+
 	obj := g.bucketHandle.Object(path.Join(g.opts.Prefix, name))
 	reader, err := obj.NewRangeReader(ctx, start, length)
 	if err != nil {
@@ -226,5 +230,13 @@ func (g googleClient) getPartialObject(ctx context.Context, name string, start,
 
 		return nil, storage.GetObjError{Err: err}
 	}
-	return reader, nil
+
+	ch := buf.GetSpan()
+	_, err = io.CopyBuffer(ch, reader, buf.CpBuf)
+	if err != nil {
+		ch.Close()
+		return nil, errors.Wrap(err, "copy")
+	}
+	reader.Close()
+	return ch, nil
 }

pbm/storage/gcs/hmac_client.go

@@ -6,6 +6,7 @@ import (
 	"path"
 	"runtime"
 	"strings"
+	"time"
 
 	"github.com/minio/minio-go/v7"
 	"github.com/minio/minio-go/v7/pkg/credentials"
@@ -181,7 +182,10 @@ func (h hmacClient) copy(src, dst string) error {
 	return err
 }
 
-func (h hmacClient) getPartialObject(ctx context.Context, name string, start, length int64) (io.ReadCloser, error) {
+func (h hmacClient) getPartialObject(name string, buf *storage.Arena, start, length int64) (io.ReadCloser, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*60)
+	defer cancel()
+
 	objectName := path.Join(h.opts.Prefix, name)
 
 	opts := minio.GetObjectOptions{}
@@ -200,6 +204,14 @@ func (h hmacClient) getPartialObject(ctx context.Context, name string, start, le
 
 		return nil, storage.GetObjError{Err: err}
 	}
+	defer object.Close()
+
+	ch := buf.GetSpan()
+	_, err = io.CopyBuffer(ch, object, buf.CpBuf)
+	if err != nil {
+		ch.Close()
+		return nil, errors.Wrap(err, "copy")
+	}
 
-	return object, nil
+	return ch, nil
 }