correct propagation of context, no mutex -> use errgroup instead, one more test for 2 loki queries; query executor and loki parallelization; added todo

Author: Tofel

wasp/benchspy/TO_DO.md

@@ -0,0 +1,11 @@
+Known things to do:
+- [ ] add logger
+- [ ] add tests
+- [ ] write documentation
+- [ ] test with Docker app (focus: resources)
+- [ ] test with k8s app (focus: resources)
+- [ ] add report builder (?)
+- [ ] add wrapper function for executing some code and then creating a report
+- [ ] think what to do with errors... do we really need a slice?
+- [ ] if TestEnd is zero, default to `time.Now()`?
+- [ ] add helper method for a profile what would create a report based on all generators?

wasp/benchspy/loki.go

@@ -11,6 +11,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/smartcontractkit/chainlink-testing-framework/lib/client"
 	"github.com/smartcontractkit/chainlink-testing-framework/wasp"
+	"golang.org/x/sync/errgroup"
 )
 
 func NewLokiQuery(queries map[string]string, lokiConfig *wasp.LokiConfig) *LokiQuery {
@@ -65,7 +66,7 @@ func (l *LokiQuery) Validate() error {
 	return nil
 }
 
-func (l *LokiQuery) Execute() error {
+func (l *LokiQuery) Execute(ctx context.Context) error {
 	splitAuth := strings.Split(l.LokiConfig.BasicAuth, ":")
 	var basicAuth client.LokiBasicAuth
 	if len(splitAuth) == 2 {
@@ -76,41 +77,54 @@ func (l *LokiQuery) Execute() error {
 	}
 
 	l.QueryResults = make(map[string][]string)
+	resultCh := make(chan map[string][]string, len(l.Queries))
+	errGroup, errCtx := errgroup.WithContext(ctx)
 
-	// TODO this can also be parallelized, just remember to add a mutex for writing to results map
 	for name, query := range l.Queries {
-		queryParams := client.LokiQueryParams{
-			Query:     query,
-			StartTime: l.StartTime,
-			EndTime:   l.EndTime,
-			Limit:     1000, //TODO make this configurable
-		}
+		errGroup.Go(func() error {
+			queryParams := client.LokiQueryParams{
+				Query:     query,
+				StartTime: l.StartTime,
+				EndTime:   l.EndTime,
+				Limit:     1000, //TODO make this configurable
+			}
 
-		parsedLokiUrl, err := url.Parse(l.LokiConfig.URL)
-		if err != nil {
-			return errors.Wrapf(err, "failed to parse Loki URL %s", l.LokiConfig.URL)
-		}
+			parsedLokiUrl, err := url.Parse(l.LokiConfig.URL)
+			if err != nil {
+				return errors.Wrapf(err, "failed to parse Loki URL %s", l.LokiConfig.URL)
+			}
 
-		lokiUrl := parsedLokiUrl.Scheme + "://" + parsedLokiUrl.Host
-		lokiClient := client.NewLokiClient(lokiUrl, l.LokiConfig.TenantID, basicAuth, queryParams)
+			lokiUrl := parsedLokiUrl.Scheme + "://" + parsedLokiUrl.Host
+			lokiClient := client.NewLokiClient(lokiUrl, l.LokiConfig.TenantID, basicAuth, queryParams)
 
-		ctx, cancelFn := context.WithTimeout(context.Background(), l.LokiConfig.Timeout)
-		rawLogs, err := lokiClient.QueryLogs(ctx)
-		if err != nil {
-			l.Errors = append(l.Errors, err)
-			cancelFn()
-			continue
-		}
+			rawLogs, err := lokiClient.QueryLogs(errCtx)
+			if err != nil {
+				return errors.Wrapf(err, "failed to query logs for %s", name)
+			}
 
-		cancelFn()
-		l.QueryResults[name] = []string{}
-		for _, log := range rawLogs {
-			l.QueryResults[name] = append(l.QueryResults[name], log.Log)
-		}
+			resultMap := make(map[string][]string)
+			for _, log := range rawLogs {
+				resultMap[name] = append(resultMap[name], log.Log)
+			}
+
+			select {
+			case resultCh <- resultMap:
+				return nil
+			case <-errCtx.Done():
+				return errCtx.Err() // Allows goroutine to exit if timeout occurs
+			}
+		})
 	}
 
-	if len(l.Errors) > 0 {
-		return errors.New("there were errors while fetching the results. Please check the errors and try again")
+	if err := errGroup.Wait(); err != nil {
+		return errors.Wrap(err, "failed to execute Loki queries")
+	}
+
+	for i := 0; i < len(l.Queries); i++ {
+		result := <-resultCh
+		for name, logs := range result {
+			l.QueryResults[name] = logs
+		}
 	}
 
 	return nil

wasp/benchspy/report.go

@@ -1,8 +1,11 @@
 package benchspy
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
+
+	"golang.org/x/sync/errgroup"
 )
 
 // StandardReport is a report that contains all the necessary data for a performance test
@@ -21,26 +24,35 @@ func (b *StandardReport) Load() error {
 	return b.LocalReportStorage.Load(b.TestName, b.CommitOrTag, b)
 }
 
-func (b *StandardReport) Fetch() error {
+func (b *StandardReport) Fetch(ctx context.Context) error {
 	basicErr := b.BasicData.Validate()
 	if basicErr != nil {
 		return basicErr
 	}
 
-	// TODO parallelize it
+	errGroup, errCtx := errgroup.WithContext(ctx)
+
 	for _, queryExecutor := range b.QueryExecutors {
-		queryExecutor.TimeRange(b.TestStart, b.TestEnd)
+		errGroup.Go(func() error {
+			queryExecutor.TimeRange(b.TestStart, b.TestEnd)
 
-		if validateErr := queryExecutor.Validate(); validateErr != nil {
-			return validateErr
-		}
+			if validateErr := queryExecutor.Validate(); validateErr != nil {
+				return validateErr
+			}
 
-		if execErr := queryExecutor.Execute(); execErr != nil {
-			return execErr
-		}
+			if execErr := queryExecutor.Execute(errCtx); execErr != nil {
+				return execErr
+			}
+
+			return nil
+		})
+	}
+
+	if err := errGroup.Wait(); err != nil {
+		return err
 	}
 
-	resourceErr := b.FetchResources()
+	resourceErr := b.FetchResources(ctx)
 	if resourceErr != nil {
 		return resourceErr
 	}

wasp/benchspy/resources.go

@@ -5,8 +5,6 @@ import (
 	"fmt"
 	"os"
 	"regexp"
-	"sync"
-	"time"
 
 	"github.com/docker/docker/api/types/container"
 	"github.com/pkg/errors"
@@ -49,14 +47,14 @@ type PodResources struct {
 	LimitsMemory   int64
 }
 
-func (r *ResourceReporter) FetchResources() error {
+func (r *ResourceReporter) FetchResources(ctx context.Context) error {
 	if r.ExecutionEnvironment == ExecutionEnvironment_Docker {
-		err := r.fetchDockerResources()
+		err := r.fetchDockerResources(ctx)
 		if err != nil {
 			return err
 		}
 	} else {
-		err := r.fetchK8sResources()
+		err := r.fetchK8sResources(ctx)
 		if err != nil {
 			return err
 		}
@@ -65,7 +63,7 @@ func (r *ResourceReporter) FetchResources() error {
 	return nil
 }
 
-func (r *ResourceReporter) fetchK8sResources() error {
+func (r *ResourceReporter) fetchK8sResources(ctx context.Context) error {
 	config, err := rest.InClusterConfig()
 	if err != nil {
 		return errors.Wrapf(err, "failed to get in-cluster config, are you sure this is running in a k8s cluster?")
@@ -82,9 +80,9 @@ func (r *ResourceReporter) fetchK8sResources() error {
 		return errors.Wrapf(err, "failed to read namespace file %s", namespaceFile)
 	}
 
-	pods, err := clientset.CoreV1().Pods(string(namespace)).List(context.TODO(), metav1.ListOptions{})
+	pods, err := clientset.CoreV1().Pods(string(namespace)).List(ctx, metav1.ListOptions{})
 	if err != nil {
-		panic(err)
+		return errors.Wrapf(err, "failed to list pods in namespace %s", namespace)
 	}
 
 	r.PodsResources = make(map[string]*PodResources)
@@ -101,22 +99,21 @@ func (r *ResourceReporter) fetchK8sResources() error {
 	return nil
 }
 
-func (r *ResourceReporter) fetchDockerResources() error {
+func (r *ResourceReporter) fetchDockerResources(ctx context.Context) error {
 	provider, err := tc.NewDockerProvider()
 	if err != nil {
 		return fmt.Errorf("failed to create Docker provider: %w", err)
 	}
 
-	containers, err := provider.Client().ContainerList(context.Background(), container.ListOptions{All: true})
+	containers, err := provider.Client().ContainerList(ctx, container.ListOptions{All: true})
 	if err != nil {
 		return fmt.Errorf("failed to list Docker containers: %w", err)
 	}
 
-	eg := &errgroup.Group{}
+	eg, errCtx := errgroup.WithContext(ctx)
 	pattern := regexp.MustCompile(r.ResourceSelectionPattern)
 
-	var dockerResources = make(map[string]*DockerResources)
-	resourceMutex := sync.Mutex{}
+	resultCh := make(chan map[string]*DockerResources, len(containers))
 
 	for _, containerInfo := range containers {
 		eg.Go(func() error {
@@ -125,32 +122,40 @@ func (r *ResourceReporter) fetchDockerResources() error {
 				return nil
 			}
 
-			ctx, cancelFn := context.WithTimeout(context.Background(), 30*time.Second)
-			info, err := provider.Client().ContainerInspect(ctx, containerInfo.ID)
+			info, err := provider.Client().ContainerInspect(errCtx, containerInfo.ID)
 			if err != nil {
-				cancelFn()
 				return errors.Wrapf(err, "failed to inspect container %s", containerName)
 			}
 
-			cancelFn()
-			resourceMutex.Lock()
-			dockerResources[containerName] = &DockerResources{
+			result := make(map[string]*DockerResources)
+			result[containerName] = &DockerResources{
 				NanoCPUs:   info.HostConfig.NanoCPUs,
 				CpuShares:  info.HostConfig.CPUShares,
 				Memory:     info.HostConfig.Memory,
 				MemorySwap: info.HostConfig.MemorySwap,
 			}
-			resourceMutex.Unlock()
 
-			return nil
+			select {
+			case resultCh <- result:
+				return nil
+			case <-errCtx.Done():
+				return errCtx.Err() // Allows goroutine to exit if timeout occurs
+			}
 		})
 	}
 
 	if err := eg.Wait(); err != nil {
 		return errors.Wrapf(err, "failed to fetch Docker resources")
 	}
 
-	r.ContainerResources = dockerResources
+	r.ContainerResources = make(map[string]*DockerResources)
+
+	for i := 0; i < len(containers); i++ {
+		result := <-resultCh
+		for name, res := range result {
+			r.ContainerResources[name] = res
+		}
+	}
 
 	return nil
 }

wasp/benchspy/types.go

@@ -1,6 +1,9 @@
 package benchspy
 
-import "time"
+import (
+	"context"
+	"time"
+)
 
 type Report interface {
 	// Store stores the report in a persistent storage and returns the path to it, or an error
@@ -17,7 +20,7 @@ type QueryExecutor interface {
 	// Validate checks if the QueryExecutor has all the necessary data and configuration to execute the queries
 	Validate() error
 	// Execute executes the queries and populates the QueryExecutor with the results
-	Execute() error
+	Execute(ctx context.Context) error
 	// Results returns the results of the queries, where key is the name of the query and value is the result
 	Results() map[string][]string
 	// IsComparable checks whether both QueryExecutors can be compared (e.g. they have the same type, queries are the same, etc.), and returns an error (if any difference is found)