From 3f2b615885031de3bf35ff4b6a931021e9ba96ee Mon Sep 17 00:00:00 2001
From: Dimitri Koshkin <dimitri.koshkin@nutanix.com>
Date: Thu, 15 Feb 2024 17:49:08 -0800
Subject: [PATCH] fix: restart Containerd when setting mirror config

---
 .../generic/mutation/mirrors/inject.go        |  4 ++++
 .../generic/mutation/mirrors/mirror.go        | 16 +++++++++++++
 .../mirrors/templates/containerd-restart.sh   | 17 +++++++++++++
 .../mirrors/tests/generate_patches.go         | 24 +++++++++++++++----
 4 files changed, 57 insertions(+), 4 deletions(-)
 create mode 100644 pkg/handlers/generic/mutation/mirrors/templates/containerd-restart.sh

diff --git a/pkg/handlers/generic/mutation/mirrors/inject.go b/pkg/handlers/generic/mutation/mirrors/inject.go
index 0ef5c66a2..741dba59e 100644
--- a/pkg/handlers/generic/mutation/mirrors/inject.go
+++ b/pkg/handlers/generic/mutation/mirrors/inject.go
@@ -191,6 +191,10 @@ func generateFilesAndCommands(
 	}
 	files = append(files, applyPatchesFile...)
 	commands = append(commands, applyPatchesCommand)
+	// generate Containerd restart script and command
+	restartFile, restartCommand := generateContainerdRestartScript()
+	files = append(files, restartFile...)
+	commands = append(commands, restartCommand)
 
 	return files, commands, err
 }
diff --git a/pkg/handlers/generic/mutation/mirrors/mirror.go b/pkg/handlers/generic/mutation/mirrors/mirror.go
index 70c6d459b..53297bd4a 100644
--- a/pkg/handlers/generic/mutation/mirrors/mirror.go
+++ b/pkg/handlers/generic/mutation/mirrors/mirror.go
@@ -28,6 +28,9 @@ const (
 	containerdPatchesDirOnRemote                = "/etc/containerd/cre.d"
 	containerdApplyPatchesScriptOnRemote        = "/etc/containerd/apply-patches.sh"
 	containerdApplyPatchesScriptOnRemoteCommand = "/bin/bash " + containerdApplyPatchesScriptOnRemote
+
+	containerdRestartScriptOnRemote        = "/etc/containerd/restart.sh"
+	containerdRestartScriptOnRemoteCommand = "/bin/bash " + containerdRestartScriptOnRemote
 )
 
 var (
@@ -47,6 +50,9 @@ var (
 
 	//go:embed templates/containerd-apply-patches.sh.gotmpl
 	containerdApplyConfigPatchesScript []byte
+
+	//go:embed templates/containerd-restart.sh
+	containerdRestartScript []byte
 )
 
 type mirrorConfig struct {
@@ -220,3 +226,13 @@ func generateContainerdApplyPatchesScript() ([]cabpkv1.File, string, error) {
 		},
 	}, containerdApplyPatchesScriptOnRemoteCommand, nil
 }
+
+func generateContainerdRestartScript() ([]cabpkv1.File, string) {
+	return []cabpkv1.File{
+		{
+			Path:        containerdRestartScriptOnRemote,
+			Content:     string(containerdRestartScript),
+			Permissions: "0700",
+		},
+	}, containerdRestartScriptOnRemoteCommand
+}
diff --git a/pkg/handlers/generic/mutation/mirrors/templates/containerd-restart.sh b/pkg/handlers/generic/mutation/mirrors/templates/containerd-restart.sh
new file mode 100644
index 000000000..fd5ad53df
--- /dev/null
+++ b/pkg/handlers/generic/mutation/mirrors/templates/containerd-restart.sh
@@ -0,0 +1,17 @@
+#!/bin/bash
+systemctl restart containerd
+
+if ! command -v crictl; then
+  echo "Command crictl is not available, will not wait for Containerd to be running"
+  exit
+fi
+
+SECONDS=0
+until crictl info; do
+  if ((SECONDS > 60)); then
+    echo "Containerd is not running. Giving up..."
+    exit 1
+  fi
+  echo "Containerd is not running yet. Waiting..."
+  sleep 5
+done
diff --git a/pkg/handlers/generic/mutation/mirrors/tests/generate_patches.go b/pkg/handlers/generic/mutation/mirrors/tests/generate_patches.go
index 117c37a3d..1654df3da 100644
--- a/pkg/handlers/generic/mutation/mirrors/tests/generate_patches.go
+++ b/pkg/handlers/generic/mutation/mirrors/tests/generate_patches.go
@@ -75,13 +75,17 @@ func TestGeneratePatches(
 						gomega.HaveKeyWithValue(
 							"path", "/etc/containerd/apply-patches.sh",
 						),
+						gomega.HaveKeyWithValue(
+							"path", "/etc/containerd/restart.sh",
+						),
 					),
 				},
 				{
 					Operation: "add",
 					Path:      "/spec/template/spec/kubeadmConfigSpec/preKubeadmCommands",
-					ValueMatcher: gomega.ContainElement(
+					ValueMatcher: gomega.ContainElements(
 						"/bin/bash /etc/containerd/apply-patches.sh",
+						"/bin/bash /etc/containerd/restart.sh",
 					),
 				},
 			},
@@ -120,13 +124,17 @@ func TestGeneratePatches(
 						gomega.HaveKeyWithValue(
 							"path", "/etc/containerd/apply-patches.sh",
 						),
+						gomega.HaveKeyWithValue(
+							"path", "/etc/containerd/restart.sh",
+						),
 					),
 				},
 				{
 					Operation: "add",
 					Path:      "/spec/template/spec/kubeadmConfigSpec/preKubeadmCommands",
-					ValueMatcher: gomega.ContainElement(
+					ValueMatcher: gomega.ContainElements(
 						"/bin/bash /etc/containerd/apply-patches.sh",
+						"/bin/bash /etc/containerd/restart.sh",
 					),
 				},
 			},
@@ -165,13 +173,17 @@ func TestGeneratePatches(
 						gomega.HaveKeyWithValue(
 							"path", "/etc/containerd/apply-patches.sh",
 						),
+						gomega.HaveKeyWithValue(
+							"path", "/etc/containerd/restart.sh",
+						),
 					),
 				},
 				{
 					Operation: "add",
 					Path:      "/spec/template/spec/preKubeadmCommands",
-					ValueMatcher: gomega.ContainElement(
+					ValueMatcher: gomega.ContainElements(
 						"/bin/bash /etc/containerd/apply-patches.sh",
+						"/bin/bash /etc/containerd/restart.sh",
 					),
 				},
 			},
@@ -218,13 +230,17 @@ func TestGeneratePatches(
 						gomega.HaveKeyWithValue(
 							"path", "/etc/containerd/apply-patches.sh",
 						),
+						gomega.HaveKeyWithValue(
+							"path", "/etc/containerd/restart.sh",
+						),
 					),
 				},
 				{
 					Operation: "add",
 					Path:      "/spec/template/spec/preKubeadmCommands",
-					ValueMatcher: gomega.ContainElement(
+					ValueMatcher: gomega.ContainElements(
 						"/bin/bash /etc/containerd/apply-patches.sh",
+						"/bin/bash /etc/containerd/restart.sh",
 					),
 				},
 			},