feat(prose/golang): add more spans to track latency

dettanym · Mar 14, 2024 · 155492d · 155492d
1 parent df87ef3
commit 155492d
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 6 deletions.
diff --git a/privacy-profile-composer/pkg/envoyfilter/filter.go b/privacy-profile-composer/pkg/envoyfilter/filter.go
@@ -130,8 +130,10 @@ func (f *Filter) DecodeData(buffer api.BufferInstance, endStream bool) api.Statu
 		return api.StopAndBuffer
 	}
 
+	ctx := common.AddTracerToContext(context.Background(), f.tracer)
+
 	span, ctx := f.tracer.StartSpanFromContext(
-		context.Background(),
+		ctx,
 		"DecodeData",
 		zipkin.Parent(f.parentSpanContext),
 	)
@@ -218,8 +220,10 @@ func (f *Filter) EncodeData(buffer api.BufferInstance, endStream bool) api.Statu
 		return api.StopAndBuffer
 	}
 
+	ctx := common.AddTracerToContext(context.Background(), f.tracer)
+
 	span, ctx := f.tracer.StartSpanFromContext(
-		context.Background(),
+		ctx,
 		"EncodeData",
 		zipkin.Parent(f.parentSpanContext),
 	)
@@ -271,13 +275,13 @@ func (f *Filter) processBody(ctx context.Context, body string, isDecode bool) (s
 
 	proseTags[PROSE_SIDECAR_DIRECTION] = string(f.config.direction)
 
-	jsonBody, err := common.GetJSONBody(f.headerMetadata, body)
+	jsonBody, err := common.GetJSONBody(ctx, f.headerMetadata, body)
 	if err != nil {
 		return false, err, proseTags
 	}
 
 	// Run Presidio and add tags for PII types or an error from Presidio
-	piiTypes, err := common.PiiAnalysis(f.config.presidioUrl, f.headerMetadata.SvcName, jsonBody)
+	piiTypes, err := common.PiiAnalysis(ctx, f.config.presidioUrl, f.headerMetadata.SvcName, jsonBody)
 	if err != nil {
 		proseTags[PROSE_PRESIDIO_ERROR] = fmt.Sprintf("%s", err)
 		return false, err, proseTags
@@ -295,6 +299,9 @@ func (f *Filter) processBody(ctx context.Context, body string, isDecode bool) (s
 }
 
 func (f *Filter) runOPA(ctx context.Context, isDecode bool, dataItems []string) (sendLocalReply bool, err error, proseTags map[string]string) {
+	span, ctx := f.tracer.StartSpanFromContext(ctx, "runOPA")
+	defer span.Finish()
+
 	proseTags = map[string]string{}
 
 	log.Printf("direction: '%v'; isDecode: '%v', host: '%v', thirdPartyUrl: '%v'", f.config.direction, isDecode, f.headerMetadata.Host, f.thirdPartyURL)

diff --git a/privacy-profile-composer/pkg/envoyfilter/internal/common/packet_processing.go b/privacy-profile-composer/pkg/envoyfilter/internal/common/packet_processing.go
@@ -1,6 +1,7 @@
 package common
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"log"
@@ -129,7 +130,10 @@ func GetDirection(callbacks api.FilterCallbackHandler) (SidecarDirection, error)
 		"check the Envoy docs for the range of values for this key", directionInt)
 }
 
-func GetJSONBody(headerMetadata HeaderMetadata, body string) (interface{}, error) {
+func GetJSONBody(ctx context.Context, headerMetadata HeaderMetadata, body string) (interface{}, error) {
+	span, ctx := TracerFromContext(ctx).StartSpanFromContext(ctx, "getJSONBody")
+	defer span.Finish()
+
 	if headerMetadata.ContentType == nil {
 		return nil, fmt.Errorf("cannot analyze body, since 'ContentType' header is not set")
 	}

diff --git a/privacy-profile-composer/pkg/envoyfilter/internal/common/pii_analysis.go b/privacy-profile-composer/pkg/envoyfilter/internal/common/pii_analysis.go
@@ -2,6 +2,7 @@ package common
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -13,7 +14,10 @@ type PresidioDataFormat struct {
 	DerivePurpose string      `json:"derive_purpose,omitempty"`
 }
 
-func PiiAnalysis(presidioSvcURL string, svcName string, bufferBytes interface{}) ([]string, error) {
+func PiiAnalysis(ctx context.Context, presidioSvcURL string, svcName string, bufferBytes interface{}) ([]string, error) {
+	span, ctx := TracerFromContext(ctx).StartSpanFromContext(ctx, "PiiAnalysis")
+	defer span.Finish()
+
 	empty := []string{}
 
 	msgString, err := json.Marshal(

diff --git a/privacy-profile-composer/pkg/envoyfilter/internal/common/zipkin.go b/privacy-profile-composer/pkg/envoyfilter/internal/common/zipkin.go
@@ -1,6 +1,7 @@
 package common
 
 import (
+	"context"
 	"fmt"
 	"log"
 
@@ -88,3 +89,18 @@ func (t ZipkinTracer) Extract(h envoyapi.HeaderMap) model.SpanContext {
 	})
 
 }
+
+func TracerFromContext(ctx context.Context) *ZipkinTracer {
+	if t, ok := ctx.Value(tracerKey).(*ZipkinTracer); ok {
+		return t
+	}
+	return nil
+}
+
+func AddTracerToContext(ctx context.Context, tracer *ZipkinTracer) context.Context {
+	return context.WithValue(ctx, tracerKey, tracer)
+}
+
+type tracerCtxKey struct{}
+
+var tracerKey = tracerCtxKey{}