gocommon/metric/metric.go

package metric

import (
	"crypto/md5"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"math"
	"os"
	"path"
	"strings"
	"sync/atomic"

	"repositories.action2quare.com/ayo/gocommon/logger"
)

const (
	METRIC_HEAD_INLINE = byte(14)
	METRIC_TAIL_INLINE = byte(15)
)

type MetricType int

const (
	MetricCounter   = MetricType(1)
	MetricGuage     = MetricType(2)
	metric_key_size = 8
)

type MetricDescription struct {
	Key         string
	Type        MetricType
	Name        string            `json:",omitempty"`
	Help        string            `json:",omitempty"`
	ConstLabels map[string]string `json:",omitempty"`
}

type writeRequest struct {
	key     string
	valfunc func() float64
}

type metricCollection struct {
	writerChan chan *writeRequest
}

var mc = metricCollection{
	writerChan: make(chan *writeRequest, 100),
}

type MetricWriter interface {
	Add(int64)
	Set(int64)
}

type metric_empty struct{}

func (mw *metric_empty) Set(newval int64) {}
func (mw *metric_empty) Add(inc int64)    {}

var MetricWriterNil = MetricWriter(&metric_empty{})

type metric_int64 struct {
	valptr     *int64
	key        string
	writerChan chan *writeRequest
}

func (mw *metric_int64) requestMetricWrite() {
	mw.writerChan <- &writeRequest{
		key:     mw.key,
		valfunc: func() float64 { return float64(atomic.LoadInt64(mw.valptr)) },
	}
}

func (mw *metric_int64) Set(newval int64) {
	atomic.StoreInt64(mw.valptr, newval)
	mw.requestMetricWrite()
}

func (mw *metric_int64) Add(inc int64) {
	atomic.AddInt64(mw.valptr, inc)
	mw.requestMetricWrite()
}

func (mc *metricCollection) metricWriter() {
	defer func() {
		r := recover()
		if r != nil {
			logger.Error(r)
		}
	}()

	// head + metric_key_size + 8byte + tail + cr = 19
	var buff [19]byte
	buff[0] = METRIC_HEAD_INLINE
	buff[17] = METRIC_TAIL_INLINE
	buff[18] = '\n'

	for req := range mc.writerChan {
		copy(buff[1:], []byte(req.key))
		binary.BigEndian.PutUint64(buff[9:], math.Float64bits(req.valfunc()))
		os.Stdout.Write(buff[:])
	}
}

var NewMetric func(MetricType, string, string, map[string]string) MetricWriter
var ConcurrentUser = MetricWriterNil

func init() {
	NewMetric = func(MetricType, string, string, map[string]string) MetricWriter {
		return &metric_empty{}
	}

	if path.Base(os.Args[0]) == "houston" {
		logger.Println("metrics are going to be generated for myself(houston)")
		go mc.metricWriter()
		NewMetric = newMetricImpl
		return
	}

	ppid := os.Getppid()
	if parent, _ := os.FindProcess(ppid); parent != nil {
		filename := fmt.Sprintf(`/proc/%d/stat`, os.Getppid())
		if fn, err := os.ReadFile(filename); err == nil {
			stats := strings.SplitN(string(fn), " ", 3)
			parentname := strings.Trim(stats[1], "()")

			if path.Base(parentname) == "houston" {
				logger.Println("metrics are going to be generated for houston")
				go mc.metricWriter()
				NewMetric = newMetricImpl
				ConcurrentUser = NewMetric(MetricGuage, "concurrent_user", "concurrent user count", nil)
			} else {
				logger.Println("metrics are NOT going to be generated. parent is not houston :", filename, string(fn))
			}
		} else {
			logger.Println("metrics are NOT going to be generated. ppid proc is missing :", filename)
		}
	} else {
		logger.Println("metrics are NOT going to be generated. parent process is missing. ppid :", ppid)
	}
}

func newMetricImpl(mt MetricType, name string, help string, constLabels map[string]string) (writer MetricWriter) {
	hash := md5.New()
	hash.Write([]byte(name))

	if constLabels == nil {
		constLabels = make(map[string]string)
	}

	hn, _ := os.Hostname()
	sn := path.Base(os.Args[0])
	constLabels["hostname"] = hn
	constLabels["service"] = sn

	key := hex.EncodeToString(hash.Sum(nil))[:metric_key_size]
	temp, _ := json.Marshal(MetricDescription{
		Key:         key,
		Type:        mt,
		Name:        name,
		Help:        help,
		ConstLabels: constLabels,
	})

	writer = &metric_int64{
		valptr:     new(int64),
		key:        key,
		writerChan: mc.writerChan,
	}

	output := append([]byte{METRIC_HEAD_INLINE}, temp...)
	output = append(output, METRIC_TAIL_INLINE, '\n')
	os.Stdout.Write(output)

	return
}