Support resource requirements by annotations

controller/webhook.go

@@ -20,29 +20,35 @@ import (
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
+	"log"
+	"net/http"
+	"os"
+	"strings"
+
 	"k8s.io/api/admission/v1beta1"
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/client-go/kubernetes"
 	corev1Types "k8s.io/client-go/kubernetes/typed/core/v1"
-	"log"
-	"net/http"
-	"os"
-	"strings"
 )
 
 const (
-	signingProxyWebhookAnnotationHostKey    = "sidecar.aws.signing-proxy/host"
-	signingProxyWebhookAnnotationInjectKey  = "sidecar.aws.signing-proxy/inject"
-	signingProxyWebhookAnnotationNameKey    = "sidecar.aws.signing-proxy/name"
-	signingProxyWebhookAnnotationRegionKey  = "sidecar.aws.signing-proxy/region"
-	signingProxyWebhookAnnotationRoleArnKey = "sidecar.aws.signing-proxy/role-arn"
-	signingProxyWebhookAnnotationStatusKey  = "sidecar.aws.signing-proxy/status"
-	signingProxyWebhookLabelHostKey         = "sidecar-host"
-	signingProxyWebhookLabelNameKey         = "sidecar-name"
-	signingProxyWebhookLabelRegionKey       = "sidecar-region"
-	signingProxyWebhookLabelRoleArnKey      = "sidecar-role-arn"
+	signingProxyWebhookAnnotationHostKey       = "sidecar.aws.signing-proxy/host"
+	signingProxyWebhookAnnotationInjectKey     = "sidecar.aws.signing-proxy/inject"
+	signingProxyWebhookAnnotationNameKey       = "sidecar.aws.signing-proxy/name"
+	signingProxyWebhookAnnotationRegionKey     = "sidecar.aws.signing-proxy/region"
+	signingProxyWebhookAnnotationRoleArnKey    = "sidecar.aws.signing-proxy/role-arn"
+	signingProxyWebhookAnnotationStatusKey     = "sidecar.aws.signing-proxy/status"
+	signingProxyWebhookAnnotationCPURequestKey = "sidecar.aws.signing-proxy/cpu"
+	signingProxyWebhookAnnotationCPULimitKey   = "sidecar.aws.signing-proxy/cpu-limit"
+	signingProxyWebhookAnnotationMemRequestKey = "sidecar.aws.signing-proxy/memory"
+	signingProxyWebhookAnnotationMemLimitKey   = "sidecar.aws.signing-proxy/memory-limit"
+	signingProxyWebhookLabelHostKey            = "sidecar-host"
+	signingProxyWebhookLabelNameKey            = "sidecar-name"
+	signingProxyWebhookLabelRegionKey          = "sidecar-region"
+	signingProxyWebhookLabelRoleArnKey         = "sidecar-role-arn"
 )
 
 var (
@@ -176,6 +182,11 @@ func (whsvr *WebhookServer) mutate(ctx context.Context, admissionReview *v1beta1
 		Args: sidecarArgs,
 	}}
 
+	resources := whsvr.getResourceRequirements(&pod.ObjectMeta)
+	if resources != nil {
+		sidecarContainer[0].Resources = *resources
+	}
+
 	patchOperations = append(patchOperations, addContainers(pod.Spec.Containers, sidecarContainer, "/spec/containers")...)
 
 	annotations := map[string]string{signingProxyWebhookAnnotationStatusKey: "injected"}
@@ -297,6 +308,68 @@ func extractParameters(host string, name string, region string) (string, string,
 	return host, name, region
 }
 
