Adding correlation IDs to reconcile loops

Signed-off-by: Aaron Schlesinger <[email protected]>

controllers/azuremachine_controller.go

@@ -137,7 +137,8 @@ func (r *AzureMachineReconciler) Reconcile(ctx context.Context, req ctrl.Request
 			attribute.String("namespace", req.Namespace),
 			attribute.String("name", req.Name),
 			attribute.String("kind", "AzureMachine"),
-		))
+		),
+	)
 	defer span.End()
 
 	// Fetch the AzureMachine VM.

exp/controllers/azuremachinepool_controller.go

@@ -158,18 +158,18 @@ func (ampr *AzureMachinePoolReconciler) SetupWithManager(ctx context.Context, mg
 
 // Reconcile idempotently gets, creates, and updates a machine pool.
 func (ampr *AzureMachinePoolReconciler) Reconcile(ctx context.Context, req ctrl.Request) (_ ctrl.Result, reterr error) {
-	ctx, cancel := context.WithTimeout(ctx, reconciler.DefaultedLoopTimeout(ampr.ReconcileTimeout))
-	defer cancel()
-
-	logger := ampr.Log.WithValues("namespace", req.Namespace, "azureMachinePool", req.Name)
-
 	ctx, span := tele.Tracer().Start(ctx, "controllers.AzureMachinePoolReconciler.Reconcile",
 		trace.WithAttributes(
 			attribute.String("namespace", req.Namespace),
 			attribute.String("name", req.Name),
 			attribute.String("kind", "AzureMachinePool"),
-		))
+		),
+	)
 	defer span.End()
+	ctx, cancel := context.WithTimeout(ctx, reconciler.DefaultedLoopTimeout(ampr.ReconcileTimeout))
+	defer cancel()
+
+	logger := ampr.Log.WithValues("namespace", req.Namespace, "azureMachinePool", req.Name)
 
 	azMachinePool := &infrav1exp.AzureMachinePool{}
 	err := ampr.Get(ctx, req.NamespacedName, azMachinePool)

exp/controllers/azuremachinepoolmachine_controller.go

@@ -127,17 +127,17 @@ func (ampmr *AzureMachinePoolMachineController) SetupWithManager(ctx context.Con
 
 // Reconcile idempotently gets, creates, and updates a machine pool.
 func (ampmr *AzureMachinePoolMachineController) Reconcile(ctx context.Context, req ctrl.Request) (_ ctrl.Result, reterr error) {
-	ctx, cancel := context.WithTimeout(ctx, reconciler.DefaultedLoopTimeout(ampmr.ReconcileTimeout))
-	defer cancel()
-	logger := ampmr.Log.WithValues("namespace", req.Namespace, "azureMachinePoolMachine", req.Name)
-
 	ctx, span := tele.Tracer().Start(ctx, "controllers.AzureMachinePoolMachineController.Reconcile",
 		trace.WithAttributes(
 			attribute.String("namespace", req.Namespace),
 			attribute.String("name", req.Name),
 			attribute.String("kind", "AzureMachinePoolMachine"),
-		))
+		),
+	)
 	defer span.End()
+	ctx, cancel := context.WithTimeout(ctx, reconciler.DefaultedLoopTimeout(ampmr.ReconcileTimeout))
+	defer cancel()
+	logger := ampmr.Log.WithValues("namespace", req.Namespace, "azureMachinePoolMachine", req.Name)
 
 	machine := &infrav1exp.AzureMachinePoolMachine{}
 	err := ampmr.Get(ctx, req.NamespacedName, machine)

util/tele/corr_id.go

@@ -0,0 +1,78 @@
+/*
+Copyright 2021 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package tele
+
+import (
+	"context"
+
+	"github.com/google/uuid"
+)
+
+type corrIDKey string
+
+// CorrID is a correlation ID that the cluster API provider
+// sends with all API requests to Azure. Do not create one
+// of these manually. Instead, use the CtxWithCorrelationID function
+// to create one of these within a context.Context.
+type CorrID string
+
+const corrIDKeyVal corrIDKey = "x-ms-correlation-id"
+
+// ctxWithCorrID creates a CorrID and creates a new context.Context
+// with the new CorrID in it. It returns the _new_ context and the
+// newly created CorrID. If there was a problem creating the correlation
+// ID, the new context will not have the correlation ID in it and the
+// returned CorrID will be the empty string.After you call this function, prefer to
+// use the newly created context over the old one. Common usage is
+// below:
+//
+// 	ctx := context.Background()
+//	ctx, newCorrID := CtxWithCorrID(ctx)
+//	fmt.Println("new corr ID: ", newCorrID)
+//	doSomething(ctx)
+func ctxWithCorrID(ctx context.Context) (context.Context, CorrID) {
+	currentCorrIDIface := ctx.Value(corrIDKeyVal)
+	if currentCorrIDIface != nil {
+		currentCorrID, ok := currentCorrIDIface.(CorrID)
+		if ok {
+			return ctx, currentCorrID
+		}
+	}
+	uid, err := uuid.NewRandom()
+	if err != nil {
+		return nil, CorrID("")
+	}
+	newCorrID := CorrID(uid.String())
+	ctx = context.WithValue(ctx, corrIDKeyVal, newCorrID)
+	return ctx, newCorrID
+}
+
+// CorrIDFromCtx attempts to fetch a correlation ID from the given
+// context.Context. If none exists, returns an empty CorrID and false.
+// Otherwise returns the CorrID value and true.
+func CorrIDFromCtx(ctx context.Context) (CorrID, bool) {
+	currentCorrIDIface := ctx.Value(corrIDKeyVal)
+	if currentCorrIDIface == nil {
+		return CorrID(""), false
+	}
+
+	if corrID, ok := currentCorrIDIface.(CorrID); ok {
+		return corrID, ok
+	}
+
+	return CorrID(""), false
+}

util/tele/tele.go

@@ -5,7 +5,7 @@ Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
 
-    http://www.apache.org/licenses/LICENSE-2.0
+	http://www.apache.org/licenses/LICENSE-2.0
 
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
@@ -17,11 +17,39 @@ limitations under the License.
 package tele
 
 import (
+	"context"
+
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/trace"
 )
 
-// Tracer returns the default opentelemetry tracer.
+type tracer struct {
+	trace.Tracer
+}
+
+// Start creates a new context with a new Azure correlation ID, then
+// creates a new trace.Span with that new context. This function then
+// returns the new Context and Span.
+func (t tracer) Start(
+	ctx context.Context,
+	op string,
+	opts ...trace.SpanOption,
+) (context.Context, trace.Span) {
+	ctx, _ = ctxWithCorrID(ctx)
+	return t.Tracer.Start(ctx, op, opts...)
+}
+
+// Tracer returns an OpenTelemetry Tracer implementation to be used
+// to create spans. Use this implementation instead of the "raw" one that
+// you could otherwise get from calling `otel.Tracer("whatever")`.
+//
+// Example usage:
+//
+//	ctx, span := tele.Tracer().Start(ctx, "myFunction")
+//	defer span.End()
+//	// use the span and context here
 func Tracer() trace.Tracer {
-	return otel.Tracer("capz")
+	return tracer{
+		Tracer: otel.Tracer("capz"),
+	}
 }