Add support for scaling kubernetes stateful sets

cmd/bb_autoscaler/main.go

@@ -106,17 +106,20 @@ func main() {
 				}
 			case *bb_autoscaler.NodeGroupConfiguration_KubernetesDeployment:
 				if kubernetesClientset == nil {
-					config, err := rest.InClusterConfig()
-					if err != nil {
-						return util.StatusWrap(err, "Failed to create Kubernetes client configuration")
+					kubernetesClientset, err = createKubernetesClient()
+					if nil != err {
+						return err
 					}
-					kubernetesClientset, err = kubernetes.NewForConfig(config)
-					if err != nil {
-						return util.StatusWrap(err, "Failed to create Kubernetes client")
+				}
+			case *bb_autoscaler.NodeGroupConfiguration_KubernetesStatefulSet:
+				if kubernetesClientset == nil {
+					kubernetesClientset, err = createKubernetesClient()
+					if nil != err {
+						return err
 					}
 				}
 			default:
-				return status.Error(codes.InvalidArgument, "No ASG, EKS managed node group, or Kubernetes deployment name specified")
+				return status.Error(codes.InvalidArgument, "No ASG, EKS managed node group, Kubernetes deployment, or Kubernetes stateful set specified")
 			}
 		}
 
@@ -250,6 +253,38 @@ func main() {
 								}); err != nil {
 							return util.StatusWrapf(err, "Failed to change number of replicas of Kubernetes deployment %#v in namespace %#v", name, namespace)
 						}
+					case *bb_autoscaler.NodeGroupConfiguration_KubernetesStatefulSet:
+						namespace := kind.KubernetesStatefulSet.Namespace
+						name := kind.KubernetesStatefulSet.Name
+						metaKind := "StatefulSet"
+						metaAPIVersion := "apps/v1"
+						if _, err := kubernetesClientset.
+							AppsV1().
+							StatefulSets(namespace).
+							Apply(
+								ctx,
+								&appsv1_apply.StatefulSetApplyConfiguration{
+									TypeMetaApplyConfiguration: metav1_apply.TypeMetaApplyConfiguration{
+										Kind:       &metaKind,
+										APIVersion: &metaAPIVersion,
+									},
+									ObjectMetaApplyConfiguration: &metav1_apply.ObjectMetaApplyConfiguration{
+										Name:      &name,
+										Namespace: &namespace,
+										Annotations: map[string]string{
+											"kubernetes.io/change-cause": "replicas updated by bb_autoscaler",
+										},
+									},
+									Spec: &appsv1_apply.StatefulSetSpecApplyConfiguration{
+										Replicas: &newDesiredCapacity,
+									},
+								},
+								metav1.ApplyOptions{
+									FieldManager: "bb_autoscaler",
+									Force:        true,
+								}); err != nil {
+							return util.StatusWrapf(err, "Failed to change number of replicas of Kubernetes stateful set %#v in namespace %#v", name, namespace)
+						}
 					default:
 						panic("Incomplete switch on node group kind")
 					}
@@ -261,3 +296,16 @@ func main() {
 		return nil
 	})
 }
+
+func createKubernetesClient() (*kubernetes.Clientset, error) {
+	config, err := rest.InClusterConfig()
+	if err != nil {
+		return nil, util.StatusWrap(err, "Failed to create Kubernetes client configuration")
+	}
+	kubernetesClientset, err := kubernetes.NewForConfig(config)
+	if err != nil {
+		return nil, util.StatusWrap(err, "Failed to create Kubernetes client")
+	}
+
+	return kubernetesClientset, nil
+}

pkg/proto/configuration/bb_autoscaler/bb_autoscaler.proto

@@ -59,7 +59,7 @@ message EKSManagedNodeGroupConfiguration {
   string node_group_name = 2;
 }
 
-message KubernetesDeploymentConfiguration {
+message KubernetesReplicaSetConfiguration {
   // Namespace containing the deployment whose replicas count should be
   // adjusted.
   string namespace = 1;
@@ -101,7 +101,10 @@ message NodeGroupConfiguration {
     EKSManagedNodeGroupConfiguration eks_managed_node_group = 5;
 
     // Kubernetes deployment whose replicas count should be adjusted.
-    KubernetesDeploymentConfiguration kubernetes_deployment = 7;
+    KubernetesReplicaSetConfiguration kubernetes_deployment = 7;
+
+    // Kubernetes stateful set whose replicas count should be adjusted.
+    KubernetesReplicaSetConfiguration kubernetes_stateful_set = 9;
   }
 
   // The number of workers that constitute to a single unit of capacity.