cli: azure scale set poller: check for power state of every instance (#78)

cli/internal/azure/client/api.go

@@ -51,6 +51,12 @@ type scaleSetsAPI interface {
 		*runtime.Poller[armcomputev2.VirtualMachineScaleSetsClientCreateOrUpdateResponse], error)
 }
 
+type virtualMachineScaleSetVMsAPI interface {
+	GetInstanceView(ctx context.Context, resourceGroupName string, vmScaleSetName string, instanceID string,
+		options *armcomputev2.VirtualMachineScaleSetVMsClientGetInstanceViewOptions,
+	) (armcomputev2.VirtualMachineScaleSetVMsClientGetInstanceViewResponse, error)
+}
+
 type publicIPAddressesAPI interface {
 	NewListVirtualMachineScaleSetVMPublicIPAddressesPager(
 		resourceGroupName string, virtualMachineScaleSetName string,

cli/internal/azure/client/client.go

@@ -37,6 +37,7 @@ type Client struct {
 	networkSecurityGroupsAPI
 	resourceAPI
 	scaleSetsAPI
+	virtualMachineScaleSetVMsAPI
 	publicIPAddressesAPI
 	networkInterfacesAPI
 	loadBalancersAPI
@@ -91,6 +92,10 @@ func NewFromDefault(subscriptionID, tenantID string) (*Client, error) {
 	if err != nil {
 		return nil, err
 	}
+	virtualMachineScaleSetVMsAPI, err := armcomputev2.NewVirtualMachineScaleSetVMsClient(subscriptionID, cred, nil)
+	if err != nil {
+		return nil, err
+	}
 	publicIPAddressesAPI, err := armnetwork.NewPublicIPAddressesClient(subscriptionID, cred, nil)
 	if err != nil {
 		return nil, err
@@ -119,22 +124,23 @@ func NewFromDefault(subscriptionID, tenantID string) (*Client, error) {
 	roleAssignmentsAPI.Authorizer = managementAuthorizer
 
 	return &Client{
-		networksAPI:              netAPI,
+		networksAPI:                  netAPI,
-		networkSecurityGroupsAPI: netSecGrpAPI,
+		networkSecurityGroupsAPI:     netSecGrpAPI,
-		resourceAPI:              resourceAPI,
+		resourceAPI:                  resourceAPI,
-		scaleSetsAPI:             scaleSetAPI,
+		scaleSetsAPI:                 scaleSetAPI,
-		publicIPAddressesAPI:     publicIPAddressesAPI,
+		virtualMachineScaleSetVMsAPI: virtualMachineScaleSetVMsAPI,
-		networkInterfacesAPI:     networkInterfacesAPI,
+		publicIPAddressesAPI:         publicIPAddressesAPI,
-		loadBalancersAPI:         loadBalancersAPI,
+		networkInterfacesAPI:         networkInterfacesAPI,
-		applicationsAPI:          applicationsAPI,
+		loadBalancersAPI:             loadBalancersAPI,
-		servicePrincipalsAPI:     servicePrincipalsAPI,
+		applicationsAPI:              applicationsAPI,
-		roleAssignmentsAPI:       roleAssignmentsAPI,
+		servicePrincipalsAPI:         servicePrincipalsAPI,
-		applicationInsightsAPI:   applicationInsightsAPI,
+		roleAssignmentsAPI:           roleAssignmentsAPI,
-		subscriptionID:           subscriptionID,
+		applicationInsightsAPI:       applicationInsightsAPI,
-		tenantID:                 tenantID,
+		subscriptionID:               subscriptionID,
-		workers:                  cloudtypes.Instances{},
+		tenantID:                     tenantID,
-		controlPlanes:            cloudtypes.Instances{},
+		workers:                      cloudtypes.Instances{},
-		pollFrequency:            time.Second * 5,
+		controlPlanes:                cloudtypes.Instances{},
+		pollFrequency:                time.Second * 5,
 	}, nil
 }
 

cli/internal/azure/client/compute.go

@@ -16,13 +16,18 @@ import (
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	armcomputev2 "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v2"
 	"github.com/edgelesssys/constellation/cli/internal/azure"
 	"github.com/edgelesssys/constellation/cli/internal/azure/internal/poller"
 	"github.com/edgelesssys/constellation/internal/cloud/cloudtypes"
 )
 
-// scaleSetCreateTimeout maximum timeout to wait for scale set creation.
+const (
-const scaleSetCreateTimeout = 5 * time.Minute
+	// scaleSetCreateTimeout maximum timeout to wait for scale set creation.
+	scaleSetCreateTimeout = 5 * time.Minute
+	powerStateStarting    = "PowerState/starting"
+	powerStateRunning     = "PowerState/running"
+)
 
 func (c *Client) CreateInstances(ctx context.Context, input CreateInstancesInput) error {
 	// Create worker scale set
@@ -150,9 +155,10 @@ func (c *Client) createScaleSet(ctx context.Context, input CreateScaleSetInput)
 	// use custom poller to wait for resource creation but skip waiting for OS provisioning.
 	// OS provisioning does not work reliably without the azure guest agent installed.
 	poller := poller.New[bool](&scaleSetCreationPollingHandler{
-		resourceGroup: c.resourceGroup,
+		resourceGroup:                c.resourceGroup,
-		scaleSet:      input.Name,
+		scaleSet:                     input.Name,
-		scaleSetsAPI:  c.scaleSetsAPI,
+		scaleSetsAPI:                 c.scaleSetsAPI,
+		virtualMachineScaleSetVMsAPI: c.virtualMachineScaleSetVMsAPI,
 	})
 
 	pollCtx, cancel := context.WithTimeout(ctx, scaleSetCreateTimeout)
@@ -220,10 +226,12 @@ type CreateScaleSetInput struct {
 
 // scaleSetCreationPollingHandler is a custom poller used to check if a scale set was created successfully.
 type scaleSetCreationPollingHandler struct {
-	done          bool
+	done                         bool
-	resourceGroup string
+	instanceIDOffset             int
-	scaleSet      string
+	resourceGroup                string
-	scaleSetsAPI  scaleSetsAPI
+	scaleSet                     string
+	scaleSetsAPI                 scaleSetsAPI
+	virtualMachineScaleSetVMsAPI virtualMachineScaleSetVMsAPI
 }
 
 // Done returns true if the condition is met.
@@ -231,19 +239,29 @@ func (h *scaleSetCreationPollingHandler) Done() bool {
 	return h.done
 }
 
-// Poll checks if the scale set resource was created successfully.
+// Poll checks if the scale set resource was created successfully and every VM is starting or running.
 func (h *scaleSetCreationPollingHandler) Poll(ctx context.Context) error {
-	_, err := h.scaleSetsAPI.Get(ctx, h.resourceGroup, h.scaleSet, nil)
+	// check if scale set can be retrieved from API
-	if err == nil {
+	scaleSet, err := h.scaleSetsAPI.Get(ctx, h.resourceGroup, h.scaleSet, nil)
-		h.done = true
+	if err != nil {
-		return nil
+		return ignoreNotFoundError(err)
 	}
-	var respErr *azcore.ResponseError
+	if scaleSet.SKU == nil || scaleSet.SKU.Capacity == nil {
-	if errors.As(err, &respErr) && respErr.StatusCode == http.StatusNotFound {
+		return errors.New("invalid scale set capacity")
-		// resource does not exist yet - retry later
-		return nil
 	}
-	return err
+	// check if every VM in the scale set has power state starting or running
+	for i := h.instanceIDOffset; i < int(*scaleSet.SKU.Capacity); i++ {
+		instanceView, err := h.virtualMachineScaleSetVMsAPI.GetInstanceView(ctx, h.resourceGroup, h.scaleSet, strconv.Itoa(i), nil)
+		if err != nil {
+			return ignoreNotFoundError(err)
+		}
+		if !vmIsStartingOrRunning(instanceView.Statuses) {
+			return nil
+		}
+		h.instanceIDOffset = i + 1 // skip this VM in the next Poll() invocation
+	}
+	h.done = true
+	return nil
 }
 
 // Result returns the result of the poller if the condition is met.
@@ -255,3 +273,27 @@ func (h *scaleSetCreationPollingHandler) Result(ctx context.Context, out *bool)
 	*out = h.done
 	return nil
 }
+
+func ignoreNotFoundError(err error) error {
+	var respErr *azcore.ResponseError
+	if errors.As(err, &respErr) && respErr.StatusCode == http.StatusNotFound {
+		// resource does not exist yet - retry later
+		return nil
+	}
+	return err
+}
+
+func vmIsStartingOrRunning(statuses []*armcomputev2.InstanceViewStatus) bool {
+	for _, status := range statuses {
+		if status == nil || status.Code == nil {
+			continue
+		}
+		switch *status.Code {
+		case powerStateStarting:
+			return true
+		case powerStateRunning:
+			return true
+		}
+	}
+	return false
+}

cli/internal/azure/client/compute_test.go

@@ -30,8 +30,8 @@ func TestCreateInstances(t *testing.T) {
 			publicIPAddressesAPI: stubPublicIPAddressesAPI{},
 			networkInterfacesAPI: stubNetworkInterfacesAPI{},
 			scaleSetsAPI: stubScaleSetsAPI{
-				stubResponse: armcomputev2.VirtualMachineScaleSetsClientCreateOrUpdateResponse{
+				getResponse: armcomputev2.VirtualMachineScaleSet{
-					VirtualMachineScaleSet: armcomputev2.VirtualMachineScaleSet{Identity: &armcomputev2.VirtualMachineScaleSetIdentity{PrincipalID: to.Ptr("principal-id")}},
+					Identity: &armcomputev2.VirtualMachineScaleSetIdentity{PrincipalID: to.Ptr("principal-id")}, SKU: &armcomputev2.SKU{Capacity: to.Ptr[int64](0)},
 				},
 			},
 			createInstancesInput: CreateInstancesInput{