diff --git a/pkg/admission/admission_controller.go b/pkg/admission/admission_controller.go index 12e5a6881..a2b8c1b88 100644 --- a/pkg/admission/admission_controller.go +++ b/pkg/admission/admission_controller.go @@ -423,9 +423,42 @@ func (c *AdmissionController) updateApplicationInfo(namespace string, pod *v1.Po zap.Any("labels", pod.Labels), zap.Any("annotations", pod.Annotations)) - annotations := getAnnotationsForApplicationInfoUpdate(pod, namespace, c.conf.GetGenerateUniqueAppIds(), c.conf.GetDefaultQueueName()) + newLabels, newAnnotations := getNewApplicationInfo(pod, namespace, c.conf.GetGenerateUniqueAppIds(), c.conf.GetDefaultQueueName()) - // check for an existing patch on annotations and update it + patch = updateLabelPatch(pod, newLabels, patch) + patch = updateAnnotationPatch(pod, newAnnotations, patch) + + return patch +} + +func updateLabelPatch(pod *v1.Pod, labels map[string]string, patch []common.PatchOperation) []common.PatchOperation { + // if found an existing patch on labels, add new labels to it + for _, p := range patch { + if p.Op == "add" && p.Path == "/metadata/labels" { + if existingLabels, ok := p.Value.(map[string]string); ok { + for k, v := range labels { + existingLabels[k] = v + } + return patch + } + } + } + // Add a new label patch + if len(labels) != 0 { + // combine new labels with existing labels in pod to create first patch + existingLabels := pod.Labels + labels = utils.MergeMaps(existingLabels, labels) + patch = append(patch, common.PatchOperation{ + Op: "add", + Path: "/metadata/labels", + Value: labels, + }) + } + return patch +} + +func updateAnnotationPatch(pod *v1.Pod, annotations map[string]string, patch []common.PatchOperation) []common.PatchOperation { + // if found an existing patch on annotations, add new annotations to it for _, p := range patch { if p.Op == "add" && p.Path == "/metadata/annotations" { if existingAnnotations, ok := p.Value.(map[string]string); ok { @@ -436,15 +469,17 @@ func (c *AdmissionController) updateApplicationInfo(namespace string, pod *v1.Po } } } - + // Add a new annotation patch if len(annotations) != 0 { + // combine new annotations with existing annotations in pod to create first patch + existingAnnotations := pod.Annotations + annotations = utils.MergeMaps(existingAnnotations, annotations) patch = append(patch, common.PatchOperation{ Op: "add", Path: "/metadata/annotations", Value: annotations, }) } - return patch } diff --git a/pkg/admission/admission_controller_test.go b/pkg/admission/admission_controller_test.go index f38d37200..5a0cbb76e 100644 --- a/pkg/admission/admission_controller_test.go +++ b/pkg/admission/admission_controller_test.go @@ -53,7 +53,7 @@ const ( // nolint: funlen func TestUpdateApplicationInfo(t *testing.T) { // verify when appId/queue are not given, - // we patch it correctly + // we patch it correctly. (Add appId/queueName/disableStateAware to labels and annotations) var patch []common.PatchOperation pod := &v1.Pod{ @@ -66,31 +66,30 @@ func TestUpdateApplicationInfo(t *testing.T) { Namespace: "default", UID: "7f5fd6c5d5", ResourceVersion: "10654", - Labels: map[string]string{ - "random": "random", - }, }, Spec: v1.PodSpec{}, Status: v1.PodStatus{}, } c := createAdmissionControllerForTest() - patch = c.updateApplicationInfo("default", pod, patch) - assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") - assert.Equal(t, patch[0].Path, "/metadata/annotations") - if updatedMap, ok := patch[0].Value.(map[string]string); ok { - assert.Equal(t, len(updatedMap), 3) - assert.Equal(t, updatedMap["yunikorn.apache.org/queue"], "root.default") - assert.Equal(t, updatedMap["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, strings.HasPrefix(updatedMap["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("patch info content is not as expected") - } + patch = c.updateApplicationInfo("default", pod, patch) - // verify if applicationId is given in the labels, - // we won't modify it + labelsInPatch := fetchPatchValues("add", "/metadata/labels", patch) + annotationsInPatch := fetchPatchValues("add", "/metadata/annotations", patch) + + assert.Equal(t, len(patch), 2) + assert.Equal(t, len(labelsInPatch), 3) + assert.Equal(t, labelsInPatch["queue"], "root.default") + assert.Equal(t, labelsInPatch["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(labelsInPatch["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(annotationsInPatch), 3) + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(annotationsInPatch["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) + + // verify if applicationId is given in the pod's labels, + // we won't modify it in admission controller, and will patch the value to annotations patch = make([]common.PatchOperation, 0) pod = &v1.Pod{ @@ -104,7 +103,6 @@ func TestUpdateApplicationInfo(t *testing.T) { UID: "7f5fd6c5d5", ResourceVersion: "10654", Labels: map[string]string{ - "random": "random", "applicationId": "app-0001", }, }, @@ -113,19 +111,20 @@ func TestUpdateApplicationInfo(t *testing.T) { } patch = c.updateApplicationInfo("default", pod, patch) - assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") - assert.Equal(t, patch[0].Path, "/metadata/annotations") //nolint:gosec - if updatedMap, ok := patch[0].Value.(map[string]string); ok { - assert.Equal(t, len(updatedMap), 2) - assert.Equal(t, updatedMap["yunikorn.apache.org/queue"], "root.default") - assert.Equal(t, updatedMap["yunikorn.apache.org/app-id"], "app-0001") - } else { - t.Fatal("patch info content is not as expected") - } + labelsInPatch = fetchPatchValues("add", "/metadata/labels", patch) + annotationsInPatch = fetchPatchValues("add", "/metadata/annotations", patch) - // verify if queue is given in the labels, - // we won't modify it + assert.Equal(t, len(patch), 2) + assert.Equal(t, len(labelsInPatch), 2) + assert.Equal(t, labelsInPatch["queue"], "root.default") + assert.Equal(t, labelsInPatch["applicationId"], "app-0001") + assert.Equal(t, len(patch), 2) + assert.Equal(t, len(annotationsInPatch), 2) + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/app-id"], "app-0001") + + // verify if queue is given in the pod's labels, + // we won't modify it in admission controller, and will patch the value to annotations patch = make([]common.PatchOperation, 0) pod = &v1.Pod{ @@ -139,8 +138,7 @@ func TestUpdateApplicationInfo(t *testing.T) { UID: "7f5fd6c5d5", ResourceVersion: "10654", Labels: map[string]string{ - "random": "random", - "queue": "root.abc", + "queue": "root.abc", }, }, Spec: v1.PodSpec{}, @@ -149,20 +147,21 @@ func TestUpdateApplicationInfo(t *testing.T) { patch = c.updateApplicationInfo("default", pod, patch) - assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") - assert.Equal(t, patch[0].Path, "/metadata/annotations") //nolint:gosec - if updatedMap, ok := patch[0].Value.(map[string]string); ok { - assert.Equal(t, len(updatedMap), 3) - assert.Equal(t, updatedMap["yunikorn.apache.org/queue"], "root.abc") - assert.Equal(t, updatedMap["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, strings.HasPrefix(updatedMap["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("patch info content is not as expected") - } - - // namespace might be empty - // labels might be empty + labelsInPatch = fetchPatchValues("add", "/metadata/labels", patch) + annotationsInPatch = fetchPatchValues("add", "/metadata/annotations", patch) + + assert.Equal(t, len(patch), 2) + assert.Equal(t, len(labelsInPatch), 3) + assert.Equal(t, labelsInPatch["queue"], "root.abc") + assert.Equal(t, labelsInPatch["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(labelsInPatch["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(annotationsInPatch), 3) + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/queue"], "root.abc") + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(annotationsInPatch["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) + + // verify if applicationId is given in the pod's annotations, + // we won't modify it in admission controller, and will patch the value to labels patch = make([]common.PatchOperation, 0) pod = &v1.Pod{ @@ -172,116 +171,170 @@ func TestUpdateApplicationInfo(t *testing.T) { }, ObjectMeta: metav1.ObjectMeta{ Name: "a-test-pod", + Namespace: "default", UID: "7f5fd6c5d5", ResourceVersion: "10654", + Annotations: map[string]string{ + "yunikorn.apache.org/app-id": "app-0001", + }, }, Spec: v1.PodSpec{}, Status: v1.PodStatus{}, } - patch = c.updateApplicationInfo("default", pod, patch) - assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") - assert.Equal(t, patch[0].Path, "/metadata/annotations") //nolint:gosec - if updatedMap, ok := patch[0].Value.(map[string]string); ok { - assert.Equal(t, len(updatedMap), 3) - assert.Equal(t, updatedMap["yunikorn.apache.org/queue"], "root.default") - assert.Equal(t, updatedMap["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, strings.HasPrefix(updatedMap["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("patch info content is not as expected") - } + labelsInPatch = fetchPatchValues("add", "/metadata/labels", patch) + annotationsInPatch = fetchPatchValues("add", "/metadata/annotations", patch) - // pod name might be empty, it can comes from generatedName - patch = make([]common.PatchOperation, 0) + assert.Equal(t, len(patch), 2) + assert.Equal(t, len(labelsInPatch), 2) + assert.Equal(t, labelsInPatch["queue"], "root.default") + assert.Equal(t, labelsInPatch["applicationId"], "app-0001") + assert.Equal(t, len(patch), 2) + assert.Equal(t, len(annotationsInPatch), 2) + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/app-id"], "app-0001") - pod = &v1.Pod{ - TypeMeta: metav1.TypeMeta{ - Kind: "Pod", - APIVersion: "v1", + // if both annotation and labels patch existing + // we won't create new patch and will add new labels/annotation to the existing patch. + patch = make([]common.PatchOperation, 0) + patch = append(patch, common.PatchOperation{ + Op: "add", + Path: "/metadata/labels", + Value: map[string]string{ + "existingLabelKey": "existingLabelValue", }, - ObjectMeta: metav1.ObjectMeta{ - GenerateName: "some-pod-", + }) + patch = append(patch, common.PatchOperation{ + Op: "add", + Path: "/metadata/annotations", + Value: map[string]string{ + "existingAnnotationKey": "existingAnnotationValue", }, - Spec: v1.PodSpec{}, - Status: v1.PodStatus{}, - } + }) + pod = &v1.Pod{} patch = c.updateApplicationInfo("default", pod, patch) - assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") - assert.Equal(t, patch[0].Path, "/metadata/annotations") //nolint:gosec - if updatedMap, ok := patch[0].Value.(map[string]string); ok { - assert.Equal(t, len(updatedMap), 3) - assert.Equal(t, updatedMap["yunikorn.apache.org/queue"], "root.default") - assert.Equal(t, updatedMap["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, strings.HasPrefix(updatedMap["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("patch info content is not as expected") + labelsInPatch = fetchPatchValues("add", "/metadata/labels", patch) + annotationsInPatch = fetchPatchValues("add", "/metadata/annotations", patch) + + assert.Equal(t, len(labelsInPatch), 4) + assert.Equal(t, labelsInPatch["existingLabelKey"], "existingLabelValue") + assert.Equal(t, labelsInPatch["queue"], "root.default") + assert.Equal(t, labelsInPatch["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(labelsInPatch["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(annotationsInPatch), 4) + assert.Equal(t, annotationsInPatch["existingAnnotationKey"], "existingAnnotationValue") + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(annotationsInPatch["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) +} + +func TestUpdateLabelPatch(t *testing.T) { + pod := &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + "label_key_in_pod": "label_value_in_pod", + }, + }, } - // pod name and generate name could be both empty - patch = make([]common.PatchOperation, 0) + // Verify if no patch on labels in patch list, + // create a new patch with pods's existing labels and put new labels into the patch + patch := make([]common.PatchOperation, 0) + newLabels := map[string]string{ + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + } + patch = updateLabelPatch(pod, newLabels, patch) + result := fetchPatchValues("add", "/metadata/labels", patch) - pod = &v1.Pod{ - TypeMeta: metav1.TypeMeta{ - Kind: "Pod", - APIVersion: "v1", + assert.Equal(t, len(result), 3) + assert.Equal(t, result["label_key_in_pod"], "label_value_in_pod") + assert.Equal(t, result["new_key_1"], "new_value_1") + assert.Equal(t, result["new_key_2"], "new_value_2") + + // Verify if have patch on labels in patch list, + // won't crete new patch and will add new labels into the patch + patch = make([]common.PatchOperation, 0) + patch = append(patch, common.PatchOperation{ + Op: "add", + Path: "/metadata/labels", + Value: map[string]string{ + "label_key_in_patch": "label_value_in_patch", }, - ObjectMeta: metav1.ObjectMeta{}, - Spec: v1.PodSpec{}, - Status: v1.PodStatus{}, + }) + newLabels = map[string]string{ + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", } - patch = c.updateApplicationInfo("default", pod, patch) + patch = updateLabelPatch(pod, newLabels, patch) + result = fetchPatchValues("add", "/metadata/labels", patch) - assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") - assert.Equal(t, patch[0].Path, "/metadata/annotations") //nolint:gosec - if updatedMap, ok := patch[0].Value.(map[string]string); ok { - assert.Equal(t, len(updatedMap), 3) - assert.Equal(t, updatedMap["yunikorn.apache.org/queue"], "root.default") - assert.Equal(t, updatedMap["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, strings.HasPrefix(updatedMap["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("patch info content is not as expected") + assert.Equal(t, len(result), 3) + assert.Equal(t, result["label_key_in_patch"], "label_value_in_patch") + assert.Equal(t, result["new_key_1"], "new_value_1") + assert.Equal(t, result["new_key_2"], "new_value_2") +} + +func TestUpdateAnnotationPatch(t *testing.T) { + pod := &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Annotations: map[string]string{ + "annotation_key_in_pod": "annotation_value_in_pod", + }, + }, + } + + // Verify if no patch on annotations in patch list, + // create a new patch with pods's existing annotations and put new annotation into the patch + patch := make([]common.PatchOperation, 0) + newAnnotations := map[string]string{ + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", } + patch = updateAnnotationPatch(pod, newAnnotations, patch) + result := fetchPatchValues("add", "/metadata/annotations", patch) - // should combine with the existing patch, which has the 'add' operation and the '/metadata/annotations' path. + assert.Equal(t, len(result), 3) + assert.Equal(t, result["annotation_key_in_pod"], "annotation_value_in_pod") + assert.Equal(t, result["new_key_1"], "new_value_1") + assert.Equal(t, result["new_key_2"], "new_value_2") + + // Verify if have patch on annotation in patch list, + // won't crete new patch and will add new annotation into the patch patch = make([]common.PatchOperation, 0) patch = append(patch, common.PatchOperation{ Op: "add", Path: "/metadata/annotations", Value: map[string]string{ - "existingAnnotationKey": "existingAnnotationValue", + "annotation_key_in_patch": "annotation_value_in_patch", }, }) - pod = &v1.Pod{ - TypeMeta: metav1.TypeMeta{ - Kind: "Pod", - APIVersion: "v1", - }, - ObjectMeta: metav1.ObjectMeta{}, - Spec: v1.PodSpec{}, - Status: v1.PodStatus{}, + newAnnotations = map[string]string{ + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", } - patch = c.updateApplicationInfo("default", pod, patch) + patch = updateAnnotationPatch(pod, newAnnotations, patch) + result = fetchPatchValues("add", "/metadata/annotations", patch) - assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") //nolint:gosec - assert.Equal(t, patch[0].Path, "/metadata/annotations") //nolint:gosec - if updatedMap, ok := patch[0].Value.(map[string]string); ok { //nolint:gosec - assert.Equal(t, len(updatedMap), 4) - assert.Equal(t, updatedMap["existingAnnotationKey"], "existingAnnotationValue") - assert.Equal(t, updatedMap["yunikorn.apache.org/queue"], "root.default") - assert.Equal(t, updatedMap["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, strings.HasPrefix(updatedMap["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("patch info content is not as expected") + assert.Equal(t, len(result), 3) + assert.Equal(t, result["annotation_key_in_patch"], "annotation_value_in_patch") + assert.Equal(t, result["new_key_1"], "new_value_1") + assert.Equal(t, result["new_key_2"], "new_value_2") +} + +func fetchPatchValues(op string, path string, patch []common.PatchOperation) map[string]string { + // fetch an existing patch by op and path + for _, p := range patch { + if p.Op == op && p.Path == path { + return p.Value.(map[string]string) + } } + return make(map[string]string) } func TestUpdateSchedulerName(t *testing.T) { @@ -882,6 +935,16 @@ func schedulerName(t *testing.T, patch []byte) string { return "" } +func labels(t *testing.T, patch []byte) map[string]interface{} { + ops := parsePatch(t, patch) + for _, op := range ops { + if op.Path == "/metadata/labels" { + return op.Value.(map[string]interface{}) + } + } + return make(map[string]interface{}) +} + func annotations(t *testing.T, patch []byte) map[string]interface{} { ops := parsePatch(t, patch) for _, op := range ops { diff --git a/pkg/admission/util.go b/pkg/admission/util.go index d70f521e3..dc2b0347a 100644 --- a/pkg/admission/util.go +++ b/pkg/admission/util.go @@ -31,13 +31,17 @@ import ( "github.com/apache/yunikorn-k8shim/pkg/log" ) -func getAnnotationsForApplicationInfoUpdate(pod *v1.Pod, namespace string, generateUniqueAppIds bool, defaultQueueName string) map[string]string { - result := make(map[string]string) - existingAnnotations := pod.Annotations - for k, v := range existingAnnotations { - result[k] = v - } +func getNewApplicationInfo(pod *v1.Pod, namespace string, generateUniqueAppIds bool, defaultQueueName string) (map[string]string, map[string]string) { + newLabels := make(map[string]string) + newAnnotations := make(map[string]string) + appID := getApplicationIDFromPod(pod) + disableStateAware := getDisableStateAwareFromPod(pod) + + // for undefined configuration, am_conf will add 'root.default' as default queue name + // if a custom name is configured for default queue, it will be used instead of root.default + queueName := utils.GetQueueNameFromPod(pod, defaultQueueName) + if appID == "" { // if app id not exist, generate one // for each namespace, we group unnamed pods to one single app - if GenerateUniqueAppId is not set @@ -46,34 +50,60 @@ func getAnnotationsForApplicationInfoUpdate(pod *v1.Pod, namespace string, gener // else // application ID convention: ${AUTO_GEN_PREFIX}-${NAMESPACE}-${AUTO_GEN_SUFFIX} appID = generateAppID(namespace, generateUniqueAppIds) - result[constants.AnnotationApplicationID] = appID // if we generate an app ID, disable state-aware scheduling for this app - disableStateAware := "true" - if value := utils.GetPodLabelValue(pod, constants.LabelDisableStateAware); value != "" { - disableStateAware = value - } - result[constants.AnnotationDisableStateAware] = disableStateAware - } else { - // if appID was not in pod annotations, add it to annotations. - if value := utils.GetPodAnnotationValue(pod, constants.AnnotationApplicationID); value == "" { - result[constants.AnnotationApplicationID] = appID + // skip it if disableStateAware has already been set in the pod + if disableStateAware == "" { + disableStateAware = "true" } } - // for undefined configuration, am_conf will add 'root.default' as default queue name - // if a custom name is configured for default queue, it will be used instead of root.default - queueName := utils.GetQueueNameFromPod(pod, defaultQueueName) + // we're looking forward to deprecate the labels and move everything to annotations + // but for backforward, we still add to labels + newLabels = updateLabelIfNotPresentInPod(pod, newLabels, constants.LabelApplicationID, appID) + newAnnotations = updateAnnotationIfNotPresentInPod(pod, newAnnotations, constants.AnnotationApplicationID, appID) - // if queue name was not in pod annotations, add it to annotations. - if value := utils.GetPodAnnotationValue(pod, constants.AnnotationQueueName); value == "" { - // if defaultQueueName is "", skip adding default queue name to the pod annotation - if queueName != "" { - result[constants.AnnotationQueueName] = queueName - } + // skip adding empty disableStateAware to the pod + if disableStateAware != "" { + newLabels = updateLabelIfNotPresentInPod(pod, newLabels, constants.LabelDisableStateAware, disableStateAware) + newAnnotations = updateAnnotationIfNotPresentInPod(pod, newAnnotations, constants.AnnotationDisableStateAware, disableStateAware) } - return result + // skip adding empty queue name to the pod + if queueName != "" { + newLabels = updateLabelIfNotPresentInPod(pod, newLabels, constants.LabelQueueName, queueName) + newAnnotations = updateAnnotationIfNotPresentInPod(pod, newAnnotations, constants.AnnotationQueueName, queueName) + } + + return newLabels, newAnnotations +} + +func updateLabelIfNotPresentInPod(pod *v1.Pod, labelMap map[string]string, label string, value string) map[string]string { + existingLabels := pod.Labels + + if existingLabels == nil { + labelMap[label] = value + } + + // skip update label if it has already been set in the pod + if _, ok := existingLabels[label]; !ok { + labelMap[label] = value + } + return labelMap +} + +func updateAnnotationIfNotPresentInPod(pod *v1.Pod, annotationMap map[string]string, annotation string, value string) map[string]string { + existingAnnotations := pod.Annotations + + if existingAnnotations == nil { + annotationMap[annotation] = value + } + + // skip update annotation if it has already been set in the pod + if _, ok := existingAnnotations[annotation]; !ok { + annotationMap[annotation] = value + } + return annotationMap } func updatePodAnnotation(pod *v1.Pod, key string, value string) map[string]string { @@ -127,3 +157,13 @@ func getApplicationIDFromPod(pod *v1.Pod) string { } return "" } + +func getDisableStateAwareFromPod(pod *v1.Pod) string { + // if existing annotation exist, it takes priority over everything else + if value := utils.GetPodAnnotationValue(pod, constants.AnnotationDisableStateAware); value != "" { + return value + } else if value := utils.GetPodLabelValue(pod, constants.LabelDisableStateAware); value != "" { + return value + } + return "" +} diff --git a/pkg/admission/util_test.go b/pkg/admission/util_test.go index 947fa0519..3cc065782 100644 --- a/pkg/admission/util_test.go +++ b/pkg/admission/util_test.go @@ -107,6 +107,15 @@ func createTestingPodWithAnnotationQueue() *v1.Pod { return pod } +func createTestingPodWithLabelEnableStateAware() *v1.Pod { + pod := createTestingPodWithMeta() + pod.ObjectMeta.Labels = map[string]string{ + constants.LabelDisableStateAware: "false", + } + + return pod +} + func createTestingPodNoNamespaceAndLabels() *v1.Pod { pod := createMinimalTestingPod() pod.ObjectMeta = @@ -118,147 +127,223 @@ func createTestingPodNoNamespaceAndLabels() *v1.Pod { return pod } -func TestGetAnnotationsForApplicationInfoUpdate(t *testing.T) { - // verify when appId/queue are not given, +func TestGetNewApplicationInfo(t *testing.T) { + // verify when appId/queue are not given, will generate it pod := createTestingPodWithMeta() + newLabels, newAnnotations := getNewApplicationInfo(pod, "default", false, "root.default") - if result := getAnnotationsForApplicationInfoUpdate(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["yunikorn.apache.org/queue"], "root.default") - assert.Equal(t, result["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, strings.HasPrefix(result["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("getAnnotationsForApplicationInfoUpdate is not as expected") - } + assert.Equal(t, len(newLabels), 3) + assert.Equal(t, newLabels["queue"], "root.default") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) // verify if applicationId is given in the labels, - // we'll write to annotation + // will copy applicationId to annotation and won't change existing value pod = createTestingPodWithLabelAppId() + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") - if result := getAnnotationsForApplicationInfoUpdate(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 2) - assert.Equal(t, result["yunikorn.apache.org/queue"], "root.default") - assert.Equal(t, result["yunikorn.apache.org/app-id"], "app-0001") - } else { - t.Fatal("getAnnotationsForApplicationInfoUpdate is not as expected") - } - - // if appId existing in pod annotation, will keep it - pod = createTestingPodWithAnnotationAppId() - if result := getAnnotationsForApplicationInfoUpdate(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 2) - assert.Equal(t, result["yunikorn.apache.org/queue"], "root.default") - assert.Equal(t, result["yunikorn.apache.org/app-id"], "app-0001") - } else { - t.Fatal("getAnnotationsForApplicationInfoUpdate is not as expected") - } + assert.Equal(t, len(newLabels), 1) + assert.Equal(t, newLabels["queue"], "root.default") + assert.Equal(t, newLabels["disableStateAware"], "") + assert.Equal(t, newLabels["applicationId"], "") + assert.Equal(t, len(newAnnotations), 2) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, newAnnotations["yunikorn.apache.org/app-id"], "app-0001") // verify if queue is given in the labels, - // we'll write to annotation + // will copy queue to annotation and won't change existing value pod = createTestingPodWithLabelQueue() - if result := getAnnotationsForApplicationInfoUpdate(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["yunikorn.apache.org/queue"], "root.abc") - assert.Equal(t, result["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, strings.HasPrefix(result["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("getAnnotationsForApplicationInfoUpdate is not as expected") - } - - // if queue existing in pod annotation, will keep it + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") + + assert.Equal(t, len(newLabels), 2) + assert.Equal(t, newLabels["queue"], "") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.abc") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) + + // verify if applicationId is given in the annotation, + // will copy to label and won't change existing value + pod = createTestingPodWithAnnotationAppId() + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") + assert.Equal(t, len(newLabels), 2) + assert.Equal(t, newLabels["applicationId"], "app-0001") + assert.Equal(t, newLabels["queue"], "root.default") + assert.Equal(t, newLabels["disableStateAware"], "") + assert.Equal(t, len(newAnnotations), 1) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, newAnnotations["yunikorn.apache.org/app-id"], "") + + // verify if queue is given in the annotation + // will copy to labels and won't change existing value pod = createTestingPodWithAnnotationQueue() - if result := getAnnotationsForApplicationInfoUpdate(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["yunikorn.apache.org/queue"], "root.abc") - assert.Equal(t, result["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, strings.HasPrefix(result["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("getAnnotationsForApplicationInfoUpdate is not as expected") - } - - // namespace might be empty - // labels might be empty - pod = createTestingPodNoNamespaceAndLabels() - - if result := getAnnotationsForApplicationInfoUpdate(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["yunikorn.apache.org/queue"], "root.default") - assert.Equal(t, result["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, strings.HasPrefix(result["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("getAnnotationsForApplicationInfoUpdate is not as expected") - } - - // pod name might be empty, it can comes from generatedName + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") + assert.Equal(t, len(newLabels), 3) + assert.Equal(t, newLabels["queue"], "root.abc") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(newAnnotations), 2) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) + + // verify if state aware label set in pod, + // won't change state aware setting and will copy to annotation + pod = createTestingPodWithLabelEnableStateAware() + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") + + assert.Equal(t, len(newLabels), 2) + assert.Equal(t, newLabels["queue"], "root.default") + assert.Equal(t, newLabels["disableStateAware"], "") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "false") + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) + + // verify empty pod info and default queue name is "" + // won't add queue name to pod + pod = createTestingPodWithMeta() + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "") + + assert.Equal(t, len(newLabels), 2) + assert.Equal(t, newLabels["queue"], "") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(newAnnotations), 2) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) + + // pod name might be empty, applicationId can come from generatedName pod = createTestingPodWithGenerateName() - if result := getAnnotationsForApplicationInfoUpdate(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["yunikorn.apache.org/queue"], "root.default") - assert.Equal(t, result["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, strings.HasPrefix(result["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("getAnnotationsForApplicationInfoUpdate is not as expected") - } + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) pod = createMinimalTestingPod() - if result := getAnnotationsForApplicationInfoUpdate(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["yunikorn.apache.org/queue"], "root.default") - assert.Equal(t, result["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, strings.HasPrefix(result["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("getAnnotationsForApplicationInfoUpdate is not as expected") - } + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") + assert.Equal(t, len(newLabels), 3) + assert.Equal(t, newLabels["queue"], "root.default") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) } func TestDefaultQueueName(t *testing.T) { defaultConf := createConfig() pod := createTestingPodWithMeta() - if result := getAnnotationsForApplicationInfoUpdate(pod, defaultConf.GetNamespace(), defaultConf.GetGenerateUniqueAppIds(), defaultConf.GetDefaultQueueName()); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["yunikorn.apache.org/app-id"], "yunikorn-default-autogen") - assert.Equal(t, result["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, result["yunikorn.apache.org/queue"], "root.default") - } else { - t.Fatal("getAnnotationsForApplicationInfoUpdate is not as expected") - } + newLabels, newAnnotations := getNewApplicationInfo(pod, defaultConf.GetNamespace(), defaultConf.GetGenerateUniqueAppIds(), defaultConf.GetDefaultQueueName()) + + assert.Equal(t, len(newLabels), 3) + assert.Equal(t, newLabels["applicationId"], "yunikorn-default-autogen") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, newLabels["queue"], "root.default") + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/app-id"], "yunikorn-default-autogen") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.default") queueNameEmptyConf := createConfigWithOverrides(map[string]string{ conf.AMFilteringDefaultQueueName: "", }) - if result := getAnnotationsForApplicationInfoUpdate(pod, queueNameEmptyConf.GetNamespace(), queueNameEmptyConf.GetGenerateUniqueAppIds(), queueNameEmptyConf.GetDefaultQueueName()); result != nil { - assert.Equal(t, len(result), 2) - assert.Equal(t, result["yunikorn.apache.org/app-id"], "yunikorn-default-autogen") - assert.Equal(t, result["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, result["yunikorn.apache.org/queue"], "") - } else { - t.Fatal("getAnnotationsForApplicationInfoUpdate is not as expected") - } + newLabels, newAnnotations = getNewApplicationInfo(pod, queueNameEmptyConf.GetNamespace(), queueNameEmptyConf.GetGenerateUniqueAppIds(), queueNameEmptyConf.GetDefaultQueueName()) + assert.Equal(t, len(newLabels), 2) + assert.Equal(t, newLabels["applicationId"], "yunikorn-default-autogen") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, newLabels["queue"], "") + assert.Equal(t, len(newAnnotations), 2) + assert.Equal(t, newAnnotations["yunikorn.apache.org/app-id"], "yunikorn-default-autogen") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "") customQueueNameConf := createConfigWithOverrides(map[string]string{ conf.AMFilteringDefaultQueueName: "yunikorn", }) - if result := getAnnotationsForApplicationInfoUpdate(pod, customQueueNameConf.GetNamespace(), customQueueNameConf.GetGenerateUniqueAppIds(), customQueueNameConf.GetDefaultQueueName()); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["yunikorn.apache.org/app-id"], "yunikorn-default-autogen") - assert.Equal(t, result["yunikorn.apache.org/disable-state-aware"], "true") - assert.Assert(t, result["yunikorn.apache.org/queue"] != "yunikorn") - } else { - t.Fatal("getAnnotationsForApplicationInfoUpdate is not as expected") - } + newLabels, newAnnotations = getNewApplicationInfo(pod, customQueueNameConf.GetNamespace(), customQueueNameConf.GetGenerateUniqueAppIds(), customQueueNameConf.GetDefaultQueueName()) + assert.Equal(t, len(newLabels), 3) + assert.Equal(t, newLabels["applicationId"], "yunikorn-default-autogen") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Assert(t, newLabels["queue"] != "yunikorn") + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/app-id"], "yunikorn-default-autogen") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Assert(t, newAnnotations["yunikorn.apache.org/queue"] != "yunikorn") customValidQueueNameConf := createConfigWithOverrides(map[string]string{ conf.AMFilteringDefaultQueueName: "root.yunikorn", }) - if result := getAnnotationsForApplicationInfoUpdate(pod, customValidQueueNameConf.GetNamespace(), - customValidQueueNameConf.GetGenerateUniqueAppIds(), customValidQueueNameConf.GetDefaultQueueName()); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["yunikorn.apache.org/app-id"], "yunikorn-default-autogen") - assert.Equal(t, result["yunikorn.apache.org/disable-state-aware"], "true") - assert.Equal(t, result["yunikorn.apache.org/queue"], "root.yunikorn") - } else { - t.Fatal("getAnnotationsForApplicationInfoUpdate is not as expected") + newLabels, newAnnotations = getNewApplicationInfo(pod, customValidQueueNameConf.GetNamespace(), + customValidQueueNameConf.GetGenerateUniqueAppIds(), customValidQueueNameConf.GetDefaultQueueName()) + assert.Equal(t, len(newLabels), 3) + assert.Equal(t, newLabels["applicationId"], "yunikorn-default-autogen") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, newLabels["queue"], "root.yunikorn") + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/app-id"], "yunikorn-default-autogen") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.yunikorn") +} + +func TestUpdateLabelIfNotPresentInPod(t *testing.T) { + pod := createTestingPodWithMeta() + pod.ObjectMeta.Labels = map[string]string{ + "exist_key": "exist_value", + } + + // Verify updating label that exist in pod + result := make(map[string]string) + result = updateLabelIfNotPresentInPod(pod, result, "exist_key", "new_value") + assert.Equal(t, len(result), 0) + assert.Equal(t, result["exist_key"], "") + + // Verify updating label that not exists in pod + result = make(map[string]string) + result = updateLabelIfNotPresentInPod(pod, result, "new_key", "new_value") + assert.Equal(t, len(result), 1) + assert.Equal(t, result["new_key"], "new_value") + + // Verify if no label in pod + pod.ObjectMeta.Labels = nil + result = make(map[string]string) + result = updateLabelIfNotPresentInPod(pod, result, "new_key", "new_value") + assert.Equal(t, len(result), 1) + assert.Equal(t, result["new_key"], "new_value") +} + +func TestUpdateAnnotationIfNotPresentInPod(t *testing.T) { + pod := createTestingPodWithMeta() + pod.ObjectMeta.Annotations = map[string]string{ + "exist_key": "exist_value", } + + // Verify updating annotation that exist in pod + result := make(map[string]string) + result = updateAnnotationIfNotPresentInPod(pod, result, "exist_key", "new_value") + assert.Equal(t, len(result), 0) + assert.Equal(t, result["exist_key"], "") + + // Verify updating annotation that not exists in pod + result = make(map[string]string) + result = updateAnnotationIfNotPresentInPod(pod, result, "new_key", "new_value") + assert.Equal(t, len(result), 1) + assert.Equal(t, result["new_key"], "new_value") + + // Verify if no annotation in pod + pod.ObjectMeta.Labels = nil + result = make(map[string]string) + result = updateAnnotationIfNotPresentInPod(pod, result, "new_key", "new_value") + assert.Equal(t, len(result), 1) + assert.Equal(t, result["new_key"], "new_value") } func TestGenerateAppID(t *testing.T) { @@ -316,11 +401,11 @@ func TestGetApplicationIDFromPod(t *testing.T) { appName := []string{"app00001", "app00002"} // test empty pod - pod1 := &v1.Pod{} - assert.Equal(t, getApplicationIDFromPod(pod1), "") + pod := &v1.Pod{} + assert.Equal(t, getApplicationIDFromPod(pod), "") // test annotation take precedence over label - pod2 := &v1.Pod{ + pod = &v1.Pod{ ObjectMeta: metav1.ObjectMeta{ Labels: map[string]string{ constants.LabelApplicationID: appName[0], @@ -330,28 +415,44 @@ func TestGetApplicationIDFromPod(t *testing.T) { }, }, } - assert.Equal(t, getApplicationIDFromPod(pod2), appName[1]) + assert.Equal(t, getApplicationIDFromPod(pod), appName[1]) // test pod with label only - pod3 := &v1.Pod{ + pod = &v1.Pod{ ObjectMeta: metav1.ObjectMeta{ Labels: map[string]string{ constants.LabelApplicationID: appName[0], }, }, } - assert.Equal(t, getApplicationIDFromPod(pod3), appName[0]) + assert.Equal(t, getApplicationIDFromPod(pod), appName[0]) +} - // test pod with empty applicationID annotation - pod4 := &v1.Pod{ +func TestGetDisableStateAwareFromPod(t *testing.T) { + // test empty pod + pod := &v1.Pod{} + assert.Equal(t, getDisableStateAwareFromPod(pod), "") + + // test annotation take precedence over label + pod = &v1.Pod{ ObjectMeta: metav1.ObjectMeta{ Labels: map[string]string{ - constants.LabelApplicationID: appName[0], + constants.LabelDisableStateAware: "true", }, Annotations: map[string]string{ - constants.AnnotationApplicationID: "", + constants.AnnotationDisableStateAware: "false", + }, + }, + } + assert.Equal(t, getDisableStateAwareFromPod(pod), "false") + + // test pod with label only + pod = &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + constants.LabelDisableStateAware: "true", }, }, } - assert.Equal(t, getApplicationIDFromPod(pod4), appName[0]) + assert.Equal(t, getDisableStateAwareFromPod(pod), "true") }