feat: Add sandbox claim label synchronization between Pod and Sandbox

Signed-off-by: 柏存 <guoxiongfeng.gxf@alibaba-inc.com>

Author: 柏存

api/v1alpha1/sandboxset_types.go

@@ -23,8 +23,11 @@ import (
 const (
 	InternalPrefix = "agents.kruise.io/"
 
-	LabelSandboxPool  = InternalPrefix + "sandbox-pool"
-	LabelTemplateHash = InternalPrefix + "template-hash"
+	// LabelSandboxPool identifies which SandboxSet generated the sandbox
+	LabelSandboxPool = InternalPrefix + "sandbox-pool"
+	// LabelSandboxIsClaimed indicates whether the sandbox has been claimed by user
+	LabelSandboxIsClaimed = InternalPrefix + "sandbox-claimed"
+	LabelTemplateHash     = InternalPrefix + "template-hash"
 
 	AnnotationLock      = InternalPrefix + "lock"
 	AnnotationOwner     = InternalPrefix + "owner"

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/google/uuid v1.6.0
 	github.com/onsi/ginkgo/v2 v2.27.3
 	github.com/onsi/gomega v1.38.2
+	github.com/prometheus/client_golang v1.22.0
 	github.com/spf13/pflag v1.0.6
 	github.com/stretchr/testify v1.11.1
 	google.golang.org/grpc v1.75.1
@@ -68,7 +69,6 @@ require (
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
-	github.com/prometheus/client_golang v1.22.0 // indirect
 	github.com/prometheus/client_model v0.6.1 // indirect
 	github.com/prometheus/common v0.62.0 // indirect
 	github.com/prometheus/procfs v0.16.0 // indirect

pkg/controller/sandbox/core/common_control.go

@@ -18,11 +18,13 @@ package core
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/tools/record"
 	"k8s.io/klog/v2"
 	"k8s.io/utils/ptr"
@@ -104,6 +106,11 @@ func (r *commonControl) EnsureSandboxUpdated(ctx context.Context, args EnsureFun
 		PodUID:   pod.UID,
 	}
 	logger.Info("sandbox newStatus", "newStatus", utils.DumpJson(newStatus))
+
+	// update pod label
+	if err := r.handleClaimSandbox(ctx, args); err != nil {
+		return err
+	}
 	// inplace update
 	done, err := r.handleInplaceUpdateSandbox(ctx, args)
 	if err != nil {
@@ -258,6 +265,15 @@ func (r *commonControl) createPod(ctx context.Context, box *agentsv1alpha1.Sandb
 	}
 	// todo, when resume, create Pod based on the revision from the paused state.
 	pod.Labels[agentsv1alpha1.PodLabelTemplateHash] = newStatus.UpdateRevision
+	// Ensure SandboxSet can retrieve these pods using label selector
+	if box.GetLabels() != nil {
+		if value, exists := box.GetLabels()[agentsv1alpha1.LabelSandboxPool]; exists {
+			pod.Labels[agentsv1alpha1.LabelSandboxPool] = value
+		}
+		if value, exists := box.GetLabels()[agentsv1alpha1.LabelSandboxIsClaimed]; exists {
+			pod.Labels[agentsv1alpha1.LabelSandboxIsClaimed] = value
+		}
+	}
 
 	volumes := make([]corev1.Volume, 0, len(box.Spec.VolumeClaimTemplates))
 	for _, template := range box.Spec.VolumeClaimTemplates {
@@ -287,6 +303,120 @@ func (r *commonControl) createPod(ctx context.Context, box *agentsv1alpha1.Sandb
 	return pod, nil
 }
 
+// handleClaimSandbox synchronizes specific labels between sandbox and pod.
+// 1. If sandbox has a label but pod doesn't, add it to pod
+// 2. If sandbox doesn't have a label but pod does, delete it from pod
+// 3. If both have the label but values differ, update pod's value to match sandbox
+func (r *commonControl) handleClaimSandbox(ctx context.Context, args EnsureFuncArgs) error {
+	pod, box := args.Pod, args.Box
+	logger := logf.FromContext(ctx).WithValues("sandbox", klog.KObj(box))
+
+	// If pod doesn't exist, no action needed
+	if pod == nil {
+		return nil
+	}
+
+	// Define label keys to sync
+	labelsToSync := []string{
+		agentsv1alpha1.LabelSandboxPool,
+		agentsv1alpha1.LabelSandboxIsClaimed,
+	}
+
+	// Get labels from both sides (safe nil handling)
+	boxLabels := box.GetLabels()
+	podLabels := pod.GetLabels()
+
+	// Initialize pod labels map if nil
+	if podLabels == nil {
+		podLabels = make(map[string]string)
+	}
+
+	// Collect labels to update (including add and modify)
+	labelsToUpdate := make(map[string]string)
+	// Collect labels to delete
+	labelsToDelete := []string{}
+
+	// Iterate through label keys that need to be synced
+	for _, key := range labelsToSync {
+		boxValue, boxExists := "", false
+		if boxLabels != nil {
+			boxValue, boxExists = boxLabels[key]
+		}
+
+		podValue, podExists := podLabels[key]
+
+		// Case 1: sandbox has it, pod doesn't -> add to pod
+		if boxExists && !podExists {
+			labelsToUpdate[key] = boxValue
+			logger.Info("label missing in pod, will add",
+				"key", key,
+				"boxValue", boxValue)
+		} else if !boxExists && podExists {
+			// Case 2: sandbox doesn't have it, pod does -> delete from pod
+			labelsToDelete = append(labelsToDelete, key)
+			logger.Info("label not in box, will delete from pod",
+				"key", key,
+				"podValue", podValue)
+		} else if boxExists && podExists && boxValue != podValue {
+			// Case 3: both have it but values differ -> update pod's value
+			labelsToUpdate[key] = boxValue
+			logger.Info("label value mismatch, will update pod",
+				"key", key,
+				"boxValue", boxValue,
+				"podValue", podValue)
+		}
+		// Case 4: both don't have it or both have same value -> no action needed
+	}
+
+	// If nothing needs to be changed, return early
+	if len(labelsToUpdate) == 0 && len(labelsToDelete) == 0 {
+		logger.V(utils.DebugLogLevel).Info("pod labels already in sync, no action needed")
+		return nil
+	}
+
+	// Build patch data using Strategic Merge Patch:
+	// - Add/Update: directly set key: value
+	// - Delete: set key: null
+	labels := make(map[string]interface{})
+
+	// Add/Update operations
+	for k, v := range labelsToUpdate {
+		labels[k] = v
+	}
+
+	// Delete operations: set to null
+	for _, k := range labelsToDelete {
+		labels[k] = nil
+	}
+
+	patchData := map[string]interface{}{
+		"metadata": map[string]interface{}{
+			"labels": labels,
+		},
+	}
+
+	patchBytes, err := json.Marshal(patchData)
+	if err != nil {
+		logger.Error(err, "failed to marshal patch data")
+		return err
+	}
+
+	// Execute patch operation
+	// Note: using null to delete only affects specified keys, won't impact labels managed by other components
+	podCopy := pod.DeepCopy()
+	if err = r.Patch(ctx, podCopy, client.RawPatch(types.StrategicMergePatchType, patchBytes)); err != nil {
+		logger.Error(err, "failed to patch pod labels",
+			"labelsToUpdate", labelsToUpdate,
+			"labelsToDelete", labelsToDelete)
+		return err
+	}
+
+	logger.Info("successfully synced pod labels",
+		"labelsUpdated", labelsToUpdate,
+		"labelsDeleted", labelsToDelete)
+	return nil
+}
+
 func (r *commonControl) handleInplaceUpdateSandbox(ctx context.Context, args EnsureFuncArgs) (bool, error) {
 	pod, box, newStatus := args.Pod, args.Box, args.NewStatus
 	logger := logf.FromContext(ctx).WithValues("sandbox", klog.KObj(box))