PMM-9471 Context deadline (#420)

* PMM-9471 Cover case with limit collections -1 (auto).

* PMM-9471 Logic, nil fixies. Reflect timeout from VM.

* PMM-9471 Revert dynamic timeout for now.

* Revert "PMM-9471 Revert dynamic timeout for now."

This reverts commit 02cb9b7.

* PMM-9471 New default timeout.

* PMM-9471 Fix missing flag.

* PMM-9471 support of collect[] query parameter

* PMM-9471 Another changes in exporter. Context background.

Co-authored-by: Nurlan Moldomurov <[email protected]>

Author: JiriCtvrtka

exporter/diagnostic_data_collector.go

@@ -53,6 +53,10 @@ func (d *diagnosticDataCollector) Collect(ch chan<- prometheus.Metric) {
 		d.logger.Errorf("cannot run getDiagnosticData: %s", err)
 	}
 
+	if m == nil || m["data"] == nil {
+		d.logger.Error("cannot run getDiagnosticData: response is empty")
+	}
+
 	m, ok := m["data"].(bson.M)
 	if !ok {
 		err := errors.Wrapf(errUnexpectedDataType, "%T for data field", m["data"])

exporter/exporter.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"strconv"
 	"sync"
 	"time"
 
@@ -96,7 +97,7 @@ func New(opts *Opts) *Exporter {
 		opts:                  opts,
 		webListenAddress:      opts.WebListenAddress,
 		lock:                  &sync.Mutex{},
-		totalCollectionsCount: -1, // not calculated yet. waiting the db connection.
+		totalCollectionsCount: -1, // Not calculated yet. waiting the db connection.
 	}
 	// Try initial connect. Connection will be retried with every scrape.
 	go func() {
@@ -116,7 +117,7 @@ func (e *Exporter) getTotalCollectionsCount() int {
 	return e.totalCollectionsCount
 }
 
-func (e *Exporter) makeRegistry(ctx context.Context, client *mongo.Client, topologyInfo labelsGetter) *prometheus.Registry {
+func (e *Exporter) makeRegistry(ctx context.Context, client *mongo.Client, topologyInfo labelsGetter, requestOpts Opts) *prometheus.Registry {
 	registry := prometheus.NewRegistry()
 
 	gc := generalCollector{
@@ -135,11 +136,11 @@ func (e *Exporter) makeRegistry(ctx context.Context, client *mongo.Client, topol
 		e.logger.Errorf("Cannot get node type to check if this is a mongos: %s", err)
 	}
 
-	// enable collectors like collstats and indexstats depending on the number of collections
+	// Enable collectors like collstats and indexstats depending on the number of collections
 	// present in the database.
 	limitsOk := false
-	if e.opts.CollStatsLimit == 0 || // Unlimited
-		(e.getTotalCollectionsCount() > 0 && e.getTotalCollectionsCount() < e.opts.CollStatsLimit) {
+	if e.opts.CollStatsLimit <= 0 || // Unlimited
+		e.getTotalCollectionsCount() <= e.opts.CollStatsLimit {
 		limitsOk = true
 	}
 
@@ -149,12 +150,14 @@ func (e *Exporter) makeRegistry(ctx context.Context, client *mongo.Client, topol
 		}
 		e.opts.EnableDiagnosticData = true
 		e.opts.EnableDBStats = true
+		e.opts.EnableCollStats = true
 		e.opts.EnableTopMetrics = true
 		e.opts.EnableReplicasetStatus = true
+		e.opts.EnableIndexStats = true
 	}
 
-	// if we manually set the collection names we want or auto discovery is set
-	if (len(e.opts.CollStatsNamespaces) > 0 || e.opts.DiscoveringMode) && e.opts.EnableCollStats && limitsOk {
+	// If we manually set the collection names we want or auto discovery is set.
+	if (len(e.opts.CollStatsNamespaces) > 0 || e.opts.DiscoveringMode) && e.opts.EnableCollStats && limitsOk && requestOpts.EnableCollStats {
 		cc := collstatsCollector{
 			ctx:             ctx,
 			client:          client,
@@ -167,8 +170,8 @@ func (e *Exporter) makeRegistry(ctx context.Context, client *mongo.Client, topol
 		registry.MustRegister(&cc)
 	}
 
-	// if we manually set the collection names we want or auto discovery is set
-	if (len(e.opts.IndexStatsCollections) > 0 || e.opts.DiscoveringMode) && e.opts.EnableIndexStats && limitsOk {
+	// If we manually set the collection names we want or auto discovery is set.
+	if (len(e.opts.IndexStatsCollections) > 0 || e.opts.DiscoveringMode) && e.opts.EnableIndexStats && limitsOk && requestOpts.EnableIndexStats {
 		ic := indexstatsCollector{
 			ctx:             ctx,
 			client:          client,
@@ -180,7 +183,7 @@ func (e *Exporter) makeRegistry(ctx context.Context, client *mongo.Client, topol
 		registry.MustRegister(&ic)
 	}
 
-	if e.opts.EnableDiagnosticData {
+	if e.opts.EnableDiagnosticData && requestOpts.EnableDiagnosticData {
 		ddc := diagnosticDataCollector{
 			ctx:            ctx,
 			client:         client,
@@ -191,7 +194,7 @@ func (e *Exporter) makeRegistry(ctx context.Context, client *mongo.Client, topol
 		registry.MustRegister(&ddc)
 	}
 
-	if e.opts.EnableDBStats {
+	if e.opts.EnableDBStats && requestOpts.EnableDBStats {
 		cc := dbstatsCollector{
 			ctx:            ctx,
 			client:         client,
@@ -202,7 +205,7 @@ func (e *Exporter) makeRegistry(ctx context.Context, client *mongo.Client, topol
 		registry.MustRegister(&cc)
 	}
 
-	if e.opts.EnableTopMetrics && nodeType != typeMongos {
+	if e.opts.EnableTopMetrics && nodeType != typeMongos && requestOpts.EnableTopMetrics {
 		tc := topCollector{
 			ctx:            ctx,
 			client:         client,
@@ -213,8 +216,8 @@ func (e *Exporter) makeRegistry(ctx context.Context, client *mongo.Client, topol
 		registry.MustRegister(&tc)
 	}
 
-	// replSetGetStatus is not supported through mongos
-	if e.opts.EnableReplicasetStatus && nodeType != typeMongos {
+	// replSetGetStatus is not supported through mongos.
+	if e.opts.EnableReplicasetStatus && nodeType != typeMongos && requestOpts.EnableReplicasetStatus {
 		rsgsc := replSetGetStatusCollector{
 			ctx:            ctx,
 			client:         client,
@@ -230,17 +233,17 @@ func (e *Exporter) makeRegistry(ctx context.Context, client *mongo.Client, topol
 
 func (e *Exporter) getClient(ctx context.Context) (*mongo.Client, error) {
 	if e.opts.GlobalConnPool {
-		// get global client. Maybe it must be initialized first.
+		// Get global client. Maybe it must be initialized first.
 		// Initialization is retried with every scrape until it succeeds once.
 		e.clientMu.Lock()
 		defer e.clientMu.Unlock()
 
-		// if client is already initialized, return it
+		// If client is already initialized, return it.
 		if e.client != nil {
 			return e.client, nil
 		}
 
-		client, err := connect(ctx, e.opts.URI, e.opts.DirectConnect)
+		client, err := connect(context.Background(), e.opts.URI, e.opts.DirectConnect)
 		if err != nil {
 			return nil, err
 		}
@@ -249,7 +252,7 @@ func (e *Exporter) getClient(ctx context.Context) (*mongo.Client, error) {
 		return client, nil
 	}
 
-	// !e.opts.GlobalConnPool: create new client for every scrape
+	// !e.opts.GlobalConnPool: create new client for every scrape.
 	client, err := connect(ctx, e.opts.URI, e.opts.DirectConnect)
 	if err != nil {
 		return nil, err
@@ -259,19 +262,50 @@ func (e *Exporter) getClient(ctx context.Context) (*mongo.Client, error) {
 }
 
 // Handler returns an http.Handler that serves metrics. Can be used instead of
-// Run for hooking up custom HTTP servers.
+// run for hooking up custom HTTP servers.
 func (e *Exporter) Handler() http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		seconds, err := strconv.Atoi(r.Header.Get("X-Prometheus-Scrape-Timeout-Seconds"))
+		// To support also older ones vmagents.
+		if err != nil {
+			seconds = 10
+		}
+
 		var client *mongo.Client
-		ctx, cancel := context.WithTimeout(r.Context(), time.Second)
+		ctx, cancel := context.WithTimeout(r.Context(), time.Duration(seconds)*time.Second)
 		defer cancel()
 
-		client, err := e.getClient(ctx)
+		filters := r.URL.Query()["collect[]"]
+
+		requestOpts := Opts{}
+
+		if len(filters) == 0 {
+			requestOpts = *e.opts
+		}
+
+		for _, filter := range filters {
+			switch filter {
+			case "diagnosticdata":
+				requestOpts.EnableDiagnosticData = true
+			case "replicasetstatus":
+				requestOpts.EnableReplicasetStatus = true
+			case "dbstats":
+				requestOpts.EnableDBStats = true
+			case "topmetrics":
+				requestOpts.EnableTopMetrics = true
+			case "indexstats":
+				requestOpts.EnableIndexStats = true
+			case "collstats":
+				requestOpts.EnableCollStats = true
+			}
+		}
+
+		client, err = e.getClient(ctx)
 		if err != nil {
 			e.logger.Errorf("Cannot connect to MongoDB: %v", err)
 		}
 
-		if client != nil && e.getTotalCollectionsCount() < 0 {
+		if client != nil && e.getTotalCollectionsCount() <= 0 {
 			count, err := nonSystemCollectionsCount(ctx, client, nil, nil)
 			if err == nil {
 				e.lock.Lock()
@@ -280,7 +314,7 @@ func (e *Exporter) Handler() http.Handler {
 			}
 		}
 
-		// Close client after usage
+		// Close client after usage.
 		if !e.opts.GlobalConnPool {
 			defer func() {
 				if client != nil {
@@ -292,7 +326,7 @@ func (e *Exporter) Handler() http.Handler {
 			}()
 		}
 
-		// topology can change between requests, so we need to get it every time
+		// Topology can change between requests, so we need to get it every time.
 		var ti *topologyInfo
 		if client != nil {
 			ti, err = newTopologyInfo(ctx, client)
@@ -306,7 +340,7 @@ func (e *Exporter) Handler() http.Handler {
 			}
 		}
 
-		registry := e.makeRegistry(ctx, client, ti)
+		registry := e.makeRegistry(ctx, client, ti, requestOpts)
 
 		var gatherers prometheus.Gatherers
 

exporter/exporter_test.go

@@ -187,7 +187,7 @@ func TestMongoS(t *testing.T) {
 			topologyInfo:   new(labelsGetterMock),
 		}
 
-		r := e.makeRegistry(ctx, client, new(labelsGetterMock))
+		r := e.makeRegistry(ctx, client, new(labelsGetterMock), *e.opts)
 
 		res := r.Unregister(&rsgsc)
 		assert.Equal(t, test.want, res, fmt.Sprintf("Port: %v", test.port))
@@ -218,7 +218,7 @@ func TestMongoUp(t *testing.T) {
 		logger: e.opts.Logger,
 	}
 
-	r := e.makeRegistry(ctx, client, new(labelsGetterMock))
+	r := e.makeRegistry(ctx, client, new(labelsGetterMock), *e.opts)
 
 	res := r.Unregister(&gc)
 	assert.Equal(t, true, res)

main.go

@@ -127,6 +127,7 @@ func buildExporter(opts GlobalFlags) *exporter.Exporter {
 		EnableCollStats:        opts.EnableCollStats,
 
 		CollStatsLimit: opts.CollStatsLimit,
+		CollectAll:     opts.CollectAll,
 	}
 
 	e := exporter.New(exporterOpts)