Use timeout flag for context cancelation

Author: lucacome

client/nginx.go

@@ -47,10 +47,7 @@ func NewNginxClient(httpClient *http.Client, apiEndpoint string) *NginxClient {
 }
 
 // GetStubStats fetches the stub_status metrics.
-func (client *NginxClient) GetStubStats() (*StubStats, error) {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
+func (client *NginxClient) GetStubStats(ctx context.Context) (*StubStats, error) {
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, client.apiEndpoint, nil)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create a get request: %w", err)

collector/nginx.go

@@ -1,8 +1,10 @@
 package collector
 
 import (
+	"context"
 	"log/slog"
 	"sync"
+	"time"
 
 	"github.com/nginxinc/nginx-prometheus-exporter/client"
 	"github.com/prometheus/client_golang/prometheus"
@@ -14,14 +16,16 @@ type NginxCollector struct {
 	logger      *slog.Logger
 	nginxClient *client.NginxClient
 	metrics     map[string]*prometheus.Desc
+	timeout     time.Duration
 	mutex       sync.Mutex
 }
 
 // NewNginxCollector creates an NginxCollector.
-func NewNginxCollector(nginxClient *client.NginxClient, namespace string, constLabels map[string]string, logger *slog.Logger) *NginxCollector {
+func NewNginxCollector(nginxClient *client.NginxClient, namespace string, constLabels map[string]string, logger *slog.Logger, timeout time.Duration) *NginxCollector {
 	return &NginxCollector{
 		nginxClient: nginxClient,
 		logger:      logger,
+		timeout:     timeout,
 		metrics: map[string]*prometheus.Desc{
 			"connections_active":   newGlobalMetric(namespace, "connections_active", "Active client connections", constLabels),
 			"connections_accepted": newGlobalMetric(namespace, "connections_accepted", "Accepted client connections", constLabels),
@@ -50,7 +54,10 @@ func (c *NginxCollector) Collect(ch chan<- prometheus.Metric) {
 	c.mutex.Lock() // To protect metrics from concurrent collects
 	defer c.mutex.Unlock()
 
-	stats, err := c.nginxClient.GetStubStats()
+	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)
+	defer cancel()
+
+	stats, err := c.nginxClient.GetStubStats(ctx)
 	if err != nil {
 		c.upMetric.Set(nginxDown)
 		ch <- c.upMetric

collector/nginx_plus.go

@@ -6,6 +6,7 @@ import (
 	"log/slog"
 	"strconv"
 	"sync"
+	"time"
 
 	plusclient "github.com/nginxinc/nginx-plus-go-client/v2/client"
 	"github.com/prometheus/client_golang/prometheus"
@@ -31,11 +32,11 @@ type LabelUpdater interface {
 
 // NginxPlusCollector collects NGINX Plus metrics. It implements prometheus.Collector interface.
 type NginxPlusCollector struct {
-	upMetric                       prometheus.Gauge
+	nginxClient                    *plusclient.NginxClient
 	logger                         *slog.Logger
+	upMetric                       prometheus.Gauge
 	cacheZoneMetrics               map[string]*prometheus.Desc
 	workerMetrics                  map[string]*prometheus.Desc
-	nginxClient                    *plusclient.NginxClient
 	streamServerZoneMetrics        map[string]*prometheus.Desc
 	streamZoneSyncMetrics          map[string]*prometheus.Desc
 	streamUpstreamMetrics          map[string]*prometheus.Desc
@@ -47,16 +48,17 @@ type NginxPlusCollector struct {
 	streamLimitConnectionMetrics   map[string]*prometheus.Desc
 	upstreamServerMetrics          map[string]*prometheus.Desc
 	upstreamMetrics                map[string]*prometheus.Desc
-	streamUpstreamServerPeerLabels map[string][]string
 	serverZoneMetrics              map[string]*prometheus.Desc
+	totalMetrics                   map[string]*prometheus.Desc
+	streamUpstreamServerPeerLabels map[string][]string
 	upstreamServerLabels           map[string][]string
 	streamUpstreamServerLabels     map[string][]string
 	serverZoneLabels               map[string][]string
 	streamServerZoneLabels         map[string][]string
 	upstreamServerPeerLabels       map[string][]string
 	cacheZoneLabels                map[string][]string
-	totalMetrics                   map[string]*prometheus.Desc
 	variableLabelNames             VariableLabelNames
+	timeout                        time.Duration
 	variableLabelsMutex            sync.RWMutex
 	mutex                          sync.Mutex
 }
@@ -256,7 +258,7 @@ func NewVariableLabelNames(upstreamServerVariableLabelNames []string, serverZone
 }
 
 // NewNginxPlusCollector creates an NginxPlusCollector.
-func NewNginxPlusCollector(nginxClient *plusclient.NginxClient, namespace string, variableLabelNames VariableLabelNames, constLabels map[string]string, logger *slog.Logger) *NginxPlusCollector {
+func NewNginxPlusCollector(nginxClient *plusclient.NginxClient, namespace string, variableLabelNames VariableLabelNames, constLabels map[string]string, logger *slog.Logger, timeout time.Duration) *NginxPlusCollector {
 	upstreamServerVariableLabelNames := variableLabelNames.UpstreamServerVariableLabelNames
 	streamUpstreamServerVariableLabelNames := variableLabelNames.StreamUpstreamServerVariableLabelNames
 
@@ -273,6 +275,7 @@ func NewNginxPlusCollector(nginxClient *plusclient.NginxClient, namespace string
 		cacheZoneLabels:                make(map[string][]string),
 		nginxClient:                    nginxClient,
 		logger:                         logger,
+		timeout:                        timeout,
 		totalMetrics: map[string]*prometheus.Desc{
 			"connections_accepted":  newGlobalMetric(namespace, "connections_accepted", "Accepted client connections", constLabels),
 			"connections_dropped":   newGlobalMetric(namespace, "connections_dropped", "Dropped client connections", constLabels),
@@ -623,8 +626,10 @@ func (c *NginxPlusCollector) Collect(ch chan<- prometheus.Metric) {
 	c.mutex.Lock() // To protect metrics from concurrent collects
 	defer c.mutex.Unlock()
 
-	// FIXME: https://github.com/nginxinc/nginx-prometheus-exporter/issues/858
-	stats, err := c.nginxClient.GetStats(context.TODO())
+	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)
+	defer cancel()
+
+	stats, err := c.nginxClient.GetStats(ctx)
 	if err != nil {
 		c.upMetric.Set(nginxDown)
 		ch <- c.upMetric

exporter.go

@@ -220,9 +220,7 @@ func main() {
 	_ = srv.Shutdown(srvCtx)
 }
 
-func registerCollector(logger *slog.Logger, transport *http.Transport,
-	addr string, labels map[string]string,
-) {
+func registerCollector(logger *slog.Logger, transport *http.Transport, addr string, labels map[string]string) {
 	if strings.HasPrefix(addr, "unix:") {
 		socketPath, requestPath, err := parseUnixSocketAddress(addr)
 		if err != nil {
@@ -239,7 +237,6 @@ func registerCollector(logger *slog.Logger, transport *http.Transport,
 	userAgent := fmt.Sprintf("NGINX-Prometheus-Exporter/v%v", common_version.Version)
 
 	httpClient := &http.Client{
-		Timeout: *timeout,
 		Transport: &userAgentRoundTripper{
 			agent: userAgent,
 			rt:    transport,
@@ -253,10 +250,10 @@ func registerCollector(logger *slog.Logger, transport *http.Transport,
 			os.Exit(1)
 		}
 		variableLabelNames := collector.NewVariableLabelNames(nil, nil, nil, nil, nil, nil, nil)
-		prometheus.MustRegister(collector.NewNginxPlusCollector(plusClient, "nginxplus", variableLabelNames, labels, logger))
+		prometheus.MustRegister(collector.NewNginxPlusCollector(plusClient, "nginxplus", variableLabelNames, labels, logger, *timeout))
 	} else {
 		ossClient := client.NewNginxClient(httpClient, addr)
-		prometheus.MustRegister(collector.NewNginxCollector(ossClient, "nginx", labels, logger))
+		prometheus.MustRegister(collector.NewNginxCollector(ossClient, "nginx", labels, logger, *timeout))
 	}
 }