From 82f1f018b0c5bf17fdceddaee1eff1ac9a201e7e Mon Sep 17 00:00:00 2001
From: Jan Sykora <jan@cast.ai>
Date: Fri, 2 Aug 2024 14:06:13 +0200
Subject: [PATCH] feat: recognize karpetner spot label in GKE and AKS

---
 internal/services/providers/aks/aks.go      | 18 +++++++++++++++---
 internal/services/providers/aks/aks_test.go | 16 ++++++++++++++++
 internal/services/providers/gke/gke.go      |  4 ++++
 internal/services/providers/gke/gke_test.go |  5 +++++
 4 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/internal/services/providers/aks/aks.go b/internal/services/providers/aks/aks.go
index a05dc83d..7de4b8ef 100644
--- a/internal/services/providers/aks/aks.go
+++ b/internal/services/providers/aks/aks.go
@@ -11,6 +11,7 @@ import (
 	"castai-agent/internal/config"
 	"castai-agent/internal/services/providers/aks/metadata"
 	"castai-agent/internal/services/providers/types"
+	"castai-agent/pkg/labels"
 )
 
 type Provider struct {
@@ -92,16 +93,27 @@ func (p *Provider) FilterSpot(_ context.Context, nodes []*corev1.Node) ([]*corev
 	var ret []*corev1.Node
 
 	for _, node := range nodes {
-		if val, ok := node.ObjectMeta.Labels["kubernetes.azure.com/scalesetpriority"]; !ok || val != "spot" {
-			continue
+		if isSpot(node) {
+			ret = append(ret, node)
 		}
 
-		ret = append(ret, node)
 	}
 
 	return ret, nil
 }
 
+func isSpot(node *corev1.Node) bool {
+	if val, ok := node.ObjectMeta.Labels["kubernetes.azure.com/scalesetpriority"]; ok && val == "spot" {
+		return true
+	}
+
+	if val, ok := node.ObjectMeta.Labels[labels.KarpenterCapacityType]; ok && val == labels.ValueKarpenterCapacityTypeSpot {
+		return true
+	}
+
+	return false
+}
+
 func (p *Provider) Name() string {
 	return Name
 }
diff --git a/internal/services/providers/aks/aks_test.go b/internal/services/providers/aks/aks_test.go
index 9af34813..791f3fa1 100644
--- a/internal/services/providers/aks/aks_test.go
+++ b/internal/services/providers/aks/aks_test.go
@@ -15,6 +15,7 @@ import (
 	"castai-agent/internal/castai"
 	mock_castai "castai-agent/internal/castai/mock"
 	"castai-agent/internal/services/providers/types"
+	"castai-agent/pkg/labels"
 )
 
 func TestProvider_RegisterCluster(t *testing.T) {
@@ -69,4 +70,19 @@ func TestProvider_IsSpot(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, []*v1.Node{node}, got)
 	})
+
+	t.Run("spot instance karpenter label", func(t *testing.T) {
+		p := &Provider{
+			log: logrus.New(),
+		}
+
+		node := &v1.Node{
+			ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{labels.KarpenterCapacityType: labels.ValueKarpenterCapacityTypeSpot}},
+		}
+
+		got, err := p.FilterSpot(context.Background(), []*v1.Node{node})
+
+		require.NoError(t, err)
+		require.Equal(t, []*v1.Node{node}, got)
+	})
 }
diff --git a/internal/services/providers/gke/gke.go b/internal/services/providers/gke/gke.go
index 75de1f4d..920b26d8 100644
--- a/internal/services/providers/gke/gke.go
+++ b/internal/services/providers/gke/gke.go
@@ -112,6 +112,10 @@ func isSpot(node *corev1.Node) bool {
 		return true
 	}
 
+	if val, ok := node.Labels[labels.KarpenterCapacityType]; ok && val == labels.ValueKarpenterCapacityTypeSpot {
+		return true
+	}
+
 	return false
 }
 
diff --git a/internal/services/providers/gke/gke_test.go b/internal/services/providers/gke/gke_test.go
index 3f12217b..feed1975 100644
--- a/internal/services/providers/gke/gke_test.go
+++ b/internal/services/providers/gke/gke_test.go
@@ -99,6 +99,11 @@ func TestProvider_IsSpot(t *testing.T) {
 			node:     &v1.Node{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{LabelSpot: "true"}}},
 			expected: true,
 		},
+		{
+			name:     "gke karpenter spot node",
+			node:     &v1.Node{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{labels.KarpenterCapacityType: labels.ValueKarpenterCapacityTypeSpot}}},
+			expected: true,
+		},
 		{
 			name:     "on demand node",
 			node:     &v1.Node{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{}}},