Repository: clouditor/clouditor
Pull request #1162 (2023-07-11 11:24:52)

Add context to hooks
This pr adds context to the hook functions. 


Patch:
diff --git a/api/assessment/assessment.go b/api/assessment/assessment.go
index d13b29c7..3e7e4e52 100644
--- a/api/assessment/assessment.go
+++ b/api/assessment/assessment.go
@@ -26,12 +26,13 @@
 package assessment
 
 import (
+	"context"
 	"errors"
 
 	"google.golang.org/protobuf/proto"
 )
 
-type ResultHookFunc func(result *AssessmentResult, err error)
+type ResultHookFunc func(ctx context.Context, result *AssessmentResult, err error)
 
 var (
 	ErrMetricConfigurationMissing            = errors.New("metric configuration in assessment result is missing")
diff --git a/api/evidence/evidence.go b/api/evidence/evidence.go
index 7a597b59..d907dd94 100644
--- a/api/evidence/evidence.go
+++ b/api/evidence/evidence.go
@@ -26,6 +26,7 @@
 package evidence
 
 import (
+	"context"
 	"errors"
 	"fmt"
 
@@ -33,7 +34,7 @@ import (
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
-type EvidenceHookFunc func(result *Evidence, err error)
+type EvidenceHookFunc func(ctx context.Context, evidence *Evidence, err error)
 
 var (
 	ErrResourceNotStruct             = errors.New("resource in evidence is not struct value")
diff --git a/service/assessment/assessment.go b/service/assessment/assessment.go
index 66c56673..609a3b41 100644
--- a/service/assessment/assessment.go
+++ b/service/assessment/assessment.go
@@ -224,7 +224,7 @@ func (svc *Service) AssessEvidence(ctx context.Context, req *assessment.AssessEv
 	}
 
 	// Assess evidence
-	_, err = svc.handleEvidence(req.Evidence, resourceId)
+	_, err = svc.handleEvidence(ctx, req.Evidence, resourceId)
 	if err != nil {
 		err = fmt.Errorf("error while handling evidence: %v", err)
 		log.Error(err)
@@ -290,7 +290,7 @@ func (svc *Service) AssessEvidences(stream assessment.Assessment_AssessEvidences
 }
 
 // handleEvidence is the helper method for the actual assessment used by AssessEvidence and AssessEvidences
-func (svc *Service) handleEvidence(ev *evidence.Evidence, resourceId string) (results []*assessment.AssessmentResult, err error) {
+func (svc *Service) handleEvidence(ctx context.Context, ev *evidence.Evidence, resourceId string) (results []*assessment.AssessmentResult, err error) {
 	var types []string
 
 	log.Debugf("Evaluating evidence %s (%s) collected by %s at %s", ev.Id, resourceId, ev.ToolId, ev.Timestamp.AsTime())
@@ -300,7 +300,7 @@ func (svc *Service) handleEvidence(ev *evidence.Evidence, resourceId string) (re
 	if err != nil {
 		newError := fmt.Errorf("could not evaluate evidence: %w", err)
 
-		go svc.informHooks(nil, newError)
+		go svc.informHooks(ctx, nil, newError)
 
 		return nil, newError
 	}
@@ -312,7 +312,7 @@ func (svc *Service) handleEvidence(ev *evidence.Evidence, resourceId string) (re
 		if err != nil {
 			err = fmt.Errorf("could not get stream to evidence store (%s): %w", svc.evidenceStore.Target, err)
 
-			go svc.informHooks(nil, err)
+			go svc.informHooks(ctx, nil, err)
 
 			return nil, err
 		}
@@ -324,7 +324,7 @@ func (svc *Service) handleEvidence(ev *evidence.Evidence, resourceId string) (re
 	if err != nil {
 		err = fmt.Errorf("could not get stream to orchestrator (%s): %w", svc.orchestrator.Target, err)
 
-		go svc.informHooks(nil, err)
+		go svc.informHooks(ctx, nil, err)
 
 		return nil, err
 	}
@@ -360,7 +360,7 @@ func (svc *Service) handleEvidence(ev *evidence.Evidence, resourceId string) (re
 		}
 
 		// Inform hooks about new assessment result
-		go svc.informHooks(result, nil)
+		go svc.informHooks(ctx, result, nil)
 
 		// Send assessment result in orchestratorChannel
 		channelOrchestrator.Send(&orchestrator.StoreAssessmentResultRequest{Result: result})
@@ -372,7 +372,7 @@ func (svc *Service) handleEvidence(ev *evidence.Evidence, resourceId string) (re
 }
 
 // informHooks informs the registered hook functions
-func (svc *Service) informHooks(result *assessment.AssessmentResult, err error) {
+func (svc *Service) informHooks(ctx context.Context, result *assessment.AssessmentResult, err error) {
 	svc.hookMutex.RLock()
 	hooks := svc.resultHooks
 	defer svc.hookMutex.RUnlock()
@@ -381,12 +381,12 @@ func (svc *Service) informHooks(result *assessment.AssessmentResult, err error)
 	if len(hooks) > 0 {
 		for _, hook := range hooks {
 			// We could do hook concurrent again (assuming different hooks don't interfere with each other)
-			hook(result, err)
+			hook(ctx, result, err)
 		}
 	}
 }
 
-func (svc *Service) RegisterAssessmentResultHook(assessmentResultsHook func(result *assessment.AssessmentResult, err error)) {
+func (svc *Service) RegisterAssessmentResultHook(assessmentResultsHook func(ctx context.Context, result *assessment.AssessmentResult, err error)) {
 	svc.hookMutex.Lock()
 	defer svc.hookMutex.Unlock()
 	svc.resultHooks = append(svc.resultHooks, assessmentResultsHook)
diff --git a/service/evidence/evidence_store.go b/service/evidence/evidence_store.go
index b758b89d..feaeba73 100644
--- a/service/evidence/evidence_store.go
+++ b/service/evidence/evidence_store.go
@@ -29,11 +29,12 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"golang.org/x/exp/slices"
 	"io"
 	"strings"
 	"sync"
 
+	"golang.org/x/exp/slices"
+
 	"clouditor.io/clouditor/api/evidence"
 	"clouditor.io/clouditor/internal/logging"
 	"clouditor.io/clouditor/persistence"
@@ -119,7 +120,7 @@ func (svc *Service) StoreEvidence(ctx context.Context, req *evidence.StoreEviden
 		return nil, status.Errorf(codes.Internal, "database error: %v", err)
 	}
 
-	go svc.informHooks(req.Evidence, nil)
+	go svc.informHooks(ctx, req.Evidence, nil)
 
 	res = &evidence.StoreEvidenceResponse{}
 
@@ -278,7 +279,7 @@ func (svc *Service) RegisterEvidenceHook(evidenceHook evidence.EvidenceHookFunc)
 	svc.evidenceHooks = append(svc.evidenceHooks, evidenceHook)
 }
 
-func (svc *Service) informHooks(result *evidence.Evidence, err error) {
+func (svc *Service) informHooks(ctx context.Context, result *evidence.Evidence, err error) {
 	svc.mu.Lock()
 	defer svc.mu.Unlock()
 
@@ -286,7 +287,7 @@ func (svc *Service) informHooks(result *evidence.Evidence, err error) {
 	if svc.evidenceHooks != nil {
 		for _, hook := range svc.evidenceHooks {
 			// TODO(all): We could do hook concurrent again (assuming different hooks don't interfere with each other)
-			hook(result, err)
+			hook(ctx, result, err)
 		}
 	}
 }


Test patch:
diff --git a/service/assessment/assessment_test.go b/service/assessment/assessment_test.go
index 0da19fa3..769fd077 100644
--- a/service/assessment/assessment_test.go
+++ b/service/assessment/assessment_test.go
@@ -584,13 +584,13 @@ func TestService_AssessmentResultHooks(t *testing.T) {
 
 	wg.Add(hookCounts)
 
-	firstHookFunction := func(assessmentResult *assessment.AssessmentResult, err error) {
+	firstHookFunction := func(ctx context.Context, assessmentResult *assessment.AssessmentResult, err error) {
 		hookCallCounter++
 		log.Println("Hello from inside the firstHookFunction")
 		wg.Done()
 	}
 
-	secondHookFunction := func(assessmentResult *assessment.AssessmentResult, err error) {
+	secondHookFunction := func(ctx context.Context, assessmentResult *assessment.AssessmentResult, err error) {
 		hookCallCounter++
 		log.Println("Hello from inside the secondHookFunction")
 		wg.Done()
@@ -957,7 +957,7 @@ func TestService_HandleEvidence(t *testing.T) {
 			}
 
 			// Two tests: 1st) wantErr function. 2nd) if wantErr false then check if the results are valid
-			results, err := s.handleEvidence(tt.args.evidence, tt.args.resourceId)
+			results, err := s.handleEvidence(context.Background(), tt.args.evidence, tt.args.resourceId)
 			if !tt.wantErr(t, err, fmt.Sprintf("handleEvidence(%v, %v)", tt.args.evidence, tt.args.resourceId)) {
 				assert.NotEmpty(t, results)
 				// Check the result by validation
diff --git a/service/evidence/evidence_store_test.go b/service/evidence/evidence_store_test.go
index 9b94905d..1843189f 100644
--- a/service/evidence/evidence_store_test.go
+++ b/service/evidence/evidence_store_test.go
@@ -529,14 +529,14 @@ func TestService_EvidenceHook(t *testing.T) {
 	)
 	wg.Add(2)
 
-	firstHookFunction := func(evidence *evidence.Evidence, err error) {
+	firstHookFunction := func(ctx context.Context, evidence *evidence.Evidence, err error) {
 		hookCallCounter++
 		log.Println("Hello from inside the firstHookFunction")
 
 		wg.Done()
 	}
 
-	secondHookFunction := func(evidence *evidence.Evidence, err error) {
+	secondHookFunction := func(ctx context.Context, evidence *evidence.Evidence, err error) {
 		hookCallCounter++
 		log.Println("Hello from inside the secondHookFunction")