From 059f66f22f991aa560deb02696c43d127a36efce Mon Sep 17 00:00:00 2001 From: Scott Leggett Date: Fri, 12 Jul 2024 16:16:17 +0800 Subject: [PATCH] feat: support the new idle labels and annotations The idle labels and annotations can now have the lagoon.sh suffix as amazee.io. This change adds support for both suffixes, with idling.lagoon.sh taking priority over idling.amazee.io. --- go.mod | 2 + go.sum | 4 ++ internal/k8s/client.go | 2 +- internal/k8s/exec.go | 83 +++++++++++++++++------- internal/k8s/exec_test.go | 130 +++++++++++++++++++++++++++++++++++++- 5 files changed, 196 insertions(+), 25 deletions(-) diff --git a/go.mod b/go.mod index ad68b414..2aa44391 100644 --- a/go.mod +++ b/go.mod @@ -37,6 +37,7 @@ require ( github.com/cespare/xxhash/v2 v2.2.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/emicklei/go-restful/v3 v3.11.0 // indirect + github.com/evanphx/json-patch v4.12.0+incompatible // indirect github.com/go-jose/go-jose/v4 v4.0.2 // indirect github.com/go-logr/logr v1.4.1 // indirect github.com/go-logr/stdr v1.2.2 // indirect @@ -62,6 +63,7 @@ require ( github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f // indirect github.com/nats-io/nkeys v0.4.7 // indirect github.com/nats-io/nuid v1.0.1 // indirect + github.com/pkg/errors v0.9.1 // indirect github.com/prometheus/client_model v0.5.0 // indirect github.com/prometheus/common v0.48.0 // indirect github.com/prometheus/procfs v0.12.0 // indirect diff --git a/go.sum b/go.sum index f4355058..a301e886 100644 --- a/go.sum +++ b/go.sum @@ -26,6 +26,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/emicklei/go-restful/v3 v3.11.0 h1:rAQeMHw1c7zTmncogyy8VvRZwtkmkZ4FxERmMY4rD+g= github.com/emicklei/go-restful/v3 v3.11.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc= +github.com/evanphx/json-patch v4.12.0+incompatible h1:4onqiflcdA9EOZ4RxV643DvftH5pOlLGNtQ5lPWQu84= +github.com/evanphx/json-patch v4.12.0+incompatible/go.mod h1:50XU6AFN0ol/bzJsmQLiYLvXMP4fmwYFNcr97nuDLSk= github.com/gliderlabs/ssh v0.3.7 h1:iV3Bqi942d9huXnzEF2Mt+CY9gLu8DNM4Obd+8bODRE= github.com/gliderlabs/ssh v0.3.7/go.mod h1:zpHEXBstFnQYtGnB8k8kQLol82umzn/2/snG7alWVD8= github.com/go-chi/chi/v5 v5.0.12 h1:9euLV5sTrTNTRUU9POmDUvfxyj6LAABLUcEWO+JJb4s= @@ -122,6 +124,8 @@ github.com/onsi/ginkgo/v2 v2.15.0 h1:79HwNRBAZHOEwrczrgSOPy+eFTTlIGELKy5as+ClttY github.com/onsi/ginkgo/v2 v2.15.0/go.mod h1:HlxMHtYF57y6Dpf+mc5529KKmSq9h2FpCF+/ZkwUxKM= github.com/onsi/gomega v1.31.0 h1:54UJxxj6cPInHS3a35wm6BK/F9nHYueZ1NVujHDrnXE= github.com/onsi/gomega v1.31.0/go.mod h1:DW9aCi7U6Yi40wNVAvT6kzFnEVEI5n3DloYBiKiT6zk= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prometheus/client_golang v1.19.1 h1:wZWJDwK+NameRJuPGDhlnFgx8e8HN3XHQeLaYJFJBOE= diff --git a/internal/k8s/client.go b/internal/k8s/client.go index bc656b3e..b68b828a 100644 --- a/internal/k8s/client.go +++ b/internal/k8s/client.go @@ -22,7 +22,7 @@ var timeoutSeconds = int64(timeout / time.Second) // Client is a k8s client. type Client struct { config *rest.Config - clientset *kubernetes.Clientset + clientset kubernetes.Interface logStreamIDs sync.Map } diff --git a/internal/k8s/exec.go b/internal/k8s/exec.go index 3cd296ad..62544994 100644 --- a/internal/k8s/exec.go +++ b/internal/k8s/exec.go @@ -18,8 +18,23 @@ import ( "k8s.io/client-go/tools/remotecommand" ) -const ( - idleAnnotation = "idling.amazee.io/unidle-replicas" +var ( + // idleReplicaAnnotations are the annotations that will be used to determine + // how many replicas to set when scaling up a deployment from idle. The + // annotations are in priority order from high to low. The first annotation + // found will be used. + idleReplicaAnnotations = []string{ + "idling.lagoon.sh/unidle-replicas", + "idling.amazee.io/unidle-replicas", + } + // idleWatchLabels are the labels that will be used to determine which + // deployments to scale when unidling an environment. The labels are in + // priority order from high to low. The first annotation found on any + // deployment will be used. + idleWatchLabels = []string{ + "idling.lagoon.sh/watch=true", + "idling.amazee.io/watch=true", + } ) // podContainer returns the first pod and first container inside that pod for @@ -68,33 +83,57 @@ 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]. +// unidleReplicas checks the idleReplicaAnnotations for the number of replicas +// to restore. If the labels cannot be found 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[idleAnnotation] - if !ok { - return 1 - } - r, err := strconv.Atoi(rs) - if err != nil || r < 1 { - return 1 + for _, ra := range idleReplicaAnnotations { + rs, ok := deploy.Annotations[ra] + if !ok { + continue + } + r, err := strconv.Atoi(rs) + if err != nil || r < 1 { + return 1 + } + if r > 16 { + return 16 + } + return r } - if r > 16 { - return 16 + return 1 +} + +// idledDeploys returns the DeploymentList of idled deployments in the given +// namespace. +func (c *Client) idledDeploys(ctx context.Context, namespace string) ( + *appsv1.DeploymentList, error, +) { + var deploys *appsv1.DeploymentList + for _, selector := range idleWatchLabels { + deploys, err := c.clientset.AppsV1().Deployments(namespace).List(ctx, + metav1.ListOptions{ + LabelSelector: selector, + }) + if err != nil { + return nil, fmt.Errorf("couldn't select deploys by label: %v", err) + } + if deploys != nil && len(deploys.Items) > 0 { + return deploys, nil + } } - return r + return deploys, nil } -// 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. +// unidleNamespace scales all deployments with the idleWatchLabels up to the +// number of replicas in the idleReplicaAnnotations. 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", - }) + deploys, err := c.idledDeploys(ctx, namespace) if err != nil { - return fmt.Errorf("couldn't select deploys by label: %v", err) + return fmt.Errorf("couldn't get idled deploys: %v", err) + } + if deploys == nil { + return nil // no deploys to unidle } for _, deploy := range deploys.Items { // check if idled diff --git a/internal/k8s/exec_test.go b/internal/k8s/exec_test.go index 3dba5d04..e92daeab 100644 --- a/internal/k8s/exec_test.go +++ b/internal/k8s/exec_test.go @@ -1,14 +1,16 @@ package k8s import ( + "context" "testing" "github.com/alecthomas/assert/v2" appsv1 "k8s.io/api/apps/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/client-go/kubernetes/fake" ) -func TestUnidleReplicas(t *testing.T) { +func TestUnidleReplicasParsing(t *testing.T) { var testCases = map[string]struct { input string expect int @@ -28,10 +30,134 @@ func TestUnidleReplicas(t *testing.T) { t.Run(name, func(tt *testing.T) { deploy := appsv1.Deployment{ ObjectMeta: metav1.ObjectMeta{ - Annotations: map[string]string{idleAnnotation: tc.input}, + Annotations: map[string]string{idleReplicaAnnotations[0]: tc.input}, }, } assert.Equal(tt, tc.expect, unidleReplicas(deploy), name) }) } } + +func TestUnidleReplicasLabels(t *testing.T) { + for _, ra := range idleReplicaAnnotations { + t.Run(ra, func(tt *testing.T) { + deploy := appsv1.Deployment{ + ObjectMeta: metav1.ObjectMeta{ + Annotations: map[string]string{ra: "9"}, + }, + } + assert.Equal(tt, 9, unidleReplicas(deploy), ra) + }) + } +} + +func deployNames(deploys *appsv1.DeploymentList) []string { + var names []string + if deploys == nil { + return names // no deploys to unidle + } + for _, deploy := range deploys.Items { + names = append(names, deploy.Name) + } + return names +} + +func TestIdledDeployLabels(t *testing.T) { + testNS := "testns" + var testCases = map[string]struct { + deploys *appsv1.DeploymentList + expect []string + }{ + "prefer lagoon.sh": { + deploys: &appsv1.DeploymentList{ + Items: []appsv1.Deployment{ + { + ObjectMeta: metav1.ObjectMeta{ + Name: "one", + Namespace: testNS, + Labels: map[string]string{ + "idling.lagoon.sh/watch": "true", + }, + }, + }, + { + ObjectMeta: metav1.ObjectMeta{ + Name: "two", + Namespace: testNS, + Labels: map[string]string{ + "idling.amazee.io/watch": "true", + }, + }, + }, + }, + }, + expect: []string{"one"}, + }, + "fall back to amazee.io": { + deploys: &appsv1.DeploymentList{ + Items: []appsv1.Deployment{ + { + ObjectMeta: metav1.ObjectMeta{ + Name: "one", + Namespace: testNS, + Labels: map[string]string{ + "idling.amazee.io/watch": "true", + }, + }, + }, + { + ObjectMeta: metav1.ObjectMeta{ + Name: "two", + Namespace: testNS, + Labels: map[string]string{ + "idling.amazee.io/watch": "true", + }, + }, + }, + }, + }, + expect: []string{"one", "two"}, + }, + "ignore mislabelled deploys": { + deploys: &appsv1.DeploymentList{ + Items: []appsv1.Deployment{ + { + ObjectMeta: metav1.ObjectMeta{ + Name: "one", + Namespace: testNS, + Labels: map[string]string{ + "idling.foo/watch": "true", + }, + }, + }, + }, + }, + }, + "ignore other namespaces": { + deploys: &appsv1.DeploymentList{ + Items: []appsv1.Deployment{ + { + ObjectMeta: metav1.ObjectMeta{ + Name: "one", + Namespace: "wrongns", + Labels: map[string]string{ + "idling.lagoon.sh/watch": "true", + }, + }, + }, + }, + }, + }, + } + for name, tc := range testCases { + t.Run(name, func(tt *testing.T) { + // create fake Kubernetes client with test deploys + c := &Client{ + clientset: fake.NewSimpleClientset(tc.deploys), + } + deploys, err := c.idledDeploys(context.Background(), testNS) + assert.NoError(tt, err, name) + assert.Equal(tt, tc.expect, deployNames(deploys), name) + }) + } +}