From ccf382219f332e2a36444ce14d3c84eeb46725a7 Mon Sep 17 00:00:00 2001
From: Yaroslava Serdiuk <yaroslava@google.com>
Date: Tue, 24 May 2022 14:24:01 +0000
Subject: [PATCH] Add ephemeral storage price to PodPrice

---
 .../cloudprovider/gce/gce_price_model.go      |  6 +++
 .../cloudprovider/gce/gce_price_model_test.go | 10 +++--
 cluster-autoscaler/utils/test/test_utils.go   | 38 +++++++++++++++++++
 3 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/cluster-autoscaler/cloudprovider/gce/gce_price_model.go b/cluster-autoscaler/cloudprovider/gce/gce_price_model.go
index 5239c8297904..0f59b9ec7132 100644
--- a/cluster-autoscaler/cloudprovider/gce/gce_price_model.go
+++ b/cluster-autoscaler/cloudprovider/gce/gce_price_model.go
@@ -185,6 +185,12 @@ func (model *GcePriceModel) getBasePrice(resources apiv1.ResourceList, instanceT
 	}
 	price += float64(mem.Value()) / float64(units.GiB) * memPrice * hours
 
+	if model.EphemeralStorageSupport {
+		ephemeralStorage := resources[apiv1.ResourceEphemeralStorage]
+		ephemeralStoragePrice := model.PriceInfo.LocalSsdPricePerHour()
+		price += float64(ephemeralStorage.Value()) / float64(units.GiB) * ephemeralStoragePrice * hours
+	}
+
 	return price
 }
 
diff --git a/cluster-autoscaler/cloudprovider/gce/gce_price_model_test.go b/cluster-autoscaler/cloudprovider/gce/gce_price_model_test.go
index 14efffbbc9d3..fbd42fb69318 100644
--- a/cluster-autoscaler/cloudprovider/gce/gce_price_model_test.go
+++ b/cluster-autoscaler/cloudprovider/gce/gce_price_model_test.go
@@ -239,16 +239,20 @@ func TestGetNodePrice(t *testing.T) {
 }
 
 func TestGetPodPrice(t *testing.T) {
-	pod1 := BuildTestPod("a1", 100, 500*units.MiB)
-	pod2 := BuildTestPod("a2", 2*100, 2*500*units.MiB)
+	pod1 := BuildTestPodWithEphemeralStorage("a1", 100, 500*units.MiB, 100*units.GiB)
+	pod2 := BuildTestPodWithEphemeralStorage("a2", 2*100, 2*500*units.MiB, 2*100*units.GiB)
+	pod3 := BuildTestPodWithEphemeralStorage("a2", 2*100, 2*500*units.MiB, 100*units.GiB)
 
-	model := NewGcePriceModel(NewGcePriceInfo(), false)
+	model := NewGcePriceModel(NewGcePriceInfo(), true)
 	now := time.Now()
 
 	price1, err := model.PodPrice(pod1, now, now.Add(time.Hour))
 	assert.NoError(t, err)
 	price2, err := model.PodPrice(pod2, now, now.Add(time.Hour))
 	assert.NoError(t, err)
+	price3, err := model.PodPrice(pod3, now, now.Add(time.Hour))
+	assert.NoError(t, err)
 	// 2 times bigger pod should cost twice as much.
 	assert.True(t, math.Abs(price1*2-price2) < 0.001)
+	assert.True(t, price2 > price3)
 }
diff --git a/cluster-autoscaler/utils/test/test_utils.go b/cluster-autoscaler/utils/test/test_utils.go
index 969c25fb39af..adf9aa860e52 100644
--- a/cluster-autoscaler/utils/test/test_utils.go
+++ b/cluster-autoscaler/utils/test/test_utils.go
@@ -65,6 +65,44 @@ func BuildTestPod(name string, cpu int64, mem int64) *apiv1.Pod {
 	return pod
 }
 
+// BuildTestPodWithEphemeralStorage creates a pod with cpu, memory and ephemeral storage resources.
+func BuildTestPodWithEphemeralStorage(name string, cpu, mem, ephemeralStorage int64) *apiv1.Pod {
+	startTime := metav1.Unix(0, 0)
+	pod := &apiv1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			UID:         types.UID(name),
+			Namespace:   "default",
+			Name:        name,
+			SelfLink:    fmt.Sprintf("/api/v1/namespaces/default/pods/%s", name),
+			Annotations: map[string]string{},
+		},
+		Spec: apiv1.PodSpec{
+			Containers: []apiv1.Container{
+				{
+					Resources: apiv1.ResourceRequirements{
+						Requests: apiv1.ResourceList{},
+					},
+				},
+			},
+		},
+		Status: apiv1.PodStatus{
+			StartTime: &startTime,
+		},
+	}
+
+	if cpu >= 0 {
+		pod.Spec.Containers[0].Resources.Requests[apiv1.ResourceCPU] = *resource.NewMilliQuantity(cpu, resource.DecimalSI)
+	}
+	if mem >= 0 {
+		pod.Spec.Containers[0].Resources.Requests[apiv1.ResourceMemory] = *resource.NewQuantity(mem, resource.DecimalSI)
+	}
+	if ephemeralStorage >= 0 {
+		pod.Spec.Containers[0].Resources.Requests[apiv1.ResourceEphemeralStorage] = *resource.NewQuantity(ephemeralStorage, resource.DecimalSI)
+	}
+
+	return pod
+}
+
 // BuildServiceTokenProjectedVolumeSource returns a ProjectedVolumeSource with SA token
 // projection
 func BuildServiceTokenProjectedVolumeSource(path string) *apiv1.ProjectedVolumeSource {