Merge pull request #95578 from marwanad/add-deallocation-methods

add methods to deallocate and starts vms in a scale set

staging/src/k8s.io/legacy-cloud-providers/azure/clients/vmssclient/azure_vmssclient.go

@@ -272,6 +272,38 @@ func (c *Client) CreateOrUpdateAsync(ctx context.Context, resourceGroupName stri
 	return future, nil
 }
 
+// WaitForCreateOrUpdateResult waits for the response of the create or update request
+func (c *Client) WaitForCreateOrUpdateResult(ctx context.Context, future *azure.Future, resourceGroupName string) (*http.Response, error) {
+	mc := metrics.NewMetricContext("vmss", "wait_for_create_or_update_result", resourceGroupName, c.subscriptionID, "")
+	res, err := c.armClient.WaitForAsyncOperationResult(ctx, future, "VMSSWaitForCreateOrUpdateResult")
+	mc.Observe(err)
+	return res, err
+}
+
+// WaitForDeleteInstancesResult waits for the response of the delete instance request
+func (c *Client) WaitForDeleteInstancesResult(ctx context.Context, future *azure.Future, resourceGroupName string) (*http.Response, error) {
+	mc := metrics.NewMetricContext("vmss", "wait_for_delete_instances_result", resourceGroupName, c.subscriptionID, "")
+	res, err := c.armClient.WaitForAsyncOperationResult(ctx, future, "VMSSWaitForDeleteInstancesResult")
+	mc.Observe(err)
+	return res, err
+}
+
+// WaitForDeallocateInstancesResult waits for the response of the delete instance request
+func (c *Client) WaitForDeallocateInstancesResult(ctx context.Context, future *azure.Future, resourceGroupName string) (*http.Response, error) {
+	mc := metrics.NewMetricContext("vmss", "wait_for_deallocate_instances_result", resourceGroupName, c.subscriptionID, "")
+	res, err := c.armClient.WaitForAsyncOperationResult(ctx, future, "VMSSWaitForDeallocateInstancesResult")
+	mc.Observe(err)
+	return res, err
+}
+
+// WaitForStartInstancesResult waits for the response of the delete instance request
+func (c *Client) WaitForStartInstancesResult(ctx context.Context, future *azure.Future, resourceGroupName string) (*http.Response, error) {
+	mc := metrics.NewMetricContext("vmss", "wait_for_start_instances_result", resourceGroupName, c.subscriptionID, "")
+	res, err := c.armClient.WaitForAsyncOperationResult(ctx, future, "VMSSWaitForStartInstancesResult")
+	mc.Observe(err)
+	return res, err
+}
+
 // WaitForAsyncOperationResult waits for the response of the request
 func (c *Client) WaitForAsyncOperationResult(ctx context.Context, future *azure.Future) (*http.Response, error) {
 	return c.armClient.WaitForAsyncOperationResult(ctx, future, "VMSSWaitForAsyncOperationResult")
@@ -473,6 +505,7 @@ func (c *Client) DeleteInstancesAsync(ctx context.Context, resourceGroupName str
 	}
 
 	future, err := azure.NewFutureFromResponse(response)
+	mc.Observe(err)
 	if err != nil {
 		klog.V(5).Infof("Received error in %s: resourceID: %s, error: %s", "vmss.deletevms.future", resourceID, rerr.Error())
 		return nil, retry.NewError(false, err)
@@ -481,6 +514,102 @@ func (c *Client) DeleteInstancesAsync(ctx context.Context, resourceGroupName str
 	return &future, nil
 }
 
+// DeallocateInstancesAsync sends the deallocate request to ARM client and DOEST NOT wait on the future
+func (c *Client) DeallocateInstancesAsync(ctx context.Context, resourceGroupName string, vmScaleSetName string, vmInstanceIDs compute.VirtualMachineScaleSetVMInstanceRequiredIDs) (*azure.Future, *retry.Error) {
+	mc := metrics.NewMetricContext("vmss", "deallocate_instances_async", resourceGroupName, c.subscriptionID, "")
+
+	// Report errors if the client is rate limited.
+	if !c.rateLimiterWriter.TryAccept() {
+		mc.RateLimitedCount()
+		return nil, retry.GetRateLimitError(true, "VMSSDeallocateInstancesAsync")
+	}
+
+	// Report errors if the client is throttled.
+	if c.RetryAfterWriter.After(time.Now()) {
+		mc.ThrottledCount()
+		rerr := retry.GetThrottlingError("VMSSDeallocateInstancesAsync", "client throttled", c.RetryAfterWriter)
+		return nil, rerr
+	}
+
+	resourceID := armclient.GetResourceID(
+		c.subscriptionID,
+		resourceGroupName,
+		"Microsoft.Compute/virtualMachineScaleSets",
+		vmScaleSetName,
+	)
+
+	response, rerr := c.armClient.PostResource(ctx, resourceID, "deallocate", vmInstanceIDs)
+	defer c.armClient.CloseResponse(ctx, response)
+
+	if rerr != nil {
+		klog.V(5).Infof("Received error in %s: resourceID: %s, error: %s", "vmss.deallocatevms.request", resourceID, rerr.Error())
+		return nil, rerr
+	}
+
+	err := autorest.Respond(response, azure.WithErrorUnlessStatusCode(http.StatusOK, http.StatusAccepted))
+	if err != nil {
+		klog.V(5).Infof("Received error in %s: resourceID: %s, error: %s", "vmss.deallocatevms.respond", resourceID, rerr.Error())
+		return nil, retry.GetError(response, err)
+	}
+
+	future, err := azure.NewFutureFromResponse(response)
+	mc.Observe(err)
+	if err != nil {
+		klog.V(5).Infof("Received error in %s: resourceID: %s, error: %s", "vmss.deallocatevms.future", resourceID, rerr.Error())
+		return nil, retry.NewError(false, err)
+	}
+
+	return &future, nil
+}
+
+// StartInstancesAsync sends the start request to ARM client and DOEST NOT wait on the future
+func (c *Client) StartInstancesAsync(ctx context.Context, resourceGroupName string, vmScaleSetName string, vmInstanceIDs compute.VirtualMachineScaleSetVMInstanceRequiredIDs) (*azure.Future, *retry.Error) {
+	mc := metrics.NewMetricContext("vmss", "start_instances_async", resourceGroupName, c.subscriptionID, "")
+
+	// Report errors if the client is rate limited.
+	if !c.rateLimiterWriter.TryAccept() {
+		mc.RateLimitedCount()
+		return nil, retry.GetRateLimitError(true, "VMSSStartInstancesAsync")
+	}
+
+	// Report errors if the client is throttled.
+	if c.RetryAfterWriter.After(time.Now()) {
+		mc.ThrottledCount()
+		rerr := retry.GetThrottlingError("VMSSStartInstancesAsync", "client throttled", c.RetryAfterWriter)
+		return nil, rerr
+	}
+
+	resourceID := armclient.GetResourceID(
+		c.subscriptionID,
+		resourceGroupName,
+		"Microsoft.Compute/virtualMachineScaleSets",
+		vmScaleSetName,
+	)
+
+	response, rerr := c.armClient.PostResource(ctx, resourceID, "start", vmInstanceIDs)
+	defer c.armClient.CloseResponse(ctx, response)
+
+	if rerr != nil {
+		klog.V(5).Infof("Received error in %s: resourceID: %s, error: %s", "vmss.startvms.request", resourceID, rerr.Error())
+		return nil, rerr
+	}
+
+	err := autorest.Respond(response, azure.WithErrorUnlessStatusCode(http.StatusOK, http.StatusAccepted))
+	if err != nil {
+		klog.V(5).Infof("Received error in %s: resourceID: %s, error: %s", "vmss.startvms.respond", resourceID, rerr.Error())
+		return nil, retry.GetError(response, err)
+	}
+
+	future, err := azure.NewFutureFromResponse(response)
+	mc.Observe(err)
+	if err != nil {
+		klog.V(5).Infof("Received error in %s: resourceID: %s, error: %s", "vmss.startvms.future", resourceID, rerr.Error())
+		return nil, retry.NewError(false, err)
+	}
+
+	return &future, nil
+}
+
 // deleteVMSSInstances deletes the instances for a VirtualMachineScaleSet.
 func (c *Client) deleteVMSSInstances(ctx context.Context, resourceGroupName string, vmScaleSetName string, vmInstanceIDs compute.VirtualMachineScaleSetVMInstanceRequiredIDs) *retry.Error {
 	resourceID := armclient.GetResourceID(

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

@@ -57,4 +57,22 @@ type Interface interface {
 
 	// DeleteInstancesAsync sends the delete request to the ARM client and DOEST NOT wait on the future
 	DeleteInstancesAsync(ctx context.Context, resourceGroupName string, vmScaleSetName string, vmInstanceIDs compute.VirtualMachineScaleSetVMInstanceRequiredIDs) (*azure.Future, *retry.Error)
+
+	// WaitForCreateOrUpdateResult waits for the response of the create or update request
+	WaitForCreateOrUpdateResult(ctx context.Context, future *azure.Future, resourceGroupName string) (*http.Response, error)
+
+	// WaitForDeleteInstancesResult waits for the response of the delete instances request
+	WaitForDeleteInstancesResult(ctx context.Context, future *azure.Future, resourceGroupName string) (*http.Response, error)
+
+	// DeallocateInstances sends the deallocate request to the ARM client and DOEST NOT wait on the future
+	DeallocateInstancesAsync(ctx context.Context, resourceGroupName string, vmScaleSetName string, vmInstanceIDs compute.VirtualMachineScaleSetVMInstanceRequiredIDs) (*azure.Future, *retry.Error)
+
+	// WaitForDeallocateInstancesResult waits for the response of the deallocate instances request
+	WaitForDeallocateInstancesResult(ctx context.Context, future *azure.Future, resourceGroupName string) (*http.Response, error)
+
+	// StartInstancesAsync starts the instances for a VirtualMachineScaleSet.
+	StartInstancesAsync(ctx context.Context, resourceGroupName string, vmScaleSetName string, vmInstanceIDs compute.VirtualMachineScaleSetVMInstanceRequiredIDs) (*azure.Future, *retry.Error)
+
+	// WaitForStartInstancesResult waits for the response of the start instances request
+	WaitForStartInstancesResult(ctx context.Context, future *azure.Future, resourceGroupName string) (*http.Response, error)
 }

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

@@ -126,6 +126,66 @@ func (mr *MockInterfaceMockRecorder) WaitForAsyncOperationResult(ctx, future int
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WaitForAsyncOperationResult", reflect.TypeOf((*MockInterface)(nil).WaitForAsyncOperationResult), ctx, future)
 }
 
+// WaitForCreateOrUpdateResult mocks base method
+func (m *MockInterface) WaitForCreateOrUpdateResult(ctx context.Context, future *azure.Future, resourceGroupName string) (*http.Response, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "WaitForCreateOrUpdateResult", ctx, future, resourceGroupName)
+	ret0, _ := ret[0].(*http.Response)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// WaitForCreateOrUpdateResult indicates an expected call of WaitForCreateOrUpdateResult
+func (mr *MockInterfaceMockRecorder) WaitForCreateOrUpdateResult(ctx, future interface{}, resourceGroupName string) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WaitForCreateOrUpdateResult", reflect.TypeOf((*MockInterface)(nil).WaitForCreateOrUpdateResult), ctx, future, resourceGroupName)
+}
+
+// WaitForDeleteInstancesResult mocks base method
+func (m *MockInterface) WaitForDeleteInstancesResult(ctx context.Context, future *azure.Future, resourceGroupName string) (*http.Response, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "WaitForDeleteInstancesResult", ctx, future, resourceGroupName)
+	ret0, _ := ret[0].(*http.Response)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// WaitForDeleteInstancesResult indicates an expected call of WaitForDeleteInstancesResult
+func (mr *MockInterfaceMockRecorder) WaitForDeleteInstancesResult(ctx, future interface{}, resourceGroupName string) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WaitForDeleteInstancesResult", reflect.TypeOf((*MockInterface)(nil).WaitForDeleteInstancesResult), ctx, future, resourceGroupName)
+}
+
+// WaitForDeallocateInstancesResult mocks base method
+func (m *MockInterface) WaitForDeallocateInstancesResult(ctx context.Context, future *azure.Future, resourceGroupName string) (*http.Response, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "WaitForDeallocateInstancesResult", ctx, future, resourceGroupName)
+	ret0, _ := ret[0].(*http.Response)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// WaitForDeallocateInstancesResult indicates an expected call of WaitForDeallocateInstancesResult
+func (mr *MockInterfaceMockRecorder) WaitForDeallocateInstancesResult(ctx, future interface{}, resourceGroupName string) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WaitForDeallocateInstancesResult", reflect.TypeOf((*MockInterface)(nil).WaitForDeallocateInstancesResult), ctx, future, resourceGroupName)
+}
+
+// WaitForStartInstancesResult mocks base method
+func (m *MockInterface) WaitForStartInstancesResult(ctx context.Context, future *azure.Future, resourceGroupName string) (*http.Response, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "WaitForStartInstancesResult", ctx, future, resourceGroupName)
+	ret0, _ := ret[0].(*http.Response)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// WaitForStartInstancesResult indicates an expected call of WaitForStartInstancesResult
+func (mr *MockInterfaceMockRecorder) WaitForStartInstancesResult(ctx, future interface{}, resourceGroupName string) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WaitForStartInstancesResult", reflect.TypeOf((*MockInterface)(nil).WaitForStartInstancesResult), ctx, future, resourceGroupName)
+}
+
 // DeleteInstances mocks base method
 func (m *MockInterface) DeleteInstances(ctx context.Context, resourceGroupName, vmScaleSetName string, vmInstanceIDs compute.VirtualMachineScaleSetVMInstanceRequiredIDs) *retry.Error {
 	m.ctrl.T.Helper()
@@ -154,3 +214,33 @@ func (mr *MockInterfaceMockRecorder) DeleteInstancesAsync(ctx, resourceGroupName
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteInstancesAsync", reflect.TypeOf((*MockInterface)(nil).DeleteInstancesAsync), ctx, resourceGroupName, VMScaleSetName, vmInstanceIDs)
 }
+
+// DeallocateInstancesAsync mocks base method
+func (m *MockInterface) DeallocateInstancesAsync(ctx context.Context, resourceGroupName, VMScaleSetName string, vmInstanceIDs compute.VirtualMachineScaleSetVMInstanceRequiredIDs) (*azure.Future, *retry.Error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "DeallocateInstancesAsync", ctx, resourceGroupName, VMScaleSetName, vmInstanceIDs)
+	ret0, _ := ret[0].(*azure.Future)
+	ret1, _ := ret[1].(*retry.Error)
+	return ret0, ret1
+}
+
+// DeallocateInstancesAsync indicates an expected call of DeleteInstancesAsync
+func (mr *MockInterfaceMockRecorder) DeallocateInstancesAsync(ctx, resourceGroupName, VMScaleSetName, vmInstanceIDs interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeallocateInstancesAsync", reflect.TypeOf((*MockInterface)(nil).DeallocateInstancesAsync), ctx, resourceGroupName, VMScaleSetName, vmInstanceIDs)
+}
+
+// StartInstancesAsync mocks base method
+func (m *MockInterface) StartInstancesAsync(ctx context.Context, resourceGroupName, VMScaleSetName string, vmInstanceIDs compute.VirtualMachineScaleSetVMInstanceRequiredIDs) (*azure.Future, *retry.Error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "StartInstancesAsync", ctx, resourceGroupName, VMScaleSetName, vmInstanceIDs)
+	ret0, _ := ret[0].(*azure.Future)
+	ret1, _ := ret[1].(*retry.Error)
+	return ret0, ret1
+}
+
+// StartInstancesAsync indicates an expected call of DeleteInstancesAsync
+func (mr *MockInterfaceMockRecorder) StartInstancesAsync(ctx, resourceGroupName, VMScaleSetName, vmInstanceIDs interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StartInstancesAsync", reflect.TypeOf((*MockInterface)(nil).StartInstancesAsync), ctx, resourceGroupName, VMScaleSetName, vmInstanceIDs)
+}