fix: support pod affinity (#308)

aquasecurity · Mar 11, 2024 · 8e927ab · 8e927ab
1 parent e511c6f
commit 8e927ab
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 4 deletions.
diff --git a/pkg/jobs/builder.go b/pkg/jobs/builder.go
@@ -51,6 +51,12 @@ func WithAnnotation(annotations map[string]string) JobOption {
 	}
 }
 
+func WithAffinity(affinity *corev1.Affinity) JobOption {
+	return func(j *JobBuilder) {
+		j.affinity = affinity
+	}
+}
+
 func WithTolerations(tolerations []corev1.Toleration) JobOption {
 	return func(j *JobBuilder) {
 		j.tolerations = tolerations
@@ -140,6 +146,7 @@ type JobBuilder struct {
 	podSecurityContext   *corev1.PodSecurityContext
 	securityContext      *corev1.SecurityContext
 	annotations          map[string]string
+	affinity             *corev1.Affinity
 	tolerations          []corev1.Toleration
 	priorityClassName    string
 	volumes              []corev1.Volume
@@ -191,6 +198,9 @@ func (b *JobBuilder) build() (*batchv1.Job, error) {
 	if len(b.serviceAccount) > 0 {
 		job.Spec.Template.Spec.ServiceAccountName = b.serviceAccount
 	}
+	if b.affinity != nil {
+		job.Spec.Template.Spec.Affinity = b.affinity
+	}
 	if len(b.tolerations) > 0 {
 		job.Spec.Template.Spec.Tolerations = b.tolerations
 	}

diff --git a/pkg/jobs/collector.go b/pkg/jobs/collector.go
@@ -45,6 +45,7 @@ type jobCollector struct {
 	podSecurityContext   *corev1.PodSecurityContext
 	securityContext      *corev1.SecurityContext
 	imageRef             string
+	affinity             *corev1.Affinity
 	tolerations          []corev1.Toleration
 	volumes              []corev1.Volume
 	volumeMounts         []corev1.VolumeMount
@@ -92,6 +93,12 @@ func WithPodPriorityClassName(priorityClassName string) CollectorOption {
 	}
 }
 
+func WithJobAffinity(affinity *corev1.Affinity) CollectorOption {
+	return func(jc *jobCollector) {
+		jc.affinity = affinity
+	}
+}
+
 func WithJobTolerations(tolerations []corev1.Toleration) CollectorOption {
 	return func(jc *jobCollector) {
 		jc.tolerations = tolerations
@@ -246,6 +253,7 @@ func (jb *jobCollector) ApplyAndCollect(ctx context.Context, nodeName string) (s
 		withSecurityContext(jb.securityContext),
 		withPodSecurityContext(jb.podSecurityContext),
 		WithNodeCollectorImageRef(jb.imageRef),
+		WithAffinity(jb.affinity),
 		WithTolerations(jb.tolerations),
 		WithPodVolumes(jb.volumes),
 		WithImagePullSecrets(jb.imagePullSecrets),
@@ -312,6 +320,7 @@ func (jb *jobCollector) Apply(ctx context.Context, nodeName string) (*batchv1.Jo
 		WithLabels(jb.labels),
 		withPodSecurityContext(jb.podSecurityContext),
 		withSecurityContext(jb.securityContext),
+		WithAffinity(jb.affinity),
 		WithTolerations(jb.tolerations),
 		WithJobServiceAccount(jb.serviceAccount),
 		WithJobTimeout(jb.collectorTimeout),

diff --git a/pkg/trivyk8s/trivyk8s.go b/pkg/trivyk8s/trivyk8s.go
@@ -169,17 +169,24 @@ func (c *client) ListArtifacts(ctx context.Context) ([]*artifacts.Artifact, erro
 }
 
 type scanJobParams struct {
-	toleration       []corev1.Toleration
+	affinity         *corev1.Affinity
+	tolerations      []corev1.Toleration
 	ignoreLabels     map[string]string
 	scanJobNamespace string
 	imageRef         string
 }
 
 type NodeCollectorOption func(*client)
 
-func WithTolerations(toleration []corev1.Toleration) NodeCollectorOption {
+func WithAffinity(affinity *corev1.Affinity) NodeCollectorOption {
 	return func(c *client) {
-		c.scanJobParams.toleration = toleration
+		c.scanJobParams.affinity = affinity
+	}
+}
+
+func WithTolerations(tolerations []corev1.Toleration) NodeCollectorOption {
+	return func(c *client) {
+		c.scanJobParams.tolerations = tolerations
 	}
 }
 
@@ -229,7 +236,8 @@ func (c *client) ListArtifactAndNodeInfo(ctx context.Context,
 		jobs.WithJobNamespace(c.scanJobParams.scanJobNamespace),
 		jobs.WithJobLabels(labels),
 		jobs.WithImageRef(c.scanJobParams.imageRef),
-		jobs.WithJobTolerations(c.scanJobParams.toleration),
+		jobs.WithJobAffinity(c.scanJobParams.affinity),
+		jobs.WithJobTolerations(c.scanJobParams.tolerations),
 		jobs.WithNodeConfig(c.nodeConfig),
 	)
 	// delete trivy namespace