浏览代码

Merge branch 'develop' into patch-2

Matt Ray 3 年之前
父节点
当前提交
308a494380
共有 2 个文件被更改,包括 88 次插入2 次删除
  1. 22 2
      pkg/cloud/awsprovider.go
  2. 66 0
      pkg/cloud/awsprovider_test.go

+ 22 - 2
pkg/cloud/awsprovider.go

@@ -55,7 +55,9 @@ const (
 	InUseState    = "in-use"
 	AttachedState = "attached"
 
-	AWSHourlyPublicIPCost = 0.005
+	AWSHourlyPublicIPCost    = 0.005
+	EKSCapacityTypeLabel     = "eks.amazonaws.com/capacityType"
+	EKSCapacitySpotTypeValue = "SPOT"
 )
 
 var (
@@ -636,6 +638,9 @@ func (k *awsKey) ID() string {
 	return ""
 }
 
+// Features will return a comma seperated list of features for the given node
+// If the node has a spot label, it will be included in the list
+// Otherwise, the list include instance type, operating system, and the region
 func (k *awsKey) Features() string {
 
 	instanceType, _ := util.GetInstanceType(k.Labels)
@@ -643,7 +648,7 @@ func (k *awsKey) Features() string {
 	region, _ := util.GetRegion(k.Labels)
 
 	key := region + "," + instanceType + "," + operatingSystem
-	usageType := PreemptibleType
+	usageType := k.getUsageType(k.Labels)
 	spotKey := key + "," + usageType
 	if l, ok := k.Labels["lifecycle"]; ok && l == "EC2Spot" {
 		return spotKey
@@ -651,9 +656,23 @@ func (k *awsKey) Features() string {
 	if l, ok := k.Labels[k.SpotLabelName]; ok && l == k.SpotLabelValue {
 		return spotKey
 	}
+	if usageType == PreemptibleType {
+		return spotKey
+	}
 	return key
 }
 
+// getUsageType returns the usage type of the instance
+// If the instance is a spot instance, it will return PreemptibleType
+// Otherwise returns an empty string
+func (k *awsKey) getUsageType(labels map[string]string) string {
+	if label, ok := labels[EKSCapacityTypeLabel]; ok && label == EKSCapacitySpotTypeValue {
+		// We currently write out spot instances as "preemptible" in the pricing data, so these need to match
+		return PreemptibleType
+	}
+	return ""
+}
+
 func (aws *AWS) PVPricing(pvk PVKey) (*PV, error) {
 	pricing, ok := aws.Pricing[pvk.Features()]
 	if !ok {
@@ -821,6 +840,7 @@ func (aws *AWS) DownloadPricingData() error {
 
 	inputkeys := make(map[string]bool)
 	for _, n := range nodeList {
+
 		if _, ok := n.Labels["eks.amazonaws.com/nodegroup"]; ok {
 			aws.clusterManagementPrice = 0.10
 			aws.clusterProvisioner = "EKS"

+ 66 - 0
pkg/cloud/awsprovider_test.go

@@ -0,0 +1,66 @@
+package cloud
+
+import "testing"
+
+func Test_awsKey_getUsageType(t *testing.T) {
+	type fields struct {
+		Labels     map[string]string
+		ProviderID string
+	}
+	type args struct {
+		labels map[string]string
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+		want   string
+	}{
+		{
+			// test with no labels should return false
+			name: "Label does not have the capacityType label associated with it",
+			args: args{
+				labels: map[string]string{},
+			},
+			want: "",
+		},
+		{
+			name: "label with a capacityType set to empty string should return empty string",
+			args: args{
+				labels: map[string]string{
+					EKSCapacityTypeLabel: "",
+				},
+			},
+			want: "",
+		},
+		{
+			name: "label with capacityType set to a random value should return empty string",
+			args: args{
+				labels: map[string]string{
+					EKSCapacityTypeLabel: "TEST_ME",
+				},
+			},
+			want: "",
+		},
+		{
+			name: "label with capacityType set to spot should return spot",
+			args: args{
+				labels: map[string]string{
+					EKSCapacityTypeLabel: EKSCapacitySpotTypeValue,
+				},
+			},
+			want: PreemptibleType,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			k := &awsKey{
+				Labels:     tt.fields.Labels,
+				ProviderID: tt.fields.ProviderID,
+			}
+			if got := k.getUsageType(tt.args.labels); got != tt.want {
+				t.Errorf("getUsageType() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}