platform/azure: support additional data disks

Add support for attaching additional data disks to instances created in
Azure. Disks are defined through the machines options as the Size in GB
and the 'sku', or storage type, e.g. '["100G:sku=UltraSSD_LRS"]' for
NVMe disks.

Author: marmijo

mantle/cmd/kola/kola.go

@@ -300,13 +300,14 @@ func writeProps() error {
 		InstanceType string `json:"type"`
 	}
 	type Azure struct {
-		DiskURI   string `json:"diskUri"`
-		Publisher string `json:"publisher"`
-		Offer     string `json:"offer"`
-		Sku       string `json:"sku"`
-		Version   string `json:"version"`
-		Location  string `json:"location"`
-		Size      string `json:"size"`
+		DiskURI          string `json:"diskUri"`
+		Publisher        string `json:"publisher"`
+		Offer            string `json:"offer"`
+		Sku              string `json:"sku"`
+		Version          string `json:"version"`
+		Location         string `json:"location"`
+		Size             string `json:"size"`
+		AvailabilityZone string `json:"availability_zone"`
 	}
 	type DO struct {
 		Region string `json:"region"`
@@ -355,13 +356,14 @@ func writeProps() error {
 			InstanceType: kola.AWSOptions.InstanceType,
 		},
 		Azure: Azure{
-			DiskURI:   kola.AzureOptions.DiskURI,
-			Publisher: kola.AzureOptions.Publisher,
-			Offer:     kola.AzureOptions.Offer,
-			Sku:       kola.AzureOptions.Sku,
-			Version:   kola.AzureOptions.Version,
-			Location:  kola.AzureOptions.Location,
-			Size:      kola.AzureOptions.Size,
+			DiskURI:          kola.AzureOptions.DiskURI,
+			Publisher:        kola.AzureOptions.Publisher,
+			Offer:            kola.AzureOptions.Offer,
+			Sku:              kola.AzureOptions.Sku,
+			Version:          kola.AzureOptions.Version,
+			Location:         kola.AzureOptions.Location,
+			Size:             kola.AzureOptions.Size,
+			AvailabilityZone: kola.AzureOptions.AvailabilityZone,
 		},
 		DO: DO{
 			Region: kola.DOOptions.Region,

mantle/cmd/kola/options.go

@@ -100,6 +100,7 @@ func init() {
 	sv(&kola.AzureOptions.Version, "azure-version", "", "Azure image version")
 	sv(&kola.AzureOptions.Location, "azure-location", "westus", "Azure location (default \"westus\"")
 	sv(&kola.AzureOptions.Size, "azure-size", "Standard_D2_v2", "Azure machine size (default \"Standard_D2_v2\")")
+	sv(&kola.AzureOptions.AvailabilityZone, "azure-availability-zone", "1", "Azure Availability Zone (default \"1\")")
 
 	// do-specific options
 	sv(&kola.DOOptions.ConfigPath, "do-config-file", "", "DigitalOcean config file (default \"~/"+auth.DOConfigPath+"\")")

mantle/platform/api/azure/disk.go

@@ -0,0 +1,105 @@
+// Copyright 2025 Red Hat
+// Copyright 2016 CoreOS, Inc.
+//
+// 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 azure
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute"
+
+	"github.com/coreos/coreos-assembler/mantle/util"
+)
+
+// CreateDisk provisions a new managed disk in the specified Azure resource group using
+// the given name, size (in GiB), and SKU (e.g., Premium_LRS). The disk is created in
+// the location and availability zone specified in the API options.
+func (a *API) CreateDisk(name, resourceGroup string, sizeGB int32, sku armcompute.DiskStorageAccountTypes) (string, error) {
+	ctx := context.Background()
+	poller, err := a.diskClient.BeginCreateOrUpdate(ctx, resourceGroup, name, armcompute.Disk{
+		Location: &a.opts.Location,
+		Zones:    []*string{&a.opts.AvailabilityZone},
+		Tags: map[string]*string{
+			"createdBy": to.Ptr("mantle"),
+		},
+		SKU: &armcompute.DiskSKU{
+			Name: to.Ptr(sku),
+		},
+		Properties: &armcompute.DiskProperties{
+			DiskSizeGB: to.Ptr(sizeGB),
+			CreationData: &armcompute.CreationData{
+				CreateOption: to.Ptr(armcompute.DiskCreateOptionEmpty),
+			},
+		},
+	}, nil)
+
+	if err != nil {
+		return "", fmt.Errorf("failed to create azure disk %v", err)
+	}
+
+	diskResponse, err := poller.PollUntilDone(context.Background(), nil)
+	if err != nil {
+		return "", err
+	}
+
+	if diskResponse.Disk.ID == nil {
+		return "", fmt.Errorf("failed to get azure disk id")
+	}
+
+	return *diskResponse.Disk.ID, nil
+}
+
+// DeleteDisk deletes a managed disk by name from the specified Azure resource group.
+func (a *API) DeleteDisk(name, resourceGroup string) error {
+	ctx := context.Background()
+	poller, err := a.diskClient.BeginDelete(ctx, resourceGroup, name, nil)
+	if err != nil {
+		return err
+	}
+	_, err = poller.PollUntilDone(ctx, nil)
+	return err
+}
+
+// example: ["10G:sku=UltraSSD_LRS"], the default disk type is "Standard_LRS"
+func (a *API) ParseDisk(spec string) (int64, armcompute.DiskStorageAccountTypes, error) {
+	sku := armcompute.DiskStorageAccountTypes(armcompute.DiskStorageAccountTypesStandardLRS)
+	size, diskmap, err := util.ParseDiskSpec(spec, false)
+	if err != nil {
+		return size, sku, fmt.Errorf("failed to parse disk spec %q: %w", spec, err)
+	}
+	for key, value := range diskmap {
+		switch key {
+		case "sku":
+			normalizedSku := strings.ToUpper(value)
+			foundSku := false
+			for _, validSku := range armcompute.PossibleDiskStorageAccountTypesValues() {
+				if strings.EqualFold(normalizedSku, string(validSku)) {
+					sku = validSku
+					foundSku = true
+					break
+				}
+			}
+			if !foundSku {
+				return size, sku, fmt.Errorf("unsupported disk sku: %s", value)
+			}
+		default:
+			return size, sku, fmt.Errorf("invalid key: %s", key)
+		}
+	}
+	return size, sku, nil
+}

mantle/platform/api/azure/instance.go

@@ -30,6 +30,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork"
 
+	"github.com/coreos/coreos-assembler/mantle/platform"
 	"github.com/coreos/coreos-assembler/mantle/util"
 )
 
@@ -108,6 +109,7 @@ func (a *API) getVMParameters(name, userdata, sshkey, storageAccountURI string,
 	return armcompute.VirtualMachine{
 		Name:     &name,
 		Location: &a.opts.Location,
+		Zones:    []*string{&a.opts.AvailabilityZone},
 		Tags: map[string]*string{
 			"createdBy": to.Ptr("mantle"),
 		},
@@ -142,7 +144,7 @@ func (a *API) getVMParameters(name, userdata, sshkey, storageAccountURI string,
 	}
 }
 
-func (a *API) CreateInstance(name, userdata, sshkey, resourceGroup, storageAccount string) (*Machine, error) {
+func (a *API) CreateInstance(name, userdata, sshkey, resourceGroup, storageAccount string, opts platform.MachineOptions) (*Machine, error) {
 	subnet, err := a.getSubnet(resourceGroup)
 	if err != nil {
 		return nil, fmt.Errorf("preparing network resources: %v", err)
@@ -206,6 +208,18 @@ func (a *API) CreateInstance(name, userdata, sshkey, resourceGroup, storageAccou
 		return nil, fmt.Errorf("couldn't get VM ID")
 	}
 
+	for i, spec := range opts.AdditionalDisks {
+		size, sku, err := a.ParseDisk(spec)
+		if err != nil {
+			return nil, fmt.Errorf("error parsing additional disk: %v", err)
+		}
+		diskName := util.RandomName(fmt.Sprintf("disk-%d", i))
+		err = a.AttachDiskToInstance(*vm.Name, resourceGroup, diskName, sku, int32(size), int32(i))
+		if err != nil {
+			return nil, fmt.Errorf("failed to attach disk to vm: %v", err)
+		}
+	}
+
 	publicaddr, privaddr, err := a.GetIPAddresses(*nic.Name, *ip.Name, resourceGroup)
 	if err != nil {
 		return nil, err
@@ -272,3 +286,45 @@ func (a *API) GetConsoleOutput(name, resourceGroup, storageAccount string) ([]by
 
 	return io.ReadAll(data)
 }
+
+func (a *API) AttachDiskToInstance(instanceName, resourceGroup, diskName string, sku armcompute.DiskStorageAccountTypes, sizeGB int32, lun int32) error {
+	ctx := context.Background()
+
+	vm, err := a.getInstance(instanceName, resourceGroup)
+	if err != nil {
+		return err
+	}
+
+	diskID, err := a.CreateDisk(diskName, resourceGroup, sizeGB, sku)
+	if err != nil {
+		return err
+	}
+
+	newDisk := armcompute.DataDisk{
+		Lun:          to.Ptr(lun),
+		Name:         to.Ptr(diskName),
+		CreateOption: to.Ptr(armcompute.DiskCreateOptionTypesAttach),
+		ManagedDisk: &armcompute.ManagedDiskParameters{
+			ID: to.Ptr(diskID),
+		},
+	}
+
+	if vm.Properties.StorageProfile.DataDisks == nil {
+		vm.Properties.StorageProfile.DataDisks = []*armcompute.DataDisk{}
+	}
+	vm.Properties.StorageProfile.DataDisks = append(vm.Properties.StorageProfile.DataDisks, &newDisk)
+
+	poller, err := a.compClient.BeginUpdate(ctx, resourceGroup, instanceName, armcompute.VirtualMachineUpdate{
+		Properties: &armcompute.VirtualMachineProperties{
+			StorageProfile: &armcompute.StorageProfile{
+				DataDisks: vm.Properties.StorageProfile.DataDisks,
+			},
+		},
+	}, nil)
+	if err != nil {
+		return fmt.Errorf("failed to attach disk to VM %s: %v", instanceName, err)
+	}
+
+	_, err = poller.PollUntilDone(ctx, nil)
+	return err
+}

mantle/platform/api/azure/options.go

@@ -25,13 +25,14 @@ type Options struct {
 	AzureCredentials  string
 	AzureSubscription string
 
-	DiskURI   string
-	Publisher string
-	Offer     string
-	Sku       string
-	Version   string
-	Size      string
-	Location  string
+	DiskURI          string
+	Publisher        string
+	Offer            string
+	Sku              string
+	Version          string
+	Size             string
+	Location         string
+	AvailabilityZone string
 
 	SubscriptionName string
 	SubscriptionID   string

mantle/platform/machine/azure/cluster.go

@@ -46,9 +46,6 @@ func (ac *cluster) NewMachine(userdata *conf.UserData) (platform.Machine, error)
 }
 
 func (ac *cluster) NewMachineWithOptions(userdata *conf.UserData, options platform.MachineOptions) (platform.Machine, error) {
-	if len(options.AdditionalDisks) > 0 {
-		return nil, errors.New("platform azure does not yet support additional disks")
-	}
 	if options.MultiPathDisk {
 		return nil, errors.New("platform azure does not support multipathed disks")
 	}
@@ -69,7 +66,7 @@ func (ac *cluster) NewMachineWithOptions(userdata *conf.UserData, options platfo
 		return nil, err
 	}
 
-	instance, err := ac.flight.api.CreateInstance(ac.vmname(), conf.String(), ac.sshKey, ac.ResourceGroup, ac.StorageAccount)
+	instance, err := ac.flight.api.CreateInstance(ac.vmname(), conf.String(), ac.sshKey, ac.ResourceGroup, ac.StorageAccount, options)
 	if err != nil {
 		return nil, err
 	}