From ca4684c5b5b3faa2e88a562abebf1202b4cc8c3e Mon Sep 17 00:00:00 2001
From: Patrick Schork <354473+pschork@users.noreply.github.com>
Date: Thu, 6 Jun 2024 22:27:04 -0700
Subject: [PATCH] Adds urfave/cli

---
 tools/batchgen/Makefile       |  2 +-
 tools/batchgen/cmd/main.go    | 24 ++++++++++--------------
 tools/batchgen/flags/flags.go | 23 +++++++++++++++--------
 3 files changed, 26 insertions(+), 23 deletions(-)

diff --git a/tools/batchgen/Makefile b/tools/batchgen/Makefile
index e27e4763dc..7afd620ad4 100644
--- a/tools/batchgen/Makefile
+++ b/tools/batchgen/Makefile
@@ -6,5 +6,5 @@ clean:
 	rm -rf ./bin
 
 run: build
-	./bin/batchgen 1000 10 localhost:32005
+	./bin/batchgen --host localhost:32005
 
diff --git a/tools/batchgen/cmd/main.go b/tools/batchgen/cmd/main.go
index c8f92b7ded..8e8f7ccade 100644
--- a/tools/batchgen/cmd/main.go
+++ b/tools/batchgen/cmd/main.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"log"
 	"os"
-	"strconv"
 
 	commonpb "github.com/Layr-Labs/eigenda/api/grpc/common"
 	pb "github.com/Layr-Labs/eigenda/api/grpc/node"
@@ -38,32 +37,26 @@ func main() {
 	app.Flags = flags.Flags
 	app.Action = batchGeneratorMain
 	if err := app.Run(os.Args); err != nil {
-		log.Fatalf("runtime failure: %v", err)
+		log.Fatal(err)
 	}
 }
 
-func batchGeneratorMain() {
-	threads, err := strconv.Atoi(os.Args[2])
-	if err != nil {
-		log.Fatalf("failed to convert threads to int:", err)
-	}
-
-	factor, err := strconv.Atoi(os.Args[1])
-	if err != nil {
-		log.Fatalf("failed to convert batch scaling factor to int:", err)
-	}
+func batchGeneratorMain(ctx *cli.Context) {
+	threads := int(ctx.GlobalUint(flags.NumThreadsFlag.Name))
+	factor := int(ctx.GlobalUint(flags.ScalingFactorFlag.Name))
+	hosts := ctx.GlobalStringSlice(flags.HostsFlag.Name)
 
-	hosts := os.Args[3:]
 	results := make(chan *pb.StoreChunksReply, len(hosts))
 	errors := make(chan error, len(hosts))
 
 	req, _, _, _, _ := makeStoreChunksRequest(100, 10, factor)
 	reqSizeMB := float64(proto.Size(req)) / 1024.0 / 1024.0
 	totalSizeMB := reqSizeMB * float64(threads) * float64(len(hosts))
-	log.Println("request size", reqSizeMB, "total data sent", totalSizeMB, "MB")
 
 	for _, host := range hosts {
+		log.Printf("host: %v", host)
 		for i := 0; i < threads; i++ {
+			log.Printf("i: %v", i)
 			go func(host string, i int) {
 				conn, err := grpc.Dial(host, grpc.WithTransportCredentials(insecure.NewCredentials()))
 				if err != nil {
@@ -93,6 +86,9 @@ func batchGeneratorMain() {
 			log.Println("ERR", err)
 		}
 	}
+
+	log.Println("request size", reqSizeMB, "total data sent", totalSizeMB, "MB")
+	os.Exit(0)
 }
 
 func makeStoreChunksRequest(quorumThreshold, adversaryThreshold uint8, factor int) (*pb.StoreChunksRequest, [32]byte, [32]byte, []*core.BlobHeader, []*pb.BlobHeader) {
diff --git a/tools/batchgen/flags/flags.go b/tools/batchgen/flags/flags.go
index 84c092b70f..26411f8d0f 100644
--- a/tools/batchgen/flags/flags.go
+++ b/tools/batchgen/flags/flags.go
@@ -6,7 +6,7 @@ import (
 )
 
 const (
-	FlagPrefix = "batchgen"
+	FlagPrefix = ""
 	envPrefix  = "BATCHGEN"
 )
 
@@ -14,28 +14,35 @@ var (
 	/* Required Flags */
 
 	NumThreadsFlag = cli.UintFlag{
-		Name:     "host-threads",
+		Name:     common.PrefixFlag(FlagPrefix, "threads"),
 		Usage:    "Number of host threads in parallel",
 		Required: false,
 		Value:    1,
-		EnvVar:   "HOST_THREADS",
+		EnvVar:   common.PrefixEnvVar(envPrefix, "HOST_THREADS"),
 	}
 	ScalingFactorFlag = cli.UintFlag{
-		Name:     "scaling-factor",
+		Name:     common.PrefixFlag(FlagPrefix, "scaling-factor"),
 		Usage:    "Scaling factor applied to default batch size",
 		Required: false,
-		Value:    10,
-		EnvVar:   "SCALING_FACTOR",
+		Value:    100,
+		EnvVar:   common.PrefixEnvVar(envPrefix, "SCALING_FACTOR"),
+	}
+	HostsFlag = cli.StringSliceFlag{
+		Name:     common.PrefixFlag(FlagPrefix, "host"),
+		Usage:    "host:port",
+		Required: true,
 	}
 )
 
 var requiredFlags = []cli.Flag{
+	HostsFlag,
+}
+
+var optionalFlags = []cli.Flag{
 	NumThreadsFlag,
 	ScalingFactorFlag,
 }
 
-var optionalFlags = []cli.Flag{}
-
 // Flags contains the list of configuration options available to the binary.
 var Flags []cli.Flag