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

     1  package netsample
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"io"
     7  	"os"
     8  	"strconv"
     9  	"time"
    10  
    11  	"github.com/c2h5oh/datasize"
    12  	"github.com/pkg/errors"
    13  	"github.com/spf13/afero"
    14  	"github.com/yandex/pandora/core"
    15  	"github.com/yandex/pandora/core/coreutil"
    16  )
    17  
    18  type PhoutConfig struct {
    19  	Destination     string                    // Destination file name
    20  	ID              bool                      // Print ammo ids if true.
    21  	FlushTime       time.Duration             `config:"flush-time"`
    22  	SampleQueueSize int                       `config:"sample-queue-size"`
    23  	Buffer          coreutil.BufferSizeConfig `config:",squash"`
    24  }
    25  
    26  func DefaultPhoutConfig() PhoutConfig {
    27  	return PhoutConfig{
    28  		FlushTime:       time.Second,
    29  		SampleQueueSize: 256 * 1024,
    30  		Buffer: coreutil.BufferSizeConfig{
    31  			BufferSize: 8 * datasize.MB,
    32  		},
    33  	}
    34  }
    35  
    36  func NewPhout(fs afero.Fs, conf PhoutConfig) (a Aggregator, err error) {
    37  	filename := conf.Destination
    38  	var file afero.File = os.Stdout
    39  	if filename != "" {
    40  		file, err = fs.Create(conf.Destination)
    41  	}
    42  	if err != nil {
    43  		err = errors.Wrap(err, "phout output file open failed")
    44  		return
    45  	}
    46  	a = &phoutAggregator{
    47  		config: conf,
    48  		sink:   make(chan *Sample, conf.SampleQueueSize),
    49  		writer: bufio.NewWriterSize(file, conf.Buffer.BufferSizeOrDefault()),
    50  		buf:    make([]byte, 0, 1024),
    51  		file:   file,
    52  	}
    53  	return
    54  }
    55  
    56  type phoutAggregator struct {
    57  	config PhoutConfig
    58  	sink   chan *Sample
    59  	writer *bufio.Writer
    60  	buf    []byte
    61  	file   io.Closer
    62  }
    63  
    64  func (a *phoutAggregator) Report(s *Sample) { a.sink <- s }
    65  
    66  func (a *phoutAggregator) Run(ctx context.Context, _ core.AggregatorDeps) error {
    67  	shouldFlush := time.NewTicker(1 * time.Second)
    68  	defer func() {
    69  		_ = a.writer.Flush()
    70  		_ = a.file.Close()
    71  		shouldFlush.Stop()
    72  	}()
    73  loop:
    74  	for {
    75  		select {
    76  		case r := <-a.sink:
    77  			if err := a.handle(r); err != nil {
    78  				return err
    79  			}
    80  			select {
    81  			case <-shouldFlush.C:
    82  				_ = a.writer.Flush()
    83  			default:
    84  			}
    85  		case <-time.After(1 * time.Second):
    86  			_ = a.writer.Flush()
    87  		case <-ctx.Done():
    88  			// Context is done, but we should read all data from sink
    89  			for {
    90  				select {
    91  				case r := <-a.sink:
    92  					if err := a.handle(r); err != nil {
    93  						return err
    94  					}
    95  				default:
    96  					break loop
    97  				}
    98  			}
    99  		}
   100  	}
   101  	return nil
   102  }
   103  
   104  func (a *phoutAggregator) handle(s *Sample) error {
   105  	a.buf = appendPhout(s, a.buf, a.config.ID)
   106  	a.buf = append(a.buf, '\n')
   107  	_, err := a.writer.Write(a.buf)
   108  	a.buf = a.buf[:0]
   109  	releaseSample(s)
   110  	return err
   111  }
   112  
   113  const phoutDelimiter = '\t'
   114  
   115  func appendPhout(s *Sample, dst []byte, id bool) []byte {
   116  	dst = appendTimestamp(s.timeStamp, dst)
   117  	dst = append(dst, phoutDelimiter)
   118  	dst = append(dst, s.tags...)
   119  	if id {
   120  		dst = append(dst, '#')
   121  		dst = strconv.AppendInt(dst, int64(s.ID()), 10)
   122  	}
   123  	for _, v := range s.fields {
   124  		dst = append(dst, phoutDelimiter)
   125  		dst = strconv.AppendInt(dst, int64(v), 10)
   126  	}
   127  	return dst
   128  }
   129  
   130  func appendTimestamp(ts time.Time, dst []byte) []byte {
   131  	// Append time stamp in phout format. Example: 1335524833.562
   132  	// Algorithm: append milliseconds string, than insert dot in right place.
   133  	dst = strconv.AppendInt(dst, ts.UnixNano()/1e6, 10)
   134  	dotIndex := len(dst) - 3
   135  	dst = append(dst, 0)
   136  	for i := len(dst) - 1; i > dotIndex; i-- {
   137  		dst[i] = dst[i-1]
   138  	}
   139  	dst[dotIndex] = '.'
   140  	return dst
   141  }