Support custom quorums in traffic generator (#567)

Layr-Labs · May 16, 2024 · 9748a38 · 9748a38
1 parent 4d5d8b5
commit 9748a38
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 1 deletion.
diff --git a/tools/traffic/config.go b/tools/traffic/config.go
@@ -1,6 +1,7 @@
 package traffic
 
 import (
+	"errors"
 	"time"
 
 	"github.com/Layr-Labs/eigenda/clients"
@@ -18,13 +19,22 @@ type Config struct {
 	LoggingConfig          common.LoggerConfig
 	RandomizeBlobs         bool
 	InstanceLaunchInterval time.Duration
+	CustomQuorums          []uint8
 }
 
 func NewConfig(ctx *cli.Context) (*Config, error) {
 	loggerConfig, err := common.ReadLoggerCLIConfig(ctx, flags.FlagPrefix)
 	if err != nil {
 		return nil, err
 	}
+	customQuorums := ctx.GlobalIntSlice(flags.CustomQuorumNumbersFlag.Name)
+	customQuorumsUint8 := make([]uint8, len(customQuorums))
+	for i, q := range customQuorums {
+		if q < 0 || q > 255 {
+			return nil, errors.New("invalid custom quorum number")
+		}
+		customQuorumsUint8[i] = uint8(q)
+	}
 	return &Config{
 		Config: *clients.NewConfig(
 			ctx.GlobalString(flags.HostnameFlag.Name),
@@ -38,5 +48,6 @@ func NewConfig(ctx *cli.Context) (*Config, error) {
 		LoggingConfig:          *loggerConfig,
 		RandomizeBlobs:         ctx.GlobalBool(flags.RandomizeBlobsFlag.Name),
 		InstanceLaunchInterval: ctx.Duration(flags.InstanceLaunchIntervalFlag.Name),
+		CustomQuorums:          customQuorumsUint8,
 	}, nil
 }
diff --git a/tools/traffic/flags/flags.go b/tools/traffic/flags/flags.go
@@ -72,6 +72,12 @@ var (
 		Required: false,
 		EnvVar:   common.PrefixEnvVar(envPrefix, "USE_SECURE_GRPC"),
 	}
+	CustomQuorumNumbersFlag = cli.IntSliceFlag{
+		Name:     common.PrefixFlag(FlagPrefix, "custom-quorum-numbers"),
+		Usage:    "Custom quorum numbers to use for the traffic generator",
+		Required: false,
+		EnvVar:   common.PrefixEnvVar(envPrefix, "CUSTOM_QUORUM_NUMBERS"),
+	}
 )
 
 var requiredFlags = []cli.Flag{
@@ -87,6 +93,7 @@ var optionalFlags = []cli.Flag{
 	RandomizeBlobsFlag,
 	InstanceLaunchIntervalFlag,
 	UseSecureGrpcFlag,
+	CustomQuorumNumbersFlag,
 }
 
 // Flags contains the list of configuration options available to the binary.

diff --git a/tools/traffic/generator.go b/tools/traffic/generator.go
@@ -97,7 +97,7 @@ func (g *TrafficGenerator) StartTraffic(ctx context.Context) error {
 func (g *TrafficGenerator) sendRequest(ctx context.Context, data []byte) error {
 	ctxTimeout, cancel := context.WithTimeout(ctx, g.Config.Timeout)
 	defer cancel()
-	blobStatus, key, err := g.DisperserClient.DisperseBlob(ctxTimeout, data, []uint8{})
+	blobStatus, key, err := g.DisperserClient.DisperseBlob(ctxTimeout, data, g.Config.CustomQuorums)
 	if err != nil {
 		return err
 	}