implement parallel queries

Signed-off-by: Markus Blaschke <mblaschke82@gmail.com>

Author: mblaschke

config/opts.go

@@ -21,6 +21,7 @@ type (
 
 		Loganalytics struct {
 			Workspace []string `long:"loganalytics.workspace"   env:"LOGANALYTICS_WORKSPACE"  env-delim:" " description:"Loganalytics workspace IDs" required:"true"`
+			Parallel  int      `long:"loganalytics.parallel"   env:"LOGANALYTICS_PARALLEL"                        description:"Specifies how many workspaces should be queried in parallel" default:"5"`
 		}
 
 		// config

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/prometheus/client_golang v1.10.0
 	github.com/prometheus/common v0.23.0 // indirect
+	github.com/remeh/sizedwaitgroup v1.0.0
 	github.com/sirupsen/logrus v1.8.1
 	github.com/webdevops/azure-resourcegraph-exporter v0.0.0-20210505184535-837efac736c6
 	golang.org/x/sys v0.0.0-20210503173754-0981d6026fa6 // indirect

go.sum

@@ -281,6 +281,8 @@ github.com/prometheus/procfs v0.2.0/go.mod h1:lV6e/gmhEcM9IjHGsFOCxxuZ+z1YqCvr4O
 github.com/prometheus/procfs v0.6.0 h1:mxy4L2jP6qMonqmq+aTtOx1ifVWUgG/TAmntgbh3xv4=
 github.com/prometheus/procfs v0.6.0/go.mod h1:cz+aTbrPOrUb4q7XlbU9ygM+/jj0fzG6c1xBZuNvfVA=
 github.com/rcrowley/go-metrics v0.0.0-20181016184325-3113b8401b8a/go.mod h1:bCqnVzQkZxMG4s8nGwiZ5l3QUCyqpo9Y+/ZMZ9VjZe4=
+github.com/remeh/sizedwaitgroup v1.0.0 h1:VNGGFwNo/R5+MJBf6yrsr110p0m4/OX4S3DCy7Kyl5E=
+github.com/remeh/sizedwaitgroup v1.0.0/go.mod h1:3j2R4OIe/SeS6YDhICBy22RWjJC5eNCJ1V+9+NVNYlo=
 github.com/rogpeppe/fastuuid v0.0.0-20150106093220-6724a57986af/go.mod h1:XWv6SoW27p1b0cqNHllgS5HIMJraePCO15w5zCzIWYg=
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
 github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=

loganalytics.go

@@ -0,0 +1,92 @@
+package main
+
+import (
+	"context"
+	operationalinsights "github.com/Azure/azure-sdk-for-go/services/operationalinsights/v1/operationalinsights"
+	"github.com/Azure/go-autorest/autorest"
+	"github.com/Azure/go-autorest/autorest/to"
+	log "github.com/sirupsen/logrus"
+	"github.com/webdevops/azure-resourcegraph-exporter/kusto"
+	"net/http"
+)
+
+const (
+	OPINSIGHTS_URL_SUFFIX = "/v1"
+)
+
+type (
+	probeResult struct {
+		Name    string
+		Metrics []kusto.MetricRow
+	}
+)
+
+func NewLoganalyticsQueryClient() operationalinsights.QueryClient {
+	// Create and authorize a operationalinsights client
+	client := operationalinsights.NewQueryClientWithBaseURI(AzureEnvironment.ResourceIdentifiers.OperationalInsights + OPINSIGHTS_URL_SUFFIX)
+	client.Authorizer = OpInsightsAuthorizer
+	client.ResponseInspector = respondDecorator()
+	return client
+}
+
+func SendQueryToLoganalyticsWorkspace(ctx context.Context, logger *log.Entry, workspaceId string, queryClient operationalinsights.QueryClient, queryConfig kusto.ConfigQuery, result chan<- probeResult) {
+	workspaceLogger := logger.WithField("workspaceId", workspaceId)
+
+	// Set options
+	workspaces := []string{}
+	queryBody := operationalinsights.QueryBody{
+		Query:      &queryConfig.Query,
+		Timespan:   queryConfig.Timespan,
+		Workspaces: &workspaces,
+	}
+
+	workspaceLogger.WithField("query", queryConfig.Query).Debug("send query to loganaltyics workspace")
+	var results, queryErr = queryClient.Execute(ctx, workspaceId, queryBody)
+
+	if queryErr == nil {
+		workspaceLogger.Debug("fetched query result")
+		resultTables := *results.Tables
+
+		if len(resultTables) >= 1 {
+			for _, table := range resultTables {
+				if table.Rows == nil || table.Columns == nil {
+					// no results found, skip table
+					continue
+				}
+
+				for _, v := range *table.Rows {
+					resultRow := map[string]interface{}{}
+
+					for colNum, colName := range *resultTables[0].Columns {
+						resultRow[to.String(colName.Name)] = v[colNum]
+					}
+
+					for metricName, metric := range kusto.BuildPrometheusMetricList(queryConfig.Metric, queryConfig.MetricConfig, resultRow) {
+						// inject workspaceId
+						for num := range metric {
+							metric[num].Labels["workspaceTable"] = to.String(table.Name)
+							metric[num].Labels["workspaceID"] = workspaceId
+						}
+
+						result <- probeResult{
+							Name:    metricName,
+							Metrics: metric,
+						}
+					}
+				}
+			}
+		}
+
+		workspaceLogger.Debug("metrics parsed")
+	} else {
+		workspaceLogger.Panic(queryErr.Error())
+	}
+}
+
+func respondDecorator() autorest.RespondDecorator {
+	return func(p autorest.Responder) autorest.Responder {
+		return autorest.ResponderFunc(func(r *http.Response) error {
+			return nil
+		})
+	}
+}

misc.go

@@ -0,0 +1,16 @@
+package main
+
+import (
+	"github.com/remeh/sizedwaitgroup"
+	"net/http"
+	"sync"
+)
+
+func NewWaitGroup() sync.WaitGroup {
+	return sync.WaitGroup{}
+}
+
+func NewWaitGroupWithSize(r *http.Request) sizedwaitgroup.SizedWaitGroup {
+	size := opts.Loganalytics.Parallel
+	return sizedwaitgroup.New(size)
+}

probe.go

@@ -4,9 +4,6 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	operationalinsights "github.com/Azure/azure-sdk-for-go/services/operationalinsights/v1/operationalinsights"
-	"github.com/Azure/go-autorest/autorest"
-	"github.com/Azure/go-autorest/autorest/to"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	log "github.com/sirupsen/logrus"
@@ -15,11 +12,15 @@ import (
 	"time"
 )
 
-const (
-	OPINSIGHTS_URL_SUFFIX = "/v1"
-)
-
 func handleProbeRequest(w http.ResponseWriter, r *http.Request) {
+	defer func() {
+		if r := recover(); r != nil {
+			log.Errorf(fmt.Sprintf("%v", r))
+			http.Error(w, fmt.Sprintf("%v", r), http.StatusBadRequest)
+			return
+		}
+	}()
+
 	registry := prometheus.NewRegistry()
 
 	requestTime := time.Now()
@@ -42,11 +43,7 @@ func handleProbeRequest(w http.ResponseWriter, r *http.Request) {
 	}
 
 	ctx := context.Background()
-
-	// Create and authorize a operationalinsights client
-	queryClient := operationalinsights.NewQueryClientWithBaseURI(AzureEnvironment.ResourceIdentifiers.OperationalInsights + OPINSIGHTS_URL_SUFFIX)
-	queryClient.Authorizer = OpInsightsAuthorizer
-	queryClient.ResponseInspector = respondDecorator()
+	queryClient := NewLoganalyticsQueryClient()
 
 	metricList := kusto.MetricList{}
 	metricList.Init()
@@ -69,7 +66,9 @@ func handleProbeRequest(w http.ResponseWriter, r *http.Request) {
 
 	if executeQuery {
 		w.Header().Add("X-metrics-cached", "false")
-		for _, queryConfig := range Config.Queries {
+		for _, queryRow := range Config.Queries {
+			queryConfig := queryRow
+
 			// check if query matches module name
 			if queryConfig.Module != moduleName {
 				continue
@@ -88,63 +87,46 @@ func handleProbeRequest(w http.ResponseWriter, r *http.Request) {
 			contextLogger.Debug("starting query")
 
 			resultTotalRecords := 0
-			for _, workspaceId := range opts.Loganalytics.Workspace {
-				workspaceLogger := contextLogger.WithField("workspaceId", workspaceId)
-
-				// Set options
-				workspaces := []string{}
-				queryBody := operationalinsights.QueryBody{
-					Query:      &queryConfig.Query,
-					Timespan:   queryConfig.Timespan,
-					Workspaces: &workspaces,
+
+			resultChannel := make(chan probeResult)
+			wgProbes := NewWaitGroupWithSize(r)
+			wgProcess := NewWaitGroup()
+
+			// collect metrics
+			wgProcess.Add(1)
+			go func() {
+				defer wgProcess.Done()
+				for result := range resultChannel {
+					resultTotalRecords++
+					metricList.Add(result.Name, result.Metrics...)
 				}
+			}()
 
+			// query workspaces
+			for _, row := range opts.Loganalytics.Workspace {
+				workspaceId := row
 				// Run the query and get the results
 				prometheusQueryRequests.With(prometheus.Labels{"workspace": workspaceId, "module": moduleName, "metric": queryConfig.Metric}).Inc()
 
-				var results, queryErr = queryClient.Execute(ctx, workspaceId, queryBody)
-				resultTotalRecords = 1
-
-				if queryErr == nil {
-					contextLogger.Debug("parsing result")
-					resultTables := *results.Tables
-
-					if len(resultTables) >= 1 {
-						for _, table := range resultTables {
-							if table.Rows == nil || table.Columns == nil {
-								// no results found, skip table
-								continue
-							}
-
-							for _, v := range *table.Rows {
-								resultTotalRecords++
-								resultRow := map[string]interface{}{}
-
-								for colNum, colName := range *resultTables[0].Columns {
-									resultRow[to.String(colName.Name)] = v[colNum]
-								}
-
-								for metricName, metric := range kusto.BuildPrometheusMetricList(queryConfig.Metric, queryConfig.MetricConfig, resultRow) {
-									// inject workspaceId
-									for num := range metric {
-										metric[num].Labels["workspaceTable"] = to.String(table.Name)
-										metric[num].Labels["workspaceID"] = workspaceId
-									}
-
-									metricList.Add(metricName, metric...)
-								}
-							}
-						}
-					}
-
-					workspaceLogger.Debug("metrics parsed")
-				} else {
-					workspaceLogger.Errorln(queryErr.Error())
-					http.Error(w, queryErr.Error(), http.StatusBadRequest)
-					return
-				}
+				wgProbes.Add()
+				go func() {
+					defer wgProbes.Done()
+					SendQueryToLoganalyticsWorkspace(
+						ctx,
+						contextLogger,
+						workspaceId,
+						queryClient,
+						queryConfig,
+						resultChannel,
+					)
+				}()
 			}
 
+			// wait until queries are done for closing channel and waiting for result process
+			wgProbes.Wait()
+			close(resultChannel)
+			wgProcess.Wait()
+
 			elapsedTime := time.Since(startTime)
 			contextLogger.WithField("results", resultTotalRecords).Debugf("fetched %v results", resultTotalRecords)
 			prometheusQueryTime.With(prometheus.Labels{"module": moduleName, "metric": queryConfig.Metric}).Observe(elapsedTime.Seconds())
@@ -186,11 +168,3 @@ func handleProbeRequest(w http.ResponseWriter, r *http.Request) {
 	h := promhttp.HandlerFor(registry, promhttp.HandlerOpts{})
 	h.ServeHTTP(w, r)
 }
-
-func respondDecorator() autorest.RespondDecorator {
-	return func(p autorest.Responder) autorest.Responder {
-		return autorest.ResponderFunc(func(r *http.Response) error {
-			return nil
-		})
-	}
-}