From e227934c880c4b95e1d9394962303aaae5a421a1 Mon Sep 17 00:00:00 2001
From: Abubakr-Sadik Nii Nai Davis <dwa2pac@gmail.com>
Date: Sun, 15 Oct 2017 12:39:29 +0000
Subject: [PATCH] Add function to get unit files for kubernetes components.

---
 cfg/config.yaml |  4 ++++
 cmd/common.go   |  2 ++
 cmd/util.go     | 31 +++++++++++++++++++++++++++++++
 3 files changed, 37 insertions(+)

diff --git a/cfg/config.yaml b/cfg/config.yaml
index e7c2e4f..8512c75 100644
--- a/cfg/config.yaml
+++ b/cfg/config.yaml
@@ -99,6 +99,10 @@ node:
       - /etc/kubernetes/kubelet.conf
       - /etc/kubernetes/kubelet 
     defaultconf: "/etc/kubernetes/kubelet.conf"
+
+    unitfiles:
+     - /etc/systemd/system/kubelet.service.d/10-kubeadm.conf
+    defaultunitfile: /etc/systemd/system/kubelet.service.d/10-kubeadm.conf
   
   proxy:
     bins:
diff --git a/cmd/common.go b/cmd/common.go
index e9cc610..2c22c9d 100644
--- a/cmd/common.go
+++ b/cmd/common.go
@@ -68,6 +68,7 @@ func runChecks(t check.NodeType) {
 	binmap := getBinaries(typeConf)
 	confmap := getConfigFiles(typeConf)
 	podspecmap := getPodSpecFiles(typeConf)
+	unitfilemap := getUnitFiles(typeConf)
 
 	switch t {
 	case check.MASTER:
@@ -90,6 +91,7 @@ func runChecks(t check.NodeType) {
 	s = makeSubstitutions(s, "bin", binmap)
 	s = makeSubstitutions(s, "conf", confmap)
 	s = makeSubstitutions(s, "podspec", podspecmap)
+	s = makeSubstitutions(s, "unitfile", unitfilemap)
 
 	glog.V(1).Info(fmt.Sprintf("Using config file: %s\n", viper.ConfigFileUsed()))
 	glog.V(1).Info(fmt.Sprintf("Using benchmark file: %s\n", path))
diff --git a/cmd/util.go b/cmd/util.go
index 5dd2303..1de375c 100644
--- a/cmd/util.go
+++ b/cmd/util.go
@@ -178,6 +178,37 @@ func getPodSpecFiles(v *viper.Viper) map[string]string {
 	return podspecmap
 }
 
+// getUnitFiles finds which of the set of candidate unit files exist
+func getUnitFiles(v *viper.Viper) map[string]string {
+	unitfilemap := make(map[string]string)
+
+	for _, component := range v.GetStringSlice("components") {
+		s := v.Sub(component)
+		if s == nil {
+			continue
+		}
+
+		// See if any of the candidate podspec files exist
+		unitfile := findConfigFile(s.GetStringSlice("unitfiles"))
+		if unitfile == "" {
+			if s.IsSet("defaultunitfile") {
+				unitfile = s.GetString("defaultunitfile")
+				glog.V(2).Info(fmt.Sprintf("Using default unit file name '%s' for component %s", unitfile, component))
+			} else {
+				// Default the config file name that we'll substitute to the name of the component
+				printlnWarn(fmt.Sprintf("Missing unit file for %s", component))
+				unitfile = component
+			}
+		} else {
+			glog.V(2).Info(fmt.Sprintf("Component %s uses unit file '%s'", component, unitfile))
+		}
+
+		unitfilemap[component] = unitfile
+	}
+
+	return unitfilemap
+}
+
 // verifyBin checks that the binary specified is running
 func verifyBin(bin string) bool {