fix(storage): modify the callback of mrd to return length of data read instead of limit. (#11687)

Currently MRD takes a callback in each Add call.
Current call back returns the offset, length and error on that add call making the callback function look like, func(int64,int64,error) .
In some scenario data read can be different from limit set.
For example: We just have 50 bytes to read and ask for 100 bytes(length/limit) in this scenario we will return just 50 bytes.

Hence changed the callback to return amount of data read instead of data limit set.

Added the go doc for callback.

Author: shubham-diwakar

storage/grpc_client.go

@@ -1062,12 +1062,13 @@ func contextMetadataFromBidiReadObject(req *storagepb.BidiReadObjectRequest) []s
 }
 
 type rangeSpec struct {
-	readID       int64
-	writer       io.Writer
-	offset       int64
-	limit        int64
-	bytesWritten int64
-	callback     func(int64, int64, error)
+	readID              int64
+	writer              io.Writer
+	offset              int64
+	limit               int64
+	currentBytesWritten int64
+	totalBytesWritten   int64
+	callback            func(int64, int64, error)
 }
 
 func (c *grpcStorageClient) NewMultiRangeDownloader(ctx context.Context, params *newMultiRangeDownloaderParams, opts ...storageOption) (mr *MultiRangeDownloader, err error) {
@@ -1202,7 +1203,7 @@ func (c *grpcStorageClient) NewMultiRangeDownloader(ctx context.Context, params
 				rr.mu.Lock()
 				if len(rr.mp) != 0 {
 					for key := range rr.mp {
-						rr.mp[key].callback(rr.mp[key].offset, rr.mp[key].limit, fmt.Errorf("stream closed early"))
+						rr.mp[key].callback(rr.mp[key].offset, rr.mp[key].totalBytesWritten, fmt.Errorf("stream closed early"))
 						delete(rr.mp, key)
 					}
 				}
@@ -1295,21 +1296,22 @@ func (c *grpcStorageClient) NewMultiRangeDownloader(ctx context.Context, params
 						}
 						_, err = rr.mp[id].writer.Write(val.GetChecksummedData().GetContent())
 						if err != nil {
-							rr.mp[id].callback(rr.mp[id].offset, rr.mp[id].limit, err)
+							rr.mp[id].callback(rr.mp[id].offset, rr.mp[id].totalBytesWritten, err)
 							rr.activeTask--
 							delete(rr.mp, id)
 						} else {
 							rr.mp[id] = rangeSpec{
-								readID:       rr.mp[id].readID,
-								writer:       rr.mp[id].writer,
-								offset:       rr.mp[id].offset,
-								limit:        rr.mp[id].limit,
-								bytesWritten: rr.mp[id].bytesWritten + int64(len(val.GetChecksummedData().GetContent())),
-								callback:     rr.mp[id].callback,
+								readID:              rr.mp[id].readID,
+								writer:              rr.mp[id].writer,
+								offset:              rr.mp[id].offset,
+								limit:               rr.mp[id].limit,
+								currentBytesWritten: rr.mp[id].currentBytesWritten + int64(len(val.GetChecksummedData().GetContent())),
+								totalBytesWritten:   rr.mp[id].totalBytesWritten + int64(len(val.GetChecksummedData().GetContent())),
+								callback:            rr.mp[id].callback,
 							}
 						}
 						if val.GetRangeEnd() {
-							rr.mp[id].callback(rr.mp[id].offset, rr.mp[id].limit, nil)
+							rr.mp[id].callback(rr.mp[id].offset, rr.mp[id].totalBytesWritten, nil)
 							rr.activeTask--
 							delete(rr.mp, id)
 						}
@@ -1340,7 +1342,7 @@ func (c *grpcStorageClient) NewMultiRangeDownloader(ctx context.Context, params
 		if err != nil {
 			rr.mu.Lock()
 			for key := range rr.mp {
-				rr.mp[key].callback(rr.mp[key].offset, rr.mp[key].limit, err)
+				rr.mp[key].callback(rr.mp[key].offset, rr.mp[key].totalBytesWritten, err)
 				delete(rr.mp, key)
 			}
 			// In case we hit an permanent error, delete entries from map and remove active tasks.
@@ -1388,12 +1390,13 @@ func getActiveRange(r *gRPCBidiReader) []rangeSpec {
 	var activeRange []rangeSpec
 	for k, v := range r.mp {
 		activeRange = append(activeRange, rangeSpec{
-			readID:       k,
-			writer:       v.writer,
-			offset:       (v.offset + v.bytesWritten),
-			limit:        v.limit - v.bytesWritten,
-			callback:     v.callback,
-			bytesWritten: 0,
+			readID:              k,
+			writer:              v.writer,
+			offset:              (v.offset + v.currentBytesWritten),
+			limit:               v.limit - v.currentBytesWritten,
+			callback:            v.callback,
+			currentBytesWritten: 0,
+			totalBytesWritten:   v.totalBytesWritten,
 		})
 		r.mp[k] = activeRange[len(activeRange)-1]
 	}
@@ -1443,22 +1446,22 @@ func (mr *gRPCBidiReader) add(output io.Writer, offset, limit int64, callback fu
 	mr.mu.Unlock()
 
 	if offset > objectSize {
-		callback(offset, limit, fmt.Errorf("offset larger than size of object: %v", objectSize))
+		callback(offset, 0, fmt.Errorf("offset larger than size of object: %v", objectSize))
 		return
 	}
 	if limit < 0 {
-		callback(offset, limit, fmt.Errorf("limit can't be negative"))
+		callback(offset, 0, fmt.Errorf("limit can't be negative"))
 		return
 	}
 	mr.mu.Lock()
 	currentID := (*mr).readID
 	(*mr).readID++
 	if !mr.done {
-		spec := rangeSpec{readID: currentID, writer: output, offset: offset, limit: limit, bytesWritten: 0, callback: callback}
+		spec := rangeSpec{readID: currentID, writer: output, offset: offset, limit: limit, currentBytesWritten: 0, totalBytesWritten: 0, callback: callback}
 		mr.activeTask++
 		mr.data <- []rangeSpec{spec}
 	} else {
-		callback(offset, limit, fmt.Errorf("stream is closed, can't add range"))
+		callback(offset, 0, fmt.Errorf("stream is closed, can't add range"))
 	}
 	mr.mu.Unlock()
 }

storage/integration_test.go

@@ -402,6 +402,55 @@ func TestIntegration_MultiRangeDownloader(t *testing.T) {
 	})
 }
 
+// TestIntegration_MRDCallbackReturnsDataLength tests if the callback returns the correct data
+// read length or not.
+func TestIntegration_MRDCallbackReturnsDataLength(t *testing.T) {
+	multiTransportTest(skipHTTP("gRPC implementation specific test"), t, func(t *testing.T, ctx context.Context, bucket string, _ string, client *Client) {
+		content := make([]byte, 1000)
+		rand.New(rand.NewSource(0)).Read(content)
+		objName := "MRDCallback"
+
+		// Upload test data.
+		obj := client.Bucket(bucket).Object(objName)
+		if err := writeObject(ctx, obj, "text/plain", content); err != nil {
+			t.Fatal(err)
+		}
+		defer func() {
+			if err := obj.Delete(ctx); err != nil {
+				log.Printf("failed to delete test object: %v", err)
+			}
+		}()
+		reader, err := obj.NewMultiRangeDownloader(ctx)
+		if err != nil {
+			t.Fatalf("NewMultiRangeDownloader: %v", err)
+		}
+		var res multiRangeDownloaderOutput
+		callback := func(x, y int64, err error) {
+			res.offset = x
+			res.limit = y
+			res.err = err
+		}
+		// Read All At Once.
+		offset := 0
+		limit := 10000
+		reader.Add(&res.buf, int64(offset), int64(limit), callback)
+		reader.Wait()
+		if res.limit != 1000 {
+			t.Errorf("Error in callback want data length 1000, got: %v", res.limit)
+		}
+		if !bytes.Equal(res.buf.Bytes(), content) {
+			t.Errorf("Error in read range offset %v, limit %v, got: %v; want: %v",
+				offset, limit, res.buf.Bytes(), content)
+		}
+		if res.err != nil {
+			t.Errorf("read range %v to %v : %v", res.offset, 10000, res.err)
+		}
+		if err = reader.Close(); err != nil {
+			t.Fatalf("Error while closing reader %v", err)
+		}
+	})
+}
+
 // TestIntegration_ReadSameFileConcurrentlyUsingMultiRangeDownloader tests for potential deadlocks
 // or race conditions when multiple goroutines call Add() concurrently on the same MRD multiple times.
 func TestIntegration_ReadSameFileConcurrentlyUsingMultiRangeDownloader(t *testing.T) {

storage/reader.go

@@ -411,6 +411,10 @@ type multiRangeDownloader interface {
 // This will initiate the read range but is non-blocking; call callback to
 // process the result. Add is thread-safe and can be called simultaneously
 // from different goroutines.
+//
+// Callback will be called with the offset, length of data read, and error
+// of the read. Note that the length of the data read may be less than the
+// requested length if the end of the object is reached.
 func (mrd *MultiRangeDownloader) Add(output io.Writer, offset, length int64, callback func(int64, int64, error)) {
 	mrd.reader.add(output, offset, length, callback)
 }