chore: correlation for worker logger

Signed-off-by: Lukas Zapletal <[email protected]>

internal/db/connection.go

@@ -74,16 +74,31 @@ func Initialize(ctx context.Context, schema string) error {
 
 		if logLevel > 0 {
 			zeroLogger := pgxlog.NewLogger(log.Logger,
-				pgxlog.WithContextFunc(func(ctx context.Context, logWith zerolog.Context) zerolog.Context {
+				pgxlog.WithContextFunc(func(ctx context.Context, zx zerolog.Context) zerolog.Context {
+					jobId := logging.JobId(ctx)
+					if jobId != "" {
+						zx = zx.Str("job_id", jobId)
+					}
+					reservationId := logging.ReservationId(ctx)
+					if reservationId != 0 {
+						zx = zx.Int64("reservation_id", reservationId)
+					}
 					traceId := logging.TraceId(ctx)
 					if traceId != "" {
-						logWith = logWith.Str("trace_id", traceId)
+						zx = zx.Str("trace_id", traceId)
+					}
+					requestId := logging.EdgeRequestId(ctx)
+					if requestId != "" {
+						zx = zx.Str("request_id", requestId)
 					}
 					accountId := identity.AccountIdOrNil(ctx)
 					if accountId != 0 {
-						logWith = logWith.Int64("account_id", accountId)
+						zx = zx.Int64("account_id", accountId)
 					}
-					return logWith
+					principal := identity.Identity(ctx)
+					zx = zx.Str("org_id", principal.Identity.OrgID)
+					zx = zx.Str("account_number", principal.Identity.AccountNumber)
+					return zx
 				}))
 			poolConfig.ConnConfig.Tracer = &tracelog.TraceLog{
 				Logger:   zeroLogger,

internal/identity/identity.go

@@ -13,7 +13,11 @@ type Principal = identity.XRHID
 
 // Identity returns identity header struct or nil when not set.
 func Identity(ctx context.Context) Principal {
-	return identity.Get(ctx)
+	val := ctx.Value(identity.Key)
+	if val == nil {
+		return Principal{}
+	}
+	return val.(Principal)
 }
 
 // IdentityHeader returns identity header (base64-encoded JSON)

internal/jobs/ctx.go

@@ -5,6 +5,9 @@ import (
 
 	"github.com/RHEnVision/provisioning-backend/internal/identity"
 	"github.com/RHEnVision/provisioning-backend/internal/logging"
+	"github.com/RHEnVision/provisioning-backend/internal/ptr"
+	"github.com/google/uuid"
+	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 )
 
@@ -18,3 +21,25 @@ func copyContext(ctx context.Context) context.Context {
 	nCtx = identity.WithAccountId(nCtx, identity.AccountId(ctx))
 	return nCtx
 }
+
+func jobContextLogger(origCtx context.Context, principal identity.Principal, traceID, edgeID string, jobID uuid.UUID, accountID, reservationID int64) (context.Context, *zerolog.Logger) {
+	ctx := logging.WithJobId(origCtx, jobID.String())
+	ctx = logging.WithReservationId(ctx, reservationID)
+	ctx = identity.WithIdentity(ctx, principal)
+	ctx = logging.WithTraceId(ctx, traceID)
+	ctx = logging.WithEdgeRequestId(ctx, edgeID)
+	ctx = identity.WithAccountId(ctx, accountID)
+
+	logger := zerolog.Ctx(ctx)
+	logger = ptr.To(logger.With().
+		Int64("account_id", accountID).
+		Str("org_id", principal.Identity.OrgID).
+		Str("account_number", principal.Identity.AccountNumber).
+		Str("trace_id", traceID).
+		Str("request_id", edgeID).
+		Int64("reservation_id", reservationID).
+		Str("job_id", jobID.String()).
+		Logger())
+	ctx = logger.WithContext(ctx)
+	return ctx, logger
+}

internal/jobs/launch_instance_aws.go

@@ -5,18 +5,16 @@ import (
 	"errors"
 	"fmt"
 
-	"go.opentelemetry.io/otel"
-	"go.opentelemetry.io/otel/codes"
-
 	"github.com/RHEnVision/provisioning-backend/internal/clients"
 	"github.com/RHEnVision/provisioning-backend/internal/clients/http"
 	"github.com/RHEnVision/provisioning-backend/internal/dao"
 	"github.com/RHEnVision/provisioning-backend/internal/models"
-	"github.com/RHEnVision/provisioning-backend/internal/ptr"
 	"github.com/RHEnVision/provisioning-backend/internal/userdata"
 	"github.com/RHEnVision/provisioning-backend/pkg/worker"
 	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
 	"github.com/rs/zerolog"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/codes"
 )
 
 type LaunchInstanceAWSTaskArgs struct {
@@ -60,9 +58,11 @@ func HandleLaunchInstanceAWS(ctx context.Context, job *worker.Job) {
 		return
 	}
 
+	// context and logger
+	ctx, logger = jobContextLogger(ctx, job.Identity, job.TraceID, job.EdgeID, job.ID, job.AccountID, args.ReservationID)
+	logger.Info().Msg("Started launch instance AWS job")
+
 	// ensure panic finishes the job
-	logger = ptr.To(logger.With().Int64("reservation_id", args.ReservationID).Logger())
-	ctx = logger.WithContext(ctx)
 	defer func() {
 		if r := recover(); r != nil {
 			panicErr := fmt.Errorf("%w: %s", ErrPanicInJob, r)

internal/jobs/launch_instance_azure.go

@@ -64,22 +64,23 @@ func HandleLaunchInstanceAzure(ctx context.Context, job *worker.Job) {
 		return
 	}
 
+	// context and logger
+	ctx, logger = jobContextLogger(ctx, job.Identity, job.TraceID, job.EdgeID, job.ID, job.AccountID, args.ReservationID)
+	logger.Info().Msg("Started launch instance Azure job")
+
 	if args.ResourceGroupName == "" {
 		logger.Debug().Msg("Resource group has not been set, defaulting to 'redhat-deployed'")
 		args.ResourceGroupName = DefaultAzureResourceGroupName
 	}
 
 	// ensure panic finishes the job
-	logger = ptr.To(logger.With().Int64("reservation_id", args.ReservationID).Logger())
-	ctx = logger.WithContext(ctx)
 	defer func() {
 		if r := recover(); r != nil {
 			panicErr := fmt.Errorf("%w: %s", ErrPanicInJob, r)
 			finishWithError(ctx, args.ReservationID, panicErr)
 		}
 	}()
 
-	logger.Info().Msg("Started launch instance Azure job")
 	ctx, span := otel.Tracer(TraceName).Start(ctx, "LaunchInstanceAzureJob")
 	defer span.End()
 

internal/jobs/launch_instance_gcp.go

@@ -11,7 +11,6 @@ import (
 	_ "github.com/RHEnVision/provisioning-backend/internal/clients/http/gcp"
 	"github.com/RHEnVision/provisioning-backend/internal/dao"
 	"github.com/RHEnVision/provisioning-backend/internal/models"
-	"github.com/RHEnVision/provisioning-backend/internal/ptr"
 	"github.com/RHEnVision/provisioning-backend/internal/userdata"
 	"github.com/RHEnVision/provisioning-backend/pkg/worker"
 	"github.com/rs/zerolog"
@@ -56,9 +55,11 @@ func HandleLaunchInstanceGCP(ctx context.Context, job *worker.Job) {
 		return
 	}
 
+	// context and logger
+	ctx, logger = jobContextLogger(ctx, job.Identity, job.TraceID, job.EdgeID, job.ID, job.AccountID, args.ReservationID)
+	logger.Info().Msg("Started launch instance GCP job")
+
 	// ensure panic finishes the job
-	logger = ptr.To(logger.With().Int64("reservation_id", args.ReservationID).Logger())
-	ctx = logger.WithContext(ctx)
 	defer func() {
 		if r := recover(); r != nil {
 			panicErr := fmt.Errorf("%w: %s", ErrPanicInJob, r)

internal/jobs/noop_job.go

@@ -33,8 +33,9 @@ func HandleNoop(ctx context.Context, job *worker.Job) {
 		return
 	}
 
-	logger := zerolog.Ctx(ctx).With().Int64("reservation_id", args.ReservationID).Logger()
-	ctx = logger.WithContext(ctx)
+	// context and logger
+	ctx, _ = jobContextLogger(ctx, job.Identity, job.TraceID, job.EdgeID, job.ID, job.AccountID, args.ReservationID)
+
 	nc := notifications.GetNotificationClient(ctx)
 
 	jobErr := DoNoop(ctx, &args)

internal/logging/ctx.go

@@ -10,6 +10,8 @@ const (
 	requestIdCtxKey     commonKeyId = iota
 	edgeRequestIdCtxKey commonKeyId = iota
 	correlationCtxKey   commonKeyId = iota
+	jobIdCtxKey         commonKeyId = iota
+	reservationIdCtxKey commonKeyId = iota
 )
 
 // CorrelationId returns UI correlation id or an empty string when not set.
@@ -53,3 +55,31 @@ func TraceId(ctx context.Context) string {
 func WithTraceId(ctx context.Context, id string) context.Context {
 	return context.WithValue(ctx, requestIdCtxKey, id)
 }
+
+// JobId returns request id or an empty string when not set.
+func JobId(ctx context.Context) string {
+	value := ctx.Value(jobIdCtxKey)
+	if value == nil {
+		return ""
+	}
+	return value.(string)
+}
+
+// WithJobId returns context copy with trace id value.
+func WithJobId(ctx context.Context, id string) context.Context {
+	return context.WithValue(ctx, jobIdCtxKey, id)
+}
+
+// ReservationId returns request id or an empty string when not set.
+func ReservationId(ctx context.Context) int64 {
+	value := ctx.Value(reservationIdCtxKey)
+	if value == nil {
+		return 0
+	}
+	return value.(int64)
+}
+
+// WithReservationId returns context copy with trace id value.
+func WithReservationId(ctx context.Context, id int64) context.Context {
+	return context.WithValue(ctx, reservationIdCtxKey, id)
+}

internal/services/aws_reservation_service.go

@@ -5,6 +5,8 @@ import (
 	"net/http"
 	"strings"
 
+	"github.com/RHEnVision/provisioning-backend/internal/logging"
+
 	"github.com/RHEnVision/provisioning-backend/internal/clients"
 	_ "github.com/RHEnVision/provisioning-backend/internal/clients/http/image_builder"
 	"github.com/RHEnVision/provisioning-backend/internal/config"
@@ -148,6 +150,8 @@ func CreateAWSReservation(w http.ResponseWriter, r *http.Request) {
 	launchJob := worker.Job{
 		Type:      jobs.TypeLaunchInstanceAws,
 		Identity:  id,
+		TraceID:   logging.TraceId(r.Context()),
+		EdgeID:    logging.EdgeRequestId(r.Context()),
 		AccountID: accountId,
 		Args: jobs.LaunchInstanceAWSTaskArgs{
 			ReservationID:    reservation.ID,

internal/services/azure_reservation_service.go

@@ -5,6 +5,8 @@ import (
 	"net/http"
 	"strings"
 
+	"github.com/RHEnVision/provisioning-backend/internal/logging"
+
 	"github.com/RHEnVision/provisioning-backend/internal/clients"
 	"github.com/RHEnVision/provisioning-backend/internal/config"
 	"github.com/RHEnVision/provisioning-backend/internal/dao"
@@ -142,6 +144,8 @@ func CreateAzureReservation(w http.ResponseWriter, r *http.Request) {
 	launchJob := worker.Job{
 		Type:      jobs.TypeLaunchInstanceAzure,
 		Identity:  identity.Identity(r.Context()),
+		TraceID:   logging.TraceId(r.Context()),
+		EdgeID:    logging.EdgeRequestId(r.Context()),
 		AccountID: identity.AccountId(r.Context()),
 		Args: jobs.LaunchInstanceAzureTaskArgs{
 			ReservationID:     reservation.ID,

internal/services/gcp_reservation_service.go

@@ -6,6 +6,8 @@ import (
 	"regexp"
 	"strings"
 
+	"github.com/RHEnVision/provisioning-backend/internal/logging"
+
 	"github.com/RHEnVision/provisioning-backend/internal/preload"
 	"github.com/google/uuid"
 	"github.com/rs/zerolog"
@@ -140,6 +142,8 @@ func CreateGCPReservation(w http.ResponseWriter, r *http.Request) {
 	launchJob := worker.Job{
 		Type:      jobs.TypeLaunchInstanceGcp,
 		AccountID: accountId,
+		TraceID:   logging.TraceId(r.Context()),
+		EdgeID:    logging.EdgeRequestId(r.Context()),
 		Identity:  id,
 		Args: jobs.LaunchInstanceGCPTaskArgs{
 			ReservationID:    reservation.ID,

internal/services/noop_reservation_service.go

@@ -3,6 +3,8 @@ package services
 import (
 	"net/http"
 
+	"github.com/RHEnVision/provisioning-backend/internal/logging"
+
 	"github.com/RHEnVision/provisioning-backend/internal/dao"
 	"github.com/RHEnVision/provisioning-backend/internal/identity"
 	"github.com/RHEnVision/provisioning-backend/internal/jobs"
@@ -44,6 +46,8 @@ func CreateNoopReservation(w http.ResponseWriter, r *http.Request) {
 		Type:      jobs.TypeNoop,
 		AccountID: accountId,
 		Identity:  identity,
+		TraceID:   logging.TraceId(r.Context()),
+		EdgeID:    logging.EdgeRequestId(r.Context()),
 		Args: jobs.NoopJobArgs{
 			ReservationID: reservation.ID,
 		},

pkg/worker/job.go

@@ -31,6 +31,12 @@ type Job struct {
 	// Associated identity
 	Identity identity.Principal
 
+	// For logging purposes
+	TraceID string
+
+	// For logging purposes
+	EdgeID string
+
 	// Job arguments.
 	Args any
 }