Prometheus metrics

Author: cjimti

README.md

@@ -87,6 +87,17 @@ curl -w "\n" -d "{\"generic\": \"$RANDOM\"}" -X POST http://localhost:8080/rx/me
  time for i in {1..1000}; do curl -w "\n" -d "{\"generic\": \"$RANDOM\"}" -X POST http://localhost:8080/rx/me/generic_data/generic/test/data; done
 ```
 
+### Profile
+
+```bash
+go build ./rxtx.go && time ./rxtx --path=./data/ --cpuprofile=rxtxcpu.prof --memprofile=rxtxmem.prof
+```
+
+Browser-based profile viewer:
+```bash
+go tool pprof -http=:8081 rxtxcpu.prof
+```
+
 ### Building and Releasing
 
 **rxtx** uses [GORELEASER] to build binaries and [Docker] containers.

rtq/routes.go

@@ -16,10 +16,16 @@ func RxRouteHandler(c *gin.Context) {
 
 	rawData, _ := c.GetRawData()
 
+	q := c.MustGet("Q").(*rtQ)
+
 	// all data is json
 	payload := make(map[string]interface{})
 	err := json.Unmarshal(rawData, &payload)
 	if err != nil {
+
+		// increment metric msg_errors
+		q.cfg.MsgError.Inc()
+
 		c.JSON(500, gin.H{
 			"status":  "FAIL",
 			"message": fmt.Sprintf("could not unmarshal json: %s", rawData),
@@ -36,9 +42,13 @@ func RxRouteHandler(c *gin.Context) {
 	}
 
 	// write the message
-	q := c.MustGet("Q").(*rtQ)
+
 	err = q.QWrite(msg)
 	if err != nil {
+
+		// increment metric msg_errors
+		q.cfg.MsgError.Inc()
+
 		c.JSON(500, gin.H{
 			"status":  "FAIL",
 			"message": fmt.Sprintf("failed to write message: %s", err.Error()),

rtq/rtq.go

@@ -6,8 +6,11 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"os"
 	"time"
 
+	"github.com/prometheus/client_golang/prometheus"
+
 	"errors"
 
 	"github.com/coreos/bbolt"
@@ -42,6 +45,12 @@ type Config struct {
 	Logger     *zap.Logger
 	Receiver   string
 	Path       string
+	Processed  prometheus.Counter
+	Queued     prometheus.Gauge
+	TxBatches  prometheus.Counter
+	TxFail     prometheus.Counter
+	DbErr      prometheus.Counter
+	MsgError   prometheus.Counter
 }
 
 // rtQ private struct see NewQ
@@ -56,6 +65,10 @@ type rtQ struct {
 	statusError func(msg string, fields ...zapcore.Field) // status error output
 }
 
+func (rt *rtQ) GetMessageCount() int {
+	return rt.mCount
+}
+
 // NewQ returns a new rtQ
 func NewQ(name string, cfg Config) (*rtQ, error) {
 	db, err := bolt.Open(cfg.Path+name+".db", 0600, &bolt.Options{Timeout: 1 * time.Second})
@@ -103,6 +116,9 @@ func (rt *rtQ) getMessageBatch() MessageBatch {
 		// get bucket stats
 		stats := bucket.Stats()
 
+		// metric: messages_in_queue
+		rt.cfg.Queued.Set(float64(stats.KeyN))
+
 		rt.status("QueueState", zapcore.Field{
 			Key:     "TotalRecords",
 			Type:    zapcore.Int32Type,
@@ -144,6 +160,9 @@ func (rt *rtQ) getMessageBatch() MessageBatch {
 	})
 
 	if err != nil {
+		// increment metric db_errors
+		rt.cfg.DbErr.Inc()
+
 		rt.cfg.Logger.Error("bbolt db View error: " + err.Error())
 	}
 
@@ -217,6 +236,10 @@ func (rt *rtQ) tx() {
 	// try to send
 	err := rt.transmit(mb)
 	if err != nil {
+
+		// increment metric tx_fails
+		rt.cfg.TxFail.Inc()
+
 		// transmission failed
 		rt.status("Transmission", zapcore.Field{
 			Key:    "TransmissionError",
@@ -232,6 +255,9 @@ func (rt *rtQ) tx() {
 		return
 	}
 
+	// increment metric tx_batches
+	rt.cfg.TxBatches.Inc()
+
 	rt.status("TransmissionComplete", zapcore.Field{
 		Key:     "RemovingMessages",
 		Type:    zapcore.Int32Type,
@@ -256,11 +282,8 @@ func (rt *rtQ) waitTx() {
 // Write to the queue
 func (rt *rtQ) QWrite(msg Message) error {
 
-	rt.status("ReceiverStatus", zapcore.Field{
-		Key:     "KeyValues",
-		Type:    zapcore.Int32Type,
-		Integer: int64(len(msg.Payload)),
-	})
+	// increment metric
+	rt.cfg.Processed.Inc()
 
 	rt.mq <- msg
 
@@ -274,7 +297,7 @@ func messageHandler(db *bolt.DB, mq chan Message, remove chan int) {
 	for {
 		select {
 		case msg := <-mq:
-			db.Update(func(tx *bolt.Tx) error {
+			err := db.Update(func(tx *bolt.Tx) error {
 				uuidV4, _ := uuid.NewV4()
 
 				msg.Time = time.Now()
@@ -295,8 +318,12 @@ func messageHandler(db *bolt.DB, mq chan Message, remove chan int) {
 
 				return nil
 			})
+			if err != nil {
+				fmt.Println("Error updating database. FATAL")
+				os.Exit(1)
+			}
 		case rmi := <-remove:
-			db.Update(func(tx *bolt.Tx) error {
+			err := db.Update(func(tx *bolt.Tx) error {
 				bucket := tx.Bucket([]byte("mq"))
 
 				c := bucket.Cursor()
@@ -313,6 +340,10 @@ func messageHandler(db *bolt.DB, mq chan Message, remove chan int) {
 
 				return nil
 			})
+			if err != nil {
+				fmt.Println("Error updating database. FATAL")
+				os.Exit(1)
+			}
 
 		}
 	}

rxtx.go

@@ -2,28 +2,89 @@ package main
 
 import (
 	"flag"
+	"fmt"
 	"io/ioutil"
 	"net/http"
+	"os"
+	"os/signal"
+	"runtime"
+	"runtime/pprof"
 	"time"
 
-	"fmt"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 
 	"github.com/gin-contrib/zap"
+
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+
 	"github.com/gin-gonic/gin"
 	"github.com/txn2/rxtx/rtq"
 	"go.uber.org/zap"
 )
 
 func main() {
+
 	var port = flag.String("port", "8080", "Server port.")
 	var path = flag.String("path", "./", "Directory to store database.")
 	var interval = flag.Int("interval", 60, "Seconds between intervals.")
 	var batch = flag.Int("batch", 100, "Batch size.")
 	var maxq = flag.Int("maxq", 100000, "Max number of message in queue.")
 	var ingest = flag.String("ingest", "http://localhost:8081/in", "Ingest server.")
+	var verbose = flag.Bool("verbose", false, "Verbose")
+
+	// Instrumentation
+	var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to `file`")
+	var memprofile = flag.String("memprofile", "", "write memory profile to `file`")
 
 	flag.Parse()
 
+	if *cpuprofile != "" {
+		f, err := os.Create(*cpuprofile)
+		if err != nil {
+			fmt.Println(err.Error())
+			os.Exit(1)
+		}
+
+		err = pprof.StartCPUProfile(f)
+		if err != nil {
+			fmt.Println(err.Error())
+			os.Exit(1)
+		}
+	}
+
+	// Instrumentation and Signal handling
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, os.Interrupt)
+	go func() {
+		for sig := range c {
+			if sig.String() == "interrupt" {
+				fmt.Printf("Exiting on interrupt...")
+				if *cpuprofile != "" {
+					pprof.StopCPUProfile()
+				}
+
+				if *memprofile != "" {
+					f, err := os.Create(*memprofile)
+					if err != nil {
+						fmt.Println("could not create memory profile: " + err.Error())
+					}
+					runtime.GC() // get up-to-date statistics
+					if err := pprof.WriteHeapProfile(f); err != nil {
+						fmt.Println("could not write memory profile: " + err.Error())
+					}
+					err = f.Close()
+					if err != nil {
+						fmt.Printf("Can not close memory profile: %s\n", err.Error())
+						return
+					}
+				}
+
+				os.Exit(0)
+			}
+		}
+	}()
+
 	zapCfg := zap.NewProductionConfig()
 	zapCfg.DisableCaller = true
 	zapCfg.DisableStacktrace = true
@@ -41,6 +102,37 @@ func main() {
 
 	logger.Info("Starting rxtx...")
 
+	// Prometheus Metrics
+	processed := promauto.NewCounter(prometheus.CounterOpts{
+		Name: "total_messages_received",
+		Help: "Total number of messages received.",
+	})
+
+	queued := promauto.NewGauge(prometheus.GaugeOpts{
+		Name: "messages_in_queue",
+		Help: "Number os messages in the queue.",
+	})
+
+	txBatches := promauto.NewCounter(prometheus.CounterOpts{
+		Name: "tx_batches",
+		Help: "Total number of batch transmissions.",
+	})
+
+	txFail := promauto.NewCounter(prometheus.CounterOpts{
+		Name: "tx_fails",
+		Help: "Total number of transaction errors.",
+	})
+
+	dbErr := promauto.NewCounter(prometheus.CounterOpts{
+		Name: "db_errors",
+		Help: "Total number database errors.",
+	})
+
+	msgErr := promauto.NewCounter(prometheus.CounterOpts{
+		Name: "msg_errors",
+		Help: "Total number message errors.",
+	})
+
 	// database
 	q, err := rtq.NewQ("rxtx", rtq.Config{
 		Interval:   time.Duration(*interval) * time.Second,
@@ -49,6 +141,12 @@ func main() {
 		Logger:     logger,
 		Receiver:   *ingest,
 		Path:       *path,
+		Processed:  processed,
+		Queued:     queued,
+		TxBatches:  txBatches,
+		TxFail:     txFail,
+		DbErr:      dbErr,
+		MsgError:   msgErr,
 	})
 	if err != nil {
 		panic(err)
@@ -70,14 +168,20 @@ func main() {
 		c.Next()
 	})
 
-	// use zap logger
-	r.Use(ginzap.Ginzap(logger, time.RFC3339, true))
+	// use zap logger on http
+	if *verbose == true {
+		r.Use(ginzap.Ginzap(logger, time.RFC3339, true))
+	}
 
 	rxRoute := "/rx/:producer/:key/*label"
 	r.POST(rxRoute, rtq.RxRouteHandler)
 	r.OPTIONS(rxRoute, preflight)
 
+	// Prometheus Metrics
+	r.GET("/metrics", gin.WrapH(promhttp.Handler()))
+
 	logger.Info("Listening on port: " + *port)
+
 	// block on server run
 	r.Run(":" + *port)
 }