Add post processor that will add a gang id label (#265) (#4029)

* Add post processor that will add a gang id label

* lint

* whitespace

Co-authored-by: Christopher Martin <[email protected]>

internal/server/configuration/types.go

@@ -93,6 +93,8 @@ type SubmissionConfig struct {
 	MaxOversubscriptionByResourceRequest map[string]float64
 	// Enforce that an init containers requestion non-integer cpu. This is due to https://github.com/kubernetes/kubernetes/issues/112228
 	AssertInitContainersRequestFractionalCpu bool
+	// Controls whether we add the gang id annotation as a label.
+	AddGangIdLabel bool
 }
 
 // TODO: we can probably just typedef this to map[string]string

internal/server/submit/conversion/post_process.go

@@ -23,6 +23,7 @@ var (
 	msgLevelProcessors = []msgProcessor{
 		templateMeta,
 		defaultGangNodeUniformityLabel,
+		addGangIdLabel,
 	}
 	podLevelProcessors = []podProcessor{
 		defaultActiveDeadlineSeconds,
@@ -166,6 +167,24 @@ func defaultGangNodeUniformityLabel(msg *armadaevents.SubmitJob, config configur
 	}
 }
 
+// Add a gangId label if the gangId annotation is set.  We do this because labels are much faster to search on than
+// annotations and a gang may want to hit the kubeapi to find its other gang members.
+func addGangIdLabel(msg *armadaevents.SubmitJob, config configuration.SubmissionConfig) {
+	if !config.AddGangIdLabel {
+		return
+	}
+
+	gangId := msg.GetObjectMeta().GetAnnotations()[configuration.GangIdAnnotation]
+	if gangId != "" {
+		labels := msg.GetObjectMeta().GetLabels()
+		if labels == nil {
+			labels = map[string]string{}
+		}
+		labels[configuration.GangIdAnnotation] = gangId
+		msg.GetObjectMeta().Labels = labels
+	}
+}
+
 // Templates the JobId in labels and annotations. This allows users to define labels and annotations containing the string
 // {JobId} and have it populated with the actual id of the job.
 func templateMeta(msg *armadaevents.SubmitJob, _ configuration.SubmissionConfig) {

internal/server/submit/conversion/post_process_test.go

@@ -584,6 +584,56 @@ func TestDefaultTerminationGracePeriod(t *testing.T) {
 	}
 }
 
+func TestAddGangIdLabel(t *testing.T) {
+	tests := map[string]struct {
+		annotations    map[string]string
+		initialLabels  map[string]string
+		expectedLabels map[string]string
+		enabled        bool
+	}{
+		"Unchanged if no gang id set": {
+			annotations: map[string]string{},
+			enabled:     true,
+		},
+		"Label added if gang id set": {
+			annotations: map[string]string{
+				configuration.GangIdAnnotation: "foo",
+			},
+			expectedLabels: map[string]string{
+				configuration.GangIdAnnotation: "foo",
+			},
+			enabled: true,
+		},
+		"Doesn't modify existing labels": {
+			annotations: map[string]string{
+				configuration.GangIdAnnotation: "foo",
+			},
+			initialLabels: map[string]string{
+				"fish": "chips",
+			},
+			expectedLabels: map[string]string{
+				"fish":                         "chips",
+				configuration.GangIdAnnotation: "foo",
+			},
+			enabled: true,
+		},
+		"Unchanged if disabled": {
+			annotations: map[string]string{
+				configuration.GangIdAnnotation: "foo",
+			},
+			enabled: false,
+		},
+	}
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			submitMsg := submitMsgFromAnnotations(tc.annotations)
+			submitMsg.ObjectMeta.Labels = tc.initialLabels
+			addGangIdLabel(submitMsg, configuration.SubmissionConfig{AddGangIdLabel: tc.enabled})
+			assert.Equal(t, tc.expectedLabels, submitMsg.ObjectMeta.Labels)
+		})
+	}
+}
+
 func submitMsgFromAnnotations(annotations map[string]string) *armadaevents.SubmitJob {
 	return &armadaevents.SubmitJob{
 		ObjectMeta: &armadaevents.ObjectMeta{