Use a less expensive method for storage policy

pkg/services/govmomi/service.go

@@ -394,35 +394,47 @@ func (vms *VMService) reconcileStoragePolicy(ctx context.Context, virtualMachine
 	if err != nil {
 		return errors.Wrap(err, "unable to retrieve storage profile ID")
 	}
-	entities, err := pbmClient.QueryAssociatedEntity(ctx, pbmTypes.PbmProfileId{UniqueId: storageProfileID}, "virtualDiskId")
-	if err != nil {
-		return err
-	}
 
 	var changes []types.BaseVirtualDeviceConfigSpec
 	devices, err := virtualMachineCtx.Obj.Device(ctx)
 	if err != nil {
 		return err
 	}
 
+	arrayOfEntities := make([]pbmTypes.PbmServerObjectRef, 0)
+	// diskMap is just an auxiliar map so we don't need to iterate over and over disks to get their configs
+	// if we realize they are not on the right storage policy
+	diskMap := make(map[string]*types.VirtualDisk)
+
 	disks := devices.SelectByType((*types.VirtualDisk)(nil))
+
+	// We iterate over disks and create an array of entities, so we just need to make a single call
+	// against vCenter, instead of one call per disk
+	// the diskMap is an auxiliar way of, besides the arrayOfEntities, we have a "searchable" disk configuration
+	// in case we need to reconfigure a disk, to get its config
 	for _, d := range disks {
 		disk := d.(*types.VirtualDisk)
-		found := false
 		// entities associated with storage policy has key in the form <vm-ID>:<disk>
 		diskID := fmt.Sprintf("%s:%d", virtualMachineCtx.Obj.Reference().Value, disk.Key)
-		for _, e := range entities {
-			if e.Key == diskID {
-				found = true
-				break
-			}
-		}
+		diskMap[diskID] = disk
+
+		arrayOfEntities = append(arrayOfEntities, pbmTypes.PbmServerObjectRef{
+			ObjectType: string(pbmTypes.PbmObjectTypeVirtualDiskId),
+			Key:        diskID,
+		})
+	}
+
+	entities, err := QueryAssociatedProfiles(ctx, pbmClient, arrayOfEntities)
+	if err != nil {
+		return err
+	}
 
-		if !found {
-			// disk wasn't associated with storage policy, create a device change to make the association
+	for k := range entities {
+		if !isStoragePolicyIDPresent(storageProfileID, entities[k]) {
+			virtualMachineCtx.Logger.V(5).Info("storage policy not found on disk, adding for reconciliation", "disk", entities[k].Object.Key)
 			config := &types.VirtualDeviceConfigSpec{
 				Operation: types.VirtualDeviceConfigSpecOperationEdit,
-				Device:    disk,
+				Device:    diskMap[entities[k].Object.Key],
 				Profile: []types.BaseVirtualMachineProfileSpec{
 					&types.VirtualMachineDefinedProfileSpec{ProfileId: storageProfileID},
 				},

pkg/services/govmomi/service_test.go

@@ -18,11 +18,14 @@ package govmomi
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/go-logr/logr"
 	. "github.com/onsi/gomega"
 	"github.com/vmware/govmomi/find"
+	"github.com/vmware/govmomi/object"
+	pbmsimulator "github.com/vmware/govmomi/pbm/simulator"
 	"github.com/vmware/govmomi/simulator"
 	"github.com/vmware/govmomi/vim25"
 	"github.com/vmware/govmomi/vim25/types"
@@ -32,6 +35,11 @@ import (
 
 	infrav1 "sigs.k8s.io/cluster-api-provider-vsphere/apis/v1beta1"
 	capvcontext "sigs.k8s.io/cluster-api-provider-vsphere/pkg/context"
+	"sigs.k8s.io/cluster-api-provider-vsphere/pkg/session"
+)
+
+const (
+	defaultStoragePolicy = "vSAN Default Storage Policy"
 )
 
 func emptyVirtualMachineContext() *virtualMachineContext {
@@ -93,9 +101,134 @@ func Test_reconcilePCIDevices(t *testing.T) {
 			pciDevices := devices.SelectByBackingInfo(&types.VirtualPCIPassthroughDynamicBackingInfo{
 				AllowedDevice: []types.VirtualPCIPassthroughAllowedDevice{
 					{DeviceId: 1234, VendorId: 5678},
-				}})
+				},
+			})
 			g.Expect(pciDevices).To(HaveLen(2))
 			return nil
 		})
 	})
 }
