github.com/argoproj/argo-events@v1.9.1/test/stress/generator/cmd/root.go (about)

     1  package cmd
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/spf13/cobra"
    11  )
    12  
    13  const (
    14  	Success = "success"
    15  	Failure = "failure"
    16  
    17  	RootCommand = "go run ./test/stress/generator/main.go"
    18  )
    19  
    20  var (
    21  	duration         time.Duration
    22  	requestPerSecond int
    23  	totalRequests    int64
    24  )
    25  
    26  var rootCmd = &cobra.Command{
    27  	Use:   RootCommand,
    28  	Short: "Events generator for stress testing.",
    29  	Long:  ``,
    30  	Run: func(cmd *cobra.Command, args []string) {
    31  		cmd.HelpFunc()(cmd, args)
    32  	},
    33  }
    34  
    35  func Execute() {
    36  	if err := rootCmd.Execute(); err != nil {
    37  		fmt.Println(err)
    38  		os.Exit(1)
    39  	}
    40  }
    41  
    42  func init() {
    43  	defaultDuration, _ := time.ParseDuration("5m")
    44  	rootCmd.PersistentFlags().DurationVarP(&duration, "duration", "d", defaultDuration, "How long it will run, e.g. 5m, 60s")
    45  	rootCmd.PersistentFlags().IntVar(&requestPerSecond, "rps", 0, "Requests per second")
    46  	rootCmd.PersistentFlags().Int64Var(&totalRequests, "total", 0, "Total requests")
    47  	rootCmd.AddCommand(NewWebhookCommand())
    48  	rootCmd.AddCommand(NewSqsCommand())
    49  }
    50  
    51  func validateGlobalParameters(cmd *cobra.Command, args []string) {
    52  	if requestPerSecond <= 0 {
    53  		fmt.Printf("\nrps needs to be specified with a number greater than 0\n\n")
    54  		cmd.HelpFunc()(cmd, args)
    55  		os.Exit(1)
    56  	}
    57  	if totalRequests < 0 {
    58  		fmt.Printf("\ntotal needs to be a number greater than 0\n\n")
    59  		cmd.HelpFunc()(cmd, args)
    60  		os.Exit(1)
    61  	}
    62  }
    63  
    64  func stressRun(ctx context.Context, f func() error) {
    65  	var mutex = &sync.Mutex{}
    66  	results := map[string]int64{Success: 0, Failure: 0}
    67  	counter := int64(0)
    68  	startTime := time.Now()
    69  
    70  	defer func(start time.Time) {
    71  		fmt.Printf("\n----------------- Summary --------------\n")
    72  		fmt.Printf("Total generated events  : %d\n", counter)
    73  		fmt.Printf("Succeeded events        : %d\n", results[Success])
    74  		fmt.Printf("Failed events           : %d\n", results[Failure])
    75  		fmt.Printf("Time taken              : %v\n", time.Since(start))
    76  	}(startTime)
    77  
    78  	wg := &sync.WaitGroup{}
    79  loop:
    80  	for {
    81  		if time.Since(startTime).Seconds() >= duration.Seconds() {
    82  			break loop
    83  		}
    84  		for i := 0; i < requestPerSecond; i++ {
    85  			select {
    86  			case <-ctx.Done():
    87  				break loop
    88  			default:
    89  			}
    90  			wg.Add(1)
    91  			go func() {
    92  				defer wg.Done()
    93  				if err := f(); err != nil {
    94  					fmt.Printf("Error: %v\n", err)
    95  					mutex.Lock()
    96  					results[Failure]++
    97  					mutex.Unlock()
    98  				} else {
    99  					mutex.Lock()
   100  					results[Success]++
   101  					mutex.Unlock()
   102  				}
   103  			}()
   104  			counter++
   105  			if totalRequests > 0 && counter >= totalRequests {
   106  				break loop
   107  			}
   108  		}
   109  		time.Sleep(1 * time.Second)
   110  	}
   111  
   112  	wg.Wait()
   113  }