Add prometheus

2026-08-31 12:44:02 +03:00 · 2021-03-17 12:01:28 +03:00
parent 8e7207d975
commit dbbfdaf918
8 changed files with 627 additions and 23 deletions
@@ -0,0 +1,123 @@
+package stats
+
+import (
+	"context"
+	"net"
+	"net/http"
+	"time"
+
+	"github.com/9seconds/mtg/v2/events"
+	"github.com/9seconds/mtg/v2/mtglib"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+)
+
+type prometheusProcessor struct {
+	streams map[string]*streamInfo
+	factory *PrometheusFactory
+}
+
+func (p prometheusProcessor) EventStart(evt mtglib.EventStart) {
+	sInfo := &streamInfo{
+		createdAt: evt.CreatedAt,
+		clientIP:  evt.RemoteIP,
+	}
+	p.streams[evt.StreamID()] = sInfo
+
+	p.factory.metricActiveConnections.WithLabelValues(sInfo.IPType()).Inc()
+}
+
+func (p prometheusProcessor) EventFinish(evt mtglib.EventFinish) {
+	sInfo, ok := p.streams[evt.StreamID()]
+	if !ok {
+		return
+	}
+
+	defer delete(p.streams, evt.StreamID())
+
+	duration := evt.CreatedAt.Sub(sInfo.createdAt)
+
+	p.factory.metricActiveConnections.WithLabelValues(sInfo.IPType()).Dec()
+	p.factory.metricSessionDuration.Observe(float64(duration) / float64(time.Second))
+}
+
+func (p prometheusProcessor) EventConcurrencyLimited(evt mtglib.EventConcurrencyLimited) {
+	p.factory.metricConcurrencyLimited.Inc()
+}
+
+func (p prometheusProcessor) Shutdown() {
+	p.streams = make(map[string]*streamInfo)
+}
+
+type PrometheusFactory struct {
+	httpServer *http.Server
+
+	metricActiveConnections  *prometheus.GaugeVec
+	metricConcurrencyLimited prometheus.Counter
+	metricSessionDuration    prometheus.Histogram
+}
+
+func (p *PrometheusFactory) Make() events.Observer {
+	return prometheusProcessor{
+		streams: make(map[string]*streamInfo),
+		factory: p,
+	}
+}
+
+func (p *PrometheusFactory) Serve(listener net.Listener) error {
+	return p.httpServer.Serve(listener)
+}
+
+func (p *PrometheusFactory) Close() error {
+	return p.httpServer.Shutdown(context.Background())
+}
+
+func NewPrometheus(metricPrefix, httpPath string) *PrometheusFactory {
+	registry := prometheus.NewPedanticRegistry()
+	httpHandler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{
+		EnableOpenMetrics: true,
+	})
+	mux := http.NewServeMux()
+
+	mux.Handle(httpPath, httpHandler)
+
+	factory := &PrometheusFactory{
+		httpServer: &http.Server{
+			Handler: mux,
+		},
+
+		metricActiveConnections: prometheus.NewGaugeVec(prometheus.GaugeOpts{
+			Namespace: metricPrefix,
+			Name:      MetricActiveConnection,
+			Help:      "A number of connections under active processing.",
+		}, []string{TagIPType}),
+		metricSessionDuration: prometheus.NewHistogram(prometheus.HistogramOpts{
+			Namespace: metricPrefix,
+			Name:      MetricSessionDuration,
+			Help:      "Session duration.",
+			Buckets: []float64{ // per 30 seconds
+				30,
+				60,
+				90,
+				120,
+				150,
+				180,
+				210,
+				240,
+				270,
+				300,
+			},
+		}),
+		metricConcurrencyLimited: prometheus.NewCounter(prometheus.CounterOpts{
+			Namespace: metricPrefix,
+			Name:      MetricConcurrencyLimited,
+			Help:      "A number of sessions that were rejected by concurrency limiter.",
+		}),
+	}
+
+	registry.MustRegister(factory.metricActiveConnections)
+	registry.MustRegister(factory.metricSessionDuration)
+	registry.MustRegister(factory.metricConcurrencyLimited)
+
+	return factory
+}
@@ -0,0 +1,97 @@
+package stats_test
+
+import (
+	"fmt"
+	"io/ioutil"
+	"net"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/9seconds/mtg/v2/events"
+	"github.com/9seconds/mtg/v2/mtglib"
+	"github.com/9seconds/mtg/v2/stats"
+	"github.com/stretchr/testify/suite"
+)
+
+type PrometheusTestSuite struct {
+	suite.Suite
+
+	httpListener net.Listener
+	factory      *stats.PrometheusFactory
+	prometheus   events.Observer
+}
+
+func (suite *PrometheusTestSuite) Get() (string, error) {
+	addr := fmt.Sprintf("http://%s/", suite.httpListener.Addr().String())
+
+	resp, err := http.Get(addr) // nolint: noctx
+	if err != nil {
+		return "", err // nolint: wrapcheck
+	}
+
+	defer resp.Body.Close()
+
+	data, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return "", err // nolint: wrapcheck
+	}
+
+	return string(data), nil
+}
+
+func (suite *PrometheusTestSuite) SetupTest() {
+	suite.httpListener, _ = net.Listen("tcp", "127.0.0.1:0")
+	suite.factory = stats.NewPrometheus("mtg", "/")
+	suite.prometheus = suite.factory.Make()
+
+	go suite.factory.Serve(suite.httpListener) // nolint: errcheck
+}
+
+func (suite *PrometheusTestSuite) TearDownTest() {
+	suite.prometheus.Shutdown()
+	suite.NoError(suite.factory.Close())
+	suite.httpListener.Close()
+}
+
+func (suite *PrometheusTestSuite) TestEventStartFinish() {
+	suite.prometheus.EventStart(mtglib.EventStart{
+		CreatedAt: time.Now(),
+		ConnID:    "connID",
+		RemoteIP:  net.ParseIP("10.0.0.10"),
+	})
+
+	time.Sleep(100 * time.Millisecond)
+
+	data, err := suite.Get()
+	suite.NoError(err)
+	suite.Contains(data, `mtg_active_connections{ip_type="ipv4"} 1`)
+
+	suite.prometheus.EventFinish(mtglib.EventFinish{
+		CreatedAt: time.Now(),
+		ConnID:    "connID",
+	})
+
+	time.Sleep(100 * time.Millisecond)
+
+	data, err = suite.Get()
+	suite.NoError(err)
+	suite.Contains(data, `mtg_active_connections{ip_type="ipv4"} 0`)
+}
+
+func (suite *PrometheusTestSuite) TestEventConcurrencyLimited() {
+	suite.prometheus.EventConcurrencyLimited(mtglib.EventConcurrencyLimited{
+		CreatedAt: time.Now(),
+	})
+
+	time.Sleep(100 * time.Millisecond)
+
+	data, err := suite.Get()
+	suite.NoError(err)
+	suite.Contains(data, `mtg_concurrency_limited 1`)
+}
+
+func TestPrometheus(t *testing.T) {
+	t.Parallel()
+	suite.Run(t, &PrometheusTestSuite{})
+}
@@ -2,7 +2,6 @@ package stats

 import (
 	"fmt"
-	"net"
 	"strings"
 	"time"

@@ -15,45 +14,34 @@ type statsdFakeLogger struct{}

 func (s statsdFakeLogger) Printf(msg string, args ...interface{}) {}

-type statsdStreamInfo struct {
-	createdAt time.Time
-	clientIP  net.IP
-}
-
-func (s *statsdStreamInfo) ClientIPTag() statsd.Tag {
-	if s.clientIP.To4() == nil {
-		return statsd.StringTag(TagIPType, TagIPTypeIPv6)
-	} else {
-		return statsd.StringTag(TagIPType, TagIPTypeIPv4)
-	}
-}
-
 type statsdProcessor struct {
-	streams map[string]*statsdStreamInfo
+	streams map[string]*streamInfo
 	client  *statsd.Client
 }

 func (s statsdProcessor) EventStart(evt mtglib.EventStart) {
-	clientInfo := &statsdStreamInfo{
+	sInfo := &streamInfo{
 		createdAt: evt.CreatedAt,
 		clientIP:  evt.RemoteIP,
 	}
-	s.streams[evt.StreamID()] = clientInfo
+	s.streams[evt.StreamID()] = sInfo
+	ipTypeTag := statsd.StringTag(TagIPType, sInfo.IPType())

-	s.client.GaugeDelta(MetricActiveConnection, 1, clientInfo.ClientIPTag())
+	s.client.GaugeDelta(MetricActiveConnection, 1, ipTypeTag)
 }

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

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

-	duration := evt.CreatedAt.Sub(clientInfo.createdAt)
+	duration := evt.CreatedAt.Sub(sInfo.createdAt)
+	ipTypeTag := statsd.StringTag(TagIPType, sInfo.IPType())

-	s.client.GaugeDelta(MetricActiveConnection, -1, clientInfo.ClientIPTag())
+	s.client.GaugeDelta(MetricActiveConnection, -1, ipTypeTag)
 	s.client.PrecisionTiming(MetricSessionDuration, duration)
 }

@@ -88,7 +76,7 @@ func (s StatsdFactory) Close() error {
 func (s StatsdFactory) Make() events.Observer {
 	return statsdProcessor{
 		client:  s.client,
-		streams: make(map[string]*statsdStreamInfo),
+		streams: make(map[string]*streamInfo),
 	}
 }

@@ -97,6 +97,7 @@ func (suite *StatsdTestSuite) TestEventStartFinish() {
 	suite.statsd.EventStart(mtglib.EventStart{
 		CreatedAt: time.Now(),
 		ConnID:    "connID",
+		RemoteIP:  net.ParseIP("10.0.0.10"),
 	})

 	time.Sleep(2 * statsd.DefaultFlushInterval)
@@ -0,0 +1,19 @@
+package stats
+
+import (
+	"net"
+	"time"
+)
+
+type streamInfo struct {
+	createdAt time.Time
+	clientIP  net.IP
+}
+
+func (s *streamInfo) IPType() string {
+	if s.clientIP.To4() == nil {
+		return TagIPTypeIPv6
+	}
+
+	return TagIPTypeIPv4
+}