+func (whsvr *WebhookServer) getResourceRequirements(podMetadata *metav1.ObjectMeta) *corev1.ResourceRequirements {
+	annotations := podMetadata.GetAnnotations()
+
+	if annotations == nil {
+		annotations = map[string]string{}
+	}
+
+	cpuReq := annotations[signingProxyWebhookAnnotationCPURequestKey]
+	cpuLimit := annotations[signingProxyWebhookAnnotationCPULimitKey]
+	memReq := annotations[signingProxyWebhookAnnotationMemRequestKey]
+	memLimit := annotations[signingProxyWebhookAnnotationMemLimitKey]
+
+	if cpuReq == "" && cpuLimit == "" && memReq == "" && memLimit == "" {
+		return nil
+	}
+
+	requests := map[corev1.ResourceName]resource.Quantity{}
+	limits := map[corev1.ResourceName]resource.Quantity{}
+
+	if cpuReq != "" {
+		quantity, err := resource.ParseQuantity(cpuReq)
+		if err != nil {
+			log.Printf("Error parsing cpu requests: %v", err)
+		} else {
+			requests[corev1.ResourceRequestsCPU] = quantity
+		}
+	}
+	if memReq != "" {
+		quantity, err := resource.ParseQuantity(memReq)
+		if err != nil {
+			log.Printf("Error parsing memory requests: %v", err)
+		} else {
+			requests[corev1.ResourceRequestsMemory] = quantity
+
+		}
+	}
+	if cpuLimit != "" {
+		quantity, err := resource.ParseQuantity(cpuLimit)
+		if err != nil {
+			log.Printf("Error parsing cpu limit: %v", err)
+		} else {
+			limits[corev1.ResourceLimitsCPU] = quantity
+		}
+	}
+	if memLimit != "" {
+		quantity, err := resource.ParseQuantity(memLimit)
+		if err != nil {
+			log.Printf("Error parsing memory limit: %v", err)
+		} else {
+			limits[corev1.ResourceLimitsMemory] = quantity
+		}
+	}
+	resource := &corev1.ResourceRequirements{}
+	if len(requests) > 0 {
+		resource.Requests = requests
+	}
+	if len(limits) > 0 {
+		resource.Limits = limits
+	}
+	return resource
+}
+
 func (whsvr *WebhookServer) getRoleArn(nsLabels map[string]string, podMetadata *metav1.ObjectMeta) string {
 	annotations := podMetadata.GetAnnotations()
 

controller/webhook_test.go

@@ -20,6 +20,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"testing"
 )
@@ -347,3 +348,48 @@ func TestWebhookServer_getRoleArn(t *testing.T) {
 		})
 	}
 }
+
+func TestWebhookServer_getResourceRequirements(t *testing.T) {
+	var testCases = []struct {
+		name          string
+		podObjectMeta *metav1.ObjectMeta
+		labels        map[string]string
+		expected      *corev1.ResourceRequirements
+		errorMessage  string
+	}{
+		{
+			name: "TestSidecarResourceAnnotationPresent",
+			podObjectMeta: &metav1.ObjectMeta{
+				Annotations: map[string]string{
+					signingProxyWebhookAnnotationCPURequestKey: "200m",
+					signingProxyWebhookAnnotationMemRequestKey: "200Mi",
+					signingProxyWebhookAnnotationCPULimitKey:   "400m",
+					signingProxyWebhookAnnotationMemLimitKey:   "400Mi",
+				},
+			},
+			expected: &corev1.ResourceRequirements{
+				Requests: map[corev1.ResourceName]resource.Quantity{
+					corev1.ResourceRequestsCPU:    resource.MustParse("200m"),
+					corev1.ResourceRequestsMemory: resource.MustParse("200Mi"),
+				},
+				Limits: map[corev1.ResourceName]resource.Quantity{
+					corev1.ResourceLimitsCPU:    resource.MustParse("400m"),
+					corev1.ResourceLimitsMemory: resource.MustParse("400Mi"),
+				},
+			},
+			errorMessage: "Should return ResourceRequiremts value",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			whsvr := &WebhookServer{
+				server:          nil,
+				namespaceClient: nil,
+			}
+
+			r := whsvr.getResourceRequirements(tc.podObjectMeta)
+			assert.Equal(t, tc.expected, r, tc.errorMessage)
+		})
+	}
+}