feat: add OpenTelemetry Tracing to Sandbox and SandboxClaim Controllers (kubernetes-sigs#234)

* Implements OpenTelemetry tracing for the SandboxClaim and Sandbox

* Autogenerated Files

* updates metrics package to indicate that these are internal metrics

* Updates tracing Instrumenter interface by adding support for initial attributes, span events, and recording state checks.

* go mod tidy after a rebase

Author: igooch

Dockerfile

@@ -18,6 +18,7 @@ COPY api/ ./api/
 COPY cmd/ ./cmd/
 COPY controllers/ ./controllers/
 COPY extensions/ ./extensions/
+COPY internal/ ./internal/
 
 # Build the binary with optimizations
 RUN CGO_ENABLED=0 GOOS=linux GOARCH=${TARGETARCH} go build -ldflags="-s -w" -o /agent-sandbox-controller ./cmd/agent-sandbox-controller

cmd/agent-sandbox-controller/main.go

@@ -15,11 +15,13 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"net/http"
 	"net/http/pprof"
 	"os"
 	"runtime"
+	"time"
 
 	// Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.)
 	// to ensure that exec-entrypoint and run can make use of them.
@@ -34,6 +36,7 @@ import (
 	"sigs.k8s.io/agent-sandbox/controllers"
 	extensionsv1alpha1 "sigs.k8s.io/agent-sandbox/extensions/api/v1alpha1"
 	extensionscontrollers "sigs.k8s.io/agent-sandbox/extensions/controllers"
+	asmetrics "sigs.k8s.io/agent-sandbox/internal/metrics"
 	//+kubebuilder:scaffold:imports
 )
 
@@ -46,6 +49,7 @@ func main() {
 	var enableLeaderElection bool
 	var probeAddr string
 	var extensions bool
+	var enableTracing bool
 	var enablePprof bool
 	var enablePprofDebug bool
 	var pprofBlockProfileRate int
@@ -56,6 +60,7 @@ func main() {
 		"Enable leader election for controller manager. "+
 			"Enabling this will ensure there is only one active controller manager.")
 	flag.BoolVar(&extensions, "extensions", false, "Enable extensions controllers.")
+	flag.BoolVar(&enableTracing, "enable-tracing", false, "Enable OpenTelemetry tracing via OTLP.")
 	flag.BoolVar(&enablePprof, "enable-pprof", false,
 		"Enable CPU profiling endpoint (/debug/pprof/profile) on the metrics server.")
 	flag.BoolVar(&enablePprofDebug, "enable-pprof-debug", false,
@@ -74,6 +79,24 @@ func main() {
 	flag.Parse()
 
 	ctrl.SetLogger(zap.New(zap.UseFlagOptions(&opts)))
+	ctx := ctrl.SetupSignalHandler()
+
+	// Initialize Tracing Provider
+	var instrumenter asmetrics.Instrumenter = asmetrics.NewNoOp()
+	if enableTracing {
+		var cleanup func()
+		var err error
+		// Use a timeout context for initialization to prevent blocking
+		initCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+		defer cancel()
+
+		instrumenter, cleanup, err = asmetrics.SetupOTel(initCtx, "agent-sandbox-controller")
+		if err != nil {
+			setupLog.Error(err, "unable to initialize tracing")
+			os.Exit(1)
+		}
+		defer cleanup()
+	}
 
 	// Importing net/http/pprof registers handlers on the global DefaultServeMux.
 	// Reset it to avoid accidentally exposing pprof via any server that uses the default mux.
@@ -135,6 +158,7 @@ func main() {
 	if err = (&controllers.SandboxReconciler{
 		Client: mgr.GetClient(),
 		Scheme: mgr.GetScheme(),
+		Tracer: instrumenter,
 	}).SetupWithManager(mgr); err != nil {
 		setupLog.Error(err, "unable to create controller", "controller", "Sandbox")
 		os.Exit(1)
@@ -145,6 +169,7 @@ func main() {
 			Client:   mgr.GetClient(),
 			Scheme:   mgr.GetScheme(),
 			Recorder: mgr.GetEventRecorderFor("sandboxclaim-controller"),
+			Tracer:   instrumenter,
 		}).SetupWithManager(mgr); err != nil {
 			setupLog.Error(err, "unable to create controller", "controller", "SandboxClaim")
 			os.Exit(1)
@@ -170,7 +195,7 @@ func main() {
 	}
 
 	setupLog.Info("starting manager")
-	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
+	if err := mgr.Start(ctx); err != nil {
 		setupLog.Error(err, "problem running manager")
 		os.Exit(1)
 	}

controllers/sandbox_controller.go

@@ -38,6 +38,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/predicate"
 
 	sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1"
+	asmetrics "sigs.k8s.io/agent-sandbox/internal/metrics"
 )
 
 const (
@@ -60,6 +61,7 @@ func init() {
 type SandboxReconciler struct {
 	client.Client
 	Scheme *runtime.Scheme
+	Tracer asmetrics.Instrumenter
 }
 
 //+kubebuilder:rbac:groups=agents.x-k8s.io,resources=sandboxes,verbs=get;list;watch;create;update;patch;delete
@@ -89,6 +91,14 @@ func (r *SandboxReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
 		return ctrl.Result{}, err
 	}
 
+	// Start Tracing Span
+	initialAttrs := map[string]string{
+		"sandbox.name":      sandbox.Name,
+		"sandbox.namespace": sandbox.Namespace,
+	}
+	ctx, end := r.Tracer.StartSpan(ctx, sandbox, "ReconcileSandbox", initialAttrs)
+	defer end()
+
 	// If the sandbox is being deleted, do nothing
 	if !sandbox.ObjectMeta.DeletionTimestamp.IsZero() {
 		log.Info("Sandbox is being deleted")
@@ -104,6 +114,22 @@ func (r *SandboxReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
 		return ctrl.Result{}, nil
 	}
 
+	// Initialize trace ID for active resources missing an ID
+	tc := r.Tracer.GetTraceContext(ctx)
+	if tc != "" && (sandbox.Annotations == nil || sandbox.Annotations[asmetrics.TraceContextAnnotation] == "") {
+		patch := client.MergeFrom(sandbox.DeepCopy())
+		if sandbox.Annotations == nil {
+			sandbox.Annotations = make(map[string]string)
+		}
+		sandbox.Annotations[asmetrics.TraceContextAnnotation] = tc
+
+		if err := r.Patch(ctx, sandbox, patch); err != nil {
+			return ctrl.Result{}, err
+		}
+		// Return to ensure the next loop uses the persisted ID
+		return ctrl.Result{}, nil
+	}
+
 	if sandbox.Spec.Replicas == nil {
 		replicas := int32(1)
 		sandbox.Spec.Replicas = &replicas
@@ -303,6 +329,10 @@ func (r *SandboxReconciler) reconcileService(ctx context.Context, sandbox *sandb
 func (r *SandboxReconciler) reconcilePod(ctx context.Context, sandbox *sandboxv1alpha1.Sandbox, nameHash string) (*corev1.Pod, error) {
 	log := log.FromContext(ctx)
 
+	// Start a child span of ReconcileSandbox
+	ctx, end := r.Tracer.StartSpan(ctx, nil, "reconcilePod", nil)
+	defer end()
+
 	// List all pods with the pool label matching the warm pool name hash
 	// TODO: find a better way to make sure one sandbox has at most one pod
 	podList := &corev1.PodList{}
@@ -344,7 +374,7 @@ func (r *SandboxReconciler) reconcilePod(ctx context.Context, sandbox *sandboxv1
 		pod = nil
 	}
 
-	// if replicas is 0, delete the pod if it exists
+	// 1. PATH: Logic for deleting Pod when replicas is 0
 	if *sandbox.Spec.Replicas == 0 {
 		if pod != nil {
 			if pod.ObjectMeta.DeletionTimestamp.IsZero() {
@@ -372,9 +402,17 @@ func (r *SandboxReconciler) reconcilePod(ctx context.Context, sandbox *sandboxv1
 		return nil, nil
 	}
 
+	// 2. PATH: Existing Pod found (e.g., adopted from WarmPool or already exists)
 	if pod != nil {
 		log.Info("Found Pod", "Pod.Namespace", pod.Namespace, "Pod.Name", pod.Name)
 
+		if r.Tracer.IsRecording(ctx) {
+			r.Tracer.AddEvent(ctx, "ExistingPodStatusObserved", map[string]string{
+				"pod.Name":  pod.Name,
+				"pod.Phase": string(pod.Status.Phase),
+			})
+		}
+
 		if pod.Labels == nil {
 			pod.Labels = make(map[string]string)
 		}
@@ -396,6 +434,7 @@ func (r *SandboxReconciler) reconcilePod(ctx context.Context, sandbox *sandboxv1
 		return pod, nil
 	}
 
+	// 3. PATH: Create new Pod
 	log.Info("Creating a new Pod", "Pod.Namespace", sandbox.Namespace, "Pod.Name", sandbox.Name)
 	labels := map[string]string{
 		sandboxLabel: nameHash,
@@ -439,11 +478,23 @@ func (r *SandboxReconciler) reconcilePod(ctx context.Context, sandbox *sandboxv1
 		return nil, err
 	}
 
+	if r.Tracer.IsRecording(ctx) {
+		r.Tracer.AddEvent(ctx, "NewPodStatusObserved", map[string]string{
+			"pod.Name":  pod.Name,
+			"pod.Phase": string(pod.Status.Phase),
+		})
+	}
+
 	return pod, nil
 }
 
 func (r *SandboxReconciler) reconcilePVCs(ctx context.Context, sandbox *sandboxv1alpha1.Sandbox) error {
 	log := log.FromContext(ctx)
+
+	// Start a child span of ReconcileSandbox
+	ctx, end := r.Tracer.StartSpan(ctx, nil, "reconcilePVCs", nil)
+	defer end()
+
 	for _, pvcTemplate := range sandbox.Spec.VolumeClaimTemplates {
 		pvc := &corev1.PersistentVolumeClaim{}
 		pvcName := pvcTemplate.Name + "-" + sandbox.Name

controllers/sandbox_controller_test.go

@@ -29,10 +29,12 @@ import (
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/utils/ptr"
-	sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
+
+	sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1"
+	asmetrics "sigs.k8s.io/agent-sandbox/internal/metrics"
 )
 
 func newFakeClient(initialObjs ...runtime.Object) client.WithWatch {
@@ -505,6 +507,7 @@ func TestReconcile(t *testing.T) {
 			r := SandboxReconciler{
 				Client: newFakeClient(append(tc.initialObjs, sb)...),
 				Scheme: Scheme,
+				Tracer: asmetrics.NewNoOp(),
 			}
 
 			_, err := r.Reconcile(t.Context(), ctrl.Request{
@@ -861,6 +864,7 @@ func TestReconcilePod(t *testing.T) {
 			r := SandboxReconciler{
 				Client: newFakeClient(append(tc.initialObjs, tc.sandbox)...),
 				Scheme: Scheme,
+				Tracer: asmetrics.NewNoOp(),
 			}
 
 			pod, err := r.reconcilePod(t.Context(), tc.sandbox, nameHash)

extensions/controllers/sandboxclaim_controller.go

@@ -40,6 +40,7 @@ import (
 	sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1"
 	sandboxcontrollers "sigs.k8s.io/agent-sandbox/controllers"
 	extensionsv1alpha1 "sigs.k8s.io/agent-sandbox/extensions/api/v1alpha1"
+	asmetrics "sigs.k8s.io/agent-sandbox/internal/metrics"
 )
 
 // TODO: These constants should be imported from the main controller package Issue #216
@@ -55,6 +56,7 @@ type SandboxClaimReconciler struct {
 	client.Client
 	Scheme   *runtime.Scheme
 	Recorder record.EventRecorder
+	Tracer   asmetrics.Instrumenter
 }
 
 //+kubebuilder:rbac:groups=extensions.agents.x-k8s.io,resources=sandboxclaims,verbs=get;list;watch;create;update;patch;delete
@@ -77,10 +79,28 @@ func (r *SandboxClaimReconciler) Reconcile(ctx context.Context, req ctrl.Request
 		return ctrl.Result{}, fmt.Errorf("failed to get sandbox claim %q: %w", req.NamespacedName, err)
 	}
 
+	// Start Tracing Span
+	ctx, end := r.Tracer.StartSpan(ctx, claim, "ReconcileSandboxClaim", nil)
+	defer end()
+
 	if !claim.DeletionTimestamp.IsZero() {
 		return ctrl.Result{}, nil
 	}
 
+	// Initialize trace ID for active resources missing an ID
+	tc := r.Tracer.GetTraceContext(ctx)
+	if tc != "" && (claim.Annotations == nil || claim.Annotations[asmetrics.TraceContextAnnotation] == "") {
+		patch := client.MergeFrom(claim.DeepCopy())
+		if claim.Annotations == nil {
+			claim.Annotations = make(map[string]string)
+		}
+		claim.Annotations[asmetrics.TraceContextAnnotation] = tc
+		if err := r.Patch(ctx, claim, patch); err != nil {
+			return ctrl.Result{}, err
+		}
+		return ctrl.Result{}, nil
+	}
+
 	originalClaimStatus := claim.Status.DeepCopy()
 
 	// Check Expiration
@@ -411,6 +431,14 @@ func (r *SandboxClaimReconciler) createSandbox(ctx context.Context, claim *exten
 		},
 	}
 
+	// Propagate the trace context annotation to the Sandbox resource
+	if sandbox.Annotations == nil {
+		sandbox.Annotations = make(map[string]string)
+	}
+	if tc, ok := claim.Annotations[asmetrics.TraceContextAnnotation]; ok {
+		sandbox.Annotations[asmetrics.TraceContextAnnotation] = tc
+	}
+
 	template.Spec.PodTemplate.DeepCopyInto(&sandbox.Spec.PodTemplate)
 	// TODO: this is a workaround, remove replica assignment related issue #202
 	replicas := int32(1)