Allow loadtest client to take a list of grpc endpoints (#1155)

* Allow loadtest client to take a list of grpc endpoints * lint
sei-protocol · Nov 28, 2023 · 7d82755 · 7d82755
1 parent 7756513
commit 7d82755
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 16 deletions.
diff --git a/loadtest/config.json b/loadtest/config.json
@@ -1,5 +1,5 @@
 {
-    "grpc_endpoint": "127.0.0.1:9090",
+    "grpc_endpoints": "127.0.0.1:9090",
     "blockchain_endpoint": "http://localhost:26657/blockchain",
     "node_uri": "tcp://localhost:26657",
     "tls": false,

diff --git a/loadtest/loadtest_client.go b/loadtest/loadtest_client.go
@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"fmt"
+	"math/rand"
 	"os"
 	"path/filepath"
 	"strings"
@@ -25,14 +26,14 @@ import (
 type LoadTestClient struct {
 	LoadTestConfig     Config
 	TestConfig         EncodingConfig
-	TxClient           typestx.ServiceClient
+	TxClients          []typestx.ServiceClient
 	TxHashFile         *os.File
 	SignerClient       *SignerClient
 	ChainID            string
 	TxHashList         []string
 	TxResponseChan     chan *string
 	TxHashListMutex    *sync.Mutex
-	GrpcConn           *grpc.ClientConn
+	GrpcConns          []*grpc.ClientConn
 	StakingQueryClient stakingtypes.QueryClient
 	// Staking specific variables
 	Validators []stakingtypes.Validator
@@ -55,11 +56,16 @@ func NewLoadTestClient(config Config) *LoadTestClient {
 	} else {
 		dialOptions = append(dialOptions, grpc.WithInsecure())
 	}
-	grpcConn, _ := grpc.Dial(
-		config.GrpcEndpoint,
-		dialOptions...,
-	)
-	TxClient := typestx.NewServiceClient(grpcConn)
+	endpoints := strings.Split(config.GrpcEndpoints, ",")
+	TxClients := make([]typestx.ServiceClient, len(endpoints))
+	GrpcConns := make([]*grpc.ClientConn, len(endpoints))
+	for i, endpoint := range endpoints {
+		grpcConn, _ := grpc.Dial(
+			endpoint,
+			dialOptions...)
+		TxClients[i] = typestx.NewServiceClient(grpcConn)
+		GrpcConns[i] = grpcConn
+	}
 
 	// setup output files
 	userHomeDir, _ := os.UserHomeDir()
@@ -70,15 +76,15 @@ func NewLoadTestClient(config Config) *LoadTestClient {
 	return &LoadTestClient{
 		LoadTestConfig:                config,
 		TestConfig:                    TestConfig,
-		TxClient:                      TxClient,
+		TxClients:                     TxClients,
 		TxHashFile:                    outputFile,
 		SignerClient:                  NewSignerClient(config.NodeURI),
 		ChainID:                       config.ChainID,
 		TxHashList:                    []string{},
 		TxResponseChan:                make(chan *string),
 		TxHashListMutex:               &sync.Mutex{},
-		GrpcConn:                      grpcConn,
-		StakingQueryClient:            stakingtypes.NewQueryClient(grpcConn),
+		GrpcConns:                     GrpcConns,
+		StakingQueryClient:            stakingtypes.NewQueryClient(GrpcConns[0]),
 		DelegationMap:                 map[string]map[string]int{},
 		TokenFactoryDenomOwner:        map[string]string{},
 		generatedAdminMessageForBlock: false,
@@ -97,7 +103,9 @@ func (c *LoadTestClient) SetValidators() {
 }
 
 func (c *LoadTestClient) Close() {
-	c.GrpcConn.Close()
+	for _, grpcConn := range c.GrpcConns {
+		_ = grpcConn.Close()
+	}
 }
 
 func (c *LoadTestClient) AppendTxHash(txHash string) {
@@ -276,7 +284,7 @@ func (c *LoadTestClient) ValidateTxs() {
 }
 
 func (c *LoadTestClient) GetTxResponse(hash string) *types.TxResponse {
-	grpcRes, err := c.TxClient.GetTx(
+	grpcRes, err := c.GetTxClient().GetTx(
 		context.Background(),
 		&typestx.GetTxRequest{
 			Hash: hash,
@@ -289,3 +297,8 @@ func (c *LoadTestClient) GetTxResponse(hash string) *types.TxResponse {
 	}
 	return grpcRes.TxResponse
 }
+
+func (c *LoadTestClient) GetTxClient() typestx.ServiceClient {
+	rand.Seed(time.Now().Unix())
+	return c.TxClients[rand.Int()%len(c.TxClients)]
+}
diff --git a/loadtest/tx.go b/loadtest/tx.go
@@ -29,7 +29,7 @@ func SendTx(
 	loadtestClient.SignerClient.SignTx(loadtestClient.ChainID, txBuilder, key, seqDelta)
 	txBytes, _ := TestConfig.TxConfig.TxEncoder()((*txBuilder).GetTx())
 	return func() {
-		grpcRes, err := loadtestClient.TxClient.BroadcastTx(
+		grpcRes, err := loadtestClient.GetTxClient().BroadcastTx(
 			context.Background(),
 			&typestx.BroadcastTxRequest{
 				Mode:    mode,
@@ -52,7 +52,7 @@ func SendTx(
 			// retry after a second until either succeed or fail for some other reason
 			fmt.Printf("Mempool full\n")
 			time.Sleep(1 * time.Second)
-			grpcRes, err = loadtestClient.TxClient.BroadcastTx(
+			grpcRes, err = loadtestClient.GetTxClient().BroadcastTx(
 				context.Background(),
 				&typestx.BroadcastTxRequest{
 					Mode:    mode,

diff --git a/loadtest/types.go b/loadtest/types.go
@@ -31,7 +31,7 @@ const (
 
 type Config struct {
 	ChainID            string                `json:"chain_id"`
-	GrpcEndpoint       string                `json:"grpc_endpoint"`
+	GrpcEndpoints      string                `json:"grpc_endpoints"`
 	BlockchainEndpoint string                `json:"blockchain_endpoint"`
 	NodeURI            string                `json:"node_uri"`
 	TxsPerBlock        uint64                `json:"txs_per_block"`