github.com/cgcardona/r-subnet-evm@v0.1.5/cmd/simulator/load/funder.go (about)

     1  // Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package load
     5  
     6  import (
     7  	"context"
     8  	"crypto/ecdsa"
     9  	"fmt"
    10  	"math/big"
    11  
    12  	"github.com/cgcardona/r-subnet-evm/cmd/simulator/key"
    13  	"github.com/cgcardona/r-subnet-evm/cmd/simulator/txs"
    14  	"github.com/cgcardona/r-subnet-evm/core/types"
    15  	"github.com/cgcardona/r-subnet-evm/ethclient"
    16  	"github.com/cgcardona/r-subnet-evm/params"
    17  	"github.com/ethereum/go-ethereum/common"
    18  	"github.com/ethereum/go-ethereum/log"
    19  )
    20  
    21  // DistributeFunds ensures that each address in keys has at least [minFundsPerAddr] by sending funds
    22  // from the key with the highest starting balance.
    23  // This function returns a set of at least [numKeys] keys, each having a minimum balance [minFundsPerAddr].
    24  func DistributeFunds(ctx context.Context, client ethclient.Client, keys []*key.Key, numKeys int, minFundsPerAddr *big.Int) ([]*key.Key, error) {
    25  	if len(keys) < numKeys {
    26  		return nil, fmt.Errorf("insufficient number of keys %d < %d", len(keys), numKeys)
    27  	}
    28  	fundedKeys := make([]*key.Key, 0, numKeys)
    29  	// TODO: clean up fund distribution.
    30  	needFundsKeys := make([]*key.Key, 0)
    31  	needFundsAddrs := make([]common.Address, 0)
    32  
    33  	maxFundsKey := keys[0]
    34  	maxFundsBalance := common.Big0
    35  	log.Info("Checking balance of each key to distribute funds")
    36  	for _, key := range keys {
    37  		balance, err := client.BalanceAt(ctx, key.Address, nil)
    38  		if err != nil {
    39  			return nil, fmt.Errorf("failed to fetch balance for addr %s: %w", key.Address, err)
    40  		}
    41  
    42  		if balance.Cmp(minFundsPerAddr) < 0 {
    43  			needFundsKeys = append(needFundsKeys, key)
    44  			needFundsAddrs = append(needFundsAddrs, key.Address)
    45  		} else {
    46  			fundedKeys = append(fundedKeys, key)
    47  		}
    48  
    49  		if balance.Cmp(maxFundsBalance) > 0 {
    50  			maxFundsKey = key
    51  			maxFundsBalance = balance
    52  		}
    53  	}
    54  	requiredFunds := new(big.Int).Mul(minFundsPerAddr, big.NewInt(int64(numKeys)))
    55  	if maxFundsBalance.Cmp(requiredFunds) < 0 {
    56  		return nil, fmt.Errorf("insufficient funds to distribute %d < %d", maxFundsBalance, requiredFunds)
    57  	}
    58  	log.Info("Found max funded key", "address", maxFundsKey.Address, "balance", maxFundsBalance, "numFundAddrs", len(needFundsAddrs))
    59  	if len(fundedKeys) >= numKeys {
    60  		return fundedKeys[:numKeys], nil
    61  	}
    62  
    63  	// If there are not enough funded keys, cut [needFundsAddrs] to the number of keys that
    64  	// must be funded to reach [numKeys] required.
    65  	fundKeysCutLen := numKeys - len(fundedKeys)
    66  	needFundsKeys = needFundsKeys[:fundKeysCutLen]
    67  	needFundsAddrs = needFundsAddrs[:fundKeysCutLen]
    68  
    69  	chainID, err := client.ChainID(ctx)
    70  	if err != nil {
    71  		return nil, fmt.Errorf("failed to fetch chainID: %w", err)
    72  	}
    73  	gasFeeCap, err := client.EstimateBaseFee(ctx)
    74  	if err != nil {
    75  		return nil, fmt.Errorf("failed to fetch estimated base fee: %w", err)
    76  	}
    77  	gasTipCap, err := client.SuggestGasTipCap(ctx)
    78  	if err != nil {
    79  		return nil, fmt.Errorf("failed to fetch suggested gas tip: %w", err)
    80  	}
    81  	signer := types.LatestSignerForChainID(chainID)
    82  
    83  	// Generate a sequence of transactions to distribute the required funds.
    84  	log.Info("Generating distribution transactions...")
    85  	i := 0
    86  	txGenerator := func(key *ecdsa.PrivateKey, nonce uint64) (*types.Transaction, error) {
    87  		tx, err := types.SignNewTx(key, signer, &types.DynamicFeeTx{
    88  			ChainID:   chainID,
    89  			Nonce:     nonce,
    90  			GasTipCap: gasTipCap,
    91  			GasFeeCap: gasFeeCap,
    92  			Gas:       params.TxGas,
    93  			To:        &needFundsAddrs[i],
    94  			Data:      nil,
    95  			Value:     requiredFunds,
    96  		})
    97  		if err != nil {
    98  			return nil, err
    99  		}
   100  		i++
   101  		return tx, nil
   102  	}
   103  
   104  	numTxs := uint64(len(needFundsAddrs))
   105  	txSequence, err := txs.GenerateTxSequence(ctx, txGenerator, client, maxFundsKey.PrivKey, numTxs)
   106  	if err != nil {
   107  		return nil, fmt.Errorf("failed to generate fund distribution sequence from %s of length %d", maxFundsKey.Address, len(needFundsAddrs))
   108  	}
   109  	worker := NewSingleAddressTxWorker(ctx, client, maxFundsKey.Address)
   110  	txFunderAgent := txs.NewIssueNAgent[*types.Transaction](txSequence, worker, numTxs)
   111  
   112  	if err := txFunderAgent.Execute(ctx); err != nil {
   113  		return nil, err
   114  	}
   115  	for _, addr := range needFundsAddrs {
   116  		balance, err := client.BalanceAt(ctx, addr, nil)
   117  		if err != nil {
   118  			return nil, fmt.Errorf("failed to fetch balance for addr %s: %w", addr, err)
   119  		}
   120  		log.Info("Funded address has balance", "addr", addr, "balance", balance)
   121  	}
   122  	fundedKeys = append(fundedKeys, needFundsKeys...)
   123  	return fundedKeys, nil
   124  }