From 8bd27adfa2a5fba03dd567d45e9fdd4d4637d732 Mon Sep 17 00:00:00 2001
From: Will Bollock <wbollock@gmail.com>
Date: Mon, 5 Sep 2022 12:55:18 -0400
Subject: [PATCH] feat: add scrape timeout as flag

---
 nagios_exporter.go | 40 ++++++++++++++++++++++------------------
 1 file changed, 22 insertions(+), 18 deletions(-)

diff --git a/nagios_exporter.go b/nagios_exporter.go
index d0f1cac..30a2d8d 100644
--- a/nagios_exporter.go
+++ b/nagios_exporter.go
@@ -156,13 +156,15 @@ var (
 type Exporter struct {
 	nagiosEndpoint, nagiosAPIKey string
 	sslVerify                    bool
+	nagiosAPITimeout             time.Duration
 }
 
-func NewExporter(nagiosEndpoint, nagiosAPIKey string, sslVerify bool) *Exporter {
+func NewExporter(nagiosEndpoint, nagiosAPIKey string, sslVerify bool, nagiosAPITimeout time.Duration) *Exporter {
 	return &Exporter{
-		nagiosEndpoint: nagiosEndpoint,
-		nagiosAPIKey:   nagiosAPIKey,
-		sslVerify:      sslVerify,
+		nagiosEndpoint:   nagiosEndpoint,
+		nagiosAPIKey:     nagiosAPIKey,
+		sslVerify:        sslVerify,
+		nagiosAPITimeout: nagiosAPITimeout,
 	}
 }
 
@@ -190,11 +192,11 @@ func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
 	ch <- servicechecksPerformance
 }
 
-func (e *Exporter) TestNagiosConnectivity(sslVerify bool) float64 {
+func (e *Exporter) TestNagiosConnectivity(sslVerify bool, nagiosAPITimeout time.Duration) float64 {
 
 	systemStatusURL := e.nagiosEndpoint + systemstatusAPI + "?apikey=" + e.nagiosAPIKey
 
-	body := QueryAPIs(systemStatusURL, sslVerify)
+	body := QueryAPIs(systemStatusURL, sslVerify, nagiosAPITimeout)
 	log.Debug("Queried API: ", systemstatusAPI)
 
 	systemStatusObject := systemStatus{}
@@ -209,7 +211,7 @@ func (e *Exporter) TestNagiosConnectivity(sslVerify bool) float64 {
 
 func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
 
-	nagiosStatus := e.TestNagiosConnectivity(e.sslVerify)
+	nagiosStatus := e.TestNagiosConnectivity(e.sslVerify, e.nagiosAPITimeout)
 
 	if nagiosStatus == 0 {
 		log.Warn("Cannot connect to Nagios endpoint")
@@ -219,20 +221,18 @@ func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
 		up, prometheus.GaugeValue, nagiosStatus,
 	)
 
-	e.QueryAPIsAndUpdateMetrics(ch, e.sslVerify)
+	e.QueryAPIsAndUpdateMetrics(ch, e.sslVerify, e.nagiosAPITimeout)
 
 }
 
-func QueryAPIs(url string, sslVerify bool) (body []byte) {
+func QueryAPIs(url string, sslVerify bool, nagiosAPITimeout time.Duration) (body []byte) {
 
 	// https://github.com/prometheus/haproxy_exporter/blob/main/haproxy_exporter.go#L337-L345
 
 	tr := &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: !sslVerify}}
-	// I don't think timeout duration needs to be a user flag
-	var timeout time.Duration
 
 	client := http.Client{
-		Timeout:   timeout,
+		Timeout:   nagiosAPITimeout,
 		Transport: tr,
 	}
 
@@ -266,13 +266,13 @@ func QueryAPIs(url string, sslVerify bool) (body []byte) {
 	return body
 }
 
-func (e *Exporter) QueryAPIsAndUpdateMetrics(ch chan<- prometheus.Metric, sslVerify bool) {
+func (e *Exporter) QueryAPIsAndUpdateMetrics(ch chan<- prometheus.Metric, sslVerify bool, nagiosAPITimeout time.Duration) {
 
 	// get system status
 	systeminfoURL := e.nagiosEndpoint + systeminfoAPI + "?apikey=" + e.nagiosAPIKey
 	log.Debug("Queried API: ", systeminfoAPI)
 
-	body := QueryAPIs(systeminfoURL, sslVerify)
+	body := QueryAPIs(systeminfoURL, sslVerify, nagiosAPITimeout)
 
 	systemInfoObject := systemInfo{}
 	jsonErr := json.Unmarshal(body, &systemInfoObject)
@@ -287,7 +287,7 @@ func (e *Exporter) QueryAPIsAndUpdateMetrics(ch chan<- prometheus.Metric, sslVer
 	// host status
 	hoststatusURL := e.nagiosEndpoint + hoststatusAPI + "?apikey=" + e.nagiosAPIKey
 
-	body = QueryAPIs(hoststatusURL, sslVerify)
+	body = QueryAPIs(hoststatusURL, sslVerify, nagiosAPITimeout)
 	log.Debug("Queried API: ", systeminfoAPI)
 
 	hostStatusObject := hostStatus{}
@@ -372,7 +372,7 @@ func (e *Exporter) QueryAPIsAndUpdateMetrics(ch chan<- prometheus.Metric, sslVer
 	// service status
 	servicestatusURL := e.nagiosEndpoint + servicestatusAPI + "?apikey=" + e.nagiosAPIKey
 
-	body = QueryAPIs(servicestatusURL, sslVerify)
+	body = QueryAPIs(servicestatusURL, sslVerify, nagiosAPITimeout)
 	log.Debug("Queried API: ", servicestatusAPI)
 
 	serviceStatusObject := serviceStatus{}
@@ -474,7 +474,7 @@ func (e *Exporter) QueryAPIsAndUpdateMetrics(ch chan<- prometheus.Metric, sslVer
 	// service status
 	systemStatusDetailURL := e.nagiosEndpoint + systemstatusDetailAPI + "?apikey=" + e.nagiosAPIKey
 
-	body = QueryAPIs(systemStatusDetailURL, sslVerify)
+	body = QueryAPIs(systemStatusDetailURL, sslVerify, nagiosAPITimeout)
 	log.Debug("Queried API: ", systemstatusDetailAPI)
 
 	systemStatusDetailObject := systemStatusDetail{}
@@ -592,6 +592,9 @@ func main() {
 			"Nagios application address")
 		sslVerify = flag.Bool("nagios.ssl-verify", false,
 			"SSL certificate validation")
+		// I think users would rather enter `5` over `5s`, e.g int vs Duration flag
+		nagiosAPITimeout = flag.Int("nagios.timeout", 5,
+			"Timeout for querying Nagios API in seconds")
 		configPath = flag.String("config.path", "/etc/prometheus-nagios-exporter/config.toml",
 			"Config file path")
 		logLevel = flag.String("log.level", "info",
@@ -611,7 +614,8 @@ func main() {
 
 	nagiosURL := *remoteAddress + nagiosAPIVersion + apiSlug
 
-	exporter := NewExporter(nagiosURL, conf.APIKey, *sslVerify)
+	// convert timeout flag to seconds
+	exporter := NewExporter(nagiosURL, conf.APIKey, *sslVerify, time.Duration(*nagiosAPITimeout)*time.Second)
 	prometheus.MustRegister(exporter)
 
 	log.Info("Using connection endpoint: ", *remoteAddress)