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

     1  package provider
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  
     8  	"github.com/pkg/errors"
     9  	"github.com/yandex/pandora/core"
    10  	"github.com/yandex/pandora/lib/errutil"
    11  	"github.com/yandex/pandora/lib/ioutil2"
    12  	"go.uber.org/zap"
    13  )
    14  
    15  type NewAmmoDecoder func(deps core.ProviderDeps, source io.Reader) (AmmoDecoder, error)
    16  
    17  // TODO(skipo): test decoder that fills ammo with random data
    18  
    19  // AmmoEncoder MAY support only concrete type of ammo.
    20  // AmmoDecoder SHOULD NOT be used after first decode fail.
    21  type AmmoDecoder interface {
    22  	// Decode fills passed ammo with data.
    23  	// Returns non nil error on fail.
    24  	// Panics if ammo type is not supported.
    25  	Decode(ammo core.Ammo) error
    26  }
    27  
    28  type AmmoDecoderFunc func(ammo core.Ammo) error
    29  
    30  func (f AmmoDecoderFunc) Decode(ammo core.Ammo) error { return f(ammo) }
    31  
    32  // TODO(skipor): test
    33  
    34  func NewDecodeProvider(newAmmo func() core.Ammo, newDecoder NewAmmoDecoder, conf DecodeProviderConfig) *DecodeProvider {
    35  	return &DecodeProvider{
    36  		AmmoQueue:  *NewAmmoQueue(newAmmo, conf.Queue),
    37  		newDecoder: newDecoder,
    38  		conf:       conf,
    39  	}
    40  }
    41  
    42  type DecodeProviderConfig struct {
    43  	Queue  AmmoQueueConfig `config:",squash"`
    44  	Source core.DataSource `config:"source" validate:"required"`
    45  	// Limit limits total num of ammo. Unlimited if zero.
    46  	Limit int `validate:"min=0"`
    47  	// Passes limits ammo file passes. Unlimited if zero.
    48  	Passes int `validate:"min=0"`
    49  }
    50  
    51  func DefaultDecodeProviderConfig() DecodeProviderConfig {
    52  	return DecodeProviderConfig{
    53  		Queue: DefaultAmmoQueueConfig(),
    54  	}
    55  }
    56  
    57  type DecodeProvider struct {
    58  	AmmoQueue
    59  	conf       DecodeProviderConfig
    60  	newDecoder NewAmmoDecoder
    61  	core.ProviderDeps
    62  }
    63  
    64  var _ core.Provider = &DecodeProvider{}
    65  
    66  func (p *DecodeProvider) Run(ctx context.Context, deps core.ProviderDeps) (err error) {
    67  	p.ProviderDeps = deps
    68  	defer close(p.OutQueue)
    69  	source, err := p.conf.Source.OpenSource()
    70  	if err != nil {
    71  		return errors.WithMessage(err, "data source open failed")
    72  	}
    73  	defer func() {
    74  		_ = errutil.Join(err, errors.Wrap(source.Close(), "data source close failed"))
    75  	}()
    76  
    77  	// Problem: can't use decoder after io.EOF, because decoder is invalidated. But decoder recreation
    78  	// is not efficient, when we have short data source.
    79  	// Now problem solved by using MultiPassReader, but in such case decoder don't know real input
    80  	// position, so can't put this important information in decode error.
    81  	// TODO(skipor):  Let's add optional Reset(io.Reader) method, that will allow efficient Decoder reset after every pass.
    82  	multipassReader := ioutil2.NewMultiPassReader(source, p.conf.Passes)
    83  	if source == multipassReader {
    84  		p.Log.Info("Ammo data source can't sought, so will be read only once")
    85  	}
    86  	decoder, err := p.newDecoder(deps, multipassReader)
    87  
    88  	if err != nil {
    89  		return errors.WithMessage(err, "decoder construction failed")
    90  	}
    91  	var ammoNum int
    92  	for ; p.conf.Limit <= 0 || ammoNum < p.conf.Limit; ammoNum++ {
    93  		ammo := p.InputPool.Get()
    94  		err = decoder.Decode(ammo)
    95  		if err == io.EOF {
    96  			p.Log.Info("Ammo finished", zap.Int("decoded", ammoNum))
    97  			return nil
    98  		}
    99  		if err != nil {
   100  			return errors.WithMessage(err, fmt.Sprintf("ammo #%v decode failed", ammoNum))
   101  		}
   102  		select {
   103  		case p.OutQueue <- ammo:
   104  		case <-ctx.Done():
   105  			p.Log.Debug("Provider run context is Done", zap.Int("decoded", ammoNum+1))
   106  			return nil
   107  		}
   108  	}
   109  	p.Log.Info("Ammo limit is reached", zap.Int("decoded", ammoNum))
   110  	return nil
   111  }