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

     1  package ammo
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"sync/atomic"
     7  
     8  	"github.com/pkg/errors"
     9  	"github.com/spf13/afero"
    10  	"github.com/yandex/pandora/core"
    11  )
    12  
    13  func NewProvider(fs afero.Fs, fileName string, start func(ctx context.Context, file afero.File) error) Provider {
    14  	return Provider{
    15  		fs:       fs,
    16  		fileName: fileName,
    17  		start:    start,
    18  		Sink:     make(chan *Ammo, 128),
    19  		Pool:     sync.Pool{New: func() interface{} { return &Ammo{} }},
    20  		Close:    func() {},
    21  	}
    22  }
    23  
    24  type Provider struct {
    25  	fs        afero.Fs
    26  	fileName  string
    27  	start     func(ctx context.Context, file afero.File) error
    28  	Sink      chan *Ammo
    29  	Pool      sync.Pool
    30  	idCounter atomic.Uint64
    31  	Close     func()
    32  	core.ProviderDeps
    33  }
    34  
    35  func (p *Provider) Acquire() (core.Ammo, bool) {
    36  	ammo, ok := <-p.Sink
    37  	if ok {
    38  		ammo.SetID(p.idCounter.Add(1))
    39  	}
    40  	return ammo, ok
    41  }
    42  
    43  func (p *Provider) Release(a core.Ammo) {
    44  	p.Pool.Put(a)
    45  }
    46  
    47  func (p *Provider) Run(ctx context.Context, deps core.ProviderDeps) error {
    48  	defer p.Close()
    49  	p.ProviderDeps = deps
    50  	defer close(p.Sink)
    51  	file, err := p.fs.Open(p.fileName)
    52  	if err != nil {
    53  		return errors.Wrap(err, "failed to open ammo file")
    54  	}
    55  	defer file.Close()
    56  	return p.start(ctx, file)
    57  }