Throughput Testing Script (#1796)

Signed-off-by: Sam Liokumovich <[email protected]>

Author: samliok

cli/cluster_info.go

@@ -0,0 +1,101 @@
+// Copyright (C) 2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package cli
+
+import (
+	"fmt"
+	"os"
+	"strings"
+
+	"github.com/ava-labs/avalanchego/ids"
+	"gopkg.in/yaml.v2"
+
+	"github.com/ava-labs/hypersdk/utils"
+)
+
+type ClusterInfo struct {
+	ChainID  string `yaml:"CHAIN_ID"` // ids.ID requires "first and last characters to be quotes"
+	SubnetID string `yaml:"SUBNET_ID"`
+	APIs     []struct {
+		CloudID string `yaml:"CLOUD_ID"`
+		IP      string `yaml:"IP"`
+		Region  string `yaml:"REGION"`
+	} `yaml:"API"`
+	Validators []struct {
+		CloudID string `yaml:"CLOUD_ID"`
+		IP      string `yaml:"IP"`
+		Region  string `yaml:"REGION"`
+		NodeID  string `yaml:"NODE_ID"`
+	} `yaml:"VALIDATOR"`
+}
+
+func ReadClusterInfoFile(cliPath string) (ids.ID, map[string]string, error) {
+	// Load yaml file
+	yamlFile, err := os.ReadFile(cliPath)
+	if err != nil {
+		return ids.Empty, nil, err
+	}
+	var yamlContents ClusterInfo
+	if err := yaml.Unmarshal(yamlFile, &yamlContents); err != nil {
+		return ids.Empty, nil, fmt.Errorf("%w: unable to unmarshal YAML", err)
+	}
+	chainID, err := ids.FromString(yamlContents.ChainID)
+	if err != nil {
+		return ids.Empty, nil, err
+	}
+
+	// Load nodes
+	nodes := make(map[string]string)
+	for i, api := range yamlContents.APIs {
+		name := fmt.Sprintf("%s-%d (%s)", "API", i, api.Region)
+		uri := fmt.Sprintf("http://%s:9650/ext/bc/%s", api.IP, chainID)
+		nodes[name] = uri
+	}
+	for i, validator := range yamlContents.Validators {
+		name := fmt.Sprintf("%s-%d (%s)", "Validator", i, validator.Region)
+		uri := fmt.Sprintf("http://%s:9650/ext/bc/%s", validator.IP, chainID)
+		nodes[name] = uri
+	}
+	return chainID, nodes, nil
+}
+
+func (h *Handler) ImportClusterInfo(cliPath string) error {
+	oldChains, err := h.DeleteChains()
+	if err != nil {
+		return err
+	}
+	if len(oldChains) > 0 {
+		utils.Outf("{{yellow}}deleted old chains:{{/}} %+v\n", oldChains)
+	}
+
+	// Load yaml file
+	chainID, nodes, err := ReadClusterInfoFile(cliPath)
+	if err != nil {
+		return err
+	}
+	for name, uri := range nodes {
+		if err := h.StoreChain(chainID, name); err != nil {
+			return err
+		}
+		utils.Outf(
+			"{{yellow}}[%s] stored chainID:{{/}} %s {{yellow}}uri:{{/}} %s\n",
+			name,
+			chainID,
+			uri,
+		)
+	}
+	return h.StoreDefaultChain(chainID)
+}
+
+func OnlyAPIs(uris map[string]string) []string {
+	apis := make([]string, 0, len(uris))
+	for k := range uris {
+		if !strings.Contains(strings.ToLower(k), "api") {
+			continue
+		}
+
+		apis = append(apis, uris[k])
+	}
+	return apis
+}

cli/spam.go

@@ -8,13 +8,17 @@ import (
 
 	"github.com/ava-labs/hypersdk/auth"
 	"github.com/ava-labs/hypersdk/cli/prompt"
+	"github.com/ava-labs/hypersdk/codec"
 	"github.com/ava-labs/hypersdk/consts"
+	"github.com/ava-labs/hypersdk/crypto/ed25519"
 	"github.com/ava-labs/hypersdk/throughput"
 )
 
-// BuildSpammer prompts the user for the spammer parameters. If [defaults], the default values are used once the
+// BuildSpammer prompts the user for the spammer parameters.
+// If [spamKey] is provided, the user is not prompted for the root key.
+// If [defaults], the default values are used once the
 // chain and root key are selected. Otherwise, the user is prompted for all parameters.
-func (h *Handler) BuildSpammer(sh throughput.SpamHelper, defaults bool) (*throughput.Spammer, error) {
+func (h *Handler) BuildSpammer(sh throughput.SpamHelper, spamKey string, defaults bool) (*throughput.Spammer, error) {
 	// Select chain
 	chains, err := h.GetChains()
 	if err != nil {
@@ -25,20 +29,35 @@ func (h *Handler) BuildSpammer(sh throughput.SpamHelper, defaults bool) (*throug
 		return nil, err
 	}
 
-	// Select root key
-	keys, err := h.GetKeys()
-	if err != nil {
-		return nil, err
-	}
 	if err := sh.CreateClient(uris[0]); err != nil {
 		return nil, err
 	}
 
-	keyIndex, err := prompt.Choice("select root key", len(keys))
-	if err != nil {
-		return nil, err
+	var key *auth.PrivateKey
+
+	if len(spamKey) == 0 {
+		// Select root key
+		keys, err := h.GetKeys()
+		if err != nil {
+			return nil, err
+		}
+		keyIndex, err := prompt.Choice("select root key", len(keys))
+		if err != nil {
+			return nil, err
+		}
+		key = keys[keyIndex]
+	} else {
+		bytes, err := codec.LoadHex(spamKey, ed25519.PrivateKeyLen)
+		if err != nil {
+			return nil, err
+		}
+		privateKey := ed25519.PrivateKey(bytes)
+		key = &auth.PrivateKey{
+			Address: auth.NewED25519Address(privateKey.PublicKey()),
+			Bytes:   bytes,
+		}
 	}
-	key := keys[keyIndex]
+
 	// No longer using db, so we close
 	if err := h.CloseDatabase(); err != nil {
 		return nil, err
@@ -50,7 +69,7 @@ func (h *Handler) BuildSpammer(sh throughput.SpamHelper, defaults bool) (*throug
 	}
 
 	if defaults {
-		sc := throughput.NewDefaultConfig(uris, authFactory)
+		sc := throughput.NewFastConfig(uris, authFactory)
 		return throughput.NewSpammer(sc, sh)
 	}
 	// Collect parameters
@@ -102,8 +121,8 @@ func (h *Handler) BuildSpammer(sh throughput.SpamHelper, defaults bool) (*throug
 	return throughput.NewSpammer(sc, sh)
 }
 
-func (h *Handler) Spam(ctx context.Context, sh throughput.SpamHelper, defaults bool) error {
-	spammer, err := h.BuildSpammer(sh, defaults)
+func (h *Handler) Spam(ctx context.Context, sh throughput.SpamHelper, spamKey string, defaults bool) error {
+	spammer, err := h.BuildSpammer(sh, spamKey, defaults)
 	if err != nil {
 		return err
 	}

examples/morpheusvm/cmd/morpheus-cli/cmd/chain.go

@@ -28,6 +28,19 @@ var setChainCmd = &cobra.Command{
 	},
 }
 
+var importAvalancheCliChainCmd = &cobra.Command{
+	Use: "import-cli [path]",
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		if len(args) != 1 {
+			return ErrInvalidArgs
+		}
+		return nil
+	},
+	RunE: func(_ *cobra.Command, args []string) error {
+		return handler.Root().ImportClusterInfo(args[0])
+	},
+}
+
 var chainInfoCmd = &cobra.Command{
 	Use: "info",
 	RunE: func(_ *cobra.Command, _ []string) error {

examples/morpheusvm/cmd/morpheus-cli/cmd/genesis.go

@@ -64,6 +64,9 @@ var genGenesisCmd = &cobra.Command{
 		if minBlockGap >= 0 {
 			genesis.Rules.MinBlockGap = minBlockGap
 		}
+		if validityWindow >= 0 {
+			genesis.Rules.ValidityWindow = validityWindow
+		}
 
 		b, err := json.Marshal(genesis)
 		if err != nil {

examples/morpheusvm/cmd/morpheus-cli/cmd/root.go

@@ -28,16 +28,18 @@ var (
 	maxBlockUnits         []string
 	windowTargetUnits     []string
 	minBlockGap           int64
+	validityWindow        int64
 	hideTxs               bool
 	checkAllChains        bool
 	spamDefaults          bool
+	spamKey               string
+	clusterInfo           string
 	prometheusBaseURI     string
 	prometheusOpenBrowser bool
 	prometheusFile        string
 	prometheusData        string
 	startPrometheus       bool
-
-	rootCmd = &cobra.Command{
+	rootCmd               = &cobra.Command{
 		Use:        "morpheus-cli",
 		Short:      "MorpheusVM CLI",
 		SuggestFor: []string{"morpheus-cli", "morpheuscli"},
@@ -106,6 +108,12 @@ func init() {
 		-1,
 		"minimum block gap (ms)",
 	)
+	genGenesisCmd.PersistentFlags().Int64Var(
+		&validityWindow,
+		"validity-window",
+		-1,
+		"validity window (ms)",
+	)
 	genesisCmd.AddCommand(
 		genGenesisCmd,
 	)
@@ -135,6 +143,7 @@ func init() {
 		importChainCmd,
 		setChainCmd,
 		chainInfoCmd,
+		importAvalancheCliChainCmd,
 		watchChainCmd,
 	)
 
@@ -150,6 +159,20 @@ func init() {
 		"use default spam parameters",
 	)
 
+	runSpamCmd.PersistentFlags().StringVar(
+		&clusterInfo,
+		"cluster-info",
+		"",
+		"output from avalanche-cli with cluster info",
+	)
+
+	runSpamCmd.PersistentFlags().StringVar(
+		&spamKey,
+		"key",
+		"",
+		"private key used to distribute funds",
+	)
+
 	// spam
 	spamCmd.AddCommand(
 		runSpamCmd,

examples/morpheusvm/cmd/morpheus-cli/cmd/spam.go

@@ -5,11 +5,19 @@ package cmd
 
 import (
 	"context"
+	"errors"
 
 	"github.com/spf13/cobra"
 
+	"github.com/ava-labs/hypersdk/auth"
+	"github.com/ava-labs/hypersdk/cli"
+	"github.com/ava-labs/hypersdk/codec"
+	"github.com/ava-labs/hypersdk/crypto/ed25519"
 	"github.com/ava-labs/hypersdk/examples/morpheusvm/throughput"
 	"github.com/ava-labs/hypersdk/examples/morpheusvm/vm"
+	"github.com/ava-labs/hypersdk/utils"
+
+	hthroughput "github.com/ava-labs/hypersdk/throughput"
 )
 
 var spamCmd = &cobra.Command{
@@ -29,6 +37,40 @@ var runSpamCmd = &cobra.Command{
 	},
 	RunE: func(_ *cobra.Command, args []string) error {
 		ctx := context.Background()
-		return handler.Root().Spam(ctx, &throughput.SpamHelper{KeyType: args[0]}, spamDefaults)
+
+		if len(clusterInfo) > 0 {
+			_, urisFromFile, err := cli.ReadClusterInfoFile(clusterInfo)
+			if err != nil {
+				utils.Outf("{{red}} failed to read cluster info: %s \n", err)
+				return err
+			}
+			uris := cli.OnlyAPIs(urisFromFile)
+
+			if len(uris) == 0 || len(spamKey) == 0 {
+				return errors.New("uris and keyHex must be non-empty")
+			}
+			bytes, err := codec.LoadHex(spamKey, ed25519.PrivateKeyLen)
+			if err != nil {
+				return err
+			}
+			authFactory := auth.NewED25519Factory(ed25519.PrivateKey(bytes))
+
+			spamConfig, err := hthroughput.NewLongRunningConfig(uris, authFactory)
+			if err != nil {
+				return err
+			}
+
+			spamHelper := &throughput.SpamHelper{KeyType: args[0]}
+			if err := spamHelper.CreateClient(uris[0]); err != nil {
+				return err
+			}
+			spammer, err := hthroughput.NewSpammer(spamConfig, spamHelper)
+			if err != nil {
+				return err
+			}
+			return spammer.Spam(ctx, spamHelper, true, "AVAX")
+		}
+
+		return handler.Root().Spam(ctx, &throughput.SpamHelper{KeyType: args[0]}, spamKey, spamDefaults)
 	},
 }