create state disk on constellation create

cli/azure/client/compute.go

@@ -17,6 +17,7 @@ func (c *Client) CreateInstances(ctx context.Context, input CreateInstancesInput
 		NamePrefix:           c.name + "-worker-" + c.uid + "-",
 		Count:                input.Count - 1,
 		InstanceType:         input.InstanceType,
+		StateDiskSizeGB:      int32(input.StateDiskSizeGB),
 		Image:                input.Image,
 		UserAssingedIdentity: input.UserAssingedIdentity,
 	}
@@ -33,6 +34,7 @@ func (c *Client) CreateInstances(ctx context.Context, input CreateInstancesInput
 		NamePrefix:           c.name + "-control-plane-" + c.uid + "-",
 		Count:                1,
 		InstanceType:         input.InstanceType,
+		StateDiskSizeGB:      int32(input.StateDiskSizeGB),
 		Image:                input.Image,
 		UserAssingedIdentity: input.UserAssingedIdentity,
 	}
@@ -63,6 +65,7 @@ func (c *Client) CreateInstances(ctx context.Context, input CreateInstancesInput
 type CreateInstancesInput struct {
 	Count                int
 	InstanceType         string
+	StateDiskSizeGB      int
 	Image                string
 	UserAssingedIdentity string
 }
@@ -165,6 +168,7 @@ func (c *Client) createScaleSet(ctx context.Context, input CreateScaleSetInput)
 		NamePrefix:           input.NamePrefix,
 		Location:             c.location,
 		InstanceType:         input.InstanceType,
+		StateDiskSizeGB:      input.StateDiskSizeGB,
 		Count:                int64(input.Count),
 		Username:             "constellation",
 		SubnetID:             c.subnetID,
@@ -239,6 +243,7 @@ type CreateScaleSetInput struct {
 	NamePrefix           string
 	Count                int
 	InstanceType         string
+	StateDiskSizeGB      int32
 	Image                string
 	UserAssingedIdentity string
 }

cli/azure/scaleset.go

@@ -14,6 +14,7 @@ type ScaleSet struct {
 	NamePrefix           string
 	Location             string
 	InstanceType         string
+	StateDiskSizeGB      int32
 	Count                int64
 	Username             string
 	SubnetID             string
@@ -52,6 +53,13 @@ func (s ScaleSet) Azure() armcompute.VirtualMachineScaleSet {
 					ImageReference: &armcompute.ImageReference{
 						ID: to.StringPtr(s.Image),
 					},
+					DataDisks: []*armcompute.VirtualMachineScaleSetDataDisk{
+						{
+							CreateOption: armcompute.DiskCreateOptionTypesEmpty.ToPtr(),
+							DiskSizeGB:   to.Int32Ptr(s.StateDiskSizeGB),
+							Lun:          to.Int32Ptr(0),
+						},
+					},
 				},
 				NetworkProfile: &armcompute.VirtualMachineScaleSetNetworkProfile{
 					NetworkInterfaceConfigurations: []*armcompute.VirtualMachineScaleSetNetworkConfiguration{

cli/cmd/create_azure.go

@@ -105,6 +105,7 @@ func createAzure(cmd *cobra.Command, cl azureclient, fileHandler file.Handler, c
 	if err := cl.CreateInstances(cmd.Context(), client.CreateInstancesInput{
 		Count:                count,
 		InstanceType:         size,
+		StateDiskSizeGB:      *config.StateDiskSizeGB,
 		Image:                *config.Provider.Azure.Image,
 		UserAssingedIdentity: *config.Provider.Azure.UserAssignedIdentity,
 	}); err != nil {

cli/cmd/create_gcp.go

@@ -68,11 +68,12 @@ func createGCP(cmd *cobra.Command, cl gcpclient, fileHandler file.Handler, confi
 	}
 
 	createInput := client.CreateInstancesInput{
-		Count:        count,
-		ImageId:      *config.Provider.GCP.Image,
-		InstanceType: size,
-		KubeEnv:      gcp.KubeEnv,
-		DisableCVM:   *config.Provider.GCP.DisableCVM,
+		Count:           count,
+		ImageId:         *config.Provider.GCP.Image,
+		InstanceType:    size,
+		StateDiskSizeGB: *config.StateDiskSizeGB,
+		KubeEnv:         gcp.KubeEnv,
+		DisableCVM:      *config.Provider.GCP.DisableCVM,
 	}
 
 	ok, err := cmd.Flags().GetBool("yes")

cli/gcp/client/instances.go

@@ -31,6 +31,7 @@ func (c *Client) CreateInstances(ctx context.Context, input CreateInstancesInput
 		Subnetwork:                   c.subnetwork,
 		ImageId:                      input.ImageId,
 		InstanceType:                 input.InstanceType,
+		StateDiskSizeGB:              int64(input.StateDiskSizeGB),
 		KubeEnv:                      input.KubeEnv,
 		Project:                      c.project,
 		Zone:                         c.zone,
@@ -46,17 +47,18 @@ func (c *Client) CreateInstances(ctx context.Context, input CreateInstancesInput
 	c.nodeTemplate = nodeTemplateInput.Name
 
 	coordinatorTemplateInput := insertInstanceTemplateInput{
-		Name:         c.name + "-control-plane-" + c.uid,
-		Network:      c.network,
-		Subnetwork:   c.subnetwork,
-		ImageId:      input.ImageId,
-		InstanceType: input.InstanceType,
-		KubeEnv:      input.KubeEnv,
-		Project:      c.project,
-		Zone:         c.zone,
-		Region:       c.region,
-		UID:          c.uid,
-		DisableCVM:   input.DisableCVM,
+		Name:            c.name + "-control-plane-" + c.uid,
+		Network:         c.network,
+		Subnetwork:      c.subnetwork,
+		ImageId:         input.ImageId,
+		InstanceType:    input.InstanceType,
+		StateDiskSizeGB: int64(input.StateDiskSizeGB),
+		KubeEnv:         input.KubeEnv,
+		Project:         c.project,
+		Zone:            c.zone,
+		Region:          c.region,
+		UID:             c.uid,
+		DisableCVM:      input.DisableCVM,
 	}
 	op, err = c.insertInstanceTemplate(ctx, coordinatorTemplateInput)
 	if err != nil {
@@ -285,11 +287,12 @@ func (i *instanceGroupManagerInput) InsertInstanceGroupManagerRequest() computep
 
 // CreateInstancesInput is the input for a CreatInstances operation.
 type CreateInstancesInput struct {
-	Count        int
-	ImageId      string
-	InstanceType string
-	KubeEnv      string
-	DisableCVM   bool
+	Count           int
+	ImageId         string
+	InstanceType    string
+	StateDiskSizeGB int
+	KubeEnv         string
+	DisableCVM      bool
 }
 
 type insertInstanceTemplateInput struct {
@@ -299,6 +302,7 @@ type insertInstanceTemplateInput struct {
 	SecondarySubnetworkRangeName string
 	ImageId                      string
 	InstanceType                 string
+	StateDiskSizeGB              int64
 	KubeEnv                      string
 	Project                      string
 	Zone                         string
@@ -327,6 +331,15 @@ func (i insertInstanceTemplateInput) insertInstanceTemplateRequest() *computepb.
 						Boot:       proto.Bool(true),
 						Mode:       proto.String(computepb.AttachedDisk_READ_WRITE.String()),
 					},
+					{
+						InitializeParams: &computepb.AttachedDiskInitializeParams{
+							DiskSizeGb: proto.Int64(i.StateDiskSizeGB),
+						},
+						AutoDelete: proto.Bool(true),
+						DeviceName: proto.String("state-disk"),
+						Mode:       proto.String(computepb.AttachedDisk_READ_WRITE.String()),
+						Type:       proto.String(computepb.AttachedDisk_PERSISTENT.String()),
+					},
 				},
 				MachineType: proto.String(i.InstanceType),
 				Metadata: &computepb.Metadata{

internal/config/config.go

@@ -53,6 +53,7 @@ type Config struct {
 	CoordinatorPort          *string         `json:"coordinatorport,omitempty"`
 	AutoscalingNodeGroupsMin *int            `json:"autoscalingnodegroupsmin,omitempty"`
 	AutoscalingNodeGroupsMax *int            `json:"autoscalingnodegroupsmax,omitempty"`
+	StateDiskSizeGB          *int            `json:"statedisksizegb,omitempty"`
 	Provider                 *ProviderConfig `json:"provider,omitempty"`
 }
 
@@ -66,6 +67,7 @@ func Default() *Config {
 		CoordinatorPort:          proto.String(strconv.Itoa(coordinatorPort)),
 		AutoscalingNodeGroupsMin: intPtr(1),
 		AutoscalingNodeGroupsMax: intPtr(10),
+		StateDiskSizeGB:          intPtr(30),
 		Provider: &ProviderConfig{
 			EC2: &EC2Config{
 				Image: proto.String("ami-07d3864beb84157d3"),