From 50d49d9c5100651dd051f3bed9acc1b8f1555831 Mon Sep 17 00:00:00 2001
From: Yang Guo <ygg@google.com>
Date: Fri, 23 Jun 2017 11:38:17 -0700
Subject: [PATCH] Enables memcg notification in cluster/node e2e tests

---
 cluster/common.sh                   | 10 +++++++
 cluster/gce/config-test.sh          |  4 +--
 cluster/gce/gci/configure-helper.sh |  4 +--
 test/e2e_node/remote/node_e2e.go    | 44 ++++++++++++++++++-----------
 4 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/cluster/common.sh b/cluster/common.sh
index 502abfbb81..0bcc69685f 100755
--- a/cluster/common.sh
+++ b/cluster/common.sh
@@ -696,6 +696,16 @@ EOF
   if [ -n "${KUBELET_TEST_ARGS:-}" ]; then
       cat >>$file <<EOF
 KUBELET_TEST_ARGS: $(yaml-quote ${KUBELET_TEST_ARGS})
+EOF
+  fi
+  if [ -n "${NODE_KUBELET_TEST_ARGS:-}" ]; then
+      cat >>$file <<EOF
+NODE_KUBELET_TEST_ARGS: $(yaml-quote ${NODE_KUBELET_TEST_ARGS})
+EOF
+  fi
+  if [ -n "${MASTER_KUBELET_TEST_ARGS:-}" ]; then
+      cat >>$file <<EOF
+MASTER_KUBELET_TEST_ARGS: $(yaml-quote ${MASTER_KUBELET_TEST_ARGS})
 EOF
   fi
   if [ -n "${KUBELET_TEST_LOG_LEVEL:-}" ]; then
diff --git a/cluster/gce/config-test.sh b/cluster/gce/config-test.sh
index 8da62d0313..66b05a485a 100755
--- a/cluster/gce/config-test.sh
+++ b/cluster/gce/config-test.sh
@@ -150,10 +150,10 @@ TEST_CLUSTER_RESYNC_PERIOD="${TEST_CLUSTER_RESYNC_PERIOD:---min-resync-period=3m
 TEST_CLUSTER_API_CONTENT_TYPE="${TEST_CLUSTER_API_CONTENT_TYPE:-}"
 
 KUBELET_TEST_ARGS="${KUBELET_TEST_ARGS:-} --max-pods=110 --serialize-image-pulls=false --outofdisk-transition-frequency=0 ${TEST_CLUSTER_API_CONTENT_TYPE}"
-if [[ "${NODE_OS_DISTRIBUTION}" == "gci" ]]; then
+if [[ "${NODE_OS_DISTRIBUTION}" == "gci" ]] || [[ "${NODE_OS_DISTRIBUTION}" == "ubuntu" ]]; then
   NODE_KUBELET_TEST_ARGS=" --experimental-kernel-memcg-notification=true"
 fi
-if [[ "${MASTER_OS_DISTRIBUTION}" == "gci" ]]; then
+if [[ "${MASTER_OS_DISTRIBUTION}" == "gci" ]] || [[ "${MASTER_OS_DISTRIBUTION}" == "ubuntu" ]]; then
   MASTER_KUBELET_TEST_ARGS=" --experimental-kernel-memcg-notification=true"
 fi
 APISERVER_TEST_ARGS="${APISERVER_TEST_ARGS:-} --runtime-config=extensions/v1beta1 ${TEST_CLUSTER_DELETE_COLLECTION_WORKERS} ${TEST_CLUSTER_MAX_REQUESTS_INFLIGHT}"
