[API change] add client-side streaming get-batch; add test

* new API: api.GetBatchStream
  - requires req.StreamingGet == true
  - rejects zip
  - no client-side writing: caller gets the stream and is responsible to
    close it
* tools: extend `tarch` sub-package with two draining flavors
* add integration smoke-test

Signed-off-by: Alex Aizman <alex.aizman@gmail.com>

Author: alex-aizman

ais/test/moss_test.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"path"
 	"strings"
 	"sync"
 	"testing"
@@ -903,3 +904,56 @@ func validateTarStreamingWithArchive(t *testing.T, m *ioContext, req *apc.MossRe
 	tlog.Logfln("Streaming archive validation passed: %d entries, correct order (format: %s)",
 		len(callback.entries), format)
 }
+
+func TestGetBatchStream_DrainTar(t *testing.T) {
+	proxyURL := tools.GetPrimaryURL()
+	bp := tools.BaseAPIParams(proxyURL)
+
+	bck := cmn.Bck{Name: trand.String(10), Provider: apc.AIS}
+	tools.CreateBucket(t, proxyURL, bck, nil, true /*cleanup*/)
+
+	// few objects
+	names := []string{"s1", "s2", "sbig", "s3"}
+	sizes := []int64{8*cos.KiB + 123, 12*cos.KiB + 456, 3*cos.MiB + 789, 4*cos.KiB + 0xabc}
+	for i, n := range names {
+		reader, err := readers.NewRand(sizes[i], cos.ChecksumNone)
+		tassert.CheckFatal(t, err)
+
+		_, err = api.PutObject(&api.PutArgs{
+			BaseParams: bp,
+			Bck:        bck,
+			ObjName:    n,
+			Reader:     reader,
+		})
+		tassert.CheckFatal(t, err)
+	}
+
+	// Build Moss request (streaming tar).
+	in := make([]apc.MossIn, 0, len(names))
+	for _, n := range names {
+		in = append(in, apc.MossIn{ObjName: n})
+	}
+	req := &apc.MossReq{In: in, StreamingGet: true}
+
+	rc, _, err := api.GetBatchStream(bp, bck, req)
+	tassert.CheckFatal(t, err)
+	defer rc.Close()
+
+	// TAR is the default output format
+	ar, err := archive.NewReader(archive.ExtTar, rc)
+	tassert.CheckFatal(t, err)
+
+	// prep. tarch.DrainVerify
+	var (
+		wantNames = make([]string, len(in))
+		sum       int64
+	)
+	for i, n := range names {
+		wantNames[i] = path.Join(bck.Name, n) // NOTE: must be api/ml.go documented naming
+		sum += sizes[i]
+	}
+
+	drain := tarch.NewDrainVerify(t, wantNames, sizes)
+	err = ar.ReadUntil(drain, "" /*match all*/, cos.EmptyMatchAll)
+	tassert.CheckFatal(t, err)
+}

api/client.go

@@ -155,10 +155,8 @@ func (reqParams *ReqParams) doReqStr(out *string) (int, error) {
 	return resp.StatusCode, err
 }
 
