github.com/yandex/pandora@v0.5.32/core/provider/num.go (about)

     1  package provider
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/yandex/pandora/core"
     7  )
     8  
     9  // NewNum returns dummy provider, that provides 0, 1 .. n int sequence as ammo.
    10  // May be useful for test or in when Gun don't need ammo.
    11  func NewNum(limit int) core.Provider {
    12  	return &num{
    13  		limit: limit,
    14  		sink:  make(chan core.Ammo),
    15  	}
    16  }
    17  
    18  func NewNumBuffered(limit int) core.Provider {
    19  	return &num{
    20  		limit: limit,
    21  		sink:  make(chan core.Ammo, limit),
    22  	}
    23  }
    24  
    25  type NumConfig struct {
    26  	Limit int
    27  }
    28  
    29  func NewNumConf(conf NumConfig) core.Provider {
    30  	return NewNum(conf.Limit)
    31  }
    32  
    33  type num struct {
    34  	i     int
    35  	limit int
    36  	sink  chan core.Ammo
    37  }
    38  
    39  func (n *num) Run(ctx context.Context, _ core.ProviderDeps) error {
    40  	defer close(n.sink)
    41  	for ; n.limit <= 0 || n.i < n.limit; n.i++ {
    42  		select {
    43  		case n.sink <- n.i:
    44  		case <-ctx.Done():
    45  			return nil
    46  		}
    47  	}
    48  	return nil
    49  }
    50  
    51  func (n *num) Acquire() (a core.Ammo, ok bool) {
    52  	a, ok = <-n.sink
    53  	return
    54  }
    55  
    56  func (n *num) Release(core.Ammo) {}