Merge pull request #104 from uselagoon/unidle

feat: unidle the entire namespace when accessed by ssh

internal/k8s/exec.go

@@ -1,12 +1,15 @@
+// Package k8s provides an interface to Kubernetes for common Lagoon operations.
 package k8s
 
 import (
 	"context"
 	"fmt"
 	"io"
+	"strconv"
 	"time"
 
-	v1 "k8s.io/api/core/v1"
+	appsv1 "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/util/wait"
@@ -60,6 +63,56 @@ func (c *Client) hasRunningPod(ctx context.Context,
 	}
 }
 
+// unidleReplicas checks the unidle-replicas annotation for the number of
+// replicas to restore. If the label cannot be read or parsed, 1 is returned.
+// The return value is clamped to the interval [1,16].
+func unidleReplicas(deploy appsv1.Deployment) int {
+	rs, ok := deploy.Annotations["idling.amazee.io/unidle-replicas"]
+	if !ok {
+		return 1
+	}
+	r, err := strconv.Atoi(rs)
+	if err != nil || r < 1 {
+		return 1
+	}
+	if r > 16 {
+		return 16
+	}
+	return r
+}
+
+// unidleNamespace scales all deployments with the
+// "idling.amazee.io/watch=true" label up to the number of replicas in the
+// "idling.amazee.io/unidle-replicas" label.
+func (c *Client) unidleNamespace(ctx context.Context, namespace string) error {
+	deploys, err := c.clientset.AppsV1().Deployments(namespace).List(ctx, metav1.ListOptions{
+		LabelSelector: "idling.amazee.io/watch=true",
+	})
+	if err != nil {
+		return fmt.Errorf("couldn't select deploys by label: %v", err)
+	}
+	for _, deploy := range deploys.Items {
+		// check if idled
+		s, err := c.clientset.AppsV1().Deployments(namespace).
+			GetScale(ctx, deploy.Name, metav1.GetOptions{})
+		if err != nil {
+			return fmt.Errorf("couldn't get deployment scale: %v", err)
+		}
+		if s.Spec.Replicas > 0 {
+			continue
+		}
+		// scale up the deployment
+		sc := *s
+		sc.Spec.Replicas = int32(unidleReplicas(deploy))
+		_, err = c.clientset.AppsV1().Deployments(namespace).
+			UpdateScale(ctx, deploy.Name, &sc, metav1.UpdateOptions{})
+		if err != nil {
+			return fmt.Errorf("couldn't scale deployment: %v", err)
+		}
+	}
+	return nil
+}
+
 func (c *Client) ensureScaled(ctx context.Context, namespace, deployment string) error {
 	// get current scale
 	s, err := c.clientset.AppsV1().Deployments(namespace).
@@ -100,7 +153,11 @@ func (c *Client) getExecutor(ctx context.Context, namespace, deployment,
 		defer wg.Wait()
 	}
 	defer cancel()
-	// ensure the deployment has at least one replica
+	// unidle the entire namespace asynchronously
+	if err := c.unidleNamespace(ctx, namespace); err != nil {
+		return nil, fmt.Errorf("couldn't unidle namespace: %v", err)
+	}
+	// ensure the target deployment has at least one replica
 	if err := c.ensureScaled(ctx, namespace, deployment); err != nil {
 		return nil, fmt.Errorf("couldn't scale deployment: %v", err)
 	}
@@ -121,7 +178,7 @@ func (c *Client) getExecutor(ctx context.Context, namespace, deployment,
 	req := c.clientset.CoreV1().RESTClient().Post().Namespace(namespace).
 		Resource("pods").Name(firstPod).SubResource("exec")
 	req.VersionedParams(
-		&v1.PodExecOptions{
+		&corev1.PodExecOptions{
 			Stdin:     true,
 			Stdout:    true,
 			Stderr:    true,