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

     1  package grpcjson
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  
     7  	jsoniter "github.com/json-iterator/go"
     8  	"github.com/pkg/errors"
     9  	"github.com/spf13/afero"
    10  	ammo "github.com/yandex/pandora/components/providers/grpc"
    11  	"github.com/yandex/pandora/lib/confutil"
    12  	"go.uber.org/zap"
    13  )
    14  
    15  func NewProvider(fs afero.Fs, conf Config) *Provider {
    16  	var p Provider
    17  	if conf.Source.Path != "" {
    18  		conf.File = conf.Source.Path
    19  	}
    20  	p = Provider{
    21  		Provider: ammo.NewProvider(fs, conf.File, p.start),
    22  		Config:   conf,
    23  	}
    24  	return &p
    25  }
    26  
    27  type Provider struct {
    28  	ammo.Provider
    29  	Config
    30  	log *zap.Logger
    31  }
    32  
    33  type Source struct {
    34  	Type string
    35  	Path string
    36  }
    37  
    38  type Config struct {
    39  	File string //`validate:"required"`
    40  	// Limit limits total num of ammo. Unlimited if zero.
    41  	Limit int `validate:"min=0"`
    42  	// Passes limits ammo file passes. Unlimited if zero.
    43  	Passes          int `validate:"min=0"`
    44  	ContinueOnError bool
    45  	//Maximum number of byte in an ammo. Default is bufio.MaxScanTokenSize
    46  	MaxAmmoSize int
    47  	Source      Source `config:"source"`
    48  	ChosenCases []string
    49  }
    50  
    51  func (p *Provider) start(ctx context.Context, ammoFile afero.File) error {
    52  	var ammoNum, passNum int
    53  	for {
    54  		passNum++
    55  		scanner := bufio.NewScanner(ammoFile)
    56  		if p.Config.MaxAmmoSize != 0 {
    57  			var buffer []byte
    58  			scanner.Buffer(buffer, p.Config.MaxAmmoSize)
    59  		}
    60  		for line := 1; scanner.Scan() && (p.Limit == 0 || ammoNum < p.Limit); line++ {
    61  			data := scanner.Bytes()
    62  			a, err := decodeAmmo(data, p.Pool.Get().(*ammo.Ammo))
    63  			if err != nil {
    64  				if p.Config.ContinueOnError {
    65  					a.Invalidate()
    66  				} else {
    67  					return errors.Wrapf(err, "failed to decode ammo at line: %v; data: %q", line, data)
    68  				}
    69  			}
    70  			if !confutil.IsChosenCase(a.Tag, p.Config.ChosenCases) {
    71  				continue
    72  			}
    73  			ammoNum++
    74  			select {
    75  			case p.Sink <- a:
    76  			case <-ctx.Done():
    77  				return nil
    78  			}
    79  		}
    80  		err := scanner.Err()
    81  		if err != nil {
    82  			return errors.Wrap(err, "gPRC Provider scan() err")
    83  		}
    84  		if p.Passes != 0 && passNum >= p.Passes {
    85  			break
    86  		}
    87  		_, err = ammoFile.Seek(0, 0)
    88  		if err != nil {
    89  			return errors.Wrap(err, "Failed to seek ammo file")
    90  		}
    91  	}
    92  	return nil
    93  }
    94  
    95  func decodeAmmo(jsonDoc []byte, am *ammo.Ammo) (*ammo.Ammo, error) {
    96  	var ammo ammo.Ammo
    97  	err := jsoniter.Unmarshal(jsonDoc, &ammo)
    98  	if err != nil {
    99  		return am, errors.WithStack(err)
   100  	}
   101  
   102  	am.Reset(ammo.Tag, ammo.Call, ammo.Metadata, ammo.Payload)
   103  	return am, nil
   104  }