github.com/yandex/pandora@v0.5.32/components/providers/scenario/provider.go (about)

     1  package scenario
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  
     8  	"github.com/yandex/pandora/components/providers/base"
     9  	"github.com/yandex/pandora/components/providers/http/decoders"
    10  	"github.com/yandex/pandora/core"
    11  )
    12  
    13  type ProviderConfig struct {
    14  	File            string
    15  	Limit           uint
    16  	Passes          uint
    17  	ContinueOnError bool
    18  	MaxAmmoSize     int
    19  }
    20  
    21  type ProvAmmo interface {
    22  	SetID(id uint64)
    23  	Clone() ProvAmmo
    24  }
    25  
    26  type Provider[A ProvAmmo] struct {
    27  	base.ProviderBase
    28  	cfg ProviderConfig
    29  
    30  	sink  chan A
    31  	ammos []A
    32  }
    33  
    34  func (p *Provider[A]) SetConfig(conf ProviderConfig) {
    35  	p.cfg = conf
    36  }
    37  
    38  func (p *Provider[A]) SetSink(sink chan A) {
    39  	p.sink = sink
    40  }
    41  
    42  func (p *Provider[A]) SetAmmos(ammos []A) {
    43  	p.ammos = ammos
    44  }
    45  
    46  func (p *Provider[A]) Run(ctx context.Context, deps core.ProviderDeps) error {
    47  	const op = "scenario.Provider.Run"
    48  	p.Deps = deps
    49  
    50  	length := uint(len(p.ammos))
    51  	if length == 0 {
    52  		return decoders.ErrNoAmmo
    53  	}
    54  	ammoNum := uint(0)
    55  	passNum := uint(0)
    56  	for {
    57  		err := ctx.Err()
    58  		if err != nil {
    59  			if !errors.Is(err, context.Canceled) {
    60  				err = fmt.Errorf("%s error from context: %w", op, err)
    61  			}
    62  			return err
    63  		}
    64  		i := ammoNum % length
    65  		passNum = ammoNum / length
    66  		if p.cfg.Passes != 0 && passNum >= p.cfg.Passes {
    67  			return decoders.ErrPassLimit
    68  		}
    69  		if p.cfg.Limit != 0 && ammoNum >= p.cfg.Limit {
    70  			return decoders.ErrAmmoLimit
    71  		}
    72  		ammoNum++
    73  		ammo := p.ammos[i]
    74  		select {
    75  		case <-ctx.Done():
    76  			err = ctx.Err()
    77  			if err != nil && !errors.Is(err, context.Canceled) {
    78  				err = fmt.Errorf("%s error from context: %w", op, err)
    79  			}
    80  			return err
    81  		case p.sink <- ammo:
    82  		}
    83  	}
    84  }
    85  
    86  func (p *Provider[A]) Acquire() (core.Ammo, bool) {
    87  	ammo, ok := <-p.sink
    88  	if !ok {
    89  		return nil, false
    90  	}
    91  	clone := ammo.Clone()
    92  	clone.SetID(p.NextID())
    93  	return clone, true
    94  }
    95  
    96  func (p *Provider[A]) Release(_ core.Ammo) {
    97  }