github.com/badrootd/nibiru-cometbft@v0.37.5-0.20240307173500-2a75559eee9b/test/e2e/runner/load.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/badrootd/nibiru-cometbft/libs/log"
    11  	rpchttp "github.com/badrootd/nibiru-cometbft/rpc/client/http"
    12  	e2e "github.com/badrootd/nibiru-cometbft/test/e2e/pkg"
    13  	"github.com/badrootd/nibiru-cometbft/test/loadtime/payload"
    14  	"github.com/badrootd/nibiru-cometbft/types"
    15  	"github.com/google/uuid"
    16  )
    17  
    18  const workerPoolSize = 16
    19  
    20  // Load generates transactions against the network until the given context is
    21  // canceled.
    22  func Load(ctx context.Context, testnet *e2e.Testnet) error {
    23  	initialTimeout := 1 * time.Minute
    24  	stallTimeout := 30 * time.Second
    25  	chSuccess := make(chan struct{})
    26  	ctx, cancel := context.WithCancel(ctx)
    27  	defer cancel()
    28  
    29  	logger.Info("load", "msg", log.NewLazySprintf("Starting transaction load (%v workers)...", workerPoolSize))
    30  	started := time.Now()
    31  	u := [16]byte(uuid.New()) // generate run ID on startup
    32  
    33  	txCh := make(chan types.Tx)
    34  	go loadGenerate(ctx, txCh, testnet, u[:])
    35  
    36  	for _, n := range testnet.Nodes {
    37  		if n.SendNoLoad {
    38  			continue
    39  		}
    40  
    41  		for w := 0; w < testnet.LoadTxConnections; w++ {
    42  			go loadProcess(ctx, txCh, chSuccess, n)
    43  		}
    44  	}
    45  
    46  	// Monitor successful transactions, and abort on stalls.
    47  	success := 0
    48  	timeout := initialTimeout
    49  	for {
    50  		select {
    51  		case <-chSuccess:
    52  			success++
    53  			timeout = stallTimeout
    54  		case <-time.After(timeout):
    55  			return fmt.Errorf("unable to submit transactions for %v", timeout)
    56  		case <-ctx.Done():
    57  			if success == 0 {
    58  				return errors.New("failed to submit any transactions")
    59  			}
    60  			logger.Info("load", "msg", log.NewLazySprintf("Ending transaction load after %v txs (%.1f tx/s)...",
    61  				success, float64(success)/time.Since(started).Seconds()))
    62  			return nil
    63  		}
    64  	}
    65  }
    66  
    67  // loadGenerate generates jobs until the context is canceled
    68  func loadGenerate(ctx context.Context, txCh chan<- types.Tx, testnet *e2e.Testnet, id []byte) {
    69  	t := time.NewTimer(0)
    70  	defer t.Stop()
    71  	for {
    72  		select {
    73  		case <-t.C:
    74  		case <-ctx.Done():
    75  			close(txCh)
    76  			return
    77  		}
    78  		t.Reset(time.Second)
    79  
    80  		// A context with a timeout is created here to time the createTxBatch
    81  		// function out. If createTxBatch has not completed its work by the time
    82  		// the next batch is set to be sent out, then the context is canceled so that
    83  		// the current batch is halted, allowing the next batch to begin.
    84  		tctx, cf := context.WithTimeout(ctx, time.Second)
    85  		createTxBatch(tctx, txCh, testnet, id)
    86  		cf()
    87  	}
    88  }
    89  
    90  // createTxBatch creates new transactions and sends them into the txCh. createTxBatch
    91  // returns when either a full batch has been sent to the txCh or the context
    92  // is canceled.
    93  func createTxBatch(ctx context.Context, txCh chan<- types.Tx, testnet *e2e.Testnet, id []byte) {
    94  	wg := &sync.WaitGroup{}
    95  	genCh := make(chan struct{})
    96  	for i := 0; i < workerPoolSize; i++ {
    97  		wg.Add(1)
    98  		go func() {
    99  			defer wg.Done()
   100  			for range genCh {
   101  				tx, err := payload.NewBytes(&payload.Payload{
   102  					Id:          id,
   103  					Size:        uint64(testnet.LoadTxSizeBytes),
   104  					Rate:        uint64(testnet.LoadTxBatchSize),
   105  					Connections: uint64(testnet.LoadTxConnections),
   106  				})
   107  				if err != nil {
   108  					panic(fmt.Sprintf("Failed to generate tx: %v", err))
   109  				}
   110  
   111  				select {
   112  				case txCh <- tx:
   113  				case <-ctx.Done():
   114  					return
   115  				}
   116  			}
   117  		}()
   118  	}
   119  	for i := 0; i < testnet.LoadTxBatchSize; i++ {
   120  		select {
   121  		case genCh <- struct{}{}:
   122  		case <-ctx.Done():
   123  			break
   124  		}
   125  	}
   126  	close(genCh)
   127  	wg.Wait()
   128  }
   129  
   130  // loadProcess processes transactions by sending transactions received on the txCh
   131  // to the client.
   132  func loadProcess(ctx context.Context, txCh <-chan types.Tx, chSuccess chan<- struct{}, n *e2e.Node) {
   133  	var client *rpchttp.HTTP
   134  	var err error
   135  	s := struct{}{}
   136  	for tx := range txCh {
   137  		if client == nil {
   138  			client, err = n.Client()
   139  			if err != nil {
   140  				logger.Info("non-fatal error creating node client", "error", err)
   141  				continue
   142  			}
   143  		}
   144  		if _, err = client.BroadcastTxSync(ctx, tx); err != nil {
   145  			continue
   146  		}
   147  		chSuccess <- s
   148  	}
   149  }