From b388518d407e6c1579a0020cc1204453a0fe17f0 Mon Sep 17 00:00:00 2001
From: Joan Miquel Luque <joan.luque@gmail.com>
Date: Mon, 14 Apr 2025 13:36:21 +0200
Subject: [PATCH] feat(k8s): Use pod affinity when KubeScheduler is enabled
 #201 (#212)

Signed-off-by: Joan Miquel Luque Oliver <joan.luque@dynatrace.com>
---
 packages/k8s/src/k8s/index.ts | 35 ++++++++++++++++++++++++++++++-----
 1 file changed, 30 insertions(+), 5 deletions(-)

diff --git a/packages/k8s/src/k8s/index.ts b/packages/k8s/src/k8s/index.ts
index 120d3ad..d67ebf1 100644
--- a/packages/k8s/src/k8s/index.ts
+++ b/packages/k8s/src/k8s/index.ts
@@ -95,10 +95,12 @@ export async function createPod(
   appPod.spec.containers = containers
   appPod.spec.restartPolicy = 'Never'
 
-  if (!useKubeScheduler()) {
-    appPod.spec.nodeName = await getCurrentNodeName()
+  const nodeName = await getCurrentNodeName()
+  if (useKubeScheduler()) {
+    appPod.spec.affinity = await getPodAffinity(nodeName)
+  } else {
+    appPod.spec.nodeName = nodeName
   }
-
   const claimName = getVolumeClaimName()
   appPod.spec.volumes = [
     {
@@ -155,8 +157,11 @@ export async function createJob(
   job.spec.template.spec.containers = [container]
   job.spec.template.spec.restartPolicy = 'Never'
 
-  if (!useKubeScheduler()) {
-    job.spec.template.spec.nodeName = await getCurrentNodeName()
+  const nodeName = await getCurrentNodeName()
+  if (useKubeScheduler()) {
+    job.spec.template.spec.affinity = await getPodAffinity(nodeName)
+  } else {
+    job.spec.template.spec.nodeName = nodeName
   }
 
   const claimName = getVolumeClaimName()
@@ -540,6 +545,26 @@ async function getCurrentNodeName(): Promise<string> {
   return nodeName
 }
 
+async function getPodAffinity(nodeName: string): Promise<k8s.V1Affinity> {
+  const affinity = new k8s.V1Affinity()
+  affinity.nodeAffinity = new k8s.V1NodeAffinity()
+  affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution =
+    new k8s.V1NodeSelector()
+  affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms =
+    [
+      {
+        matchExpressions: [
+          {
+            key: 'kubernetes.io/hostname',
+            operator: 'In',
+            values: [nodeName]
+          }
+        ]
+      }
+    ]
+  return affinity
+}
+
 export function namespace(): string {
   if (process.env['ACTIONS_RUNNER_KUBERNETES_NAMESPACE']) {
     return process.env['ACTIONS_RUNNER_KUBERNETES_NAMESPACE']