From af59851f33f10f99c3c1cbcbf3c90c60ad6d3e98 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Oliveirinha?= <joliveirinha@cloudflare.com>
Date: Tue, 22 Nov 2022 11:30:28 +0000
Subject: [PATCH] TUN-6963: Refactor Metrics service setup

---
 cmd/cloudflared/proxydns/cmd.go |  3 ++-
 cmd/cloudflared/tunnel/cmd.go   |  7 +++++-
 metrics/metrics.go              | 41 ++++++++++++++++++++-------------
 3 files changed, 33 insertions(+), 18 deletions(-)

diff --git a/cmd/cloudflared/proxydns/cmd.go b/cmd/cloudflared/proxydns/cmd.go
index ae03d8c6..950da293 100644
--- a/cmd/cloudflared/proxydns/cmd.go
+++ b/cmd/cloudflared/proxydns/cmd.go
@@ -1,6 +1,7 @@
 package proxydns
 
 import (
+	"context"
 	"net"
 	"os"
 	"os/signal"
@@ -73,7 +74,7 @@ func Run(c *cli.Context) error {
 		log.Fatal().Err(err).Msg("Failed to open the metrics listener")
 	}
 
-	go metrics.ServeMetrics(metricsListener, nil, nil, "", nil, log)
+	go metrics.ServeMetrics(metricsListener, context.Background(), metrics.Config{}, log)
 
 	listener, err := tunneldns.CreateListener(
 		c.String("address"),
diff --git a/cmd/cloudflared/tunnel/cmd.go b/cmd/cloudflared/tunnel/cmd.go
index 6acafe4b..860fecf6 100644
--- a/cmd/cloudflared/tunnel/cmd.go
+++ b/cmd/cloudflared/tunnel/cmd.go
@@ -372,7 +372,12 @@ func StartServer(
 		defer wg.Done()
 		readinessServer := metrics.NewReadyServer(log, clientID)
 		observer.RegisterSink(readinessServer)
-		errC <- metrics.ServeMetrics(metricsListener, ctx.Done(), readinessServer, quickTunnelURL, orchestrator, log)
+		metricsConfig := metrics.Config{
+			ReadyServer:         readinessServer,
+			QuickTunnelHostname: quickTunnelURL,
+			Orchestrator:        orchestrator,
+		}
+		errC <- metrics.ServeMetrics(metricsListener, ctx, metricsConfig, log)
 	}()
 
 	reconnectCh := make(chan supervisor.ReconnectSignal, c.Int("ha-connections"))
diff --git a/metrics/metrics.go b/metrics/metrics.go
index f8b3b1bc..61d25a3b 100644
--- a/metrics/metrics.go
+++ b/metrics/metrics.go
@@ -18,18 +18,24 @@ import (
 )
 
 const (
-	shutdownTimeout = time.Second * 15
-	startupTime     = time.Millisecond * 500
+	startupTime            = time.Millisecond * 500
+	defaultShutdownTimeout = time.Second * 15
 )
 
+type Config struct {
+	ReadyServer         *ReadyServer
+	QuickTunnelHostname string
+	Orchestrator        orchestrator
+
+	ShutdownTimeout time.Duration
+}
+
 type orchestrator interface {
 	GetVersionedConfigJSON() ([]byte, error)
 }
 
 func newMetricsHandler(
-	readyServer *ReadyServer,
-	quickTunnelHostname string,
-	orchestrator orchestrator,
+	config Config,
 	log *zerolog.Logger,
 ) *mux.Router {
 	router := mux.NewRouter()
@@ -39,15 +45,15 @@ func newMetricsHandler(
 	router.HandleFunc("/healthcheck", func(w http.ResponseWriter, r *http.Request) {
 		_, _ = fmt.Fprintf(w, "OK\n")
 	})
-	if readyServer != nil {
-		router.Handle("/ready", readyServer)
+	if config.ReadyServer != nil {
+		router.Handle("/ready", config.ReadyServer)
 	}
 	router.HandleFunc("/quicktunnel", func(w http.ResponseWriter, r *http.Request) {
-		_, _ = fmt.Fprintf(w, `{"hostname":"%s"}`, quickTunnelHostname)
+		_, _ = fmt.Fprintf(w, `{"hostname":"%s"}`, config.QuickTunnelHostname)
 	})
-	if orchestrator != nil {
+	if config.Orchestrator != nil {
 		router.HandleFunc("/config", func(w http.ResponseWriter, r *http.Request) {
-			json, err := orchestrator.GetVersionedConfigJSON()
+			json, err := config.Orchestrator.GetVersionedConfigJSON()
 			if err != nil {
 				w.WriteHeader(500)
 				_, _ = fmt.Fprintf(w, "ERR: %v", err)
@@ -63,10 +69,8 @@ func newMetricsHandler(
 
 func ServeMetrics(
 	l net.Listener,
-	shutdownC <-chan struct{},
-	readyServer *ReadyServer,
-	quickTunnelHostname string,
-	orchestrator orchestrator,
+	ctx context.Context,
+	config Config,
 	log *zerolog.Logger,
 ) (err error) {
 	var wg sync.WaitGroup
@@ -74,7 +78,7 @@ func ServeMetrics(
 	trace.AuthRequest = func(*http.Request) (bool, bool) { return true, true }
 	// TODO: parameterize ReadTimeout and WriteTimeout. The maximum time we can
 	// profile CPU usage depends on WriteTimeout
-	h := newMetricsHandler(readyServer, quickTunnelHostname, orchestrator, log)
+	h := newMetricsHandler(config, log)
 	server := &http.Server{
 		ReadTimeout:  10 * time.Second,
 		WriteTimeout: 10 * time.Second,
@@ -91,7 +95,12 @@ func ServeMetrics(
 	// fully started up. So add artificial delay.
 	time.Sleep(startupTime)
 
-	<-shutdownC
+	<-ctx.Done()
+	shutdownTimeout := config.ShutdownTimeout
+	if shutdownTimeout == 0 {
+		shutdownTimeout = defaultShutdownTimeout
+	}
+
 	ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
 	_ = server.Shutdown(ctx)
 	cancel()