Use resize API to scale running clusters (#1541)

shreyas-goenka · web-flow · commit d7faf789e709 · 2022-08-24T22:47:02.000+02:00
diff --git a/clusters/acceptance/clusters_api_test.go b/clusters/acceptance/clusters_api_test.go
@@ -64,3 +64,64 @@ func TestAccListClustersIntegration(t *testing.T) {
 	assert.NoError(t, err, err)
 	assert.True(t, clusterReadInfo.State == clusters.ClusterStateRunning)
 }
+
+func TestAccListClustersResizeIntegrationTest(t *testing.T) {
+	qa.RequireAnyCloudEnv(t)
+	t.Parallel()
+	client := common.CommonEnvironmentClient()
+	ctx := context.Background()
+	clustersAPI := clusters.NewClustersAPI(ctx, client)
+	randomName := qa.RandomName()
+
+	cluster := clusters.Cluster{
+		NumWorkers:  1,
+		ClusterName: "Terraform Integration Test " + randomName,
+		SparkVersion: clustersAPI.LatestSparkVersionOrDefault(
+			clusters.SparkVersionRequest{
+				Latest:          true,
+				LongTermSupport: true,
+			}),
+		InstancePoolID:         compute.CommonInstancePoolID(),
+		IdempotencyToken:       "acc-list-" + randomName,
+		AutoterminationMinutes: 15,
+	}
+
+	clusterReadInfo, err := clustersAPI.Create(cluster)
+	require.NoError(t, err, err)
+	assert.True(t, clusterReadInfo.NumWorkers == cluster.NumWorkers)
+	assert.True(t, clusterReadInfo.ClusterName == cluster.ClusterName)
+	assert.True(t, reflect.DeepEqual(clusterReadInfo.SparkEnvVars, cluster.SparkEnvVars))
+	assert.True(t, clusterReadInfo.SparkVersion == cluster.SparkVersion)
+	assert.True(t, clusterReadInfo.AutoterminationMinutes == cluster.AutoterminationMinutes)
+	assert.True(t, clusterReadInfo.State == clusters.ClusterStateRunning)
+
+	// Resize num workers
+	clusterReadInfo, err = clustersAPI.Resize(clusters.ResizeRequest{ClusterID: clusterReadInfo.ClusterID, NumWorkers: 2})
+	require.NoError(t, err, err)
+	assert.True(t, clusterReadInfo.NumWorkers == 2)
+	assert.True(t, clusterReadInfo.State == clusters.ClusterStateRunning)
+
+	// Resize cluster to become autoscaling
+	clusterReadInfo, err = clustersAPI.Resize(clusters.ResizeRequest{ClusterID: clusterReadInfo.ClusterID, AutoScale: &clusters.AutoScale{
+		MinWorkers: 1,
+		MaxWorkers: 2,
+	}})
+	require.NoError(t, err, err)
+	assert.True(t, clusterReadInfo.AutoScale.MinWorkers == 1)
+	assert.True(t, clusterReadInfo.AutoScale.MaxWorkers == 2)
+	assert.True(t, clusterReadInfo.State == clusters.ClusterStateRunning)
+
+	// cleanup
+	err = clustersAPI.Terminate(clusterReadInfo.ClusterID)
+	assert.NoError(t, err, err)
+
+	clusterReadInfo, err = clustersAPI.Get(clusterReadInfo.ClusterID)
+	assert.NoError(t, err, err)
+	assert.True(t, clusterReadInfo.State == clusters.ClusterStateTerminated)
+
+	err = clustersAPI.Unpin(clusterReadInfo.ClusterID)
+	assert.NoError(t, err, err)
+
+	err = clustersAPI.PermanentDelete(clusterReadInfo.ClusterID)
+	assert.NoError(t, err, err)
+}
diff --git a/clusters/clusters_api.go b/clusters/clusters_api.go
@@ -311,6 +311,12 @@ type ClusterSize struct {
 	AutoScale  *AutoScale `json:"autoscale"`
 }
 
+type ResizeRequest struct {
+	ClusterID  string     `json:"cluster_id"`
+	NumWorkers int32      `json:"num_workers"`
+	AutoScale  *AutoScale `json:"autoscale,omitempty"`
+}
+
 // ResizeCause holds reason for resizing
 type ResizeCause string
 
@@ -533,6 +539,24 @@ func (a ClustersAPI) Create(cluster Cluster) (info ClusterInfo, err error) {
 	return
 }
 
+// Resize api can only be used when the cluster is in Running State
+func (a ClustersAPI) Resize(resizeRequest ResizeRequest) (info ClusterInfo, err error) {
+	info, err = a.Get(resizeRequest.ClusterID)
+	if err != nil {
+		return info, err
+	}
+	if info.State != ClusterStateRunning {
+		return info, fmt.Errorf("resize: Cluster %v is in %v state. RUNNING state required to use resize API", info.ClusterID, info.State)
+	}
+
+	err = a.client.Post(a.context, "/clusters/resize", resizeRequest, &info)
+	if err != nil {
+		return info, fmt.Errorf("resize: %w", err)
+	}
+	info, err = a.waitForClusterStatus(resizeRequest.ClusterID, ClusterStateRunning)
+	return info, err
+}
+
 // Edit edits the configuration of a cluster to match the provided attributes and size
 func (a ClustersAPI) Edit(cluster Cluster) (info ClusterInfo, err error) {
 	info, err = a.Get(cluster.ClusterID)
diff --git a/clusters/clusters_api_test.go b/clusters/clusters_api_test.go
@@ -313,6 +313,92 @@ func TestEditCluster_Pending(t *testing.T) {
 	assert.Equal(t, ClusterStateRunning, string(clusterInfo.State))
 }
 
+func TestResizeCluster_FailsForNonRunningCluster(t *testing.T) {
+	clusterStates := []ClusterState{ClusterStateUnknown,
+		ClusterStateError,
+		ClusterStatePending,
+		ClusterStateRestarting,
+		ClusterStateResizing,
+		ClusterStateTerminating,
+		ClusterStateTerminated,
+	}
+	for _, clusterState := range clusterStates {
+		t.Run(fmt.Sprintf("CLUSTER STATE %s", clusterState), func(t *testing.T) {
+			client, server, err := qa.HttpFixtureClient(t, []qa.HTTPFixture{
+				{
+					Method:   "GET",
+					Resource: "/api/2.0/clusters/get?cluster_id=abc",
+					Response: ClusterInfo{
+						State:     clusterState,
+						ClusterID: "abc",
+					},
+				},
+			})
+			require.NoError(t, err)
+
+			ctx := context.Background()
+			_, err = NewClustersAPI(ctx, client).Resize(ResizeRequest{
+				ClusterID:  "abc",
+				NumWorkers: 10,
+			})
+			require.Error(t, err)
+			assert.Contains(t, err.Error(), "resize: Cluster abc is in "+clusterState+" state. RUNNING state required to use resize API")
+			server.Close()
+		})
+	}
+}
+
+func TestResizeCluster_NormalRun(t *testing.T) {
+	client, server, err := qa.HttpFixtureClient(t, []qa.HTTPFixture{
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/clusters/get?cluster_id=abc",
+			Response: ClusterInfo{
+				State:      ClusterStateRunning,
+				ClusterID:  "abc",
+				NumWorkers: 4,
+			},
+		},
+		{
+			Method:   "POST",
+			Resource: "/api/2.0/clusters/resize",
+			ExpectedRequest: ResizeRequest{
+				ClusterID:  "abc",
+				NumWorkers: 10,
+			},
+		},
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/clusters/get?cluster_id=abc",
+			Response: ClusterInfo{
+				State:      ClusterStateResizing,
+				ClusterID:  "abc",
+				NumWorkers: 10,
+			},
+		},
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/clusters/get?cluster_id=abc",
+			Response: ClusterInfo{
+				State:      ClusterStateRunning,
+				ClusterID:  "abc",
+				NumWorkers: 10,
+			},
+		},
+	})
+	defer server.Close()
+	require.NoError(t, err)
+
+	ctx := context.Background()
+	clusterInfo, err := NewClustersAPI(ctx, client).Resize(ResizeRequest{
+		ClusterID:  "abc",
+		NumWorkers: 10,
+	})
+	require.NoError(t, err)
+	assert.Equal(t, ClusterStateRunning, string(clusterInfo.State))
+	assert.Equal(t, 10, int(clusterInfo.NumWorkers))
+}
+
 func TestEditCluster_Terminating(t *testing.T) {
 	client, server, err := qa.HttpFixtureClient(t, []qa.HTTPFixture{
 		{
diff --git a/clusters/resource_cluster.go b/clusters/resource_cluster.go
@@ -241,17 +241,76 @@ func resourceClusterUpdate(ctx context.Context, d *schema.ResourceData, c *commo
 	common.DataToStructPointer(d, clusterSchema, &cluster)
 	var clusterInfo ClusterInfo
 	var err error
+
 	if hasClusterConfigChanged(d) {
 		log.Printf("[DEBUG] Cluster state has changed!")
 		if err := cluster.Validate(); err != nil {
 			return err
 		}
 		cluster.ModifyRequestOnInstancePool()
 		fixInstancePoolChangeIfAny(d, &cluster)
-		clusterInfo, err = clusters.Edit(cluster)
+
+		// We can only call the resize api if the cluster is in the running state
+		// and only the cluster size (ie num_workers OR autoscale) is being changed
+		hasNumWorkersChanged := d.HasChange("num_workers")
+		hasAutoscaleChanged := d.HasChange("autoscale")
+		hasOnlyResizeClusterConfigChanged := true
+		for k := range clusterSchema {
+			if k == "library" ||
+				k == "is_pinned" ||
+				k == "num_workers" ||
+				k == "autoscale" {
+				continue
+			}
+			if d.HasChange(k) {
+				hasOnlyResizeClusterConfigChanged = false
+			}
+		}
+		clusterInfo, err = clusters.Get(clusterID)
+		if err != nil {
+			return err
+		}
+
+		isNumWorkersResizeForNonAutoscalingCluster := hasOnlyResizeClusterConfigChanged &&
+			hasNumWorkersChanged &&
+			!hasAutoscaleChanged &&
+			clusterInfo.State == ClusterStateRunning
+		isAutoScalingToNonAutoscalingResize := hasOnlyResizeClusterConfigChanged &&
+			hasAutoscaleChanged &&
+			hasNumWorkersChanged &&
+			cluster.Autoscale == nil &&
+			clusterInfo.State == ClusterStateRunning
+		isAutoscaleConfigResizeForAutoscalingCluster := hasOnlyResizeClusterConfigChanged &&
+			hasAutoscaleChanged &&
+			!hasNumWorkersChanged &&
+			clusterInfo.State == ClusterStateRunning
+		isNonAutoScalingToAutoscalingResize := hasOnlyResizeClusterConfigChanged &&
+			hasAutoscaleChanged &&
+			hasNumWorkersChanged &&
+			cluster.Autoscale != nil &&
+			clusterInfo.State == ClusterStateRunning
+
+		// We prefer to use the resize API in cases when only the number of
+		// workers is changed because a resizing cluster can still serve queries
+		if isNumWorkersResizeForNonAutoscalingCluster ||
+			isAutoScalingToNonAutoscalingResize {
+			clusterInfo, err = clusters.Resize(ResizeRequest{
+				ClusterID:  clusterID,
+				NumWorkers: cluster.NumWorkers,
+			})
+		} else if isAutoscaleConfigResizeForAutoscalingCluster ||
+			isNonAutoScalingToAutoscalingResize {
+			clusterInfo, err = clusters.Resize(ResizeRequest{
+				ClusterID: clusterID,
+				AutoScale: cluster.Autoscale,
+			})
+		} else {
+			clusterInfo, err = clusters.Edit(cluster)
+		}
 		if err != nil {
 			return err
 		}
+
 	} else {
 		clusterInfo, err = clusters.Get(clusterID)
 		if err != nil {
diff --git a/clusters/resource_cluster_test.go b/clusters/resource_cluster_test.go
