chore: add diskclient.Update interface

Author: andyzhangx

staging/src/k8s.io/legacy-cloud-providers/azure/clients/diskclient/azure_diskclient.go

@@ -206,6 +206,74 @@ func (c *Client) createOrUpdateResponder(resp *http.Response) (*compute.Disk, *r
 	return result, retry.GetError(resp, err)
 }
 
+// Update creates or updates a Disk.
+func (c *Client) Update(ctx context.Context, resourceGroupName string, diskName string, diskParameter compute.DiskUpdate) *retry.Error {
+	mc := metrics.NewMetricContext("disks", "update", resourceGroupName, c.subscriptionID, "")
+
+	// Report errors if the client is rate limited.
+	if !c.rateLimiterWriter.TryAccept() {
+		mc.RateLimitedCount()
+		return retry.GetRateLimitError(true, "DiskUpdate")
+	}
+
+	// Report errors if the client is throttled.
+	if c.RetryAfterWriter.After(time.Now()) {
+		mc.ThrottledCount()
+		rerr := retry.GetThrottlingError("DiskUpdate", "client throttled", c.RetryAfterWriter)
+		return rerr
+	}
+
+	rerr := c.updateDisk(ctx, resourceGroupName, diskName, diskParameter)
+	mc.Observe(rerr.Error())
+	if rerr != nil {
+		if rerr.IsThrottled() {
+			// Update RetryAfterReader so that no more requests would be sent until RetryAfter expires.
+			c.RetryAfterWriter = rerr.RetryAfter
+		}
+
+		return rerr
+	}
+
+	return nil
+}
+
+// updateDisk updates a Disk.
+func (c *Client) updateDisk(ctx context.Context, resourceGroupName string, diskName string, diskParameter compute.DiskUpdate) *retry.Error {
+	resourceID := armclient.GetResourceID(
+		c.subscriptionID,
+		resourceGroupName,
+		"Microsoft.Compute/disks",
+		diskName,
+	)
+
+	response, rerr := c.armClient.PatchResource(ctx, resourceID, diskParameter)
+	defer c.armClient.CloseResponse(ctx, response)
+	if rerr != nil {
+		klog.V(5).Infof("Received error in %s: resourceID: %s, error: %s", "disk.put.request", resourceID, rerr.Error())
+		return rerr
+	}
+
+	if response != nil && response.StatusCode != http.StatusNoContent {
+		_, rerr = c.updateResponder(response)
+		if rerr != nil {
+			klog.V(5).Infof("Received error in %s: resourceID: %s, error: %s", "disk.put.respond", resourceID, rerr.Error())
+			return rerr
+		}
+	}
+
+	return nil
+}
+
+func (c *Client) updateResponder(resp *http.Response) (*compute.Disk, *retry.Error) {
+	result := &compute.Disk{}
+	err := autorest.Respond(
+		resp,
+		azure.WithErrorUnlessStatusCode(http.StatusOK, http.StatusCreated),
+		autorest.ByUnmarshallingJSON(&result))
+	result.Response = autorest.Response{Response: resp}
+	return result, retry.GetError(resp, err)
+}
+
 // Delete deletes a Disk by name.
 func (c *Client) Delete(ctx context.Context, resourceGroupName string, diskName string) *retry.Error {
 	mc := metrics.NewMetricContext("disks", "delete", resourceGroupName, c.subscriptionID, "")

staging/src/k8s.io/legacy-cloud-providers/azure/clients/diskclient/azure_diskclient_test.go

@@ -161,6 +161,49 @@ func TestCreateOrUpdate(t *testing.T) {
 	assert.Equal(t, throttleErr, rerr)
 }
 
+func TestUpdate(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	resourceID := "/subscriptions/subscriptionID/resourceGroups/rg/providers/Microsoft.Compute/disks/disk1"
+	diskUpdate := getTestDiskUpdate()
+	armClient := mockarmclient.NewMockInterface(ctrl)
+	response := &http.Response{
+		StatusCode: http.StatusOK,
+		Body:       ioutil.NopCloser(bytes.NewReader([]byte(""))),
+	}
+	armClient.EXPECT().PatchResource(gomock.Any(), resourceID, diskUpdate).Return(response, nil).Times(1)
+	armClient.EXPECT().CloseResponse(gomock.Any(), gomock.Any()).Times(1)
+
+	diskClient := getTestDiskClient(armClient)
+	rerr := diskClient.Update(context.TODO(), "rg", "disk1", diskUpdate)
+	assert.Nil(t, rerr)
+
+	response = &http.Response{
+		StatusCode: http.StatusTooManyRequests,
+		Body:       ioutil.NopCloser(bytes.NewReader([]byte("{}"))),
+	}
+	throttleErr := &retry.Error{
+		HTTPStatusCode: http.StatusTooManyRequests,
+		RawError:       fmt.Errorf("error"),
+		Retriable:      true,
+		RetryAfter:     time.Unix(100, 0),
+	}
+
+	armClient.EXPECT().PatchResource(gomock.Any(), resourceID, diskUpdate).Return(response, throttleErr).Times(1)
+	armClient.EXPECT().CloseResponse(gomock.Any(), gomock.Any()).Times(1)
+	rerr = diskClient.Update(context.TODO(), "rg", "disk1", diskUpdate)
+	assert.Equal(t, throttleErr, rerr)
+}
+
+func getTestDiskUpdate() compute.DiskUpdate {
+	return compute.DiskUpdate{
+		DiskUpdateProperties: &compute.DiskUpdateProperties{
+			DiskSizeGB: to.Int32Ptr(100),
+		},
+	}
+}
+
 func TestDelete(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()

staging/src/k8s.io/legacy-cloud-providers/azure/clients/diskclient/interface.go

@@ -40,6 +40,9 @@ type Interface interface {
 	// CreateOrUpdate creates or updates a Disk.
 	CreateOrUpdate(ctx context.Context, resourceGroupName string, diskName string, diskParameter compute.Disk) *retry.Error
 
+	// Update updates a Disk.
+	Update(ctx context.Context, resourceGroupName string, diskName string, diskParameter compute.DiskUpdate) *retry.Error
+
 	// Delete deletes a Disk by name.
 	Delete(ctx context.Context, resourceGroupName string, diskName string) *retry.Error
 

staging/src/k8s.io/legacy-cloud-providers/azure/clients/diskclient/mockdiskclient/interface.go

@@ -79,6 +79,20 @@ func (mr *MockInterfaceMockRecorder) CreateOrUpdate(ctx, resourceGroupName, disk
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateOrUpdate", reflect.TypeOf((*MockInterface)(nil).CreateOrUpdate), ctx, resourceGroupName, diskName, diskParameter)
 }
 
+// Update mocks base method
+func (m *MockInterface) Update(ctx context.Context, resourceGroupName, diskName string, diskParameter compute.DiskUpdate) *retry.Error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Update", ctx, resourceGroupName, diskName, diskParameter)
+	ret0, _ := ret[0].(*retry.Error)
+	return ret0
+}
+
+// Update indicates an expected call of Update
+func (mr *MockInterfaceMockRecorder) Update(ctx, resourceGroupName, diskName, diskParameter interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Update", reflect.TypeOf((*MockInterface)(nil).Update), ctx, resourceGroupName, diskName, diskParameter)
+}
+
 // Delete mocks base method
 func (m *MockInterface) Delete(ctx context.Context, resourceGroupName, diskName string) *retry.Error {
 	m.ctrl.T.Helper()