feat(storage/dataflux): add dataflux interface

storage/dataflux/fast_list.go

@@ -41,16 +41,16 @@ type ListerInput struct {
 	// BucketName is the name of the bucket to list objects from. Required.
 	BucketName string
 
-	// Parallelism is number of parallel workers to use for listing. Optional.
+	// Parallelism is number of parallel workers to use for listing. Default value is 10x number of available CPU. Optional.
 	Parallelism int
 
-	// BatchSize is the number of objects to list. Optional.
+	// BatchSize is the number of objects to list. Default value returns all objects at once. Optional.
 	BatchSize int
 
-	// Query is the query to filter objects for listing. Optional.
+	// Query is the query to filter objects for listing. Default value is nil. Optional.
 	Query storage.Query
 
-	// SkipDirectoryObjects is to indicate whether to list directory objects. Optional.
+	// SkipDirectoryObjects is to indicate whether to list directory objects. Default value is false. Optional.
 	SkipDirectoryObjects bool
 }
 
@@ -100,6 +100,8 @@ func (c *Lister) NextBatch(ctx context.Context) ([]*storage.ObjectAttrs, error)
 	var results []*storage.ObjectAttrs
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
+	// Errgroup takes care of running both methods in parallel. As soon as one of the method
+	// is complete, the running method also stops.
 	g, childCtx := errgroup.WithContext(ctx)
 
 	// To start listing method is Open and runs both worksteal and sequential listing in parallel.
