From 1cbcfbcbe79a352d4d824dd6cddc2c1a8082bd9e 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      |  7 ++++
 .../cloudprovider/gce/gce_price_model_test.go | 10 +++--
 cluster-autoscaler/utils/test/test_utils.go   | 38 +++++++++++++++++++
 3 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/cluster-autoscaler/cloudprovider/gce/gce_price_model.go b/cluster-autoscaler/cloudprovider/gce/gce_price_model.go
index f1d4c784b8d5..68dae383e3db 100644
--- a/cluster-autoscaler/cloudprovider/gce/gce_price_model.go
+++ b/cluster-autoscaler/cloudprovider/gce/gce_price_model.go
@@ -197,6 +197,13 @@ func (model *GcePriceModel) getBasePrice(resources apiv1.ResourceList, instanceT
 	}
 	price += float64(mem.Value()) / float64(units.GiB) * memPrice * hours
 
+	if model.EphemeralStorageSupport {
+		ephemeralStorage := resources[apiv1.ResourceEphemeralStorage]
+		// For simplification using a fixed price for default boot disk.
+		ephemeralStoragePrice := model.PriceInfo.BootDiskPricePerHour()[DefaultBootDiskType]
+		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 b0a27b8fbb61..2d9268ce5b52 100644
--- a/cluster-autoscaler/cloudprovider/gce/gce_price_model_test.go
+++ b/cluster-autoscaler/cloudprovider/gce/gce_price_model_test.go
@@ -232,16 +232,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 18839814f3ce..71d161b7f560 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 {