fix(azure-flex): handle mismatched agentpool types in GC paths

Author: chokevin

karpenter/pkg/cloudproviders/azure/api.go

@@ -26,6 +26,16 @@ func IsNotFound(err error) bool {
 	return false
 }
 
+// IsTypeMismatch returns true if err indicates the plugin returned an object of
+// a different concrete protobuf type than the caller expected.
+func IsTypeMismatch(err error) bool {
+	if err == nil {
+		return false
+	}
+	s, ok := status.FromError(err)
+	return ok && s.Code() == codes.InvalidArgument && strings.Contains(s.Message(), "type mismatch")
+}
+
 // IsQuotaError returns true if err signals an Azure quota / capacity exhaustion.
 // We classify both HTTP 429 and the well-known Azure ARM error codes.
 func IsQuotaError(err error) bool {

karpenter/pkg/cloudproviders/azure/api_test.go

@@ -0,0 +1,55 @@
+package azure
+
+import (
+	"fmt"
+	"testing"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func TestIsTypeMismatch(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{
+			name: "exact type mismatch status",
+			err:  status.Error(codes.InvalidArgument, "type mismatch"),
+			want: true,
+		},
+		{
+			name: "wrapped type mismatch status",
+			err:  fmt.Errorf("wrap: %w", status.Error(codes.InvalidArgument, "type mismatch")),
+			want: true,
+		},
+		{
+			name: "other invalid argument",
+			err:  status.Error(codes.InvalidArgument, "bad request"),
+			want: false,
+		},
+		{
+			name: "not found",
+			err:  status.Error(codes.NotFound, "not found"),
+			want: false,
+		},
+		{
+			name: "nil",
+			err:  nil,
+			want: false,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			if got := IsTypeMismatch(tc.err); got != tc.want {
+				t.Fatalf("IsTypeMismatch(%v) = %v, want %v", tc.err, got, tc.want)
+			}
+		})
+	}
+}

karpenter/pkg/cloudproviders/azure/cloudprovider.go

