From 0ad70dba52534c4eb8baeacd49c4ea4ff7a927f9 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 |  9 ++++++++-
 cmd/lnmuxd/run.go    | 45 ++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/cmd/lnmuxd/config.go b/cmd/lnmuxd/config.go
index af67653..ca5500b 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"
+	DefaultInstrumentationAddress = "localhost:2112"
+)
 
 type Config struct {
 	// Lnd contains the configuration of the nodes.
@@ -27,6 +30,10 @@ type Config struct {
 
 	// ListenAddress is the network address that we listen on.
 	ListenAddress string `yaml:"listenAddress"`
+
+	// InstrumentationAddress is the network address that Prometheus
+	// and eventually pprof and/or other instrumentation listen on.
+	InstrumentationAddress string `yaml:"instrumentationAddress"`
 }
 
 func (c *Config) GetIdentityKey() ([32]byte, error) {
diff --git a/cmd/lnmuxd/run.go b/cmd/lnmuxd/run.go
index b5c8bcc..f3d343d 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,13 @@ func runAction(c *cli.Context) error {
 		return err
 	}
 
-	// Instantiate grpc server.
-	grpcServer := grpc.NewServer()
+	// Instantiate grpc server and enable reflection and Prometheus metrics.
+	grpcServer := grpc.NewServer(
+		grpc.StreamInterceptor(grpc_prometheus.StreamServerInterceptor),
+		grpc.UnaryInterceptor(grpc_prometheus.UnaryServerInterceptor),
+	)
 	reflection.Register(grpcServer)
+	grpc_prometheus.Register(grpcServer)
 
 	server := newServer(creator, registry)
 
@@ -126,6 +133,24 @@ func runAction(c *cli.Context) error {
 		return err
 	}
 
+	instAddress := cfg.InstrumentationAddress
+	if instAddress == "" {
+		instAddress = DefaultInstrumentationAddress
+	}
+
+	// Instantiate a new HTTP server and mux.
+	instMux := http.NewServeMux()
+	instMux.Handle("/metrics", promhttp.Handler())
+
+	instServer := &http.Server{
+		Addr:    instAddress,
+		Handler: instMux,
+
+		// Even though this server should only be exposed to trusted
+		// clients, this mitigates slowloris-like DoS attacks.
+		ReadHeaderTimeout: 10 * time.Second,
+	}
+
 	group, ctx := errgroup.WithContext(context.Background())
 
 	// Run multiplexer.
@@ -159,6 +184,22 @@ func runAction(c *cli.Context) error {
 		return nil
 	})
 
+	group.Go(func() error {
+		log.Infow("Instrumentation HTTP server starting",
+			"instrumentationAddress", instAddress)
+
+		return instServer.ListenAndServe()
+	})
+
+	group.Go(func() error {
+		<-ctx.Done()
+
+		// Stop instrumentation server
+		log.Infow("Instrumentation server stopping")
+
+		return instServer.Close()
+	})
+
 	// Run ctrl-c handler.
 	group.Go(func() error {
 		log.Infof("Press ctrl-c to exit")