@@ -109,7 +111,7 @@ func (c *Lister) NextBatch(ctx context.Context) ([]*storage.ObjectAttrs, error)
 	if c.method != worksteal {
 
 		g.Go(func() error {
-			objects, nextToken, err := sequentialListing(childCtx, *c)
+			objects, nextToken, err := c.sequentialListing(childCtx)
 			if err != nil {
 				countError++
 				return fmt.Errorf("error in running sequential listing: %w", err)
@@ -128,9 +130,11 @@ func (c *Lister) NextBatch(ctx context.Context) ([]*storage.ObjectAttrs, error)
 	// Close all functions if either sequential listing or worksteal listing is complete.
 	err := g.Wait()
 
-	// If there is not context.Canceled, then return error.
+	// If the error is not context.Canceled, then return error instead of falling back
+	// to the other method. This is so that the error can be fixed and user can take
+	//  advantage of fast-listing.
 	// As one of the listing method completes, it is expected to cancel context for the other method.
-	// If both sequential and worksteal listing fail due to context canceled, then return error.
+	// If both sequential and worksteal listing fail due to context canceled, only then return error.
 	if err != nil && (!errors.Is(err, context.Canceled) || countError > 1) {
 		return nil, fmt.Errorf("failed waiting for sequntial and work steal lister : %w", err)
 	}

storage/dataflux/integration_test.go

@@ -45,7 +45,6 @@ var (
 	//  These buckets are shared amongst download tests. They are created,
 	// populated with objects and cleaned up in TestMain.
 	httpTestBucket = downloadTestBucket{}
-	grpcTestBucket = downloadTestBucket{}
 )
 
 func TestMain(m *testing.M) {
@@ -67,7 +66,7 @@ func TestMain(m *testing.M) {
 }
 
 // Lists the all the objects in the bucket.
-func TestIntegration_NextBatch(t *testing.T) {
+func TestIntegration_NextBatch_All(t *testing.T) {
 	if testing.Short() {
 		t.Skip("Integration tests skipped in short mode")
 	}
@@ -85,14 +84,55 @@ func TestIntegration_NextBatch(t *testing.T) {
 
 	objects, err := df.NextBatch(ctx)
 	if err != nil && err != iterator.Done {
-		t.Errorf("df.NextBatch(: %v", err)
+		t.Errorf("df.NextBatch : %v", err)
 	}
 
 	if len(objects) != len(httpTestBucket.objects) {
 		t.Errorf("expected to receive %d results, got %d results", len(httpTestBucket.objects), len(objects))
 	}
 }
 
+func TestIntegration_NextBatch(t *testing.T) {
+	if testing.Short() {
+		t.Skip("Integration tests skipped in short mode")
+	}
+	const landsatBucket = "gcp-public-data-landsat"
+	const landsatPrefix = "LC08/01/001/00"
+	wantObjects := 17225
+	ctx := context.Background()
+	c, err := storage.NewClient(ctx)
+	if err != nil {
+		t.Fatalf("NewClient: %v", err)
+	}
+
+	in := &ListerInput{
+		BucketName: landsatBucket,
+		Query:      storage.Query{Prefix: landsatPrefix},
+		BatchSize:  2000,
+	}
+
+	df := NewLister(c, in)
+	defer df.Close()
+	totalObjects := 0
+	for {
+		objects, err := df.NextBatch(ctx)
+		if err != nil && err != iterator.Done {
+			t.Errorf("df.NextBatch : %v", err)
+		}
+		totalObjects += len(objects)
+		if err == iterator.Done {
+			break
+		}
+		if len(objects) > in.BatchSize {
+			t.Errorf("expected to receive %d objects in each batch, got %d objects in a batch", in.BatchSize, len(objects))
+		}
+	}
+	if totalObjects != wantObjects {
+		t.Errorf("expected to receive %d objects in results, got %d objects in results", wantObjects, totalObjects)
+
+	}
+}
+
 // generateRandomFileInGCS uploads a file with random contents to GCS and returns
 // the crc32c hash of the contents.
 func generateFileInGCS(ctx context.Context, o *storage.ObjectHandle, size int64) (uint32, error) {
@@ -257,7 +297,3 @@ func (tb *downloadTestBucket) Cleanup() error {
 
 	return b.Delete(ctx)
 }
-
-func crc32c(b []byte) uint32 {
-	return crc32.Checksum(b, crc32.MakeTable(crc32.Castagnoli))
-}

storage/dataflux/range_splitter.go

@@ -33,7 +33,6 @@ type listRange struct {
 
 // newRangeSplitter creates a new RangeSplitter with the given alphabets.
 func newRangeSplitter(alphabet string) *rangeSplitter {
-
 	return &rangeSplitter{}
 }
 

storage/dataflux/sequential.go

@@ -17,7 +17,6 @@ package dataflux
 import (
 	"context"
 	"fmt"
-	"math"
 	"strings"
 
 	"cloud.google.com/go/storage"
@@ -26,47 +25,51 @@ import (
 
 const (
 	// defaultPageSize specifies the number of object results to include on a single page.
-	defaultPageSize = 5000
+	defaultPageSize = 2000
 )
 
 // sequentialListing performs a sequential listing on the given bucket.
 // It returns a list of objects and the next token to use to continue listing.
 // If the next token is empty, then listing is complete.
-func sequentialListing(ctx context.Context, opts Lister) ([]*storage.ObjectAttrs, string, error) {
+func (c *Lister) sequentialListing(ctx context.Context) ([]*storage.ObjectAttrs, string, error) {
 	var result []*storage.ObjectAttrs
+	var objectsListed int
+	var lastToken string
+	objectIterator := c.bucket.Objects(ctx, &c.query)
+	objectIterator.PageInfo().Token = c.pageToken
+	objectIterator.PageInfo().MaxSize = defaultPageSize
 
-	objectIterator := opts.bucket.Objects(ctx, &opts.query)
-
-	var numObject int
-	if opts.batchSize < defaultPageSize {
-		numObject = defaultPageSize
-	} else {
-		numObject = int(math.Floor(float64(opts.batchSize)/float64(defaultPageSize))) * defaultPageSize
-	}
-
-	pageInfo := objectIterator.PageInfo()
-	pageInfo.MaxSize = defaultPageSize
-	pageInfo.Token = opts.pageToken
-
-	i := 0
 	for {
-		// If page size is set, then stop listing after numPageRequest.
-		if opts.batchSize > 0 && i >= numObject {
+		nextToken, err := doListing(objectIterator, &result, c.skipDirectoryObjects, &objectsListed)
+		if err != nil {
+			return nil, "", fmt.Errorf("failed while listing objects: %w", err)
+		}
+		lastToken = nextToken
+		if nextToken == "" || (c.batchSize > 0 && objectsListed >= c.batchSize) {
 			break
 		}
-		i++
+		c.pageToken = nextToken
+	}
+	return result, lastToken, nil
+}
 
+func doListing(objectIterator *storage.ObjectIterator, result *[]*storage.ObjectAttrs, skipDirectoryObjects bool, objectsListed *int) (string, error) {
+	for {
 		attrs, err := objectIterator.Next()
-		// When last item for the assigned range is listed, then stop listing.
+		*objectsListed++
+		// Stop listing when all the requested objects have been listed.
 		if err == iterator.Done {
 			break
 		}
 		if err != nil {
-			return nil, "", fmt.Errorf("iterating through objects %w", err)
+			return "", fmt.Errorf("iterating through objects %w", err)
 		}
-		if !(opts.skipDirectoryObjects && strings.HasSuffix(attrs.Name, "/")) {
-			result = append(result, attrs)
+		if !(skipDirectoryObjects && strings.HasSuffix(attrs.Name, "/")) {
+			*result = append(*result, attrs)
+		}
+		if objectIterator.PageInfo().Remaining() == 0 {
+			break
 		}
 	}
-	return result, objectIterator.PageInfo().Token, nil
+	return objectIterator.PageInfo().Token, nil
 }

storage/dataflux/worksteal.go

@@ -49,7 +49,7 @@ type worker struct {
 
 // workstealListing is the main entry point of the worksteal algorithm.
 // It performs worksteal to achieve highly dynamic object listing.
-func workstealListing(ctx context.Context, opts Lister) []*storage.ObjectAttrs {
+func (c *Lister) workstealListing(ctx context.Context) []*storage.ObjectAttrs {
 	return nil
 }