github.com/yandex/pandora@v0.5.32/core/aggregator/encoder.go (about)

     1  package aggregator
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"time"
     7  
     8  	"github.com/pkg/errors"
     9  	"github.com/yandex/pandora/core"
    10  	"github.com/yandex/pandora/core/coreutil"
    11  	"github.com/yandex/pandora/lib/errutil"
    12  )
    13  
    14  type NewSampleEncoder func(w io.Writer, onFlush func()) SampleEncoder
    15  
    16  //go:generate mockery --name=SampleEncoder --case=underscore --outpkg=aggregatemock
    17  
    18  // SampleEncoder is efficient, buffered encoder of samples.
    19  // SampleEncoder MAY support only concrete type of sample.
    20  // MAY also implement SampleEncodeCloser.
    21  type SampleEncoder interface {
    22  	// SampleEncoder SHOULD panic, if passed sample type is not supported.
    23  	Encode(s core.Sample) error
    24  	// Flush flushes internal buffer to wrapped io.Writer.
    25  	Flush() error
    26  	// Optional. Close MUST be called, if io.Closer is implemented.
    27  	// io.Closer
    28  }
    29  
    30  //go:generate mockery --name=SampleEncodeCloser --case=underscore --outpkg=aggregatemock
    31  
    32  // SampleEncoderCloser is SampleEncoder that REQUIRE Close call to finish encoding.
    33  type SampleEncodeCloser interface {
    34  	SampleEncoder
    35  	io.Closer
    36  }
    37  
    38  type EncoderAggregatorConfig struct {
    39  	Sink           core.DataSink  `config:"sink" validate:"required"`
    40  	BufferSize     int            `config:"buffer-size"`
    41  	FlushInterval  time.Duration  `config:"flush-interval"`
    42  	ReporterConfig ReporterConfig `config:",squash"`
    43  }
    44  
    45  func DefaultEncoderAggregatorConfig() EncoderAggregatorConfig {
    46  	return EncoderAggregatorConfig{
    47  		FlushInterval:  time.Second,
    48  		ReporterConfig: DefaultReporterConfig(),
    49  	}
    50  }
    51  
    52  // NewEncoderAggregator returns aggregator that use SampleEncoder to marshall samples to core.DataSink.
    53  // Handles encoder flushing and sample dropping on queue overflow.
    54  // putSample is optional func, that called on handled sample. Usually returns sample to pool.
    55  func NewEncoderAggregator(
    56  	newEncoder NewSampleEncoder,
    57  	conf EncoderAggregatorConfig,
    58  ) core.Aggregator {
    59  	return &dataSinkAggregator{
    60  		Reporter:   *NewReporter(conf.ReporterConfig),
    61  		newEncoder: newEncoder,
    62  		conf:       conf,
    63  	}
    64  }
    65  
    66  type dataSinkAggregator struct {
    67  	Reporter
    68  	core.AggregatorDeps
    69  
    70  	newEncoder NewSampleEncoder
    71  	conf       EncoderAggregatorConfig
    72  }
    73  
    74  func (a *dataSinkAggregator) Run(ctx context.Context, deps core.AggregatorDeps) (err error) {
    75  	a.AggregatorDeps = deps
    76  
    77  	sink, err := a.conf.Sink.OpenSink()
    78  	if err != nil {
    79  		return
    80  	}
    81  	defer func() {
    82  		closeErr := sink.Close()
    83  		err = errutil.Join(err, closeErr)
    84  		err = errutil.Join(err, a.DroppedErr())
    85  	}()
    86  
    87  	var flushes int
    88  	encoder := a.newEncoder(sink, func() {
    89  		flushes++
    90  	})
    91  	defer func() {
    92  		if encoder, ok := encoder.(io.Closer); ok {
    93  			closeErr := encoder.Close()
    94  			err = errutil.Join(err, errors.WithMessage(closeErr, "encoder close failed"))
    95  			return
    96  		}
    97  		flushErr := encoder.Flush()
    98  		err = errutil.Join(err, errors.WithMessage(flushErr, "final flush failed"))
    99  	}()
   100  
   101  	var flushTick <-chan time.Time
   102  	if a.conf.FlushInterval > 0 {
   103  		flushTicker := time.NewTicker(a.conf.FlushInterval)
   104  		flushTick = flushTicker.C
   105  		defer flushTicker.Stop()
   106  	}
   107  
   108  	var previousFlushes int
   109  HandleLoop:
   110  	for {
   111  		select {
   112  		case sample := <-a.Incomming:
   113  			err = a.handleSample(encoder, sample)
   114  			if err != nil {
   115  				return
   116  			}
   117  		case <-flushTick:
   118  			if previousFlushes == flushes {
   119  				a.Log.Debug("Flushing")
   120  				err = encoder.Flush()
   121  				if err != nil {
   122  					return
   123  				}
   124  			}
   125  			previousFlushes = flushes
   126  		case <-ctx.Done():
   127  			break HandleLoop // Still need to handle all queued samples.
   128  		}
   129  	}
   130  
   131  	for {
   132  		select {
   133  		case sample := <-a.Incomming:
   134  			err = a.handleSample(encoder, sample)
   135  			if err != nil {
   136  				return
   137  			}
   138  		default:
   139  			return nil
   140  		}
   141  	}
   142  }
   143  
   144  func (a *dataSinkAggregator) handleSample(enc SampleEncoder, sample core.Sample) error {
   145  	err := enc.Encode(sample)
   146  	if err != nil {
   147  		return errors.WithMessage(err, "sample encode failed")
   148  	}
   149  	coreutil.ReturnSampleIfBorrowed(sample)
   150  	return nil
   151  }