diff --git a/tools/traffic/config.go b/tools/traffic/config.go index d8ccd1e9f9..f0f797338d 100644 --- 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,6 +19,7 @@ type Config struct { LoggingConfig common.LoggerConfig RandomizeBlobs bool InstanceLaunchInterval time.Duration + CustomQuorums []uint8 } func NewConfig(ctx *cli.Context) (*Config, error) { @@ -25,6 +27,14 @@ func NewConfig(ctx *cli.Context) (*Config, error) { 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 index e3c9543722..91ffff728e 100644 --- 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 index fdbcb629da..6650fc6afd 100644 --- 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 }