From 83a1a24c1ec714163f9e71b14a06a72a1375c709 Mon Sep 17 00:00:00 2001
From: smuu <18609909+smuu@users.noreply.github.com>
Date: Fri, 13 Dec 2024 07:44:25 +0100
Subject: [PATCH] feat: add node selector to intance (#594)

Signed-off-by: Smuu <18609909+Smuu@users.noreply.github.com>
---
 pkg/instance/build.go     | 10 ++++++++++
 pkg/instance/errors.go    |  1 +
 pkg/instance/execution.go |  1 +
 pkg/k8s/pod.go            |  2 ++
 pkg/knuu/knuu.go          |  2 ++
 5 files changed, 16 insertions(+)

diff --git a/pkg/instance/build.go b/pkg/instance/build.go
index fecc13d..16714a5 100644
--- a/pkg/instance/build.go
+++ b/pkg/instance/build.go
@@ -25,6 +25,7 @@ type build struct {
 	env             map[string]string
 	imageCache      *sync.Map
 	buildDir        string
+	nodeSelector    map[string]string
 }
 
 func (i *Instance) Build() *build {
@@ -159,6 +160,15 @@ func (b *build) SetUser(user string) error {
 	return nil
 }
 
+func (b *build) SetNodeSelector(nodeSelector map[string]string) error {
+	if !b.instance.IsInState(StatePreparing, StateCommitted) {
+		return ErrSettingNodeSelectorNotAllowed.WithParams(b.instance.state.String())
+	}
+
+	b.nodeSelector = nodeSelector
+	return nil
+}
+
 // Commit commits the instance
 // This function can only be called in the state 'Preparing'
 func (b *build) Commit(ctx context.Context) error {
diff --git a/pkg/instance/errors.go b/pkg/instance/errors.go
index 6a87395..7dd6d6a 100644
--- a/pkg/instance/errors.go
+++ b/pkg/instance/errors.go
@@ -114,6 +114,7 @@ var (
 	ErrSrcDoesNotExistOrIsNotDirectory           = errors.New("SrcDoesNotExistOrIsNotDirectory", "src '%s' does not exist or is not a directory")
 	ErrCopyingFolderToInstance                   = errors.New("CopyingFolderToInstance", "error copying folder '%s' to instance '%s")
 	ErrSettingUserNotAllowed                     = errors.New("SettingUserNotAllowed", "setting user is only allowed in state 'Preparing'. Current state is '%s")
+	ErrSettingNodeSelectorNotAllowed             = errors.New("SettingNodeSelectorNotAllowed", "setting node selector is only allowed in state 'Preparing'. Current state is '%s")
 	ErrSettingUser                               = errors.New("SettingUser", "error setting user '%s' for instance '%s")
 	ErrCommittingNotAllowed                      = errors.New("CommittingNotAllowed", "committing is only allowed in state 'Preparing'. Current state is '%s")
 	ErrGettingImageRegistry                      = errors.New("GettingImageRegistry", "error getting image registry")
diff --git a/pkg/instance/execution.go b/pkg/instance/execution.go
index f8733a5..622c5fb 100644
--- a/pkg/instance/execution.go
+++ b/pkg/instance/execution.go
@@ -417,6 +417,7 @@ func (e *execution) prepareReplicaSetConfig() k8s.ReplicaSetConfig {
 		ServiceAccountName: e.instance.name,
 		ContainerConfig:    containerConfig,
 		SidecarConfigs:     sidecarConfigs,
+		NodeSelector:       e.instance.build.nodeSelector,
 	}
 
 	return k8s.ReplicaSetConfig{
diff --git a/pkg/k8s/pod.go b/pkg/k8s/pod.go
index a0ee9dd..035923a 100644
--- a/pkg/k8s/pod.go
+++ b/pkg/k8s/pod.go
@@ -62,6 +62,7 @@ type PodConfig struct {
 	ContainerConfig    ContainerConfig   // ContainerConfig for the Pod
 	SidecarConfigs     []ContainerConfig // SideCarConfigs for the Pod
 	Annotations        map[string]string // Annotations to apply to the Pod
+	NodeSelector       map[string]string // NodeSelector to apply to the Pod
 }
 
 type Volume struct {
@@ -623,6 +624,7 @@ func (c *Client) preparePodSpec(spec PodConfig, init bool) v1.PodSpec {
 		InitContainers:     c.prepareInitContainers(spec.ContainerConfig, init),
 		Containers:         []v1.Container{prepareContainer(spec.ContainerConfig)},
 		Volumes:            preparePodVolumes(spec.ContainerConfig),
+		NodeSelector:       spec.NodeSelector,
 	}
 
 	// Prepare sidecar containers and append to the pod spec
diff --git a/pkg/knuu/knuu.go b/pkg/knuu/knuu.go
index 2babcb0..4c4ac0c 100644
--- a/pkg/knuu/knuu.go
+++ b/pkg/knuu/knuu.go
@@ -109,6 +109,8 @@ func (k *Knuu) HandleStopSignal(ctx context.Context) {
 			k.Logger.Errorf("Error cleaning up resources with timeout handler: %v", err)
 		}
 		k.K8sClient.Terminate()
+		// Allow other signal handlers to run
+		signal.Reset(syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
 		os.Exit(ExitCodeSIGINT)
 	}()
 }