Fix review comments 3

Author: LiangquanLi930

controllers/awsmachinetemplate_controller.go

@@ -27,6 +27,9 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/utils/ptr"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/controller"
@@ -230,153 +233,145 @@ func (r *AWSMachineTemplateReconciler) getInstanceTypeCapacity(ctx context.Conte
 }
 
 // getNodeInfo queries node information (architecture and OS) for the AWSMachineTemplate.
-// It uses AMI ID if specified, otherwise attempts AMI lookup or falls back to instance type info.
+// It attempts to resolve nodeInfo using three strategies in order of priority:
+//  1. Directly from explicitly specified AMI ID
+//  2. From default AMI lookup (requires Kubernetes version from owner MachineDeployment/KubeadmControlPlane)
+//  3. From instance type architecture (OS cannot be determined, only architecture)
 func (r *AWSMachineTemplateReconciler) getNodeInfo(ctx context.Context, ec2Client *ec2.Client, template *infrav1.AWSMachineTemplate, instanceType string) (*infrav1.NodeInfo, error) {
-	nodeInfo := &infrav1.NodeInfo{}
-	amiID := template.Spec.Template.Spec.AMI.ID
-	if amiID != nil && *amiID != "" {
-		// AMI ID is specified, query it directly
-		arch, os, err := r.getNodeInfoFromAMI(ctx, ec2Client, *amiID)
-		if err == nil {
-			if arch != "" {
-				nodeInfo.Architecture = arch
-			}
-			if os != "" {
-				nodeInfo.OperatingSystem = os
-			}
-		}
-	} else {
-		// AMI ID is not specified, query instance type to get architecture
-		input := &ec2.DescribeInstanceTypesInput{
-			InstanceTypes: []ec2types.InstanceType{ec2types.InstanceType(instanceType)},
-		}
-
-		result, err := ec2Client.DescribeInstanceTypes(ctx, input)
+	// Strategy 1: Extract nodeInfo from the AMI if an ID is set.
+	if amiID := ptr.Deref(template.Spec.Template.Spec.AMI.ID, ""); amiID != "" {
+		result, err := ec2Client.DescribeImages(ctx, &ec2.DescribeImagesInput{
+			ImageIds: []string{amiID},
+		})
 		if err != nil {
-			return nil, errors.Wrapf(err, "failed to describe instance type %q", instanceType)
+			return nil, errors.Wrapf(err, "failed to describe AMI %q", amiID)
 		}
-
-		if len(result.InstanceTypes) == 0 {
-			return nil, errors.Errorf("no information found for instance type %q", instanceType)
+		if len(result.Images) == 0 {
+			return nil, errors.Errorf("no information found for AMI %q", amiID)
 		}
+		// Extract nodeInfo directly from the image object (no additional API call needed)
+		return r.extractNodeInfoFromImage(result.Images[0]), nil
+	}
 
-		instanceTypeInfo := result.InstanceTypes[0]
-
-		// Infer architecture from instance type
-		var architecture string
-		if instanceTypeInfo.ProcessorInfo != nil && len(instanceTypeInfo.ProcessorInfo.SupportedArchitectures) == 1 {
-			// Use the supported architecture
-			switch instanceTypeInfo.ProcessorInfo.SupportedArchitectures[0] {
-			case ec2types.ArchitectureTypeX8664:
-				architecture = ec2service.Amd64ArchitectureTag
-				nodeInfo.Architecture = infrav1.ArchitectureAmd64
-			case ec2types.ArchitectureTypeArm64:
-				architecture = ec2service.Arm64ArchitectureTag
-				nodeInfo.Architecture = infrav1.ArchitectureArm64
-			}
-		} else {
-			return nil, errors.Errorf("instance type must support exactly one architecture, got %d", len(instanceTypeInfo.ProcessorInfo.SupportedArchitectures))
-		}
+	// No explicit AMI ID specified, query instance type to determine architecture
+	// This architecture will be used to lookup default AMI (Strategy 2) or as fallback (Strategy 3)
+	result, err := ec2Client.DescribeInstanceTypes(ctx, &ec2.DescribeInstanceTypesInput{
+		InstanceTypes: []ec2types.InstanceType{ec2types.InstanceType(instanceType)},
+	})
+	if err != nil {
+		return nil, errors.Wrapf(err, "failed to describe instance type %q", instanceType)
+	}
 
-		// Attempt to get Kubernetes version from MachineDeployment
-		kubernetesVersion, versionErr := r.getKubernetesVersion(ctx, template)
-		if versionErr == nil && kubernetesVersion != "" {
-			// Try to look up AMI using the version
-			image, err := ec2service.DefaultAMILookup(
-				ec2Client,
-				template.Spec.Template.Spec.ImageLookupOrg,
-				template.Spec.Template.Spec.ImageLookupBaseOS,
-				kubernetesVersion,
-				architecture,
-				template.Spec.Template.Spec.ImageLookupFormat,
-			)
-			if err == nil && image != nil {
-				// Successfully found AMI, extract accurate nodeInfo from it
-				arch, os, _ := r.getNodeInfoFromAMI(ctx, ec2Client, *image.ImageId)
-				if arch != "" {
-					nodeInfo.Architecture = arch
-				}
-				if os != "" {
-					nodeInfo.OperatingSystem = os
-				}
-				return nodeInfo, nil
-			}
-			// AMI lookup failed, fall through to defaults
-		} else {
-			return nil, errors.Errorf("failed to get Kubernetes version: %v", versionErr)
-		}
+	if len(result.InstanceTypes) == 0 {
+		return nil, errors.Errorf("no information found for instance type %q", instanceType)
 	}
 
-	return nodeInfo, nil
-}
+	instanceTypeInfo := result.InstanceTypes[0]
 
-// getNodeInfoFromAMI queries the AMI to determine architecture and operating system.
-func (r *AWSMachineTemplateReconciler) getNodeInfoFromAMI(ctx context.Context, ec2Client *ec2.Client, amiID string) (infrav1.Architecture, infrav1.OperatingSystem, error) {
-	input := &ec2.DescribeImagesInput{
-		ImageIds: []string{amiID},
+	// Instance type must support exactly one architecture
+	if instanceTypeInfo.ProcessorInfo == nil || len(instanceTypeInfo.ProcessorInfo.SupportedArchitectures) != 1 {
+		return nil, errors.Errorf("instance type must support exactly one architecture, got %d", len(instanceTypeInfo.ProcessorInfo.SupportedArchitectures))
 	}
 
-	result, err := ec2Client.DescribeImages(ctx, input)
-	if err != nil {
-		return "", "", errors.Wrapf(err, "failed to describe AMI %q", amiID)
+	// Map EC2 architecture type to architecture tag for AMI lookup
+	var architecture string
+	var nodeInfoArch infrav1.Architecture
+	switch instanceTypeInfo.ProcessorInfo.SupportedArchitectures[0] {
+	case ec2types.ArchitectureTypeX8664:
+		architecture = ec2service.Amd64ArchitectureTag
+		nodeInfoArch = infrav1.ArchitectureAmd64
+	case ec2types.ArchitectureTypeArm64:
+		architecture = ec2service.Arm64ArchitectureTag
+		nodeInfoArch = infrav1.ArchitectureArm64
+	default:
+		return nil, errors.Errorf("unsupported architecture: %v", instanceTypeInfo.ProcessorInfo.SupportedArchitectures[0])
 	}
 
-	if len(result.Images) == 0 {
-		return "", "", errors.Errorf("no information found for AMI %q", amiID)
+	// Strategy 2: Try to get Kubernetes version and lookup default AMI
+	kubernetesVersion, err := r.getKubernetesVersion(ctx, template)
+	if err == nil && kubernetesVersion != "" {
+		// Attempt AMI lookup with the version
+		image, err := ec2service.DefaultAMILookup(
+			ec2Client,
+			template.Spec.Template.Spec.ImageLookupOrg,
+			template.Spec.Template.Spec.ImageLookupBaseOS,
+			kubernetesVersion,
+			architecture,
+			template.Spec.Template.Spec.ImageLookupFormat,
+		)
+		if err == nil && image != nil {
+			// Successfully found AMI, extract accurate nodeInfo from it
+			return r.extractNodeInfoFromImage(*image), nil
+		}
+		// AMI lookup failed, fall through to Strategy 3
 	}
 
-	image := result.Images[0]
+	// Strategy 3: Fallback to instance type architecture only
+	// Note: OS cannot be determined from instance type alone, only architecture
+	return &infrav1.NodeInfo{
+		Architecture: nodeInfoArch,
+	}, nil
+}
 
-	// Get architecture from AMI
-	var arch infrav1.Architecture
+// extractNodeInfoFromImage extracts nodeInfo (architecture and OS) from an EC2 image.
+// This is a pure function with no AWS API calls.
+func (r *AWSMachineTemplateReconciler) extractNodeInfoFromImage(image ec2types.Image) *infrav1.NodeInfo {
+	nodeInfo := &infrav1.NodeInfo{}
+
+	// Extract architecture from AMI
 	switch image.Architecture {
 	case ec2types.ArchitectureValuesX8664:
-		arch = infrav1.ArchitectureAmd64
+		nodeInfo.Architecture = infrav1.ArchitectureAmd64
 	case ec2types.ArchitectureValuesArm64:
-		arch = infrav1.ArchitectureArm64
+		nodeInfo.Architecture = infrav1.ArchitectureArm64
 	}
 
 	// Determine OS - default to Linux, change to Windows if detected
 	// Most AMIs are Linux-based, so we initialize with Linux as the default
-	os := infrav1.OperatingSystemLinux
+	nodeInfo.OperatingSystem = infrav1.OperatingSystemLinux
 
-	// 1. Check Platform field (most reliable for Windows detection)
+	// Check Platform field (most reliable for Windows detection)
 	if image.Platform == ec2types.PlatformValuesWindows {
-		os = infrav1.OperatingSystemWindows
+		nodeInfo.OperatingSystem = infrav1.OperatingSystemWindows
+		return nodeInfo
 	}
 
-	// 2. Check PlatformDetails field for Windows indication
-	if os != infrav1.OperatingSystemWindows && image.PlatformDetails != nil {
+	// Check PlatformDetails field for Windows indication
+	if image.PlatformDetails != nil {
 		platformDetails := strings.ToLower(*image.PlatformDetails)
 		if strings.Contains(platformDetails, string(infrav1.OperatingSystemWindows)) {
-			os = infrav1.OperatingSystemWindows
+			nodeInfo.OperatingSystem = infrav1.OperatingSystemWindows
 		}
 	}
 
-	return arch, os, nil
+	return nodeInfo
 }
 
 // getKubernetesVersion attempts to find the Kubernetes version by querying MachineDeployments
 // or KubeadmControlPlanes that reference this AWSMachineTemplate.
 func (r *AWSMachineTemplateReconciler) getKubernetesVersion(ctx context.Context, template *infrav1.AWSMachineTemplate) (string, error) {
+	listOpts, err := getParentListOptions(template.ObjectMeta)
+	if err != nil {
+		return "", errors.Wrap(err, "failed to get parent list options")
+	}
+
 	// Try to find version from MachineDeployment first
 	machineDeploymentList := &clusterv1.MachineDeploymentList{}
-	if err := r.List(ctx, machineDeploymentList, client.InNamespace(template.Namespace)); err != nil {
+	if err := r.List(ctx, machineDeploymentList, listOpts...); err != nil {
 		return "", errors.Wrap(err, "failed to list MachineDeployments")
 	}
 
 	// Find MachineDeployments that reference this AWSMachineTemplate
 	for _, md := range machineDeploymentList.Items {
-		if md.Spec.Template.Spec.InfrastructureRef.Kind == "AWSMachineTemplate" &&
-			md.Spec.Template.Spec.InfrastructureRef.Name == template.Name &&
-			md.Spec.Template.Spec.Version != nil {
-			return *md.Spec.Template.Spec.Version, nil
+		if version := ptr.Deref(md.Spec.Template.Spec.Version, ""); md.Spec.Template.Spec.InfrastructureRef.Kind == "AWSMachineTemplate" &&
+			md.Spec.Template.Spec.InfrastructureRef.Name == template.Name && version != "" {
+			return version, nil
 		}
 	}
 
 	// If not found in MachineDeployment, try KubeadmControlPlane
 	kcpList := &controlplanev1.KubeadmControlPlaneList{}
-	if err := r.List(ctx, kcpList, client.InNamespace(template.Namespace)); err != nil {
+	if err := r.List(ctx, kcpList, listOpts...); err != nil {
 		return "", errors.Wrap(err, "failed to list KubeadmControlPlanes")
 	}
 
@@ -391,3 +386,26 @@ func (r *AWSMachineTemplateReconciler) getKubernetesVersion(ctx context.Context,
 
 	return "", errors.New("no MachineDeployment or KubeadmControlPlane found referencing this AWSMachineTemplate with a version")
 }
+
+func getParentListOptions(obj metav1.ObjectMeta) ([]client.ListOption, error) {
+	listOpts := []client.ListOption{
+		client.InNamespace(obj.Namespace),
+	}
+
+	for _, ref := range obj.GetOwnerReferences() {
+		if ref.Kind != "Cluster" {
+			continue
+		}
+		gv, err := schema.ParseGroupVersion(ref.APIVersion)
+		if err != nil {
+			return nil, errors.WithStack(err)
+		}
+		if gv.Group == clusterv1.GroupVersion.Group {
+			listOpts = append(listOpts, client.MatchingLabels{
+				clusterv1.ClusterNameLabel: ref.Name,
+			})
+			break
+		}
+	}
+	return listOpts, nil
+}