From 54b38e7383c2e0563ea2b823e71bb270b55e1d2f Mon Sep 17 00:00:00 2001
From: SuperQ <superq@gmail.com>
Date: Sat, 1 Apr 2023 09:44:00 +0200
Subject: [PATCH] Add toolkit package

Implement a top level package to make it easy to bootstrap an exporter.
* Move the flag package to the top level.
* Add support for `--web.telemetry-path` flag. Defaults to`/metrics`.
* Add a self-check function to the web FlagConfig.

Signed-off-by: SuperQ <superq@gmail.com>
---
 web/kingpinflag/flag.go => flag.go |   6 +-
 toolkit.go                         | 153 +++++++++++++++++++++++++++++
 web/tls_config.go                  |  21 +++-
 web/tls_config_test.go             |   2 +
 4 files changed, 179 insertions(+), 3 deletions(-)
 rename web/kingpinflag/flag.go => flag.go (92%)
 create mode 100644 toolkit.go

diff --git a/web/kingpinflag/flag.go b/flag.go
similarity index 92%
rename from web/kingpinflag/flag.go
rename to flag.go
index ce6850ff..81e0fcdf 100644
--- a/web/kingpinflag/flag.go
+++ b/flag.go
@@ -10,7 +10,7 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-package kingpinflag
+package toolkit
 
 import (
 	"runtime"
@@ -31,6 +31,10 @@ func AddFlags(a *kingpin.Application, defaultAddress string) *web.FlagConfig {
 		).Bool()
 	}
 	flags := web.FlagConfig{
+		MetricsPath: a.Flag(
+			"web.telemetry-path",
+			"Path under which to expose metrics.",
+		).Default("/metrics").String(),
 		WebListenAddresses: a.Flag(
 			"web.listen-address",
 			"Addresses on which to expose metrics and web interface. Repeatable for multiple addresses.",
diff --git a/toolkit.go b/toolkit.go
new file mode 100644
index 00000000..74ac79f8
--- /dev/null
+++ b/toolkit.go
@@ -0,0 +1,153 @@
+// Copyright 2023 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package toolkit
+
+import (
+	"errors"
+	stdlog "log"
+	"net/http"
+	"os"
+
+	"github.com/alecthomas/kingpin/v2"
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/collectors"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"github.com/prometheus/common/promlog"
+	promlogflag "github.com/prometheus/common/promlog/flag"
+	"github.com/prometheus/common/version"
+	"github.com/prometheus/exporter-toolkit/web"
+)
+
+var (
+	ErrNoFlagConfig = errors.New("Missing FlagConfig")
+	ErrNoHandler    = errors.New("Missing one of MetricsHandler or MetricsHandlerFunc")
+	ErrOneHandler   = errors.New("Only one of MetricsHandler or MetricsHandlerFunc allowed")
+)
+
+type Config struct {
+	Name               string
+	Description        string
+	DefaultAddress     string
+	Logger             log.Logger
+	MetricsHandlerFunc *func(http.ResponseWriter, *http.Request)
+}
+
+type Toolkit struct {
+	Logger      log.Logger
+	MaxRequests int
+
+	flagConfig         *web.FlagConfig
+	landingConfig      web.LandingConfig
+	metricsHandler     http.Handler
+	metricsHandlerFunc *func(http.ResponseWriter, *http.Request)
+}
+
+func New(c Config) *Toolkit {
+	disableExporterMetrics := kingpin.Flag(
+		"web.disable-exporter-metrics",
+		"Exclude metrics about the exporter itself (promhttp_*, process_*, go_*).",
+	).Bool()
+	maxRequests := kingpin.Flag(
+		"web.max-requests",
+		"Maximum number of parallel scrape requests. Use 0 to disable.",
+	).Default("40").Int()
+
+	t := Toolkit{
+		flagConfig: AddFlags(kingpin.CommandLine, c.DefaultAddress),
+		landingConfig: web.LandingConfig{
+			Name:        c.Name,
+			Description: c.Description,
+			Version:     version.Info(),
+		},
+		metricsHandlerFunc: c.MetricsHandlerFunc,
+	}
+
+	promlogConfig := &promlog.Config{}
+	promlogflag.AddFlags(kingpin.CommandLine, promlogConfig)
+
+	kingpin.Version(version.Print(c.Name))
+	kingpin.HelpFlag.Short('h')
+	kingpin.Parse()
+
+	t.Logger = promlog.New(promlogConfig)
+	t.MaxRequests = *maxRequests
+
+	handlerOpts := promhttp.HandlerOpts{
+		ErrorLog:            stdlog.New(log.NewStdlibAdapter(level.Error(t.Logger)), "", 0),
+		MaxRequestsInFlight: t.MaxRequests,
+	}
+	promHandler := promhttp.InstrumentMetricHandler(
+		prometheus.DefaultRegisterer, promhttp.HandlerFor(prometheus.DefaultGatherer, handlerOpts),
+	)
+	if *disableExporterMetrics {
+		prometheus.Unregister(collectors.NewGoCollector())
+		prometheus.Unregister(collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}))
+		promHandler = promhttp.HandlerFor(prometheus.DefaultGatherer, handlerOpts)
+	}
+
+	t.metricsHandler = promHandler
+
+	return &t
+}
+
+func (t *Toolkit) SetMetricsHandler(h http.Handler) {
+	t.metricsHandler = h
+}
+
+func (t *Toolkit) SetMetricsHandlerFunc(h *func(http.ResponseWriter, *http.Request)) {
+	t.metricsHandlerFunc = h
+}
+
+func (t *Toolkit) Run() error {
+	if t.flagConfig == nil {
+		return ErrNoFlagConfig
+	}
+	err := t.flagConfig.CheckFlags()
+	if err != nil {
+		return err
+	}
+	if t.metricsHandler == nil && t.metricsHandlerFunc == nil {
+		return ErrNoHandler
+	}
+	if t.metricsHandler != nil && t.metricsHandlerFunc != nil {
+		return ErrOneHandler
+	}
+	if *t.flagConfig.MetricsPath != "" && t.metricsHandler != nil {
+		http.Handle(*t.flagConfig.MetricsPath, t.metricsHandler)
+	}
+	if *t.flagConfig.MetricsPath != "" && t.metricsHandlerFunc != nil {
+		http.HandleFunc(*t.flagConfig.MetricsPath, *t.metricsHandlerFunc)
+	}
+	if *t.flagConfig.MetricsPath != "/" && *t.flagConfig.MetricsPath != "" {
+		t.landingConfig.Links = append(t.landingConfig.Links,
+			web.LandingLinks{
+				Address: *t.flagConfig.MetricsPath,
+				Text:    "Metrics",
+			},
+		)
+		landingPage, err := web.NewLandingPage(t.landingConfig)
+		if err != nil {
+			level.Error(t.Logger).Log("err", err)
+			os.Exit(1)
+		}
+		http.Handle("/", landingPage)
+	}
+
+	level.Info(t.Logger).Log("msg", "Starting "+t.landingConfig.Name, "version", version.Info())
+	level.Info(t.Logger).Log("msg", "Build context", "build_context", version.BuildContext())
+
+	srv := &http.Server{}
+	return web.ListenAndServe(srv, t.flagConfig, t.Logger)
+}
diff --git a/web/tls_config.go b/web/tls_config.go
index b20fed52..c561fbed 100644
--- a/web/tls_config.go
+++ b/web/tls_config.go
@@ -34,6 +34,7 @@ import (
 var (
 	errNoTLSConfig = errors.New("TLS config is not present")
 	ErrNoListeners = errors.New("no web listen address or systemd socket flag specified")
+	ErrMissingFlag = errors.New("Flag config is empty")
 )
 
 type Config struct {
@@ -55,11 +56,26 @@ type TLSConfig struct {
 }
 
 type FlagConfig struct {
+	MetricsPath        *string
 	WebListenAddresses *[]string
 	WebSystemdSocket   *bool
 	WebConfigFile      *string
 }
 
+// CheckFlags validates that the FlagConfig has all required values set and has at least one listener.
+func (c *FlagConfig) CheckFlags() error {
+	if c.MetricsPath == nil {
+		return ErrMissingFlag
+	}
+	if c.WebSystemdSocket == nil && (c.WebListenAddresses == nil || len(*c.WebListenAddresses) == 0) {
+		return ErrNoListeners
+	}
+	if c.WebConfigFile == nil {
+		return ErrMissingFlag
+	}
+	return nil
+}
+
 // SetDirectory joins any relative file paths with dir.
 func (t *TLSConfig) SetDirectory(dir string) {
 	t.TLSCertPath = config_util.JoinDir(dir, t.TLSCertPath)
@@ -204,8 +220,9 @@ func ServeMultiple(listeners []net.Listener, server *http.Server, flags *FlagCon
 // WebSystemdSocket in the FlagConfig is true. The FlagConfig is also passed on
 // to ServeMultiple.
 func ListenAndServe(server *http.Server, flags *FlagConfig, logger log.Logger) error {
-	if flags.WebSystemdSocket == nil && (flags.WebListenAddresses == nil || len(*flags.WebListenAddresses) == 0) {
-		return ErrNoListeners
+	err := flags.CheckFlags()
+	if err != nil {
+		return err
 	}
 
 	if flags.WebSystemdSocket != nil && *flags.WebSystemdSocket {
diff --git a/web/tls_config_test.go b/web/tls_config_test.go
index ceb7f9b2..9f8d520d 100644
--- a/web/tls_config_test.go
+++ b/web/tls_config_test.go
@@ -387,6 +387,7 @@ func TestConfigReloading(t *testing.T) {
 			}
 		}()
 		flagsBadYAMLPath := FlagConfig{
+			MetricsPath:        "/metrics",
 			WebListenAddresses: &([]string{port}),
 			WebSystemdSocket:   OfBool(false),
 			WebConfigFile:      OfString(badYAMLPath),
@@ -461,6 +462,7 @@ func (test *TestInputs) Test(t *testing.T) {
 			}
 		}()
 		flags := FlagConfig{
+			MetricsPath:        "/metrics",
 			WebListenAddresses: &([]string{port}),
 			WebSystemdSocket:   OfBool(false),
 			WebConfigFile:      &test.YAMLConfigPath,