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

     1  package provider
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  
     8  	"github.com/yandex/pandora/components/providers/base"
     9  	httpProvider "github.com/yandex/pandora/components/providers/http/ammo"
    10  	"github.com/yandex/pandora/components/providers/http/config"
    11  	"github.com/yandex/pandora/components/providers/http/decoders"
    12  	"github.com/yandex/pandora/core"
    13  	"github.com/yandex/pandora/lib/confutil"
    14  	"go.uber.org/zap"
    15  	"golang.org/x/xerrors"
    16  )
    17  
    18  type Provider struct {
    19  	base.ProviderBase
    20  	config.Config
    21  	decoders.Decoder
    22  
    23  	Close func() error
    24  
    25  	Sink  chan decoders.DecodedAmmo
    26  	ammos []decoders.DecodedAmmo
    27  }
    28  
    29  func (p *Provider) Acquire() (core.Ammo, bool) {
    30  	ammo, ok := <-p.Sink
    31  	if !ok {
    32  		return nil, false
    33  	}
    34  	req, err := ammo.BuildRequest()
    35  	if err != nil {
    36  		p.Deps.Log.Error("http build request error", zap.Error(err))
    37  		return ammo, false
    38  	}
    39  	for _, mw := range p.Middlewares {
    40  		err := mw.UpdateRequest(req)
    41  		if err != nil {
    42  			p.Deps.Log.Error("error on Middleware.UpdateRequest", zap.Error(err))
    43  			return ammo, false
    44  		}
    45  	}
    46  	return httpProvider.NewGunAmmo(req, ammo.Tag(), p.NextID()), ok
    47  }
    48  
    49  func (p *Provider) Release(a core.Ammo) {
    50  	if p.Preload {
    51  		return
    52  	}
    53  	p.Decoder.Release(a)
    54  }
    55  
    56  func (p *Provider) Run(ctx context.Context, deps core.ProviderDeps) (err error) {
    57  	p.Deps = deps
    58  	defer func() {
    59  		close(p.Sink)
    60  		// TODO: wrap in go 1.20
    61  		// err = errors.Join(err, p.Close())
    62  		if p.Close == nil {
    63  			return
    64  		}
    65  		closeErr := p.Close()
    66  		if closeErr != nil {
    67  			if err != nil {
    68  				err = xerrors.Errorf("Multiple errors faced: %w, %w", err, closeErr)
    69  			} else {
    70  				err = closeErr
    71  			}
    72  		}
    73  	}()
    74  
    75  	for _, mw := range p.Middlewares {
    76  		if err := mw.InitMiddleware(ctx, deps.Log); err != nil {
    77  			return fmt.Errorf("cant InitMiddleware %T, err: %w", mw, err)
    78  		}
    79  	}
    80  
    81  	if p.Config.Preload {
    82  		err = p.loadAmmo(ctx)
    83  		if err == nil {
    84  			err = p.runPreloaded(ctx)
    85  		}
    86  	} else {
    87  		err = p.runFullScan(ctx)
    88  	}
    89  
    90  	return
    91  }
    92  
    93  func (p *Provider) runFullScan(ctx context.Context) error {
    94  	for {
    95  		if err := ctx.Err(); err != nil {
    96  			if !errors.Is(err, context.Canceled) {
    97  				err = xerrors.Errorf("error from context: %w", err)
    98  			}
    99  			return err
   100  		}
   101  		ammo, err := p.Decoder.Scan(ctx)
   102  		if err != nil {
   103  			if errors.Is(err, decoders.ErrAmmoLimit) || errors.Is(err, decoders.ErrPassLimit) {
   104  				err = nil
   105  			}
   106  			return err
   107  		}
   108  		if !confutil.IsChosenCase(ammo.Tag(), p.Config.ChosenCases) {
   109  			continue
   110  		}
   111  
   112  		select {
   113  		case <-ctx.Done():
   114  			err = ctx.Err()
   115  			if err != nil && !errors.Is(err, context.Canceled) {
   116  				err = xerrors.Errorf("error from context: %w", err)
   117  			}
   118  			return err
   119  		case p.Sink <- ammo:
   120  		}
   121  	}
   122  }
   123  
   124  func (p *Provider) loadAmmo(ctx context.Context) error {
   125  	ammos, err := p.Decoder.LoadAmmo(ctx)
   126  	if err != nil {
   127  		return fmt.Errorf("cant LoadAmmo, err: %w", err)
   128  	}
   129  	p.ammos = make([]decoders.DecodedAmmo, 0, len(ammos))
   130  	for _, ammo := range ammos {
   131  		if confutil.IsChosenCase(ammo.Tag(), p.Config.ChosenCases) {
   132  			p.ammos = append(p.ammos, ammo)
   133  		}
   134  	}
   135  	return nil
   136  }
   137  
   138  func (p *Provider) runPreloaded(ctx context.Context) error {
   139  	length := uint(len(p.ammos))
   140  	if length == 0 {
   141  		return decoders.ErrNoAmmo
   142  	}
   143  	ammoNum := uint(0)
   144  	passNum := uint(0)
   145  	for {
   146  		err := ctx.Err()
   147  		if err != nil {
   148  			if !errors.Is(err, context.Canceled) {
   149  				err = xerrors.Errorf("error from context: %w", err)
   150  			}
   151  			return err
   152  		}
   153  		i := ammoNum % length
   154  		passNum = ammoNum / length
   155  		if p.Passes != 0 && passNum >= p.Passes {
   156  			return decoders.ErrPassLimit
   157  		}
   158  		if p.Limit != 0 && ammoNum >= p.Limit {
   159  			return decoders.ErrAmmoLimit
   160  		}
   161  		ammoNum++
   162  		ammo := p.ammos[i]
   163  		select {
   164  		case <-ctx.Done():
   165  			err = ctx.Err()
   166  			if err != nil && !errors.Is(err, context.Canceled) {
   167  				err = xerrors.Errorf("error from context: %w", err)
   168  			}
   169  			return err
   170  		case p.Sink <- ammo:
   171  		}
   172  	}
   173  }
   174  
   175  var _ core.Provider = (*Provider)(nil)