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

     1  // (c) 2023, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package txs
     5  
     6  import (
     7  	"context"
     8  	"errors"
     9  	"fmt"
    10  )
    11  
    12  // TxSequence provides an interface to return a channel of transactions.
    13  // The sequence is responsible for closing the channel when there are no further
    14  // transactions.
    15  type TxSequence[T any] interface {
    16  	Chan() <-chan T
    17  }
    18  
    19  // Worker defines the interface for issuance and confirmation of transactions.
    20  // The caller is responsible for calling Close to cleanup resources used by the
    21  // worker at the end of the simulation.
    22  type Worker[T any] interface {
    23  	IssueTx(ctx context.Context, tx T) error
    24  	ConfirmTx(ctx context.Context, tx T) error
    25  	Close(ctx context.Context) error
    26  }
    27  
    28  // Execute the work of the given agent.
    29  type Agent[T any] interface {
    30  	Execute(ctx context.Context) error
    31  }
    32  
    33  // issueNAgent issues and confirms a batch of N transactions at a time.
    34  type issueNAgent[T any] struct {
    35  	sequence TxSequence[T]
    36  	worker   Worker[T]
    37  	n        uint64
    38  }
    39  
    40  // NewIssueNAgent creates a new issueNAgent
    41  func NewIssueNAgent[T any](sequence TxSequence[T], worker Worker[T], n uint64) Agent[T] {
    42  	return &issueNAgent[T]{
    43  		sequence: sequence,
    44  		worker:   worker,
    45  		n:        n,
    46  	}
    47  }
    48  
    49  // Execute issues txs in batches of N and waits for them to confirm
    50  func (a issueNAgent[T]) Execute(ctx context.Context) error {
    51  	if a.n == 0 {
    52  		return errors.New("batch size n cannot be equal to 0")
    53  	}
    54  
    55  	txChan := a.sequence.Chan()
    56  
    57  	for {
    58  		var (
    59  			txs = make([]T, 0, a.n)
    60  			tx  T
    61  			ok  bool
    62  		)
    63  		for i := uint64(0); i < a.n; i++ {
    64  			select {
    65  			case tx, ok = <-txChan:
    66  				if !ok {
    67  					return a.worker.Close(ctx)
    68  				}
    69  			case <-ctx.Done():
    70  				return ctx.Err()
    71  			}
    72  
    73  			if err := a.worker.IssueTx(ctx, tx); err != nil {
    74  				return fmt.Errorf("failed to issue transaction %d: %w", len(txs), err)
    75  			}
    76  			txs = append(txs, tx)
    77  		}
    78  
    79  		for i, tx := range txs {
    80  			if err := a.worker.ConfirmTx(ctx, tx); err != nil {
    81  				return fmt.Errorf("failed to await transaction %d: %w", i, err)
    82  			}
    83  		}
    84  	}
    85  }