-// Makes request via do() and uses provided writer to write `resp.Body`
-// (which is also closes)
-//
-// Returns the entire wrapped response.
+// do() and use provided writer (to write `resp.Body`)
+// return the entire wrapped response
 func (reqParams *ReqParams) doWriter(w io.Writer) (wresp *wrappedResp, err error) {
 	var resp *http.Response
 	resp, err = reqParams.do()
@@ -171,6 +169,20 @@ func (reqParams *ReqParams) doWriter(w io.Writer) (wresp *wrappedResp, err error
 	return
 }
 
+// do() and return a checked response *with* the `resp.Body` as is
+// and *without*  draining or closing the latter
+func (reqParams *ReqParams) doStream() (wresp *wrappedResp, body io.ReadCloser, err error) {
+	resp, err := reqParams.do()
+	if err != nil {
+		return nil, nil, err
+	}
+	if err = reqParams.checkResp(resp); err != nil {
+		resp.Body.Close()
+		return nil, nil, err
+	}
+	return &wrappedResp{Response: resp}, resp.Body, nil
+}
+
 // same as above except that it returns response body (as io.ReadCloser) for subsequent reading
 func (reqParams *ReqParams) doReader() (io.ReadCloser, int64, error) {
 	resp, err := reqParams.do()

api/ml.go

@@ -5,6 +5,7 @@
 package api
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -48,36 +49,83 @@ import (
 //     - but when read range is defined: MossOut.Size = (length of this range)
 
 func GetBatch(bp BaseParams, bck cmn.Bck, req *apc.MossReq, w io.Writer) (resp apc.MossResp, err error) {
-	bp.Method = http.MethodGet
-
-	q, path := _optionalBucket(&bck)
-	reqParams := AllocRp()
-	{
-		reqParams.BaseParams = bp
-		reqParams.Path = path
-		reqParams.Body = cos.MustMarshal(req)
-		reqParams.Header = http.Header{cos.HdrContentType: []string{cos.ContentJSON}}
-		reqParams.Query = q
-	}
+	rp, q := _makeMossReq(bp, bck, req)
 	if req.StreamingGet {
 		var wresp *wrappedResp
-		wresp, err = reqParams.doWriter(w)
+		wresp, err = rp.doWriter(w)
 		if err == nil {
-			ct := wresp.Header.Get(cos.HdrContentType)
-			if !strings.HasPrefix(ct, "application/") {
-				err = fmt.Errorf("unexpected Content-Type %q", ct)
-			}
+			err = _checkMossResp(wresp.Header)
 		}
 	} else {
-		_, err = reqParams.readMultipart(&resp, w)
+		_, err = rp.readMultipart(&resp, w)
 	}
-	FreeRp(reqParams)
+	FreeRp(rp)
 	if q != nil {
 		qfree(q)
 	}
 	return resp, err
 }
 
+// GetBatchStream starts a streaming GetBatch and returns the response body _as is_
+// and response headers:
+// - the returned body is forward-only (non-seekable)
+// - supported streaming formats: .tar/.tgz/.tar.lz4; zip is excepted as non-streamable
+// - it is the caller's responsibility to close the body
+// - compare with GetBatch() above
+
+const zipext = ".zip"
+
+func GetBatchStream(bp BaseParams, bck cmn.Bck, req *apc.MossReq) (io.ReadCloser, http.Header, error) {
+	if !req.StreamingGet {
+		return nil, nil, errors.New("GetBatchStream: expecting req.StreamingGet to be set")
+	}
+	if req.OutputFormat != "" {
+		if req.OutputFormat == zipext || strings.Contains(req.OutputFormat, zipext[1:]) {
+			return nil, nil, fmt.Errorf("GetBatchStream: output format %q is not streamable", req.OutputFormat)
+		}
+	}
+	rp, q := _makeMossReq(bp, bck, req)
+
+	wresp, body, err := rp.doStream()
+	FreeRp(rp)
+	if q != nil {
+		qfree(q)
+	}
+	if err != nil {
+		return nil, nil, err
+	}
+	if err := _checkMossResp(wresp.Header); err != nil {
+		body.Close()
+		return nil, nil, err
+	}
+	return body, wresp.Header, nil
+}
+
+//
+// misc. helpers
+//
+
+// not a very strict check for: application/x-tar | application/gzip (tgz) | application/x-gtar | ...
+func _checkMossResp(hdr http.Header) error {
+	ct := hdr.Get(cos.HdrContentType)
+	if !strings.HasPrefix(ct, "application/") {
+		return fmt.Errorf("unexpected Content-Type %q", ct)
+	}
+	return nil
+}
+
+func _makeMossReq(bp BaseParams, bck cmn.Bck, req *apc.MossReq) (*ReqParams, url.Values) {
+	bp.Method = http.MethodGet
+	q, path := _optionalBucket(&bck)
+	rp := AllocRp()
+	rp.BaseParams = bp
+	rp.Path = path
+	rp.Body = cos.MustMarshal(req)
+	rp.Header = http.Header{cos.HdrContentType: []string{cos.ContentJSON}}
+	rp.Query = q
+	return rp, q
+}
+
 func _optionalBucket(bck *cmn.Bck) (q url.Values, path string) {
 	if bck.IsEmpty() {
 		path = apc.URLPathML.Join(apc.Moss)

tools/tarch/archive.go

@@ -14,13 +14,15 @@ import (
 	"os"
 	"strconv"
 	"sync"
+	"testing"
 	"time"
 
 	"github.com/NVIDIA/aistore/cmn/archive"
 	"github.com/NVIDIA/aistore/cmn/cos"
 	"github.com/NVIDIA/aistore/cmn/debug"
 	"github.com/NVIDIA/aistore/cmn/mono"
 	"github.com/NVIDIA/aistore/ext/dsort/shard"
+	"github.com/NVIDIA/aistore/tools/tassert"
 	"github.com/NVIDIA/aistore/tools/trand"
 )
 
@@ -327,3 +329,60 @@ func newBuf32k() (buf []byte) {
 func freeBuf32k(buf []byte) {
 	pool32k.Put(&buf)
 }
+
+///////////
+// Drain //
+///////////
+
+type (
+	Drain struct {
+		Total int64
+		Num   int
+	}
+
+	DrainVerify struct {
+		t         *testing.T
+		wantNames []string
+		wantSizes []int64
+		i         int
+		total     int64
+	}
+)
+
+func (drain *Drain) Call(_ string, r cos.ReadCloseSizer, _ any) (bool, error) {
+	n, err := io.Copy(io.Discard, r)
+	drain.Total += n
+	_ = r.Close()
+	if err == nil {
+		drain.Num++
+	}
+	return false, err
+}
+
+func NewDrainVerify(t *testing.T, wantNames []string, wantSizes []int64) *DrainVerify {
+	return &DrainVerify{
+		t:         t,
+		wantNames: wantNames,
+		wantSizes: wantSizes,
+	}
+}
+
+func (drain *DrainVerify) Call(name string, r cos.ReadCloseSizer, hdr any) (bool, error) {
+	tarhdr, ok := hdr.(*tar.Header)
+	if !ok {
+		tassert.Fatalf(drain.t, false, "expected *tar.Header, got %T", hdr)
+	}
+	expSize := drain.wantSizes[drain.i]
+	tassert.Errorf(drain.t, tarhdr.Size == expSize, "entry[%d] size mismatch: hdr=%d exp=%d", drain.i, tarhdr.Size, expSize)
+
+	expName := drain.wantNames[drain.i]
+	tassert.Errorf(drain.t, name == expName, "entry[%d] name mismatch: got %q exp %q", drain.i, name, expName)
+
+	n, err := io.Copy(io.Discard, r)
+	drain.total += n
+	_ = r.Close()
+	tassert.Errorf(drain.t, n == tarhdr.Size, "entry[%d] drained %d bytes != hdr.Size %d", drain.i, n, tarhdr.Size)
+
+	drain.i++
+	return false, err
+}