diploma/atree/writers.go

package atree

import (
	"bytes"
	"fmt"
	"io"

	octopus "gordenko.dev/dima/diploma"
	"gordenko.dev/dima/diploma/bin"
	"gordenko.dev/dima/diploma/proto"
)

// CURRENT VALUE WRITER

type CurrentValue struct {
	MetricID  uint32
	Timestamp uint32
	Value     float64
}

type CurrentValueWriter struct {
	arr       []byte
	responder *ChunkedResponder
}

func NewCurrentValueWriter(dst io.Writer) *CurrentValueWriter {
	return &CurrentValueWriter{
		arr:       make([]byte, 16),
		responder: NewChunkedResponder(dst),
	}
}

func (s *CurrentValueWriter) BufferValue(m CurrentValue) {
	bin.PutUint32(s.arr[0:], m.MetricID)
	bin.PutUint32(s.arr[4:], m.Timestamp)
	bin.PutFloat64(s.arr[8:], m.Value)
	s.responder.BufferRecord(s.arr)
}

func (s *CurrentValueWriter) Close() error {
	return s.responder.Flush()
}

// INSTANT MEASURE WRITER

type InstantMeasure struct {
	Timestamp uint32
	Value     float64
}

type InstantMeasureWriter struct {
	arr       []byte
	responder *ChunkedResponder
}

func NewInstantMeasureWriter(dst io.Writer) *InstantMeasureWriter {
	return &InstantMeasureWriter{
		arr:       make([]byte, 12),
		responder: NewChunkedResponder(dst),
	}
}

func (s *InstantMeasureWriter) BufferMeasure(m InstantMeasure) {
	bin.PutUint32(s.arr[0:], m.Timestamp)
	bin.PutFloat64(s.arr[4:], m.Value)
	s.responder.BufferRecord(s.arr)
}

func (s *InstantMeasureWriter) WriteMeasure(m InstantMeasure) error {
	bin.PutUint32(s.arr[0:], m.Timestamp)
	bin.PutFloat64(s.arr[4:], m.Value)
	return s.responder.AppendRecord(s.arr)
}

func (s *InstantMeasureWriter) Close() error {
	return s.responder.Flush()
}

// CUMULATIVE MEASURE WRITER

type CumulativeMeasure struct {
	Timestamp uint32
	Value     float64
	Total     float64
}

type CumulativeMeasureWriter struct {
	arr       []byte
	responder *ChunkedResponder
}

func NewCumulativeMeasureWriter(dst io.Writer) *CumulativeMeasureWriter {
	return &CumulativeMeasureWriter{
		arr:       make([]byte, 20),
		responder: NewChunkedResponder(dst),
	}
}

func (s *CumulativeMeasureWriter) BufferMeasure(m CumulativeMeasure) {
	bin.PutUint32(s.arr[0:], m.Timestamp)
	bin.PutFloat64(s.arr[4:], m.Value)
	bin.PutFloat64(s.arr[12:], m.Total)
	s.responder.BufferRecord(s.arr)
}

func (s *CumulativeMeasureWriter) WriteMeasure(m CumulativeMeasure) error {
	bin.PutUint32(s.arr[0:], m.Timestamp)
	bin.PutFloat64(s.arr[4:], m.Value)
	bin.PutFloat64(s.arr[12:], m.Total)
	return s.responder.AppendRecord(s.arr)
}

func (s *CumulativeMeasureWriter) Close() error {
	return s.responder.Flush()
}

// INSTANT AGGREGATE WRITER

type InstantPeriodsWriter struct {
	aggregateFuncs byte
	arr            []byte
	responder      *ChunkedResponder
}

func NewInstantPeriodsWriter(dst io.Writer, aggregateFuncs byte) *InstantPeriodsWriter {
	var q int
	if (aggregateFuncs & octopus.AggregateMin) == octopus.AggregateMin {
		q++
	}
	if (aggregateFuncs & octopus.AggregateMax) == octopus.AggregateMax {
		q++
	}
	if (aggregateFuncs & octopus.AggregateAvg) == octopus.AggregateAvg {
		q++
	}
	return &InstantPeriodsWriter{
		aggregateFuncs: aggregateFuncs,
		arr:            make([]byte, 12+q*8),
		responder:      NewChunkedResponder(dst),
	}
}

type InstantPeriod struct {
	Period uint32
	Since  uint32
	Until  uint32
	Min    float64
	Max    float64
	Avg    float64
}

