mtg/stats/statsd.go

package stats

import (
	"fmt"
	"strings"
	"time"

	"github.com/9seconds/mtg/v2/events"
	"github.com/9seconds/mtg/v2/logger"
	"github.com/9seconds/mtg/v2/mtglib"
	statsd "github.com/smira/go-statsd"
)

type statsdProcessor struct {
	streams map[string]*streamInfo
	client  *statsd.Client
}

func (s statsdProcessor) EventStart(evt mtglib.EventStart) {
	sInfo := &streamInfo{
		createdAt: evt.CreatedAt,
		clientIP:  evt.RemoteIP,
	}
	s.streams[evt.StreamID()] = sInfo

	s.client.GaugeDelta(MetricClientConnections,
		1,
		statsd.StringTag(TagIPType, sInfo.GetClientIPType()))
}

func (s statsdProcessor) EventConnectedToDC(evt mtglib.EventConnectedToDC) {
	sInfo, ok := s.streams[evt.StreamID()]
	if !ok {
		return
	}

	sInfo.remoteIP = evt.RemoteIP
	sInfo.dc = evt.DC

	s.client.GaugeDelta(MetricTelegramConnections,
		1,
		statsd.StringTag(TagIPType, sInfo.GetRemoteIPType()),
		statsd.StringTag(TagTelegramIP, sInfo.remoteIP.String()),
		statsd.IntTag(TagDC, sInfo.dc))
}

func (s statsdProcessor) EventTraffic(evt mtglib.EventTraffic) {
	sInfo, ok := s.streams[evt.StreamID()]
	if !ok {
		return
	}

	tags := []statsd.Tag{
		statsd.StringTag(TagIPType, sInfo.GetRemoteIPType()),
		statsd.StringTag(TagTelegramIP, sInfo.remoteIP.String()),
		statsd.IntTag(TagDC, sInfo.dc),
	}

	if evt.IsRead {
		tags = append(tags, statsd.StringTag(TagDirection, TagDirectionClient))
		s.client.Incr(MetricTraffic, int64(evt.Traffic), tags...)
	} else {
		tags = append(tags, statsd.StringTag(TagDirection, TagDirectionTelegram))
		s.client.Incr(MetricTraffic, int64(evt.Traffic), tags...)
	}
}

func (s statsdProcessor) EventFinish(evt mtglib.EventFinish) {
	sInfo, ok := s.streams[evt.StreamID()]
	if !ok {
		return
	}

	defer delete(s.streams, evt.StreamID())

	s.client.GaugeDelta(MetricClientConnections,
		-1,
		statsd.StringTag(TagIPType, sInfo.GetClientIPType()))
	s.client.PrecisionTiming(MetricSessionDuration,
		evt.CreatedAt.Sub(sInfo.createdAt))

	if sInfo.remoteIP != nil {
		s.client.GaugeDelta(MetricTelegramConnections,
			-1,
			statsd.StringTag(TagIPType, sInfo.GetRemoteIPType()),
			statsd.StringTag(TagTelegramIP, sInfo.remoteIP.String()),
			statsd.IntTag(TagDC, sInfo.dc))
	}
}

func (s statsdProcessor) EventConcurrencyLimited(_ mtglib.EventConcurrencyLimited) {
	s.client.Incr(MetricConcurrencyLimited, 1)
}

func (s statsdProcessor) EventIPBlocklisted(evt mtglib.EventIPBlocklisted) {
	var tag statsd.Tag

	if evt.RemoteIP.To4() == nil {
		tag = statsd.StringTag(TagIPType, TagIPTypeIPv6)
	} else {
		tag = statsd.StringTag(TagIPType, TagIPTypeIPv4)
	}

	s.client.Incr(MetricIPBlocklisted, 1, tag)
}

func (s statsdProcessor) Shutdown() {
	now := time.Now()
	events := make([]mtglib.EventFinish, 0, len(s.streams))

	for k := range s.streams {
		events = append(events, mtglib.EventFinish{
			CreatedAt: now,
			ConnID:    k,
		})
	}

	for i := range events {
		s.EventFinish(events[i])
	}
}

type StatsdFactory struct {
	client *statsd.Client
}

func (s StatsdFactory) Close() error {
	return s.client.Close()
}

func (s StatsdFactory) Make() events.Observer {
	return statsdProcessor{
		client:  s.client,
		streams: make(map[string]*streamInfo),
	}
}

func NewStatsd(address string, log logger.StdLikeLogger,
	metricPrefix, tagFormat string) (StatsdFactory, error) {
	options := []statsd.Option{
		statsd.MetricPrefix(metricPrefix),
		statsd.Logger(log),
	}

	switch strings.ToLower(tagFormat) {
	case "datadog":
		options = append(options, statsd.TagStyle(statsd.TagFormatDatadog))
	case "influxdb":
		options = append(options, statsd.TagStyle(statsd.TagFormatInfluxDB))
	case "graphite":
		options = append(options, statsd.TagStyle(statsd.TagFormatGraphite))
	default:
		return StatsdFactory{}, fmt.Errorf("unknown tag format %s", tagFormat)
	}

	return StatsdFactory{
		client: statsd.NewClient(address, options...),
	}, nil
}