Add support for scaling kubernetes stateful sets

Author: kormide

README.md

@@ -5,7 +5,8 @@ combination with [Buildbarn Remote Execution](https://github.com/buildbarn/bb-re
 to automatically adjust the size of Amazon EC2
 [Auto Scaling Groups (ASGs)](https://docs.aws.amazon.com/autoscaling/ec2/userguide/AutoScalingGroup.html),
 EKS [Managed Node Groups](https://docs.aws.amazon.com/eks/latest/userguide/managed-node-groups.html),
-or [Kubernetes deployments](https://kubernetes.io/docs/concepts/workloads/controllers/deployment/)
+or Kubernetes [deployments](https://kubernetes.io/docs/concepts/workloads/controllers/deployment/)/
+[stateful sets](https://kubernetes.io/docs/concepts/workloads/controllers/statefulset/)
 of Buildbarn workers based on load. Load metrics are obtained by
 querying Prometheus, which in its turn extracts metrics from
 `bb_scheduler`. It relies on Prometheus to normalize the load metrics

cmd/bb_autoscaler/main.go

@@ -104,7 +104,7 @@ func main() {
 					autoScalingClient = autoscaling.NewFromConfig(cfg)
 					eksClient = eks.NewFromConfig(cfg)
 				}
-			case *bb_autoscaler.NodeGroupConfiguration_KubernetesDeployment:
+			case *bb_autoscaler.NodeGroupConfiguration_KubernetesDeployment, *bb_autoscaler.NodeGroupConfiguration_KubernetesStatefulSet:
 				if kubernetesClientset == nil {
 					config, err := rest.InClusterConfig()
 					if err != nil {
@@ -116,7 +116,7 @@ func main() {
 					}
 				}
 			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")
 			}
 		}
 
@@ -174,6 +174,9 @@ func main() {
 				case *bb_autoscaler.NodeGroupConfiguration_KubernetesDeployment:
 					minSize = kind.KubernetesDeployment.MinimumReplicas
 					maxSize = kind.KubernetesDeployment.MaximumReplicas
+				case *bb_autoscaler.NodeGroupConfiguration_KubernetesStatefulSet:
+					minSize = kind.KubernetesStatefulSet.MinimumReplicas
+					maxSize = kind.KubernetesStatefulSet.MaximumReplicas
 				default:
 					panic("Incomplete switch on node group kind")
 				}
@@ -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.pb.go

@@ -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.