github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/testing/load/suite.go (about)

     1  package load
     2  
     3  import (
     4  	"runtime"
     5  	"sync"
     6  	"time"
     7  )
     8  
     9  type StorageWriteSuite struct {
    10  	apps    []*App
    11  	sources int
    12  
    13  	interval time.Duration
    14  	period   time.Duration
    15  	from     time.Time
    16  
    17  	seed    int
    18  	writers int
    19  	writeFn func(Input)
    20  }
    21  
    22  type StorageWriteSuiteConfig struct {
    23  	Sources  int
    24  	Interval time.Duration
    25  	Period   time.Duration
    26  	From     time.Time
    27  
    28  	Seed    int
    29  	Writers int
    30  	WriteFn func(Input)
    31  }
    32  
    33  const (
    34  	defaultInterval = 10 * time.Second
    35  	defaultRandSeed = 23061912
    36  )
    37  
    38  var defaultWriters = runtime.NumCPU()
    39  
    40  func NewStorageWriteSuite(c StorageWriteSuiteConfig) *StorageWriteSuite {
    41  	s := StorageWriteSuite{
    42  		sources:  c.Sources,
    43  		period:   c.Period,
    44  		from:     c.From,
    45  		writeFn:  c.WriteFn,
    46  		interval: defaultInterval,
    47  		seed:     defaultRandSeed,
    48  		writers:  defaultWriters,
    49  	}
    50  	if s.writeFn == nil {
    51  		panic("WriteFn is required")
    52  	}
    53  	if s.period == 0 {
    54  		panic("Period duration is required")
    55  	}
    56  	if s.sources == 0 {
    57  		panic("Number of sources is required")
    58  	}
    59  	if s.from.IsZero() {
    60  		s.from = time.Now().Add(-s.period)
    61  	}
    62  	if c.Interval > 0 {
    63  		s.interval = c.Interval
    64  	}
    65  	if c.Seed > 0 {
    66  		s.seed = c.Seed
    67  	}
    68  	if c.Writers > 0 {
    69  		s.writers = c.Writers
    70  	}
    71  	return &s
    72  }
    73  
    74  func (s *StorageWriteSuite) AddApp(app *App) *StorageWriteSuite {
    75  	s.apps = append(s.apps, app)
    76  	return s
    77  }
    78  
    79  func (s *StorageWriteSuite) AddAppWithConfig(name string, c AppConfig) *StorageWriteSuite {
    80  	s.apps = append(s.apps, NewApp(s.seed, name, c))
    81  	return s
    82  }
    83  
    84  type Stats struct {
    85  	RemainingPeriod time.Duration
    86  }
    87  
    88  func (s *StorageWriteSuite) Stats() Stats {
    89  	return Stats{
    90  		RemainingPeriod: s.period,
    91  	}
    92  }
    93  
    94  func (s *StorageWriteSuite) Start() {
    95  	q := make(chan Input)
    96  	wg := new(sync.WaitGroup)
    97  	wg.Add(s.writers)
    98  	for i := 0; i < s.writers; i++ {
    99  		go func() {
   100  			defer wg.Done()
   101  			for p := range q {
   102  				s.writeFn(p)
   103  			}
   104  		}()
   105  	}
   106  	from := s.from
   107  	for s.period > 0 {
   108  		to := from.Add(s.interval)
   109  		for i := 0; i < s.sources; i++ {
   110  			a := s.apps[i%len(s.apps)]
   111  			q <- a.CreateInput(from, to)
   112  		}
   113  		from = to
   114  		s.period -= s.interval
   115  	}
   116  	close(q)
   117  	wg.Wait()
   118  }