func (s *InstantPeriodsWriter) BufferMeasure(p InstantPeriod) {
	s.pack(p)
	s.responder.BufferRecord(s.arr)
}

func (s *InstantPeriodsWriter) WritePeriod(p InstantPeriod) error {
	s.pack(p)
	return s.responder.AppendRecord(s.arr)
}

func (s *InstantPeriodsWriter) Close() error {
	return s.responder.Flush()
}

func (s *InstantPeriodsWriter) pack(p InstantPeriod) {
	bin.PutUint32(s.arr[0:], p.Period)
	bin.PutUint32(s.arr[4:], p.Since)
	bin.PutUint32(s.arr[8:], p.Until)

	pos := 12
	if (s.aggregateFuncs & octopus.AggregateMin) == octopus.AggregateMin {
		bin.PutFloat64(s.arr[pos:], p.Min)
		pos += 8
	}
	if (s.aggregateFuncs & octopus.AggregateMax) == octopus.AggregateMax {
		bin.PutFloat64(s.arr[pos:], p.Max)
		pos += 8
	}
	if (s.aggregateFuncs & octopus.AggregateAvg) == octopus.AggregateAvg {
		bin.PutFloat64(s.arr[pos:], p.Avg)
	}
}

// CUMULATIVE AGGREGATE WRITER

type CumulativePeriodsWriter struct {
	arr       []byte
	responder *ChunkedResponder
}

func NewCumulativePeriodsWriter(dst io.Writer) *CumulativePeriodsWriter {
	return &CumulativePeriodsWriter{
		arr:       make([]byte, 28),
		responder: NewChunkedResponder(dst),
	}
}

type CumulativePeriod struct {
	Period   uint32
	Since    uint32
	Until    uint32
	EndValue float64
	Total    float64
}

func (s *CumulativePeriodsWriter) BufferMeasure(p CumulativePeriod) {
	s.pack(p)
	s.responder.BufferRecord(s.arr)
}

func (s *CumulativePeriodsWriter) WritePeriod(p CumulativePeriod) error {
	s.pack(p)
	return s.responder.AppendRecord(s.arr)
}

func (s *CumulativePeriodsWriter) Close() error {
	return s.responder.Flush()
}

func (s *CumulativePeriodsWriter) pack(p CumulativePeriod) {
	bin.PutUint32(s.arr[0:], p.Period)
	bin.PutUint32(s.arr[4:], p.Since)
	bin.PutUint32(s.arr[8:], p.Until)
	bin.PutFloat64(s.arr[12:], p.EndValue)
	bin.PutFloat64(s.arr[20:], p.Total)
}

// CHUNKED RESPONDER

//const headerSize = 3

var endMsg = []byte{
	proto.RespEndOfValue, // end of stream
}

type ChunkedResponder struct {
	recordsQty int
	buf        *bytes.Buffer
	dst        io.Writer
}

func NewChunkedResponder(dst io.Writer) *ChunkedResponder {
	s := &ChunkedResponder{
		recordsQty: 0,
		buf:        bytes.NewBuffer(nil),
		dst:        dst,
	}

	s.buf.Write([]byte{
		proto.RespPartOfValue, // message type
		0, 0, 0, 0,            // records qty
	})
	return s
}

func (s *ChunkedResponder) BufferRecord(rec []byte) {
	s.buf.Write(rec)
	s.recordsQty++
}

func (s *ChunkedResponder) AppendRecord(rec []byte) error {
	s.buf.Write(rec)
	s.recordsQty++

	if s.buf.Len() < 1500 {
		return nil
	}

	if err := s.sendBuffered(); err != nil {
		return err
	}

	s.buf.Write([]byte{
		proto.RespPartOfValue, // message type
		0, 0, 0, 0,            // records qty
	})
	s.recordsQty = 0
	return nil
}

func (s *ChunkedResponder) Flush() error {
	if s.recordsQty > 0 {
		if err := s.sendBuffered(); err != nil {
			return err
		}
	}
	if _, err := s.dst.Write(endMsg); err != nil {
		return err
	}
	return nil
}

func (s *ChunkedResponder) sendBuffered() (err error) {
	msg := s.buf.Bytes()
	bin.PutUint32(msg[1:], uint32(s.recordsQty))
	n, err := s.dst.Write(msg)
	if err != nil {
		return
	}
	if n != len(msg) {
		return fmt.Errorf("incomplete write %d bytes instead of %d", n, len(msg))
	}
	s.buf.Reset()
	return
}