From 8a71f7b571a926628726e369f68b57701fdfd1f3 Mon Sep 17 00:00:00 2001
From: Harshavardhana <harsha@minio.io>
Date: Mon, 28 Oct 2024 23:42:59 -0700
Subject: [PATCH] add support for stat with --list-existing

---
 cli/get.go        | 18 ++++++------
 cli/stat.go       | 14 +++++++++
 pkg/bench/stat.go | 74 ++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 96 insertions(+), 10 deletions(-)

diff --git a/cli/get.go b/cli/get.go
index 624cea7..6c80d5b 100644
--- a/cli/get.go
+++ b/cli/get.go
@@ -25,14 +25,6 @@ import (
 )
 
 var getFlags = []cli.Flag{
-	cli.BoolFlag{
-		Name:  "list-existing",
-		Usage: "Instead of preparing the bench by PUTing some objects, only use objects already in the bucket",
-	},
-	cli.BoolFlag{
-		Name:  "list-flat",
-		Usage: "When using --list-existing, do not use recursive listing",
-	},
 	cli.IntFlag{
 		Name:  "objects",
 		Value: 2500,
@@ -62,6 +54,14 @@ var getFlags = []cli.Flag{
 		Value: 1,
 		Usage: "Number of versions to upload. If more than 1, versioned listing will be benchmarked",
 	},
+	cli.BoolFlag{
+		Name:  "list-existing",
+		Usage: "Instead of preparing the bench by PUTing some objects, only use objects already in the bucket",
+	},
+	cli.BoolFlag{
+		Name:  "list-flat",
+		Usage: "When using --list-existing, do not use recursive listing",
+	},
 }
 
 var getCmd = cli.Command{
@@ -107,7 +107,7 @@ func mainGet(ctx *cli.Context) error {
 		ListFlat:      ctx.Bool("list-flat"),
 		ListPrefix:    ctx.String("prefix"),
 	}
-	if b.ListExisting && !ctx.IsSet("objects") {
+	if b.ListExisting {
 		b.CreateObjects = 0
 	}
 	return runBench(ctx, &b)
diff --git a/cli/stat.go b/cli/stat.go
index 79eb04a..7fa46fe 100644
--- a/cli/stat.go
+++ b/cli/stat.go
@@ -40,6 +40,14 @@ var statFlags = []cli.Flag{
 		Value: 1,
 		Usage: "Number of versions to upload. If more than 1, versioned listing will be benchmarked",
 	},
+	cli.BoolFlag{
+		Name:  "list-existing",
+		Usage: "Instead of preparing the bench by PUTing some objects, only use objects already in the bucket",
+	},
+	cli.BoolFlag{
+		Name:  "list-flat",
+		Usage: "When using --list-existing, do not use recursive listing",
+	},
 }
 
 var statCmd = cli.Command{
@@ -72,6 +80,12 @@ func mainStat(ctx *cli.Context) error {
 		StatOpts: minio.StatObjectOptions{
 			ServerSideEncryption: sse,
 		},
+		ListExisting: ctx.Bool("list-existing"),
+		ListFlat:     ctx.Bool("list-flat"),
+		ListPrefix:   ctx.String("prefix"),
+	}
+	if b.ListExisting {
+		b.CreateObjects = 0
 	}
 	return runBench(ctx, &b)
 }
diff --git a/pkg/bench/stat.go b/pkg/bench/stat.go
index e822cdb..9246d5c 100644
--- a/pkg/bench/stat.go
+++ b/pkg/bench/stat.go
@@ -35,15 +35,87 @@ type Stat struct {
 	Common
 
 	// Default Stat options.
-	StatOpts      minio.StatObjectOptions
+	StatOpts   minio.StatObjectOptions
+	ListPrefix string
+
 	objects       generator.Objects
 	CreateObjects int
 	Versions      int
+
+	ListExisting bool
+	ListFlat     bool
 }
 
 // Prepare will create an empty bucket or delete any content already there
 // and upload a number of objects.
 func (g *Stat) Prepare(ctx context.Context) error {
+	// prepare the bench by listing object from the bucket
+	g.addCollector()
+	if g.ListExisting {
+		cl, done := g.Client()
+
+		// ensure the bucket exist
+		found, err := cl.BucketExists(ctx, g.Bucket)
+		if err != nil {
+			return err
+		}
+		if !found {
+			return (fmt.Errorf("bucket %s does not exist and --list-existing has been set", g.Bucket))
+		}
+
+		// list all objects
+		ctx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		objectCh := cl.ListObjects(ctx, g.Bucket, minio.ListObjectsOptions{
+			WithVersions: g.Versions > 1,
+			Prefix:       g.ListPrefix,
+			Recursive:    !g.ListFlat,
+		})
+
+		versions := map[string]int{}
+
+		for object := range objectCh {
+			if object.Err != nil {
+				return object.Err
+			}
+			if object.Size == 0 {
+				continue
+			}
+			obj := generator.Object{
+				Name: object.Key,
+				Size: object.Size,
+			}
+
+			if g.Versions > 1 {
+				if object.VersionID == "" {
+					continue
+				}
+
+				if version, found := versions[object.Key]; found {
+					if version >= g.Versions {
+						continue
+					}
+					versions[object.Key]++
+				} else {
+					versions[object.Key] = 1
+				}
+				obj.VersionID = object.VersionID
+			}
+
+			g.objects = append(g.objects, obj)
+
+			// limit to ListingMaxObjects
+			if g.CreateObjects > 0 && len(g.objects) >= g.CreateObjects {
+				break
+			}
+		}
+		if len(g.objects) == 0 {
+			return (fmt.Errorf("no objects found for bucket %s", g.Bucket))
+		}
+		done()
+		return nil
+	}
+
 	if err := g.createEmptyBucket(ctx); err != nil {
 		return err
 	}