Add size unit annotation for load balancers. (#386)

Author: wez470

ci/headers-go.sh

@@ -36,9 +36,9 @@ FILES=$(find . -name "*.go" -not -path "./vendor/*")
 EXIT=0
 
 for FILE in $FILES; do
-        if head -n 1 "$FILE" | grep -q '// +build '; then
-                # Remove +build comment and subsequent blank line.
-                HEADERS=$(tail -n +3 "$FILE" | head -n 15)
+        if head -n 1 "$FILE" | grep -q '//go:build '; then
+                # Remove go:build and +build comments as well as blank line below.
+                HEADERS=$(tail -n +4 "$FILE" | head -n 15)
         else
                 HEADERS=$(head -n 15 "$FILE")
         fi

cloud-controller-manager/do/loadbalancers.go

@@ -120,9 +120,13 @@ const (
 
 	// annDOSizeSlug is the annotation specifying the size of the LB.
 	// Options are `lb-small`, `lb-medium`, and `lb-large`.
-	// Defaults to `lb-small`.
+	// Defaults to `lb-small`. Only one of annDOSizeSlug and annDOSizeUnit can be specified.
 	annDOSizeSlug = "service.beta.kubernetes.io/do-loadbalancer-size-slug"
 
+	// annDOSizeUnit is the annotation specifying the size of the LB.
+	// Options are numbers greater than or equal to `1`. Only one of annDOSizeUnit and annDOSizeSlug can be specified.
+	annDOSizeUnit = "service.beta.kubernetes.io/do-loadbalancer-size-unit"
+
 	// annDOStickySessionsType is the annotation specifying which sticky session type
 	// DO loadbalancer should use. Options are none and cookies. Defaults
 	// to none.
@@ -665,6 +669,15 @@ func (l *loadBalancers) buildLoadBalancerRequest(ctx context.Context, service *v
 		return nil, err
 	}
 
+	sizeUnit, err := getSizeUnit(service)
+	if err != nil {
+		return nil, err
+	}
+
+	if sizeSlug != "" && sizeUnit > 0 {
+		return nil, fmt.Errorf("only one of LB size slug and size unit can be provided")
+	}
+
 	redirectHTTPToHTTPS, err := getRedirectHTTPToHTTPS(service)
 	if err != nil {
 		return nil, err
@@ -690,6 +703,7 @@ func (l *loadBalancers) buildLoadBalancerRequest(ctx context.Context, service *v
 		DropletIDs:             dropletIDs,
 		Region:                 l.region,
 		SizeSlug:               sizeSlug,
+		SizeUnit:               sizeUnit,
 		ForwardingRules:        forwardingRules,
 		HealthCheck:            healthCheck,
 		StickySessions:         stickySessions,
@@ -1095,19 +1109,39 @@ func getAlgorithm(service *v1.Service) string {
 	}
 }
 
+// getSizeSlug returns the load balancer size as a slug
 func getSizeSlug(service *v1.Service) (string, error) {
-	sizeSlug, ok := service.Annotations[annDOSizeSlug]
+	sizeSlug, _ := service.Annotations[annDOSizeSlug]
 
-	if !ok || sizeSlug == "" {
-		sizeSlug = "lb-small"
+	if sizeSlug != "" {
+		switch sizeSlug {
+		case "lb-small", "lb-medium", "lb-large":
+		default:
+			return "", fmt.Errorf("invalid LB size slug provided: %s", sizeSlug)
+		}
 	}
 
-	switch sizeSlug {
-	case "lb-small", "lb-medium", "lb-large":
-		return sizeSlug, nil
-	default:
-		return "", fmt.Errorf("invalid LB size slug provided: %s", sizeSlug)
+	return sizeSlug, nil
+}
+
+// getSizeUnit returns the load balancer size as a number
+func getSizeUnit(service *v1.Service) (uint32, error) {
+	sizeUnitStr, ok := service.Annotations[annDOSizeUnit]
+
+	if !ok || sizeUnitStr == "" {
+		return uint32(0), nil
 	}
+
+	sizeUnit, err := strconv.Atoi(sizeUnitStr)
+	if err != nil {
+		return 0, fmt.Errorf("invalid LB size unit %q provided: %s", sizeUnitStr, err)
+	}
+
+	if sizeUnit < 0 {
+		return 0, fmt.Errorf("LB size unit must be non-negative. %d provided", sizeUnit)
+	}
+
+	return uint32(sizeUnit), nil
 }
 
 // getStickySessionsType returns the sticky session type to use for