Include Spam Config Defaults (#1657)

cli/spam.go

@@ -29,42 +29,23 @@ func (h *Handler) BuildSpammer(sh throughput.SpamHelper, defaults bool) (*throug
 	if err != nil {
 		return nil, err
 	}
-	balances := make([]uint64, len(keys))
 	if err := sh.CreateClient(uris[0]); err != nil {
 		return nil, err
 	}
-	for i := 0; i < len(keys); i++ {
-		balance, err := sh.LookupBalance(keys[i].Address)
-		if err != nil {
-			return nil, err
-		}
-		balances[i] = balance
-	}
 
 	keyIndex, err := prompt.Choice("select root key", len(keys))
 	if err != nil {
 		return nil, err
 	}
 	key := keys[keyIndex]
-	balance := balances[keyIndex]
 	// No longer using db, so we close
 	if err := h.CloseDatabase(); err != nil {
 		return nil, err
 	}
 
 	if defaults {
-		return throughput.NewSpammer(
-			uris,
-			key,
-			balance,
-			1.01,
-			2.7,
-			100000,   // tx per second
-			15000,    // min tx per second
-			1000,     // tx per second step
-			10,       // num clients
-			10000000, // num accounts
-		), nil
+		sc := throughput.NewDefaultConfig(uris, key)
+		return throughput.NewSpammer(sc, sh)
 	}
 	// Collect parameters
 	numAccounts, err := prompt.Int("number of accounts", consts.MaxInt)
@@ -100,18 +81,19 @@ func (h *Handler) BuildSpammer(sh throughput.SpamHelper, defaults bool) (*throug
 		return nil, err
 	}
 
-	return throughput.NewSpammer(
+	sc := throughput.NewConfig(
 		uris,
 		key,
-		balance,
 		sZipf,
 		vZipf,
 		txsPerSecond,
 		minTxsPerSecond,
 		txsPerSecondStep,
 		numClients,
 		numAccounts,
-	), nil
+	)
+
+	return throughput.NewSpammer(sc, sh)
 }
 
 func (h *Handler) Spam(ctx context.Context, sh throughput.SpamHelper, defaults bool) error {

tests/e2e/e2e.go

@@ -124,25 +124,16 @@ var _ = ginkgo.Describe("[HyperSDK Spam Workloads]", func() {
 		tc := e2e.NewTestContext()
 		require := require.New(tc)
 		blockchainID := e2e.GetEnv(tc).GetNetwork().GetSubnet(vmName).Chains[0].ChainID
-
-		// Spam Args
 		uris := getE2EURIs(tc, blockchainID)
 		key := spamKey
-		sZipf := 1.01
-		vZipf := 2.7
-		txsPerSecond := 500
-		minTxsPerSecond := 100
-		txsPerSecondStep := 200
-		numClients := 10
-		numAccounts := 25
-
-		// run spammer
+
 		err := spamHelper.CreateClient(uris[0])
 		require.NoError(err)
-		balance, err := spamHelper.LookupBalance(key.Address)
+
+		spamConfig := throughput.NewDefaultConfig(uris, key)
+		spammer, err := throughput.NewSpammer(spamConfig, spamHelper)
 		require.NoError(err)
 
-		spammer := throughput.NewSpammer(uris, key, balance, sZipf, vZipf, txsPerSecond, minTxsPerSecond, txsPerSecondStep, numClients, numAccounts)
 		err = spammer.Spam(tc.DefaultContext(), spamHelper, true, "AVAX")
 		require.NoError(err)
 	})

throughput/config.go

@@ -0,0 +1,59 @@
+// Copyright (C) 2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package throughput
+
+import "github.com/ava-labs/hypersdk/auth"
+
+type Config struct {
+	uris             []string
+	key              *auth.PrivateKey
+	sZipf            float64
+	vZipf            float64
+	txsPerSecond     int
+	minTxsPerSecond  int
+	txsPerSecondStep int
+	numClients       int
+	numAccounts      int
+}
+
+func NewDefaultConfig(
+	uris []string,
+	key *auth.PrivateKey,
+) *Config {
+	return &Config{
+		uris:             uris,
+		key:              key,
+		sZipf:            1.01,
+		vZipf:            2.7,
+		txsPerSecond:     500,
+		minTxsPerSecond:  100,
+		txsPerSecondStep: 200,
+		numClients:       10,
+		numAccounts:      25,
+	}
+}
+
+func NewConfig(
+	uris []string,
+	key *auth.PrivateKey,
+	sZipf float64,
+	vZipf float64,
+	txsPerSecond int,
+	minTxsPerSecond int,
+	txsPerSecondStep int,
+	numClients int,
+	numAccounts int,
+) *Config {
+	return &Config{
+		uris,
+		key,
+		sZipf,
+		vZipf,
+		txsPerSecond,
+		minTxsPerSecond,
+		txsPerSecondStep,
+		numClients,
+		numAccounts,
+	}
+}

throughput/spam.go

@@ -71,29 +71,28 @@ type Spammer struct {
 	numAccounts int
 }
 
-func NewSpammer(
-	uris []string,
-	key *auth.PrivateKey,
-	balance uint64,
-	sZipf, vZipf float64,
-	txsPerSecond, minTxsPerSecond, txsPerSecondStep, numClients, numAccounts int,
-) *Spammer {
+func NewSpammer(sc *Config, sh SpamHelper) (*Spammer, error) {
 	// Log Zipf participants
 	zipfSeed := rand.New(rand.NewSource(0)) //nolint:gosec
+	balance, err := sh.LookupBalance(sc.key.Address)
+	if err != nil {
+		return nil, err
+	}
 
 	return &Spammer{
-		uris,
-		key,
-		balance,
-		zipfSeed,
-		sZipf,
-		vZipf,
-		txsPerSecond,
-		minTxsPerSecond,
-		txsPerSecondStep,
-		numClients,
-		numAccounts,
-	}
+		uris:     sc.uris,
+		key:      sc.key,
+		balance:  balance,
+		zipfSeed: zipfSeed,
+		sZipf:    sc.sZipf,
+		vZipf:    sc.vZipf,
+
+		txsPerSecond:     sc.txsPerSecond,
+		minTxsPerSecond:  sc.minTxsPerSecond,
+		txsPerSecondStep: sc.txsPerSecondStep,
+		numClients:       sc.numClients,
+		numAccounts:      sc.numAccounts,
+	}, nil
 }
 
 // Spam tests the throughput of the network by sending transactions using