From 81e58c045d43c7152466e3ef13e95d71528fb5be Mon Sep 17 00:00:00 2001
From: Alex Akselrod <alex@akselrod.org>
Date: Mon, 25 Jul 2022 21:46:50 -0700
Subject: [PATCH] initial prometheus metrics support

---
 cmd/lnmuxd/config.go |  8 +++++++-
 cmd/lnmuxd/run.go    | 33 ++++++++++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/cmd/lnmuxd/config.go b/cmd/lnmuxd/config.go
index af67653..4f46c11 100644
--- a/cmd/lnmuxd/config.go
+++ b/cmd/lnmuxd/config.go
@@ -10,7 +10,10 @@ import (
 	"gopkg.in/yaml.v2"
 )
 
-var DefaultListenAddress = "localhost:19090"
+var (
+	DefaultListenAddress     = "localhost:19090"
+	DefaultPrometheusAddress = "localhost:2112"
+)
 
 type Config struct {
 	// Lnd contains the configuration of the nodes.
@@ -27,6 +30,9 @@ type Config struct {
 
 	// ListenAddress is the network address that we listen on.
 	ListenAddress string `yaml:"listenAddress"`
+
+	// PrometheusAddress is the network address that Prometheus listens on.
+	PrometheusAddress string `yaml:"prometheusAddress"`
 }
 
 func (c *Config) GetIdentityKey() ([32]byte, error) {
diff --git a/cmd/lnmuxd/run.go b/cmd/lnmuxd/run.go
index b5c8bcc..5fc7d4c 100644
--- a/cmd/lnmuxd/run.go
+++ b/cmd/lnmuxd/run.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"net/http"
 	"os"
 	"os/signal"
 	"syscall"
@@ -17,8 +18,10 @@ import (
 	"github.com/bottlepay/lnmux/persistence"
 	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/go-pg/pg/v10"
+	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
 	"github.com/lightningnetwork/lnd/clock"
 	"github.com/lightningnetwork/lnd/keychain"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/urfave/cli/v2"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
@@ -106,9 +109,10 @@ func runAction(c *cli.Context) error {
 		return err
 	}
 
-	// Instantiate grpc server.
+	// Instantiate grpc server and enable reflection and Prometheus metrics.
 	grpcServer := grpc.NewServer()
 	reflection.Register(grpcServer)
+	grpc_prometheus.Register(grpcServer)
 
 	server := newServer(creator, registry)
 
@@ -126,6 +130,18 @@ func runAction(c *cli.Context) error {
 		return err
 	}
 
+	promAddress := cfg.PrometheusAddress
+	if promAddress == "" {
+		promAddress = DefaultPrometheusAddress
+	}
+
+	// Instantiate a new HTTP server and mux.
+	metrics := &http.Server{
+		Addr:              promAddress,
+		Handler:           promhttp.Handler(),
+		ReadHeaderTimeout: 10 * time.Second,
+	}
+
 	group, ctx := errgroup.WithContext(context.Background())
 
 	// Run multiplexer.
@@ -159,6 +175,21 @@ func runAction(c *cli.Context) error {
 		return nil
 	})
 
+	group.Go(func() error {
+		log.Infow("Prometheus server starting", "promAddress", promAddress)
+
+		return metrics.ListenAndServe()
+	})
+
+	group.Go(func() error {
+		<-ctx.Done()
+
+		// Stop Prometheus server
+		log.Infow("Prometheus server stopping")
+
+		return metrics.Close()
+	})
+
 	// Run ctrl-c handler.
 	group.Go(func() error {
 		log.Infof("Press ctrl-c to exit")