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

     1  package aggregator
     2  
     3  import (
     4  	"bufio"
     5  	"io"
     6  
     7  	jsoniter "github.com/json-iterator/go"
     8  	"github.com/yandex/pandora/core"
     9  	"github.com/yandex/pandora/core/coreutil"
    10  	"github.com/yandex/pandora/lib/ioutil2"
    11  )
    12  
    13  type JSONLineAggregatorConfig struct {
    14  	EncoderAggregatorConfig `config:",squash"`
    15  	JSONLineEncoderConfig   `config:",squash"`
    16  }
    17  
    18  type JSONLineEncoderConfig struct {
    19  	JSONIterConfig            `config:",squash"`
    20  	coreutil.BufferSizeConfig `config:",squash"`
    21  }
    22  
    23  // JSONIterConfig is subset of jsoniter.Config that may be useful to configure.
    24  type JSONIterConfig struct {
    25  	// MarshalFloatWith6Digits makes float marshalling faster.
    26  	MarshalFloatWith6Digits bool `config:"marshal-float-with-6-digits"`
    27  	// SortMapKeys useful, when sample contains map object, and you want to see them in same order.
    28  	SortMapKeys bool `config:"sort-map-keys"`
    29  }
    30  
    31  func DefaultJSONLinesAggregatorConfig() JSONLineAggregatorConfig {
    32  	return JSONLineAggregatorConfig{
    33  		EncoderAggregatorConfig: DefaultEncoderAggregatorConfig(),
    34  	}
    35  }
    36  
    37  // Aggregates samples in JSON Lines format: each output line is a Valid JSON Value of one sample.
    38  // See http://jsonlines.org/ for details.
    39  func NewJSONLinesAggregator(conf JSONLineAggregatorConfig) core.Aggregator {
    40  	var newEncoder NewSampleEncoder = func(w io.Writer, onFlush func()) SampleEncoder {
    41  		w = ioutil2.NewCallbackWriter(w, onFlush)
    42  		return NewJSONEncoder(w, conf.JSONLineEncoderConfig)
    43  	}
    44  	return NewEncoderAggregator(newEncoder, conf.EncoderAggregatorConfig)
    45  }
    46  
    47  func NewJSONEncoder(w io.Writer, conf JSONLineEncoderConfig) SampleEncoder {
    48  	apiConfig := jsoniter.Config{
    49  		SortMapKeys:             conf.JSONIterConfig.SortMapKeys,
    50  		MarshalFloatWith6Digits: conf.JSONIterConfig.MarshalFloatWith6Digits,
    51  	}
    52  
    53  	api := apiConfig.Froze()
    54  	// NOTE(skipor): internal buffering is not working really. Don't know why
    55  	// OPTIMIZE(skipor): don't wrap into buffer, if already ioutil2.ByteWriter
    56  	buf := bufio.NewWriterSize(w, conf.BufferSizeOrDefault())
    57  	stream := jsoniter.NewStream(api, buf, conf.BufferSizeOrDefault())
    58  	return &jsonEncoder{stream, buf}
    59  }
    60  
    61  type jsonEncoder struct {
    62  	*jsoniter.Stream
    63  	buf *bufio.Writer
    64  }
    65  
    66  func (e *jsonEncoder) Encode(s core.Sample) error {
    67  	e.WriteVal(s)
    68  	e.WriteRaw("\n")
    69  	return e.Error
    70  }
    71  
    72  func (e *jsonEncoder) Flush() error {
    73  	err := e.Stream.Flush()
    74  	_ = e.buf.Flush()
    75  	return err
    76  }