@@ -18,12 +18,15 @@ import (
 	"github.com/Azure/karpenter-provider-azure/pkg/utils"
 	"github.com/awslabs/operatorpkg/status"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	grpcstatus "google.golang.org/grpc/status"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	v1 "sigs.k8s.io/karpenter/pkg/apis/v1"
 	corecloudprovider "sigs.k8s.io/karpenter/pkg/cloudprovider"
 
+	pluginapi "github.com/Azure/aks-flex/plugin/api"
 	stretchhelper "github.com/Azure/aks-flex/plugin/pkg/helper"
 	stretchservices "github.com/Azure/aks-flex/plugin/pkg/services"
 	agentpoolsapi "github.com/Azure/aks-flex/plugin/pkg/services/agentpools/api"
@@ -50,6 +53,8 @@ type CloudProvider struct {
 	instanceTypeProvider *instancetype.Provider
 }
 
+var flexAgentPoolTypeURL = "type.googleapis.com/" + string((&flexvm.AgentPool{}).ProtoReflect().Descriptor().FullName())
+
 func newCloudProvider(
 	stretchPluginConn *grpc.ClientConn,
 	kubeClient client.Client,
@@ -177,11 +182,8 @@ func (c *CloudProvider) Delete(ctx context.Context, nodeClaim *v1.NodeClaim) err
 
 	// Per CloudProvider.Delete contract: signal NodeClaimNotFoundError if the
 	// remote resource is already gone (so karpenter knows it's safe to drop).
-	if _, err := stretchhelper.Get[*flexvm.AgentPool](
-		c.stretchAgentPoolsClient.Get,
-		ctx, nodeClaim.Name,
-	); err != nil {
-		if IsNotFound(err) {
+	if _, err := c.getFlexAgentPool(ctx, nodeClaim.Name); err != nil {
+		if IsNotFound(err) || IsTypeMismatch(err) {
 			return corecloudprovider.NewNodeClaimNotFoundError(err)
 		}
 		// Non-NotFound get failure: log and proceed with delete in best effort.
@@ -192,6 +194,9 @@ func (c *CloudProvider) Delete(ctx context.Context, nodeClaim *v1.NodeClaim) err
 		c.stretchAgentPoolsClient.Delete,
 		ctx, nodeClaim.Name,
 	); err != nil {
+		if IsNotFound(err) || IsTypeMismatch(err) {
+			return corecloudprovider.NewNodeClaimNotFoundError(err)
+		}
 		return fmt.Errorf("deleting azure-flex agent pool: %w", err)
 	}
 	logger.Info("deleted azure-flex agent pool", "nodeClaim", nodeClaim.Name)
@@ -203,12 +208,9 @@ func (c *CloudProvider) Get(ctx context.Context, providerID string) (*v1.NodeCla
 	if err != nil {
 		return nil, err
 	}
-	ap, err := stretchhelper.Get[*flexvm.AgentPool](
-		c.stretchAgentPoolsClient.Get,
-		ctx, name,
-	)
+	ap, err := c.getFlexAgentPool(ctx, name)
 	if err != nil {
-		if IsNotFound(err) {
+		if IsNotFound(err) || IsTypeMismatch(err) {
 			return nil, corecloudprovider.NewNodeClaimNotFoundError(err)
 		}
 		return nil, err
@@ -219,6 +221,26 @@ func (c *CloudProvider) Get(ctx context.Context, providerID string) (*v1.NodeCla
 	return agentPoolToNodeClaim(ap, nil), nil
 }
 
+func (c *CloudProvider) getFlexAgentPool(ctx context.Context, id string) (*flexvm.AgentPool, error) {
+	req := &pluginapi.GetRequest{}
+	req.SetId(id)
+	resp, err := c.stretchAgentPoolsClient.Get(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+	return flexAgentPoolFromGetResponse(resp)
+}
+
+func flexAgentPoolFromGetResponse(resp *pluginapi.GetResponse) (*flexvm.AgentPool, error) {
+	if resp == nil || resp.GetItem() == nil {
+		return nil, grpcstatus.Error(codes.NotFound, "")
+	}
+	if resp.GetItem().GetTypeUrl() != flexAgentPoolTypeURL {
+		return nil, grpcstatus.Error(codes.NotFound, "")
+	}
+	return stretchhelper.AnyTo[*flexvm.AgentPool](resp.GetItem())
+}
+
 func (c *CloudProvider) List(ctx context.Context) ([]*v1.NodeClaim, error) {
 	aps, err := stretchhelper.List[*flexvm.AgentPool](
 		c.stretchAgentPoolsClient.List,

karpenter/pkg/cloudproviders/azure/cloudprovider_test.go

@@ -0,0 +1,89 @@
+package azure
+
+import (
+	"testing"
+
+	pluginapi "github.com/Azure/aks-flex/plugin/api"
+	"github.com/Azure/aks-flex/plugin/pkg/services/agentpools/azure/flexvm"
+	"github.com/Azure/aks-flex/plugin/pkg/services/agentpools/azure/ubuntu2404vmss"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/anypb"
+)
+
+func TestFlexAgentPoolFromGetResponse(t *testing.T) {
+	t.Parallel()
+
+	mkMeta := func(id string) *pluginapi.Metadata {
+		return pluginapi.Metadata_builder{Id: proto.String(id)}.Build()
+	}
+	mkFlexResp := func(id string) *pluginapi.GetResponse {
+		item, err := anypb.New(flexvm.AgentPool_builder{
+			Metadata: mkMeta(id),
+		}.Build())
+		if err != nil {
+			t.Fatalf("building flex anypb: %v", err)
+		}
+		return pluginapi.GetResponse_builder{Item: item}.Build()
+	}
+	mkVMSSResp := func(id string) *pluginapi.GetResponse {
+		item, err := anypb.New(ubuntu2404vmss.AgentPool_builder{
+			Metadata: mkMeta(id),
+		}.Build())
+		if err != nil {
+			t.Fatalf("building vmss anypb: %v", err)
+		}
+		return pluginapi.GetResponse_builder{Item: item}.Build()
+	}
+
+	tests := []struct {
+		name    string
+		resp    *pluginapi.GetResponse
+		wantID  string
+		wantErr bool
+	}{
+		{
+			name:    "nil response is not found",
+			resp:    nil,
+			wantErr: true,
+		},
+		{
+			name:    "nil item is not found",
+			resp:    pluginapi.GetResponse_builder{}.Build(),
+			wantErr: true,
+		},
+		{
+			name:    "wrong item type is not found",
+			resp:    mkVMSSResp("node-1"),
+			wantErr: true,
+		},
+		{
+			name:   "flex agentpool item returns parsed object",
+			resp:   mkFlexResp("node-2"),
+			wantID: "node-2",
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			got, err := flexAgentPoolFromGetResponse(tc.resp)
+			if tc.wantErr {
+				if err == nil {
+					t.Fatalf("expected error, got nil")
+				}
+				if !IsNotFound(err) {
+					t.Fatalf("expected NotFound-style error, got: %v", err)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			}
+			if got.GetMetadata().GetId() != tc.wantID {
+				t.Fatalf("got id %q, want %q", got.GetMetadata().GetId(), tc.wantID)
+			}
+		})
+	}
+}