From fb2173d6d564148634f1f7d2ef6089fc8ece476e Mon Sep 17 00:00:00 2001
From: Ondrej Ezr <oezr@redhat.com>
Date: Wed, 11 Oct 2023 23:17:19 +0200
Subject: [PATCH] wip: pass spans properly

Otel needs a root spans, if that is not passed from external service,
we should start one ourselves.
Also jobs shall retrieve these spans and not start new one.
---
 cmd/pbackend/api.go              |  2 +-
 internal/logging/ctx.go          | 13 +++++-------
 internal/middleware/telemetry.go | 36 ++++++++++++++++++++++++++++++++
 internal/middleware/trace_id.go  | 35 -------------------------------
 pkg/worker/job.go                |  8 +++++++
 5 files changed, 50 insertions(+), 44 deletions(-)
 create mode 100644 internal/middleware/telemetry.go
 delete mode 100644 internal/middleware/trace_id.go

diff --git a/cmd/pbackend/api.go b/cmd/pbackend/api.go
index a293ac83..7989a556 100644
--- a/cmd/pbackend/api.go
+++ b/cmd/pbackend/api.go
@@ -100,7 +100,7 @@ func api() {
 	apiRouter.Use(telemetry.Middleware(apiRouter))
 	apiRouter.Use(m.VersionMiddleware)
 	apiRouter.Use(m.CorrelationID)
-	apiRouter.Use(m.TraceID)
+	apiRouter.Use(m.Telemetry)
 	apiRouter.Use(m.LoggerMiddleware(&log.Logger))
 
 	// Mount paths
diff --git a/internal/logging/ctx.go b/internal/logging/ctx.go
index 0bd89b04..b24063d9 100644
--- a/internal/logging/ctx.go
+++ b/internal/logging/ctx.go
@@ -2,6 +2,8 @@ package logging
 
 import (
 	"context"
+
+	"go.opentelemetry.io/otel/trace"
 )
 
 type commonKeyId int
@@ -42,14 +44,9 @@ func WithEdgeRequestId(ctx context.Context, id string) context.Context {
 
 // TraceId returns request id or an empty string when not set.
 func TraceId(ctx context.Context) string {
-	value := ctx.Value(requestIdCtxKey)
-	if value == nil {
+	value := trace.SpanContextFromContext(ctx).TraceID()
+	if !value.IsValid() {
 		return ""
 	}
-	return value.(string)
-}
-
-// WithTraceId returns context copy with trace id value.
-func WithTraceId(ctx context.Context, id string) context.Context {
-	return context.WithValue(ctx, requestIdCtxKey, id)
+	return value.String()
 }
diff --git a/internal/middleware/telemetry.go b/internal/middleware/telemetry.go
new file mode 100644
index 00000000..4dc4a7f1
--- /dev/null
+++ b/internal/middleware/telemetry.go
@@ -0,0 +1,36 @@
+package middleware
+
+import (
+	"net/http"
+
+	"github.com/RHEnVision/provisioning-backend/internal/logging"
+	"github.com/RHEnVision/provisioning-backend/internal/telemetry"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/trace"
+)
+
+const TraceName = telemetry.TracePrefix + "internal/middleware"
+
+// Telemetry middleware starts a new telemetry span for this request,
+// it tries to find the parent trace in the request,
+// if none is found, it starts new root span.
+func Telemetry(next http.Handler) http.Handler {
+	fn := func(w http.ResponseWriter, r *http.Request) {
+		var span trace.Span
+		ctx := r.Context()
+
+		// Edge request id
+		edgeId := r.Header.Get("X-Rh-Edge-Request-Id")
+		if edgeId != "" {
+			ctx = logging.WithEdgeRequestId(ctx, edgeId)
+		}
+
+		ctx, span = otel.Tracer(TraceName).Start(ctx, r.URL.Path)
+
+		// Store TraceID in response headers for easier debugging
+		w.Header().Set("X-Trace-Id", span.SpanContext().TraceID().String())
+
+		next.ServeHTTP(w, r.WithContext(ctx))
+	}
+	return http.HandlerFunc(fn)
+}
diff --git a/internal/middleware/trace_id.go b/internal/middleware/trace_id.go
deleted file mode 100644
index 3b1f941c..00000000
--- a/internal/middleware/trace_id.go
+++ /dev/null
@@ -1,35 +0,0 @@
-package middleware
-
-import (
-	"net/http"
-
-	"github.com/RHEnVision/provisioning-backend/internal/logging"
-	"github.com/RHEnVision/provisioning-backend/internal/random"
-	"go.opentelemetry.io/otel/trace"
-)
-
-func TraceID(next http.Handler) http.Handler {
-	fn := func(w http.ResponseWriter, r *http.Request) {
-		ctx := r.Context()
-
-		// Edge request id
-		edgeId := r.Header.Get("X-Rh-Edge-Request-Id")
-		if edgeId != "" {
-			ctx = logging.WithEdgeRequestId(ctx, edgeId)
-		}
-
-		// OpenTelemetry trace id
-		traceId := trace.SpanFromContext(ctx).SpanContext().TraceID()
-		if !traceId.IsValid() {
-			// OpenTelemetry library does not provide a public interface to create new IDs
-			traceId = random.TraceID()
-		}
-
-		// Store in response headers for easier debugging
-		w.Header().Set("X-Trace-Id", traceId.String())
-
-		ctx = logging.WithTraceId(ctx, traceId.String())
-		next.ServeHTTP(w, r.WithContext(ctx))
-	}
-	return http.HandlerFunc(fn)
-}
diff --git a/pkg/worker/job.go b/pkg/worker/job.go
index 7eb1e0d5..2ec44838 100644
--- a/pkg/worker/job.go
+++ b/pkg/worker/job.go
@@ -5,8 +5,11 @@ import (
 	"errors"
 
 	"github.com/RHEnVision/provisioning-backend/internal/identity"
+	"github.com/RHEnVision/provisioning-backend/internal/telemetry"
 	"github.com/google/uuid"
 	"github.com/rs/zerolog"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/trace"
 )
 
 func init() {
@@ -77,9 +80,14 @@ func contextLogger(ctx context.Context, job *Job) context.Context {
 		return ctx
 	}
 
+	// if otel enabled
+	var span trace.Span
+	ctx, span = otel.Tracer(telemetry.TracePrefix+"pkg/worker").Start(ctx, string(job.Type))
+
 	accountId := job.AccountID
 	id := job.Identity
 	logger := zerolog.Ctx(ctx).With().
+		Str("trace_id", span.SpanContext().TraceID().String()).
 		Str("job_id", job.ID.String()).
 		Int64("account_id", accountId).
 		Str("account_number", id.Identity.AccountNumber).