cc-metric-store/lineprotocol.go

package main

import (
	"bufio"
	"context"
	"errors"
	"io"
	"log"
	"math"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/nats-io/nats.go"
)

// Go's JSON encoder for floats does not support NaN (https://github.com/golang/go/issues/3480).
// This program uses NaN as a signal for missing data.
// For the HTTP JSON API to be able to handle NaN values,
// we have to use our own type which implements encoding/json.Marshaler itself.
type Float float64

var NaN Float = Float(math.NaN())

func (f Float) IsNaN() bool {
	return math.IsNaN(float64(f))
}

func (f Float) MarshalJSON() ([]byte, error) {
	if math.IsNaN(float64(f)) {
		return []byte("null"), nil
	}

	return []byte(strconv.FormatFloat(float64(f), 'f', 2, 64)), nil
}

func (f *Float) UnmarshalJSON(input []byte) error {
	s := string(input)
	if s == "null" {
		*f = NaN
		return nil
	}

	val, err := strconv.ParseFloat(s, 64)
	if err != nil {
		return err
	}
	*f = Float(val)
	return nil
}

type Metric struct {
	Name  string
	Value Float
}

// measurement: node or cpu
// tags:        host, cluster, cpu (cpu only if measurement is cpu)
// fields:      metrics...
// t:           timestamp (accuracy: seconds)
type Line struct {
	Measurement string
	Tags        map[string]string
	Fields      []Metric
	Ts          time.Time
}

// Parse a single line as string.
//
// There is performance to be gained by implementing a parser
// that directly reads from a bufio.Scanner.
func ParseLine(rawline string) (*Line, error) {
	line := &Line{}
	parts := strings.Fields(rawline)
	if len(parts) != 3 {
		return nil, errors.New("line format error")
	}

	tagsAndMeasurement := strings.Split(parts[0], ",")
	line.Measurement = tagsAndMeasurement[0]
	line.Tags = map[string]string{}
	for i := 1; i < len(tagsAndMeasurement); i++ {
		pair := strings.Split(tagsAndMeasurement[i], "=")
		if len(pair) != 2 {
			return nil, errors.New("line format error")
		}
		line.Tags[pair[0]] = pair[1]
	}

	rawfields := strings.Split(parts[1], ",")
	line.Fields = []Metric{}
	for i := 0; i < len(rawfields); i++ {
		pair := strings.Split(rawfields[i], "=")
		if len(pair) != 2 {
			return nil, errors.New("line format error")
		}
		field, err := strconv.ParseFloat(pair[1], 64)
		if err != nil {
			return nil, err
		}

		line.Fields = append(line.Fields, Metric{
			Name:  pair[0],
			Value: Float(field),
		})
	}

	unixTimestamp, err := strconv.ParseInt(parts[2], 10, 64)
	if err != nil {
		return nil, err
	}

	line.Ts = time.Unix(unixTimestamp, 0)
	return line, nil
}

func ParseLines(raw string) ([]*Line, error) {
	lines := make([]*Line, 0, 1)
	for _, line := range strings.Split(raw, "\n") {
		if len(line) == 0 {
			continue
		}

		line, err := ParseLine(line)
		if err != nil {
			return nil, err
		}
		lines = append(lines, line)
	}
	return lines, nil
}

// Listen for connections sending metric data in the line protocol format.
//
// This is a blocking function, send `true` through the channel argument to shut down the server.
// `handleLine` will be called from different go routines for different connections.
//
func ReceiveTCP(address string, handleLine func(line *Line), done chan bool) error {
	ln, err := net.Listen("tcp", address)
	if err != nil {
		return err
	}

	handleConnection := func(conn net.Conn, handleLine func(line *Line)) {
		reader := bufio.NewReader(conn)
		for {
			rawline, err := reader.ReadString('\n')
			if err == io.EOF {
				return
			}

			if err != nil {
				log.Printf("reading from connection failed: %s\n", err.Error())
				return
			}

			line, err := ParseLine(rawline)
			if err != nil {
				log.Printf("parsing line failed: %s\n", err.Error())
				return
			}

			handleLine(line)
		}
	}

	go func() {
		for {
			stop := <-done
			if stop {
				err := ln.Close()
				if err != nil {
					log.Printf("closing listener failed: %s\n", err.Error())
				}
				return
			}
		}
	}()

	for {
		conn, err := ln.Accept()
		if err != nil {
			return err
		}

		go handleConnection(conn, handleLine)
	}
}

// Connect to a nats server and subscribe to "updates". This is a blocking
// function. handleLine will be called for each line recieved via nats.
// Send `true` through the done channel for gracefull termination.
func ReceiveNats(address string, handleLine func(line *Line), workers int, ctx context.Context) error {
	nc, err := nats.Connect(nats.DefaultURL)
	if err != nil {
		return err
	}
	defer nc.Close()

	var sub *nats.Subscription

	if workers < 2 {
		sub, err = nc.Subscribe("updates", func(m *nats.Msg) {
			lines, err := ParseLines(string(m.Data))
			if err != nil {
				log.Println(err.Error())
			}

			for _, line := range lines {
				handleLine(line)
			}
		})
		if err != nil {
			return err
		}

		log.Printf("NATS subscription to 'updates' on '%s' established\n", address)

		_ = <-ctx.Done()
		err = sub.Unsubscribe()
	} else {
		msgs := make(chan *nats.Msg, 16)
		var wg sync.WaitGroup
		wg.Add(workers)

		for i := 0; i < workers; i++ {
			go func() {
				for m := range msgs {
					lines, err := ParseLines(string(m.Data))
					if err != nil {
						log.Println(err.Error())
					}

					for _, line := range lines {
						handleLine(line)
					}
				}

				wg.Done()
			}()
		}

		sub, err = nc.Subscribe("updates", func(m *nats.Msg) {
			msgs <- m
		})
		if err != nil {
			return err
		}

		log.Printf("NATS subscription to 'updates' on '%s' established\n", address)

		_ = <-ctx.Done()
		err = sub.Unsubscribe()
		close(msgs)
		wg.Wait()
	}

	if err != nil {
		return err
	}

	nc.Close()
	log.Println("NATS connection closed")
	return nil
}