refactor: trace interface again

w-h-a · Oct 13, 2024 · 165437d · 165437d
1 parent e731466
commit 165437d
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 33 deletions.
diff --git a/sidecar/custom/sidecar.go b/sidecar/custom/sidecar.go
@@ -184,9 +184,9 @@ func (s *customSidecar) UnsubscribeFromBroker(brokerId string) error {
 func (s *customSidecar) ReadFromSecretStore(ctx context.Context, secretStore string, name string) (*sidecar.Secret, error) {
 	// TODO: make sure we have a tracer
 
-	newCtx := s.options.Tracer.Start(ctx, "customSidecar.ReadFromSecretStore")
+	_, spanId := s.options.Tracer.Start(ctx, "customSidecar.ReadFromSecretStore")
 
-	s.options.Tracer.AddMetadata(newCtx, map[string]string{
+	s.options.Tracer.AddMetadata(spanId, map[string]string{
 		"secretStore": secretStore,
 		"name":        name,
 	})
@@ -195,20 +195,20 @@ func (s *customSidecar) ReadFromSecretStore(ctx context.Context, secretStore str
 	if !ok {
 		log.Warnf("secret store %s was not found", secretStore)
 		// TODO: update status of span
-		s.options.Tracer.Finish(newCtx)
+		s.options.Tracer.Finish(spanId)
 		return nil, sidecar.ErrComponentNotFound
 	}
 
 	mp, err := sc.GetSecret(name)
 	if err != nil {
 		// TODO: update status of span
-		s.options.Tracer.Finish(newCtx)
+		s.options.Tracer.Finish(spanId)
 		return nil, err
 	}
 
 	// TODO: update status of span
 
-	s.options.Tracer.Finish(newCtx)
+	s.options.Tracer.Finish(spanId)
 
 	return &sidecar.Secret{
 		Data: mp,

diff --git a/telemetry/tracev2/memory/trace.go b/telemetry/tracev2/memory/trace.go
@@ -24,7 +24,7 @@ func (t *memoryTrace) Options() tracev2.TraceOptions {
 	return t.options
 }
 
-func (t *memoryTrace) Start(ctx context.Context, name string) context.Context {
+func (t *memoryTrace) Start(ctx context.Context, name string) (context.Context, string) {
 	// TODO: make sure we're enabled
 
 	t.mtx.Lock()
@@ -42,27 +42,20 @@ func (t *memoryTrace) Start(ctx context.Context, name string) context.Context {
 
 	ctx, span := t.start(ctx, name, parentCtxCfg)
 
-	t.spans[span.SpanContext().SpanID().String()] = span
+	key := span.SpanContext().SpanID().String()
 
-	newCtx, _ := tracev2.ContextWithSpan(ctx, span.SpanContext().SpanID())
+	t.spans[key] = span
 
-	log.Infof("MY CTX %+#v", newCtx)
+	newCtx, _ := tracev2.ContextWithSpan(ctx, span.SpanContext().SpanID())
 
-	return newCtx
+	return newCtx, key
 }
 
-func (t *memoryTrace) AddMetadata(ctx context.Context, md map[string]string) {
+func (t *memoryTrace) AddMetadata(span string, md map[string]string) {
 	t.mtx.Lock()
 	defer t.mtx.Unlock()
 
-	span, found := tracev2.SpanFromContext(ctx)
-	if !found {
-		return
-	}
-
-	key := string(span[:])
-
-	if t.spans[key] == nil {
+	if t.spans[span] == nil {
 		return
 	}
 
@@ -80,23 +73,16 @@ func (t *memoryTrace) AddMetadata(ctx context.Context, md map[string]string) {
 		return
 	}
 
-	t.spans[key].SetAttributes(attrs...)
+	t.spans[span].SetAttributes(attrs...)
 }
 
-func (t *memoryTrace) Finish(ctx context.Context) {
+func (t *memoryTrace) Finish(span string) {
 	t.mtx.Lock()
 	defer t.mtx.Unlock()
 
-	span, found := tracev2.SpanFromContext(ctx)
-	if !found {
-		return
-	}
-
-	key := string(span[:])
-
-	t.spans[key].End()
+	t.spans[span].End()
 
-	delete(t.spans, key)
+	delete(t.spans, span)
 }
 
 func (t *memoryTrace) Read(opts ...tracev2.ReadOption) ([]*tracev2.SpanData, error) {

diff --git a/telemetry/tracev2/trace.go b/telemetry/tracev2/trace.go
@@ -8,9 +8,9 @@ import (
 type Trace interface {
 	Options() TraceOptions
 	// TODO: add cfg argument to check for whether tracing is enabled
-	Start(ctx context.Context, name string) context.Context
-	AddMetadata(ctx context.Context, md map[string]string)
-	Finish(ctx context.Context)
+	Start(ctx context.Context, name string) (context.Context, string)
+	AddMetadata(span string, md map[string]string)
+	Finish(span string)
 	Read(opts ...ReadOption) ([]*SpanData, error)
 	String() string
 }