diploma/client/client.go

package client

import (
	"fmt"
	"net"

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

const (
	metricKeySize = 4
)

type Error struct {
	Code    uint16
	Message string
}

func (s Error) Error() string {
	return fmt.Sprintf("%d: %s", s.Code, s.Message)
}

type Connection struct {
	conn net.Conn
	src  *bufreader.BufferedReader
}

func Connect(address string) (*Connection, error) {
	conn, err := net.Dial("tcp", address)
	if err != nil {
		return nil, err
	}

	return &Connection{
		conn: conn,
		src:  bufreader.New(conn, 1500),
	}, nil
}

func (s *Connection) String() string {
	return s.conn.LocalAddr().String()
}

func (s *Connection) Close() {
	s.conn.Close()
}

func (s *Connection) mustSuccess(reader *bufreader.BufferedReader) (err error) {
	code, err := reader.ReadByte()
	if err != nil {
		return fmt.Errorf("read response code: %s", err)
	}

	switch code {
	case proto.RespSuccess:
		return nil // ok

	case proto.RespError:
		return s.onError()

	default:
		return fmt.Errorf("unknown reponse code %d", code)
	}
}

func (s *Connection) AddMetric(req proto.AddMetricReq) error {
	arr := []byte{
		proto.TypeAddMetric,
		0, 0, 0, 0, //
		byte(req.MetricType),
		byte(req.FracDigits),
	}
	bin.PutUint32(arr[1:], req.MetricID)

	if _, err := s.conn.Write(arr); err != nil {
		return err
	}
	return s.mustSuccess(s.src)
}

func (s *Connection) GetMetric(metricID uint32) (*proto.Metric, error) {
	arr := []byte{
		proto.TypeGetMetric,
		0, 0, 0, 0,
	}
	bin.PutUint32(arr[1:], metricID)

	if _, err := s.conn.Write(arr); err != nil {
		return nil, err
	}

	code, err := s.src.ReadByte()
	if err != nil {
		return nil, fmt.Errorf("read response code: %s", err)
	}

	switch code {
	case proto.RespValue:
		arr, err := s.src.ReadN(6)
		if err != nil {
			return nil, fmt.Errorf("read body: %s", err)
		}

		return &proto.Metric{
			MetricID:   bin.GetUint32(arr),
			MetricType: diploma.MetricType(arr[4]),
			FracDigits: int(arr[5]),
		}, nil

	case proto.RespError:
		return nil, s.onMaybeError()

	default:
		return nil, fmt.Errorf("unknown reponse code %d", code)
	}
}

func (s *Connection) DeleteMetric(metricID uint32) error {
	arr := []byte{
		proto.TypeDeleteMetric,
		0, 0, 0, 0, //
	}
	bin.PutUint32(arr[1:], metricID)

	if _, err := s.conn.Write(arr); err != nil {
		return err
	}
	return s.mustSuccess(s.src)
}

func (s *Connection) AppendMeasure(req proto.AppendMeasureReq) (err error) {
	arr := []byte{
		proto.TypeAppendMeasure,
		0, 0, 0, 0, // metricID
		0, 0, 0, 0, // timestamp
		0, 0, 0, 0, 0, 0, 0, 0, // value
	}
	bin.PutUint32(arr[1:], req.MetricID)
	bin.PutUint32(arr[5:], req.Timestamp)
	bin.PutFloat64(arr[9:], req.Value)

	if _, err := s.conn.Write(arr); err != nil {
		return err
	}
	return s.mustSuccess(s.src)
}

func (s *Connection) AppendMeasures(req proto.AppendMeasuresReq) (err error) {
	if len(req.Measures) > 65535 {
		return fmt.Errorf("wrong measures qty: %d", len(req.Measures))
	}
	var (
		prefixSize = 7
		recordSize = 12
		arr        = make([]byte, prefixSize+len(req.Measures)*recordSize)
	)
	arr[0] = proto.TypeAppendMeasures
	bin.PutUint32(arr[1:], req.MetricID)
	bin.PutUint16(arr[5:], uint16(len(req.Measures)))
	pos := prefixSize
	for _, measure := range req.Measures {
		bin.PutUint32(arr[pos:], measure.Timestamp)
		bin.PutFloat64(arr[pos+4:], measure.Value)
		pos += recordSize
	}
	if _, err := s.conn.Write(arr); err != nil {
		return err
	}
	return s.mustSuccess(s.src)
}

func (s *Connection) ListAllInstantMeasures(metricID uint32) ([]proto.InstantMeasure, error) {
	arr := []byte{
		proto.TypeListAllInstantMeasures,
		0, 0, 0, 0, // metricID
	}
	bin.PutUint32(arr[1:], metricID)

	if _, err := s.conn.Write(arr); err != nil {
		return nil, err
	}

	var (
		result []proto.InstantMeasure
		tmp    = make([]byte, 12)
	)

	for {
		code, err := s.src.ReadByte()
		if err != nil {
			return nil, fmt.Errorf("read response code: %s", err)
		}

		switch code {
		case proto.RespPartOfValue:
			q, err := bin.ReadUint32(s.src)
			if err != nil {
				return nil, fmt.Errorf("read records qty: %s", err)
			}

			for i := range int(q) {
				err = bin.ReadNInto(s.src, tmp)
				if err != nil {
					return nil, fmt.Errorf("read record #%d: %s", i, err)
				}

				result = append(result, proto.InstantMeasure{
					Timestamp: bin.GetUint32(tmp),
					Value:     bin.GetFloat64(tmp[4:]),
				})
			}

		case proto.RespEndOfValue:
			return result, nil

		case proto.RespError:
			return nil, s.onError()

		default:
			return nil, fmt.Errorf("unknown reponse code %d", code)
		}
	}
}

func (s *Connection) ListInstantMeasures(req proto.ListInstantMeasuresReq) ([]proto.InstantMeasure, error) {
	arr := []byte{
		proto.TypeListInstantMeasures,
		0, 0, 0, 0, // metricID
		0, 0, 0, 0, // since
		0, 0, 0, 0, // until
	}
	bin.PutUint32(arr[1:], req.MetricID)
	bin.PutUint32(arr[5:], req.Since)
	bin.PutUint32(arr[9:], req.Until)

	if _, err := s.conn.Write(arr); err != nil {
		return nil, err
	}

	var (
		result []proto.InstantMeasure
		tmp    = make([]byte, 12)
	)

	for {
		code, err := s.src.ReadByte()
		if err != nil {
			return nil, fmt.Errorf("read response code: %s", err)
		}

		switch code {
		case proto.RespPartOfValue:
			q, err := bin.ReadUint32(s.src)
			if err != nil {
				return nil, fmt.Errorf("read records qty: %s", err)
			}

			for i := range int(q) {
				err = bin.ReadNInto(s.src, tmp)
				if err != nil {
					return nil, fmt.Errorf("read record #%d: %s", i, err)
				}

				result = append(result, proto.InstantMeasure{
					Timestamp: bin.GetUint32(tmp),
					Value:     bin.GetFloat64(tmp[4:]),
				})
			}

		case proto.RespEndOfValue:
			return result, nil

		case proto.RespError:
			return nil, s.onError()

		default:
			return nil, fmt.Errorf("unknown reponse code %d", code)
		}
	}
}

func (s *Connection) ListAllCumulativeMeasures(metricID uint32) ([]proto.CumulativeMeasure, error) {
	arr := []byte{
		proto.TypeListAllCumulativeMeasures,
		0, 0, 0, 0, // metricID
	}
	bin.PutUint32(arr[1:], metricID)

	if _, err := s.conn.Write(arr); err != nil {
		return nil, err
	}

	var (
		result []proto.CumulativeMeasure
		tmp    = make([]byte, 20)
	)

	for {
		code, err := s.src.ReadByte()
		if err != nil {
			return nil, fmt.Errorf("read response code: %s", err)
		}

		switch code {
		case proto.RespPartOfValue:
			q, err := bin.ReadUint32(s.src)
			if err != nil {
				return nil, fmt.Errorf("read records qty: %s", err)
			}

			for i := range int(q) {
				err = bin.ReadNInto(s.src, tmp)
				if err != nil {
					return nil, fmt.Errorf("read record #%d: %s", i, err)
				}

				result = append(result, proto.CumulativeMeasure{
					Timestamp: bin.GetUint32(tmp),
					Value:     bin.GetFloat64(tmp[4:]),
					Total:     bin.GetFloat64(tmp[12:]),
				})
			}

		case proto.RespEndOfValue:
			return result, nil

		case proto.RespError:
			return nil, s.onError()

		default:
			return nil, fmt.Errorf("unknown reponse code %d", code)
		}
	}
}

func (s *Connection) ListCumulativeMeasures(req proto.ListCumulativeMeasuresReq) ([]proto.CumulativeMeasure, error) {
	arr := []byte{
		proto.TypeListCumulativeMeasures,
		0, 0, 0, 0, // metricID
		0, 0, 0, 0, // since
		0, 0, 0, 0, // until
	}
	bin.PutUint32(arr[1:], req.MetricID)
	bin.PutUint32(arr[5:], req.Since)
	bin.PutUint32(arr[9:], req.Until)

	if _, err := s.conn.Write(arr); err != nil {
		return nil, err
	}

	var (
		result []proto.CumulativeMeasure
		tmp    = make([]byte, 20)
	)

	for {
		code, err := s.src.ReadByte()
		if err != nil {
			return nil, fmt.Errorf("read response code: %s", err)
		}

		switch code {
		case proto.RespPartOfValue:
			q, err := bin.ReadUint32(s.src)
			if err != nil {
				return nil, fmt.Errorf("read records qty: %s", err)
			}

			for i := range int(q) {
				err = bin.ReadNInto(s.src, tmp)
				if err != nil {
					return nil, fmt.Errorf("read record #%d: %s", i, err)
				}

				result = append(result, proto.CumulativeMeasure{
					Timestamp: bin.GetUint32(tmp),
					Value:     bin.GetFloat64(tmp[4:]),
					Total:     bin.GetFloat64(tmp[12:]),
				})
			}

		case proto.RespEndOfValue:
			return result, nil

		case proto.RespError:
			return nil, s.onError()

		default:
			return nil, fmt.Errorf("unknown reponse code %d", code)
		}
	}
}

func (s *Connection) ListInstantPeriods(req proto.ListInstantPeriodsReq) ([]proto.InstantPeriod, error) {
	arr := []byte{
		proto.TypeListInstantPeriods,
		0, 0, 0, 0, // metricID
		0, 0, 0, 0, // since
		0, 0, 0, 0, // until
		byte(req.GroupBy),
		req.AggregateFuncs,
		byte(req.FirstHourOfDay),
	}
	bin.PutUint32(arr[1:], req.MetricID)
	bin.PutUint16(arr[5:], uint16(req.Since.Year))
	arr[7] = byte(req.Since.Month)
	arr[8] = byte(req.Since.Day)
	bin.PutUint16(arr[9:], uint16(req.Until.Year))
	arr[11] = byte(req.Until.Month)
	arr[12] = byte(req.Until.Day)

	if _, err := s.conn.Write(arr); err != nil {
		return nil, err
	}

	var q int
	if (req.AggregateFuncs & diploma.AggregateMin) == diploma.AggregateMin {
		q++
	}
	if (req.AggregateFuncs & diploma.AggregateMax) == diploma.AggregateMax {
		q++
	}
	if (req.AggregateFuncs & diploma.AggregateAvg) == diploma.AggregateAvg {
		q++
	}

	var (
		result []proto.InstantPeriod
		// 12 bytes - period, since, until
		// q * 8 bytes - min, max, avg
		tmp = make([]byte, 12+q*8)
	)

	for {
		code, err := s.src.ReadByte()
		if err != nil {
			return nil, fmt.Errorf("read response code: %s", err)
		}

		switch code {
		case proto.RespPartOfValue:
			q, err := bin.ReadUint32(s.src)
			if err != nil {
				return nil, fmt.Errorf("read records qty: %s", err)
			}

			for i := range int(q) {
				err = bin.ReadNInto(s.src, tmp)
				if err != nil {
					return nil, fmt.Errorf("read record #%d: %s", i, err)
				}

				var (
					p = proto.InstantPeriod{
						Period: bin.GetUint32(tmp[0:]),
						Start:  bin.GetUint32(tmp[4:]),
						End:    bin.GetUint32(tmp[8:]),
					}
					// 12 bytes - period, since, until
					pos = 12
				)

				if (req.AggregateFuncs & diploma.AggregateMin) == diploma.AggregateMin {
					p.Min = bin.GetFloat64(tmp[pos:])
					pos += 8
				}
				if (req.AggregateFuncs & diploma.AggregateMax) == diploma.AggregateMax {
					p.Max = bin.GetFloat64(tmp[pos:])
					pos += 8
				}
				if (req.AggregateFuncs & diploma.AggregateAvg) == diploma.AggregateAvg {
					p.Avg = bin.GetFloat64(tmp[pos:])
				}
				result = append(result, p)
			}

		case proto.RespEndOfValue:
			return result, nil

		case proto.RespError:
			return nil, s.onError()

		default:
			return nil, fmt.Errorf("unknown reponse code %d", code)
		}
	}
}

func (s *Connection) ListCumulativePeriods(req proto.ListCumulativePeriodsReq) ([]proto.CumulativePeriod, error) {
	arr := []byte{
		proto.TypeListCumulativePeriods,
		0, 0, 0, 0, // metricID
		0, 0, 0, 0, // since
		0, 0, 0, 0, // until
		byte(req.GroupBy),
		byte(req.FirstHourOfDay),
	}
	bin.PutUint32(arr[1:], req.MetricID)
	bin.PutUint16(arr[5:], uint16(req.Since.Year))
	arr[7] = byte(req.Since.Month)
	arr[8] = byte(req.Since.Day)
	bin.PutUint16(arr[9:], uint16(req.Until.Year))
	arr[11] = byte(req.Until.Month)
	arr[12] = byte(req.Until.Day)

	if _, err := s.conn.Write(arr); err != nil {
		return nil, err
	}

	var (
		result []proto.CumulativePeriod
		tmp    = make([]byte, 28)
	)

	for {
		code, err := s.src.ReadByte()
		if err != nil {
			return nil, fmt.Errorf("read response code: %s", err)
		}

		switch code {
		case proto.RespPartOfValue:
			q, err := bin.ReadUint32(s.src)
			if err != nil {
				return nil, fmt.Errorf("read records qty: %s", err)
			}

			for i := range int(q) {
				err = bin.ReadNInto(s.src, tmp)
				if err != nil {
					return nil, fmt.Errorf("read record #%d: %s", i, err)
				}
				result = append(result, proto.CumulativePeriod{
					Period:     bin.GetUint32(tmp[0:]),
					Start:      bin.GetUint32(tmp[4:]),
					End:        bin.GetUint32(tmp[8:]),
					StartValue: bin.GetFloat64(tmp[12:]),
					EndValue:   bin.GetFloat64(tmp[20:]),
				})
			}

		case proto.RespEndOfValue:
			return result, nil

		case proto.RespError:
			return nil, s.onError()

		default:
			return nil, fmt.Errorf("unknown reponse code %d", code)
		}
	}
}

func (s *Connection) ListCurrentValues(metricIDs []uint32) ([]proto.CurrentValue, error) {
	arr := make([]byte, 3+metricKeySize*len(metricIDs))
	arr[0] = proto.TypeListCurrentValues

	bin.PutUint16(arr[1:], uint16(len(metricIDs)))

	off := 3
	for _, metricID := range metricIDs {
		bin.PutUint32(arr[off:], metricID)
		off += metricKeySize
	}

	if _, err := s.conn.Write(arr); err != nil {
		return nil, err
	}

	var (
		result []proto.CurrentValue
		tmp    = make([]byte, 16)
	)

	for {
		code, err := s.src.ReadByte()
		if err != nil {
			return nil, fmt.Errorf("read response code: %s", err)
		}

		switch code {
		case proto.RespPartOfValue:
			q, err := bin.ReadUint32(s.src)
			if err != nil {
				return nil, fmt.Errorf("read records qty: %s", err)
			}

			for i := range int(q) {
				err = bin.ReadNInto(s.src, tmp)
				if err != nil {
					return nil, fmt.Errorf("read record #%d: %s", i, err)
				}

				result = append(result, proto.CurrentValue{
					MetricID:  bin.GetUint32(tmp),
					Timestamp: bin.GetUint32(tmp[4:]),
					Value:     bin.GetFloat64(tmp[8:]),
				})
			}

		case proto.RespEndOfValue:
			return result, nil

		case proto.RespError:
			return nil, s.onError()

		default:
			return nil, fmt.Errorf("unknown reponse code %d", code)
		}
	}
}

func (s *Connection) DeleteMeasures(req proto.DeleteMeasuresReq) (err error) {
	arr := []byte{
		proto.TypeDeleteMeasures,
		0, 0, 0, 0, // metricID
		0, 0, 0, 0, // since
	}
	bin.PutUint32(arr[1:], req.MetricID)
	bin.PutUint32(arr[5:], req.Since)

	if _, err := s.conn.Write(arr); err != nil {
		return err
	}
	return s.mustSuccess(s.src)
}

func (s *Connection) onError() error {
	errorCode, err := bin.ReadUint16(s.src)
	if err != nil {
		return fmt.Errorf("read error code: %s", err)
	}
	return Error{
		Code:    errorCode,
		Message: proto.ErrorCodeToText(errorCode),
	}
}

func (s *Connection) onMaybeError() error {
	errorCode, err := bin.ReadUint16(s.src)
	if err != nil {
		return fmt.Errorf("read error code: %s", err)
	}
	if errorCode == proto.ErrNoMetric {
		return nil
	}
	return Error{
		Code:    errorCode,
		Message: proto.ErrorCodeToText(errorCode),
	}
}