diff --git a/cluster/gce/gci/configure-helper.sh b/cluster/gce/gci/configure-helper.sh
index b05637957b..31b0cf53f7 100644
--- a/cluster/gce/gci/configure-helper.sh
+++ b/cluster/gce/gci/configure-helper.sh
@@ -878,7 +878,7 @@ function start-kubelet {
     flags+=" --port=${KUBELET_PORT}"
   fi
   if [[ "${KUBERNETES_MASTER:-}" == "true" ]]; then
-    flags+="${MASTER_KUBELET_TEST_ARGS:-}"
+    flags+=" ${MASTER_KUBELET_TEST_ARGS:-}"
     flags+=" --enable-debugging-handlers=false"
     flags+=" --hairpin-mode=none"
     if [[ "${REGISTER_MASTER_KUBELET:-false}" == "true" ]]; then
@@ -893,7 +893,7 @@ function start-kubelet {
       flags+=" --pod-cidr=${MASTER_IP_RANGE}"
     fi
   else # For nodes
-    flags+="${NODE_KUBELET_TEST_ARGS:-}"
+    flags+=" ${NODE_KUBELET_TEST_ARGS:-}"
     flags+=" --enable-debugging-handlers=true"
     flags+=" --bootstrap-kubeconfig=/var/lib/kubelet/bootstrap-kubeconfig"
     flags+=" --require-kubeconfig"
diff --git a/test/e2e_node/remote/node_e2e.go b/test/e2e_node/remote/node_e2e.go
index 66b85deff4..3068cd5e32 100644
--- a/test/e2e_node/remote/node_e2e.go
+++ b/test/e2e_node/remote/node_e2e.go
@@ -108,19 +108,9 @@ func tarAddCOSMounter(tar string) error {
 	return nil
 }
 
-// updateCOSKubeletFlags updates kubelet flags to set gci mounter path, and enables memcg notifications. This will only take effect for
-// GCI/COS image.
-func updateCOSKubeletFlags(args, host, workspace string) (string, error) {
-	// Determine if tests will run on a GCI/COS node.
-	output, err := SSH(host, "cat", "/etc/os-release")
-	if err != nil {
-		return args, fmt.Errorf("issue detecting node's OS via node's /etc/os-release. Err: %v, Output:\n%s", err, output)
-	}
-	if !strings.Contains(output, "ID=gci") && !strings.Contains(output, "ID=cos") {
-		// This is not a GCI/COS image
-		return args, nil
-	}
-
+// prependCOSMounterFlag prepends the flag for setting the GCI mounter path to
+// args and returns the result.
+func prependCOSMounterFlag(args, host, workspace string) (string, error) {
 	// If we are testing on a GCI/COS node, we chmod 544 the mounter and specify a different mounter path in the test args.
 	// We do this here because the local var `workspace` tells us which /tmp/node-e2e-%d is relevant to the current test run.
 
@@ -140,16 +130,38 @@ func updateCOSKubeletFlags(args, host, workspace string) (string, error) {
 	// Note this implicitly requires the script to be where we expect in the tarball, so if that location changes the error
 	// here will tell us to update the remote test runner.
 	mounterPath := filepath.Join(workspace, localCOSMounterPath)
-	output, err = SSH(host, "sh", "-c", fmt.Sprintf("'chmod 544 %s'", mounterPath))
+	output, err := SSH(host, "sh", "-c", fmt.Sprintf("'chmod 544 %s'", mounterPath))
 	if err != nil {
 		return args, fmt.Errorf("unabled to chmod 544 GCI/COS mounter script. Err: %v, Output:\n%s", err, output)
 	}
 	// Insert args at beginning of test args, so any values from command line take precedence
-	args = "--kubelet-flags=--experimental-kernel-memcg-notification=true " + args
 	args = fmt.Sprintf("--kubelet-flags=--experimental-mounter-path=%s ", mounterPath) + args
 	return args, nil
 }
 
+// prependMemcgNotificationFlag prepends the flag for enabling memcg
+// notification to args and returns the result.
+func prependMemcgNotificationFlag(args string) string {
+	return "--kubelet-flags=--experimental-kernel-memcg-notification=true " + args
+}
+
+// updateOSSpecificKubeletFlags updates the Kubelet args with OS specific
+// settings.
+func updateOSSpecificKubeletFlags(args, host, workspace string) (string, error) {
+	output, err := SSH(host, "cat", "/etc/os-release")
+	if err != nil {
+		return "", fmt.Errorf("issue detecting node's OS via node's /etc/os-release. Err: %v, Output:\n%s", err, output)
+	}
+	switch {
+	case strings.Contains(output, "ID=gci"), strings.Contains(output, "ID=cos"):
+		args = prependMemcgNotificationFlag(args)
+		return prependCOSMounterFlag(args, host, workspace)
+	case strings.Contains(output, "ID=ubuntu"):
+		return prependMemcgNotificationFlag(args), nil
+	}
+	return args, nil
+}
+
 // RunTest runs test on the node.
 func (n *NodeE2ERemote) RunTest(host, workspace, results, junitFilePrefix, testArgs, ginkgoArgs string, timeout time.Duration) (string, error) {
 	// Install the cni plugins and add a basic CNI configuration.
@@ -165,7 +177,7 @@ func (n *NodeE2ERemote) RunTest(host, workspace, results, junitFilePrefix, testA
 	// Kill any running node processes
 	cleanupNodeProcesses(host)
 
-	testArgs, err := updateCOSKubeletFlags(testArgs, host, workspace)
+	testArgs, err := updateOSSpecificKubeletFlags(testArgs, host, workspace)
 	if err != nil {
 		return "", err
 	}