passing custom headers to cluster and node pool create/update requests

Author: michalno1

azure/defaults.go

@@ -84,6 +84,14 @@ const (
 	ProviderIDPrefix = "azure://"
 )
 
+const (
+	// CustomHeaderPrefix is the prefix of annotations that enable additional cluster / node pool features.
+	// Whatever follows the prefix will be passed as a header to cluster/node pool creation/update requests.
+	// E.g. add `"infrastructure.cluster.x-k8s.io/custom-header-UseGPUDedicatedVHD": "true"` annotation to
+	// AzureManagedMachinePool CR to enable creating GPU nodes by the node pool.
+	CustomHeaderPrefix = "infrastructure.cluster.x-k8s.io/custom-header-"
+)
+
 var (
 	// LinuxBootstrapExtensionCommand is the command the VM bootstrap extension will execute to verify Linux nodes bootstrap completes successfully.
 	LinuxBootstrapExtensionCommand = fmt.Sprintf("for i in $(seq 1 %d); do test -f %s && break; if [ $i -eq %d ]; then return 1; else sleep %d; fi; done", bootstrapExtensionRetries, bootstrapSentinelFile, bootstrapExtensionRetries, bootstrapExtensionSleep)

azure/scope/managedcontrolplane.go

@@ -391,6 +391,10 @@ func (s *ManagedControlPlaneScope) FailureDomains() []string {
 	return []string{}
 }
 
+func (s *ManagedControlPlaneScope) ManagedClusterAnnotations() map[string]string {
+	return s.ControlPlane.Annotations
+}
+
 // ManagedClusterSpec returns the managed cluster spec.
 func (s *ManagedControlPlaneScope) ManagedClusterSpec() (azure.ManagedClusterSpec, error) {
 	decodedSSHPublicKey, err := base64.StdEncoding.DecodeString(s.ControlPlane.Spec.SSHPublicKey)
@@ -557,6 +561,10 @@ func (s *ManagedControlPlaneScope) GetAllAgentPoolSpecs(ctx context.Context) ([]
 	return ammps, nil
 }
 
+func (s *ManagedControlPlaneScope) AgentPoolAnnotations() map[string]string {
+	return s.InfraMachinePool.Annotations
+}
+
 // AgentPoolSpec returns an azure.AgentPoolSpec for currently reconciled AzureManagedMachinePool.
 func (s *ManagedControlPlaneScope) AgentPoolSpec() azure.AgentPoolSpec {
 	return buildAgentPoolSpec(s.ControlPlane, s.MachinePool, s.InfraMachinePool)

azure/services/agentpools/agentpools.go

@@ -27,6 +27,7 @@ import (
 	"github.com/pkg/errors"
 	infrav1alpha4 "sigs.k8s.io/cluster-api-provider-azure/api/v1beta1"
 	"sigs.k8s.io/cluster-api-provider-azure/azure"
+	"sigs.k8s.io/cluster-api-provider-azure/util/maps"
 	"sigs.k8s.io/cluster-api-provider-azure/util/tele"
 )
 
@@ -35,6 +36,7 @@ type ManagedMachinePoolScope interface {
 	azure.ClusterDescriber
 
 	NodeResourceGroup() string
+	AgentPoolAnnotations() map[string]string
 	AgentPoolSpec() azure.AgentPoolSpec
 	SetAgentPoolProviderIDList([]string)
 	SetAgentPoolReplicas(int32)
@@ -64,7 +66,6 @@ func (s *Service) Reconcile(ctx context.Context) error {
 	defer done()
 
 	agentPoolSpec := s.scope.AgentPoolSpec()
-
 	profile := containerservice.AgentPool{
 		ManagedClusterAgentPoolProfileProperties: &containerservice.ManagedClusterAgentPoolProfileProperties{
 			VMSize:              &agentPoolSpec.SKU,
@@ -96,8 +97,10 @@ func (s *Service) Reconcile(ctx context.Context) error {
 	// AKS will populate defaults and read-only values, which we want
 	// to strip/clean to match what we expect.
 
+	customHeaders := maps.FilterByKeyPrefix(s.scope.AgentPoolAnnotations(), azure.CustomHeaderPrefix)
 	if isCreate := azure.ResourceNotFound(err); isCreate {
-		err = s.Client.CreateOrUpdate(ctx, agentPoolSpec.ResourceGroup, agentPoolSpec.Cluster, agentPoolSpec.Name, profile)
+		err = s.Client.CreateOrUpdate(ctx, agentPoolSpec.ResourceGroup, agentPoolSpec.Cluster, agentPoolSpec.Name,
+			profile, customHeaders)
 		if err != nil && azure.ResourceNotFound(err) {
 			return azure.WithTransientError(errors.Wrap(err, "agent pool dependent resource does not exist yet"), 20*time.Second)
 		} else if err != nil {
@@ -138,7 +141,8 @@ func (s *Service) Reconcile(ctx context.Context) error {
 		diff := cmp.Diff(normalizedProfile, existingProfile)
 		if diff != "" {
 			log.V(2).Info(fmt.Sprintf("Update required (+new -old):\n%s", diff))
-			err = s.Client.CreateOrUpdate(ctx, agentPoolSpec.ResourceGroup, agentPoolSpec.Cluster, agentPoolSpec.Name, profile)
+			err = s.Client.CreateOrUpdate(ctx, agentPoolSpec.ResourceGroup, agentPoolSpec.Cluster, agentPoolSpec.Name,
+				profile, customHeaders)
 			if err != nil {
 				return errors.Wrap(err, "failed to create or update agent pool")
 			}

azure/services/agentpools/agentpools_test.go

@@ -55,7 +55,7 @@ func TestReconcile(t *testing.T) {
 			expectedError:            "",
 			expect: func(m *mock_agentpools.MockClientMockRecorder, provisioningstate string) {
 				pv := provisioningstate
-				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agentpool", gomock.Any()).Return(nil)
+				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agentpool", gomock.Any(), gomock.Any()).Return(nil)
 				m.Get(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agentpool").Return(containerservice.AgentPool{ManagedClusterAgentPoolProfileProperties: &containerservice.ManagedClusterAgentPoolProfileProperties{
 					ProvisioningState: &pv,
 				}}, nil)
@@ -144,7 +144,7 @@ func TestReconcile(t *testing.T) {
 			},
 			expectedError: "",
 			expect: func(m *mock_agentpools.MockClientMockRecorder) {
-				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agentpool", gomock.Any()).Return(nil)
+				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agentpool", gomock.Any(), gomock.Any()).Return(nil)
 				m.Get(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agentpool").Return(containerservice.AgentPool{}, autorest.NewErrorWithResponse("", "", &http.Response{StatusCode: 404}, "Not Found"))
 			},
 		},
@@ -182,7 +182,7 @@ func TestReconcile(t *testing.T) {
 			expectedError: "",
 			expect: func(m *mock_agentpools.MockClientMockRecorder) {
 				m.Get(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agent-pool").Return(containerservice.AgentPool{}, autorest.NewErrorWithResponse("", "", &http.Response{StatusCode: 404}, "Not found"))
-				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agent-pool", gomock.AssignableToTypeOf(containerservice.AgentPool{})).Return(nil)
+				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agent-pool", gomock.AssignableToTypeOf(containerservice.AgentPool{}), gomock.Any()).Return(nil)
 			},
 		},
 		{
@@ -201,7 +201,7 @@ func TestReconcile(t *testing.T) {
 			expectedError: "failed to create or update agent pool: #: Internal Server Error: StatusCode=500",
 			expect: func(m *mock_agentpools.MockClientMockRecorder) {
 				m.Get(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agent-pool").Return(containerservice.AgentPool{}, autorest.NewErrorWithResponse("", "", &http.Response{StatusCode: 404}, "Not found"))
-				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agent-pool", gomock.AssignableToTypeOf(containerservice.AgentPool{})).Return(autorest.NewErrorWithResponse("", "", &http.Response{StatusCode: 500}, "Internal Server Error"))
+				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agent-pool", gomock.AssignableToTypeOf(containerservice.AgentPool{}), gomock.Any()).Return(autorest.NewErrorWithResponse("", "", &http.Response{StatusCode: 500}, "Internal Server Error"))
 			},
 		},
 		{
@@ -228,7 +228,7 @@ func TestReconcile(t *testing.T) {
 						ProvisioningState:   to.StringPtr("Failed"),
 					},
 				}, nil)
-				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agent-pool", gomock.AssignableToTypeOf(containerservice.AgentPool{})).Return(autorest.NewErrorWithResponse("", "", &http.Response{StatusCode: 500}, "Internal Server Error"))
+				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-cluster", "my-agent-pool", gomock.AssignableToTypeOf(containerservice.AgentPool{}), gomock.Any()).Return(autorest.NewErrorWithResponse("", "", &http.Response{StatusCode: 500}, "Internal Server Error"))
 			},
 		},
 		{

azure/services/agentpools/client.go

@@ -29,7 +29,7 @@ import (
 // Client wraps go-sdk.
 type Client interface {
 	Get(context.Context, string, string, string) (containerservice.AgentPool, error)
-	CreateOrUpdate(context.Context, string, string, string, containerservice.AgentPool) error
+	CreateOrUpdate(context.Context, string, string, string, containerservice.AgentPool, map[string]string) error
 	Delete(context.Context, string, string, string) error
 }
 
@@ -62,11 +62,20 @@ func (ac *AzureClient) Get(ctx context.Context, resourceGroupName, cluster, name
 }
 
 // CreateOrUpdate creates or updates an agent pool.
-func (ac *AzureClient) CreateOrUpdate(ctx context.Context, resourceGroupName, cluster, name string, properties containerservice.AgentPool) error {
+func (ac *AzureClient) CreateOrUpdate(ctx context.Context, resourceGroupName, cluster, name string,
+	properties containerservice.AgentPool, customHeaders map[string]string) error {
 	ctx, _, done := tele.StartSpanWithLogger(ctx, "agentpools.AzureClient.CreateOrUpdate")
 	defer done()
 
-	future, err := ac.agentpools.CreateOrUpdate(ctx, resourceGroupName, cluster, name, properties)
+	preparer, err := ac.agentpools.CreateOrUpdatePreparer(ctx, resourceGroupName, cluster, name, properties)
+	if err != nil {
+		return errors.Wrap(err, "failed to prepare operation")
+	}
+	for key, element := range customHeaders {
+		preparer.Header.Add(key, element)
+	}
+
+	future, err := ac.agentpools.CreateOrUpdateSender(preparer)
 	if err != nil {
 		return errors.Wrap(err, "failed to begin operation")
 	}

azure/services/agentpools/mock_agentpools/agentpools_mock.go

@@ -30,7 +30,7 @@ import (
 type Client interface {
 	Get(context.Context, string, string) (containerservice.ManagedCluster, error)
 	GetCredentials(context.Context, string, string) ([]byte, error)
-	CreateOrUpdate(context.Context, string, string, containerservice.ManagedCluster) (containerservice.ManagedCluster, error)
+	CreateOrUpdate(context.Context, string, string, containerservice.ManagedCluster, map[string]string) (containerservice.ManagedCluster, error)
 	Delete(context.Context, string, string) error
 }
 
@@ -78,11 +78,19 @@ func (ac *AzureClient) GetCredentials(ctx context.Context, resourceGroupName, na
 }
 
 // CreateOrUpdate creates or updates a managed cluster.
-func (ac *AzureClient) CreateOrUpdate(ctx context.Context, resourceGroupName, name string, cluster containerservice.ManagedCluster) (containerservice.ManagedCluster, error) {
+func (ac *AzureClient) CreateOrUpdate(ctx context.Context, resourceGroupName, name string, cluster containerservice.ManagedCluster, headers map[string]string) (containerservice.ManagedCluster, error) {
 	ctx, _, done := tele.StartSpanWithLogger(ctx, "managedclusters.AzureClient.CreateOrUpdate")
 	defer done()
 
-	future, err := ac.managedclusters.CreateOrUpdate(ctx, resourceGroupName, name, cluster)
+	preparer, err := ac.managedclusters.CreateOrUpdatePreparer(ctx, resourceGroupName, name, cluster)
+	if err != nil {
+		return containerservice.ManagedCluster{}, errors.Wrap(err, "failed to prepare operation")
+	}
+	for key, value := range headers {
+		preparer.Header.Add(key, value)
+	}
+
+	future, err := ac.managedclusters.CreateOrUpdateSender(preparer)
 	if err != nil {
 		return containerservice.ManagedCluster{}, errors.Wrap(err, "failed to begin operation")
 	}

azure/services/managedclusters/client.go

@@ -30,6 +30,7 @@ import (
 	"k8s.io/klog/v2"
 	infrav1alpha4 "sigs.k8s.io/cluster-api-provider-azure/api/v1beta1"
 	"sigs.k8s.io/cluster-api-provider-azure/azure"
+	"sigs.k8s.io/cluster-api-provider-azure/util/maps"
 	"sigs.k8s.io/cluster-api-provider-azure/util/tele"
 	clusterv1 "sigs.k8s.io/cluster-api/api/v1beta1"
 )
@@ -42,6 +43,7 @@ var (
 // ManagedClusterScope defines the scope interface for a managed cluster.
 type ManagedClusterScope interface {
 	azure.ClusterDescriber
+	ManagedClusterAnnotations() map[string]string
 	ManagedClusterSpec() (azure.ManagedClusterSpec, error)
 	GetAllAgentPoolSpecs(ctx context.Context) ([]azure.AgentPoolSpec, error)
 	SetControlPlaneEndpoint(clusterv1.APIEndpoint)
@@ -292,8 +294,9 @@ func (s *Service) Reconcile(ctx context.Context) error {
 		}
 	}
 
+	customHeaders := maps.FilterByKeyPrefix(s.Scope.ManagedClusterAnnotations(), azure.CustomHeaderPrefix)
 	if isCreate {
-		managedCluster, err = s.Client.CreateOrUpdate(ctx, managedClusterSpec.ResourceGroupName, managedClusterSpec.Name, managedCluster)
+		managedCluster, err = s.Client.CreateOrUpdate(ctx, managedClusterSpec.ResourceGroupName, managedClusterSpec.Name, managedCluster, customHeaders)
 		if err != nil {
 			return fmt.Errorf("failed to create managed cluster, %w", err)
 		}
@@ -322,7 +325,7 @@ func (s *Service) Reconcile(ctx context.Context) error {
 		diff := computeDiffOfNormalizedClusters(managedCluster, existingMC)
 		if diff != "" {
 			klog.V(2).Infof("Update required (+new -old):\n%s", diff)
-			managedCluster, err = s.Client.CreateOrUpdate(ctx, managedClusterSpec.ResourceGroupName, managedClusterSpec.Name, managedCluster)
+			managedCluster, err = s.Client.CreateOrUpdate(ctx, managedClusterSpec.ResourceGroupName, managedClusterSpec.Name, managedCluster, customHeaders)
 			if err != nil {
 				return fmt.Errorf("failed to update managed cluster, %w", err)
 			}

azure/services/managedclusters/managedclusters.go

@@ -43,7 +43,7 @@ func TestReconcile(t *testing.T) {
 			provisioningStatesToTest: []string{"Canceled", "Succeeded", "Failed"},
 			expectedError:            "",
 			expect: func(m *mock_managedclusters.MockClientMockRecorder, provisioningstate string, s *mock_managedclusters.MockManagedClusterScopeMockRecorder) {
-				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-managedcluster", gomock.Any()).Return(containerservice.ManagedCluster{ManagedClusterProperties: &containerservice.ManagedClusterProperties{
+				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-managedcluster", gomock.Any(), map[string]string{"myFeature": "true"}).Return(containerservice.ManagedCluster{ManagedClusterProperties: &containerservice.ManagedClusterProperties{
 					Fqdn:              pointer.String("my-managedcluster-fqdn"),
 					ProvisioningState: &provisioningstate,
 				}}, nil)
@@ -55,6 +55,9 @@ func TestReconcile(t *testing.T) {
 				m.GetCredentials(gomockinternal.AContext(), "my-rg", "my-managedcluster").Times(1)
 				s.ClusterName().AnyTimes().Return("my-managedcluster")
 				s.ResourceGroup().AnyTimes().Return("my-rg")
+				s.ManagedClusterAnnotations().Times(1).Return(map[string]string{
+					"infrastructure.cluster.x-k8s.io/custom-header-myFeature": "true",
+				})
 				s.ManagedClusterSpec().AnyTimes().Return(azure.ManagedClusterSpec{
 					Name:              "my-managedcluster",
 					ResourceGroupName: "my-rg",
@@ -73,6 +76,7 @@ func TestReconcile(t *testing.T) {
 				}}, nil)
 				s.ClusterName().AnyTimes().Return("my-managedcluster")
 				s.ResourceGroup().AnyTimes().Return("my-rg")
+				s.ManagedClusterAnnotations().Times(1).Return(map[string]string{})
 				s.ManagedClusterSpec().AnyTimes().Return(azure.ManagedClusterSpec{
 					Name:              "my-managedcluster",
 					ResourceGroupName: "my-rg",
@@ -120,11 +124,12 @@ func TestReconcile(t *testing.T) {
 			name:          "no managedcluster exists",
 			expectedError: "",
 			expect: func(m *mock_managedclusters.MockClientMockRecorder, s *mock_managedclusters.MockManagedClusterScopeMockRecorder) {
-				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-managedcluster", gomock.Any()).Return(containerservice.ManagedCluster{ManagedClusterProperties: &containerservice.ManagedClusterProperties{}}, nil)
+				m.CreateOrUpdate(gomockinternal.AContext(), "my-rg", "my-managedcluster", gomock.Any(), gomock.Any()).Return(containerservice.ManagedCluster{ManagedClusterProperties: &containerservice.ManagedClusterProperties{}}, nil)
 				m.Get(gomockinternal.AContext(), "my-rg", "my-managedcluster").Return(containerservice.ManagedCluster{}, autorest.NewErrorWithResponse("", "", &http.Response{StatusCode: 404}, "Not Found"))
 				m.GetCredentials(gomockinternal.AContext(), "my-rg", "my-managedcluster").Times(1)
 				s.ClusterName().AnyTimes().Return("my-managedcluster")
 				s.ResourceGroup().AnyTimes().Return("my-rg")
+				s.ManagedClusterAnnotations().Times(1).Return(map[string]string{})
 				s.ManagedClusterSpec().AnyTimes().Return(azure.ManagedClusterSpec{
 					Name:              "my-managedcluster",
 					ResourceGroupName: "my-rg",