Add UpdateDisks SDK method for safe disk attach/detach

go.mod

@@ -8,8 +8,8 @@ require (
 	github.com/Azure/go-autorest/autorest/date v0.3.0
 	github.com/google/go-cmp v0.7.0
 	github.com/google/uuid v1.6.0
-	github.com/microsoft/moc v0.43.3
-	google.golang.org/grpc v1.79.3
+	github.com/microsoft/moc v0.43.4
+	google.golang.org/grpc v1.82.1
 	gopkg.in/yaml.v3 v3.0.1
 	k8s.io/klog v1.0.0
 	sigs.k8s.io/controller-runtime v0.20.4
@@ -62,11 +62,11 @@ require (
 	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/subosito/gotenv v1.6.0 // indirect
 	github.com/x448/float16 v0.8.4 // indirect
-	go.opentelemetry.io/otel v1.40.0 // indirect
-	go.opentelemetry.io/otel/sdk/metric v1.40.0 // indirect
+	go.opentelemetry.io/otel v1.43.0 // indirect
+	go.opentelemetry.io/otel/sdk/metric v1.43.0 // indirect
 	golang.org/x/crypto v0.54.0 // indirect
 	golang.org/x/net v0.57.0 // indirect
-	golang.org/x/oauth2 v0.34.0 // indirect
+	golang.org/x/oauth2 v0.36.0 // indirect
 	golang.org/x/sync v0.22.0 // indirect
 	golang.org/x/term v0.45.0 // indirect
 	golang.org/x/text v0.40.0 // indirect

go.sum

@@ -130,6 +130,8 @@ github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0
 github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
 github.com/microsoft/moc v0.43.3 h1:mhYqI83KPWzPbKyB302rBUoTMkXc9DgFtGNjt3tp/Ck=
 github.com/microsoft/moc v0.43.3/go.mod h1:BMmUfVTvjLWs+zOqxIMmYJxaP7w76n6v5F/poQNHaUk=
+github.com/microsoft/moc v0.43.4 h1:hQGDD/htyrwe29yTQBwmM4nFaHCeOr5CBi7QEoSDPjs=
+github.com/microsoft/moc v0.43.4/go.mod h1:R9YDj96FLWV8daoO2kdCyoei3+MhUhiRzqh8Y/UnSFs=
 github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
 github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo=
 github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
@@ -211,14 +213,21 @@ go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ
 go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
 go.opentelemetry.io/otel v1.40.0 h1:oA5YeOcpRTXq6NN7frwmwFR0Cn3RhTVZvXsP4duvCms=
 go.opentelemetry.io/otel v1.40.0/go.mod h1:IMb+uXZUKkMXdPddhwAHm6UfOwJyh4ct1ybIlV14J0g=
+go.opentelemetry.io/otel v1.43.0 h1:mYIM03dnh5zfN7HautFE4ieIig9amkNANT+xcVxAj9I=
+go.opentelemetry.io/otel v1.43.0/go.mod h1:JuG+u74mvjvcm8vj8pI5XiHy1zDeoCS2LB1spIq7Ay0=
 go.opentelemetry.io/otel/metric v1.40.0 h1:rcZe317KPftE2rstWIBitCdVp89A2HqjkxR3c11+p9g=
 go.opentelemetry.io/otel/metric v1.40.0/go.mod h1:ib/crwQH7N3r5kfiBZQbwrTge743UDc7DTFVZrrXnqc=
+go.opentelemetry.io/otel/metric v1.43.0 h1:d7638QeInOnuwOONPp4JAOGfbCEpYb+K6DVWvdxGzgM=
 go.opentelemetry.io/otel/sdk v1.40.0 h1:KHW/jUzgo6wsPh9At46+h4upjtccTmuZCFAc9OJ71f8=
 go.opentelemetry.io/otel/sdk v1.40.0/go.mod h1:Ph7EFdYvxq72Y8Li9q8KebuYUr2KoeyHx0DRMKrYBUE=
+go.opentelemetry.io/otel/sdk v1.43.0 h1:pi5mE86i5rTeLXqoF/hhiBtUNcrAGHLKQdhg4h4V9Dg=
 go.opentelemetry.io/otel/sdk/metric v1.40.0 h1:mtmdVqgQkeRxHgRv4qhyJduP3fYJRMX4AtAlbuWdCYw=
 go.opentelemetry.io/otel/sdk/metric v1.40.0/go.mod h1:4Z2bGMf0KSK3uRjlczMOeMhKU2rhUqdWNoKcYrtcBPg=
+go.opentelemetry.io/otel/sdk/metric v1.43.0 h1:S88dyqXjJkuBNLeMcVPRFXpRw2fuwdvfCGLEo89fDkw=
+go.opentelemetry.io/otel/sdk/metric v1.43.0/go.mod h1:C/RJtwSEJ5hzTiUz5pXF1kILHStzb9zFlIEe85bhj6A=
 go.opentelemetry.io/otel/trace v1.40.0 h1:WA4etStDttCSYuhwvEa8OP8I5EWu24lkOzp+ZYblVjw=
 go.opentelemetry.io/otel/trace v1.40.0/go.mod h1:zeAhriXecNGP/s2SEG3+Y8X9ujcJOTqQ5RgdEJcawiA=
+go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09nk+3A=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
 go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
 go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
@@ -319,10 +328,13 @@ gomodules.xyz/jsonpatch/v2 v2.4.0 h1:Ci3iUJyx9UeRx7CeFN8ARgGbkESwJK+KB9lLcWxY/Zw
 gomodules.xyz/jsonpatch/v2 v2.4.0/go.mod h1:AH3dM2RI6uoBZxn3LVrfvJ3E0/9dG4cSrbuBJT4moAY=
 gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
 gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=
+gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20260511170946-3700d4141b60 h1:seT2EwLWM78plQ7wcDfuWBc/4FAEAXDDiaSol4ku4qo=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20260511170946-3700d4141b60/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8=
 google.golang.org/grpc v1.79.3 h1:sybAEdRIEtvcD68Gx7dmnwjZKlyfuc61Dyo9pGXXkKE=
 google.golang.org/grpc v1.79.3/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ=
+google.golang.org/grpc v1.82.1 h1:NnAxzGRA0677vCa4BUkOAnO5+FfQqVl9iUXeD0IqcGE=
+google.golang.org/grpc v1.82.1/go.mod h1:yzTZ1TB1Z3SG+LIYaI+WiE8D5+PZ3ArnrSp8zF3+/ZA=
 google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8=
 google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0=
 google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM=

services/compute/compute.go

@@ -149,6 +149,21 @@ type StorageProfile struct {
 	// VMConfigContainerName - Name of the storage container that hosts the VM configuration file
 	VmConfigContainerName *string `json:"vmConfigContainerName,omitempty"`
 }
+
+// VirtualMachineDiskOperation selects whether UpdateDisks attaches or detaches
+// the supplied data disks on an existing virtual machine. The values mirror the
+// wire enum moc.cloudagent.compute.VirtualMachineDiskOperation.
+type VirtualMachineDiskOperation int32
+
+const (
+	// VirtualMachineDiskOperationUnspecified is invalid and rejected by the agent.
+	VirtualMachineDiskOperationUnspecified VirtualMachineDiskOperation = 0
+	// VirtualMachineDiskOperationAttach attaches the supplied data disks.
+	VirtualMachineDiskOperationAttach VirtualMachineDiskOperation = 1
+	// VirtualMachineDiskOperationDetach detaches the supplied data disks.
+	VirtualMachineDiskOperationDetach VirtualMachineDiskOperation = 2
+)
+
 type SSHPublicKey struct {
 	// Path - Specifies the full path on the created VM where ssh public key is stored. If the file already exists, the specified key is appended to the file. Example: /home/user/.ssh/authorized_keys
 	Path *string `json:"path,omitempty"`
@@ -858,7 +873,7 @@ type AzureBlobImageProperties struct {
 	Audience    string `json:"audience,omitempty"`
 	Version     string `json:"version,omitempty"`
 	ReleaseName string `json:"releaseName,omitempty"`
-	Parts       uint32  `json:"parts,omitempty"`
+	Parts       uint32 `json:"parts,omitempty"`
 	Cloud       string `json:"cloud,omitempty"`
 	Endpoint    string `json:"endpoint,omitempty"`
 }

services/compute/virtualmachine/client.go

@@ -34,6 +34,7 @@ type Service interface {
 	GetHyperVVmId(context.Context, string, string) (*compute.VirtualMachineHyperVVmId, error)
 	GetHostNodeName(context.Context, string, string) (*compute.VirtualMachineHostNodeName, error)
 	GetHostNodeIpAddress(context.Context, string, string) (*compute.VirtualMachineHostNodeIpAddress, error)
+	UpdateDisks(context.Context, string, string, []compute.DataDisk, compute.VirtualMachineDiskOperation) (*compute.VirtualMachine, error)
 }
 
 type VirtualMachineClient struct {
@@ -195,71 +196,23 @@ func (c *VirtualMachineClient) ResizeEx(ctx context.Context, group string, vmNam
 	return
 }
 
+// DiskAttach attaches diskName to vmName. It delegates to the UpdateDisks RPC,
+// which resolves the VM and routes through Update (never Create), so a request
+// that races with a delete fails with NotFound instead of resurrecting the VM.
+// Attaching an already-attached disk is an idempotent no-op.
 func (c *VirtualMachineClient) DiskAttach(ctx context.Context, group string, vmName, diskName string) (err error) {
-	for {
-		vms, err := c.Get(ctx, group, vmName)
-		if err != nil {
-			return err
-		}
-		if vms == nil || len(*vms) == 0 {
-			return errors.Wrapf(errors.NotFound, "Unable to find Virtual Machine [%s]", vmName)
-		}
-
-		vm := (*vms)[0]
-
-		for _, disk := range *vm.StorageProfile.DataDisks {
-			if *disk.Vhd.URI == diskName {
-				return errors.Wrapf(errors.AlreadyExists, "DataDisk [%s] is already attached to the VM [%s]", diskName, vmName)
-			}
-		}
-
-		*vm.StorageProfile.DataDisks = append(*vm.StorageProfile.DataDisks, compute.DataDisk{Vhd: &compute.VirtualHardDisk{URI: &diskName}})
-
-		_, err = c.CreateOrUpdate(ctx, group, vmName, &vm)
-		if err != nil {
-			if errors.IsInvalidVersion(err) {
-				// Retry only on invalid version
-				time.Sleep(100 * time.Millisecond)
-				continue
-			}
-			return err
-		}
-		break
-	}
-	return
+	dataDisks := []compute.DataDisk{{Vhd: &compute.VirtualHardDisk{URI: &diskName}}}
+	_, err = c.UpdateDisks(ctx, group, vmName, dataDisks, compute.VirtualMachineDiskOperationAttach)
+	return err
 }
-func (c *VirtualMachineClient) DiskDetach(ctx context.Context, group string, vmName, diskName string) (err error) {
-	for {
-		vms, err := c.Get(ctx, group, vmName)
-		if err != nil {
-			return err
-		}
-		if vms == nil || len(*vms) == 0 {
-			return errors.Wrapf(errors.NotFound, "Unable to find Virtual Machine [%s]", vmName)
-		}
 
-		vm := (*vms)[0]
-
-		for i, element := range *vm.StorageProfile.DataDisks {
-			if *element.Vhd.URI == diskName {
-				*vm.StorageProfile.DataDisks = append((*vm.StorageProfile.DataDisks)[:i], (*vm.StorageProfile.DataDisks)[i+1:]...)
-				break
-			}
-		}
-
-		_, err = c.CreateOrUpdate(ctx, group, vmName, &vm)
-		if err != nil {
-			if errors.IsInvalidVersion(err) {
-				log.Printf("Retrying because of stale version\n")
-				// Retry only on invalid version
-				time.Sleep(100 * time.Millisecond)
-				continue
-			}
-			return err
-		}
-		break
-	}
-	return
+// DiskDetach detaches diskName from vmName. Like DiskAttach it delegates to the
+// UpdateDisks RPC (Update, never Create). Detaching an absent disk is an
+// idempotent no-op.
+func (c *VirtualMachineClient) DiskDetach(ctx context.Context, group string, vmName, diskName string) (err error) {
+	dataDisks := []compute.DataDisk{{Vhd: &compute.VirtualHardDisk{URI: &diskName}}}
+	_, err = c.UpdateDisks(ctx, group, vmName, dataDisks, compute.VirtualMachineDiskOperationDetach)
+	return err
 }
 
 func (c *VirtualMachineClient) NetworkInterfaceAdd(ctx context.Context, group string, vmName, nicName string) (err error) {
@@ -350,6 +303,15 @@ func (c *VirtualMachineClient) RunCommand(ctx context.Context, group, vmName str
 	return c.internal.RunCommand(ctx, group, vmName, request)
 }
 
+// UpdateDisks attaches or detaches the supplied data disks on an existing
+// virtual machine via the agent's UpdateDisks RPC. Unlike DiskAttach/DiskDetach,
+// it does not read-modify-write the full VM through CreateOrUpdate, so it will
+// not recreate a VM that was deleted concurrently: a missing VM returns NotFound.
+// Only the disks in dataDisks are acted on; every other disk is preserved.
+func (c *VirtualMachineClient) UpdateDisks(ctx context.Context, group, vmName string, dataDisks []compute.DataDisk, operation compute.VirtualMachineDiskOperation) (*compute.VirtualMachine, error) {
+	return c.internal.UpdateDisks(ctx, group, vmName, dataDisks, operation)
+}
+
 func (c *VirtualMachineClient) RepairGuestAgent(ctx context.Context, group, vmName string) (err error) {
 	return c.internal.RepairGuestAgent(ctx, group, vmName)
 }

services/compute/virtualmachine/wssd.go

@@ -217,6 +217,53 @@ func (c *client) RunCommand(ctx context.Context, group, name string, request *co
 	return
 }
 
+// UpdateDisks attaches or detaches the given data disks on an existing virtual
+// machine via the agent's UpdateDisks RPC. It never recreates the VM: a missing
+// VM results in NotFound from the agent.
+func (c *client) UpdateDisks(ctx context.Context, group, name string, dataDisks []compute.DataDisk, operation compute.VirtualMachineDiskOperation) (*compute.VirtualMachine, error) {
+	request, err := c.getVirtualMachineDiskRequest(group, name, dataDisks, operation)
+	if err != nil {
+		return nil, err
+	}
+
+	response, err := c.VirtualMachineAgentClient.UpdateDisks(ctx, request)
+	if err != nil {
+		return nil, err
+	}
+	return c.getVirtualMachineFromDiskResponse(response, group)
+}
+
+func (c *client) getVirtualMachineDiskRequest(group, name string, dataDisks []compute.DataDisk, operation compute.VirtualMachineDiskOperation) (*wssdcloudcompute.VirtualMachineDiskRequest, error) {
+	request := &wssdcloudcompute.VirtualMachineDiskRequest{
+		GroupName: group,
+		VmName:    name,
+		Operation: wssdcloudcompute.VirtualMachineDiskOperation(operation),
+		Disks:     []*wssdcloudcompute.Disk{},
+	}
+	for i := range dataDisks {
+		disk, err := c.getWssdVirtualMachineStorageConfigurationDataDisk(&dataDisks[i])
+		if err != nil {
+			return nil, err
+		}
+		request.Disks = append(request.Disks, disk)
+	}
+	return request, nil
+}
+
+func (c *client) getVirtualMachineFromDiskResponse(response *wssdcloudcompute.VirtualMachineDiskResponse, group string) (*compute.VirtualMachine, error) {
+	if response == nil {
+		return nil, errors.Wrapf(errors.Failed, "UpdateDisks returned an empty response")
+	}
+	if response.GetResult() != nil && !response.GetResult().GetValue() {
+		return nil, errors.New(response.GetError())
+	}
+	vm := response.GetVirtualMachine()
+	if vm == nil {
+		return nil, nil
+	}
+	return c.getVirtualMachine(vm), nil
+}
+
 // Get
 func (c *client) Validate(ctx context.Context, group, name string) error {
 	request, err := c.getVirtualMachineRequest(wssdcloudproto.Operation_VALIDATE, group, name, nil)