From 3490e8af56b3a2fc8c4a16ee58c079cfaf89abe8 Mon Sep 17 00:00:00 2001
From: Ifeanyi Ubah <ifeanyi.ubah@embark-studios.com>
Date: Mon, 29 Jun 2020 16:28:48 +0200
Subject: [PATCH] Add option to use the same tolerations as agent (#9)

As a shorthand, in place of manually providing the same values
in each pipeline configuration as a jsonpatch object.
---
 README.md     |  4 ++++
 hooks/command | 14 +++++++++++---
 plugin.yml    |  2 ++
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 4ef8a62..d559ef1 100644
--- a/README.md
+++ b/README.md
@@ -195,6 +195,10 @@ Sets [memory request](https://kubernetes.io/docs/concepts/configuration/manage-c
 
 Sets [memory limit](https://kubernetes.io/docs/concepts/configuration/manage-compute-resources-container/) for the build container.
 
+### `use-agent-node-affinity` (optional, boolean)
+
+If set to `true`, the spawned jobs will use the same [node affinity](https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/) and [tolerations](https://kubernetes.io/docs/concepts/scheduling-eviction/taint-and-toleration/) as the buildkite agent.
+
 ### `workdir` (optional, string)
 
 Override the working directory to run the command in, inside the container. The default is the build directory where the buildkite bootstrap and git checkout runs.
diff --git a/hooks/command b/hooks/command
index ff9840b..42cd2dc 100755
--- a/hooks/command
+++ b/hooks/command
@@ -31,13 +31,21 @@ echo "--- :kubernetes: Starting Kubernetes Job"
 
 export patchFunc=${BUILDKITE_PLUGIN_K8S_PATCH:-"function(f) f"}
 
-jsonnet \
+jobspec=$(jsonnet \
   --tla-str "jobName=${job_name}" \
   --tla-str-file "stepEnvFile=${BUILDKITE_ENV_FILE}" \
   --tla-code "agentEnv=$(jq -c -n env)" \
   --tla-code patchFunc \
-  "${basedir}/lib/job.jsonnet" \
-  | kubectl apply -f -
+  "${basedir}/lib/job.jsonnet")
+
+if [[ "$BUILDKITE_PLUGIN_K8S_USE_AGENT_NODE_AFFINITY" == 'true' ]]; then
+  for field in affinity tolerations; do
+    buildkite_agent_value=$(kubectl get pod "$(cat /etc/hostname)" -o json | jq ".spec.$field")
+    jobspec=$(echo "$jobspec" | jq ".spec.template.spec.$field=$buildkite_agent_value")
+  done
+fi
+
+echo "$jobspec" | kubectl apply -f -
 
 echo "Timeout: ${timeout}s"
 
diff --git a/plugin.yml b/plugin.yml
index beea8fc..1b07259 100644
--- a/plugin.yml
+++ b/plugin.yml
@@ -57,6 +57,8 @@ configuration:
       type: string
     resources-limit-memory:
       type: string
+    use-agent-node-affinity:
+      type: boolean
   required:
   - image
   additionalProperties: false