+
+func Test_ReconcileStoragePolicy(t *testing.T) {
+	var vmCtx *virtualMachineContext
+	var g *WithT
+	var vms *VMService
+
+	before := func() {
+		vmCtx = emptyVirtualMachineContext()
+		vmCtx.Client = fake.NewClientBuilder().Build()
+
+		vms = &VMService{}
+	}
+	t.Run("when VM has no storage policy spec", func(t *testing.T) {
+		g = NewWithT(t)
+		before()
+		vmCtx.VSphereVM = &infrav1.VSphereVM{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "vsphereVM1",
+				Namespace: "my-namespace",
+			},
+			Spec: infrav1.VSphereVMSpec{
+				VirtualMachineCloneSpec: infrav1.VirtualMachineCloneSpec{},
+			},
+		}
+		g.Expect(vms.reconcileStoragePolicy(context.Background(), vmCtx)).ToNot(HaveOccurred())
+		g.Expect(vmCtx.VSphereVM.Status.TaskRef).To(BeEmpty())
+	})
+
+	t.Run("when the requested storage policy does not exists should fail", func(t *testing.T) {
+		g = NewWithT(t)
+		before()
+		model, err := storagePolicyModel()
+		g.Expect(err).ToNot(HaveOccurred())
+
+		simulator.Run(func(ctx context.Context, c *vim25.Client) error {
+			authSession, err := getAuthSession(ctx, model.Service.Listen.Host)
+			g.Expect(err).ToNot(HaveOccurred())
+			vmCtx.Session = authSession
+			vm, err := getPoweredoffVM(ctx, c)
+			g.Expect(err).ToNot(HaveOccurred())
+
+			vmCtx.Obj = vm
+			vmCtx.VSphereVM = &infrav1.VSphereVM{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "vsphereVM1",
+					Namespace: "my-namespace",
+				},
+				Spec: infrav1.VSphereVMSpec{
+					VirtualMachineCloneSpec: infrav1.VirtualMachineCloneSpec{
+						StoragePolicyName: "non-existing-storagepolicy",
+					},
+				},
+			}
+			err = vms.reconcileStoragePolicy(context.Background(), vmCtx)
+			g.Expect(err.Error()).To(ContainSubstring("no pbm profile found with name"))
+			return nil
+		}, model)
+	})
+
+	t.Run("when the requested storage policy exists should pass", func(t *testing.T) {
+		g = NewWithT(t)
+		before()
+		model, err := storagePolicyModel()
+		g.Expect(err).ToNot(HaveOccurred())
+
+		simulator.Run(func(ctx context.Context, c *vim25.Client) error {
+			authSession, err := getAuthSession(ctx, model.Service.Listen.Host)
+			g.Expect(err).ToNot(HaveOccurred())
+			vmCtx.Session = authSession
+			vm, err := getPoweredoffVM(ctx, c)
+			g.Expect(err).ToNot(HaveOccurred())
+
+			vmCtx.Obj = vm
+			vmCtx.VSphereVM = &infrav1.VSphereVM{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "vsphereVM1",
+					Namespace: "my-namespace",
+				},
+				Spec: infrav1.VSphereVMSpec{
+					VirtualMachineCloneSpec: infrav1.VirtualMachineCloneSpec{
+						StoragePolicyName: defaultStoragePolicy,
+					},
+				},
+			}
+			err = vms.reconcileStoragePolicy(context.Background(), vmCtx)
+			g.Expect(err).ToNot(HaveOccurred())
+			return nil
+		}, model)
+	})
+}
+
+func getAuthSession(ctx context.Context, server string) (*session.Session, error) {
+	password, _ := simulator.DefaultLogin.Password()
+	return session.GetOrCreate(
+		ctx,
+		session.NewParams().
+			WithUserInfo(simulator.DefaultLogin.Username(), password).
+			WithServer(fmt.Sprintf("http://%s", server)).
+			WithDatacenter("*"))
+}
+
+func getPoweredoffVM(ctx context.Context, c *vim25.Client) (*object.VirtualMachine, error) {
+	finder := find.NewFinder(c)
+	vm, err := finder.VirtualMachine(ctx, "DC0_H0_VM0")
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = vm.PowerOff(ctx)
+	return vm, err
+}
+
+func storagePolicyModel() (*simulator.Model, error) {
+	model := simulator.VPX()
+	err := model.Create()
+	if err != nil {
+		return nil, err
+	}
+	model.Service.RegisterSDK(pbmsimulator.New())
+	model.Machine = 1
+	model.Datacenter = 1
+	model.Host = 1
+	return model, nil
+}

pkg/services/govmomi/storageprofile_util.go

@@ -0,0 +1,59 @@
+/*
+Copyright 2023 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package govmomi
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/vmware/govmomi/pbm"
+	"github.com/vmware/govmomi/pbm/methods"
+	pbmTypes "github.com/vmware/govmomi/pbm/types"
+)
+
+// QueryAssociatedProfiles is an implementation of https://github.com/vmware/govmomi/blob/6d70dadbd96b5640e25f782b7b0244a02aaedb58/pbm/client.go#L303
+// Because this is not yet released on govmomi used by CAPV, it is just a "copy" of the method
+// TODO: Remove this and use the method from pbmclient once govmomi v0.34 is released
+// and bumbed on CAPV.
+func QueryAssociatedProfiles(ctx context.Context, client *pbm.Client, entities []pbmTypes.PbmServerObjectRef) ([]pbmTypes.PbmQueryProfileResult, error) {
+	if client == nil {
+		return nil, fmt.Errorf("pbm client cannot be null")
+	}
+
+	req := pbmTypes.PbmQueryAssociatedProfiles{
+		This:     client.ServiceContent.ProfileManager,
+		Entities: entities,
+	}
+
+	res, err := methods.PbmQueryAssociatedProfiles(ctx, client, &req)
+	if err != nil {
+		return nil, err
+	}
+
+	return res.Returnval, nil
+}
+
+// isStoragePolicyIDPresent checks, given a ProfileResult if the requested storageprofileid
+// is associated at least with one entity.
+func isStoragePolicyIDPresent(storageProfileID string, result pbmTypes.PbmQueryProfileResult) bool {
+	for _, id := range result.ProfileId {
+		if id.UniqueId == storageProfileID {
+			return true
+		}
+	}
+	return false
+}