From c10af738aad1be9c4360b3c65c4c7afb5fe28968 Mon Sep 17 00:00:00 2001
From: "allen.wang" <allen.wq@alipay.com>
Date: Mon, 28 Jan 2019 20:35:32 +0800
Subject: [PATCH] support memory swap by cri

Signed-off-by: allen.wang <allen.wq@alipay.com>
---
 cri/annotations/annotations.go |  6 +++++
 cri/v1alpha2/cri_utils.go      | 25 ++++++++++++++++++
 cri/v1alpha2/cri_utils_test.go | 48 ++++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+)

diff --git a/cri/annotations/annotations.go b/cri/annotations/annotations.go
index 581aa0be46..c6cac757f1 100644
--- a/cri/annotations/annotations.go
+++ b/cri/annotations/annotations.go
@@ -22,4 +22,10 @@ const (
 
 	// LxcfsEnabled whether to enable lxcfs for a container
 	LxcfsEnabled = "io.kubernetes.lxcfs.enabled"
+
+	// ExtendAnnotationPrefix is the extend annotation prefix
+	ExtendAnnotationPrefix = "io.alibaba.pouch"
+
+	// MemorySwapExtendAnnotation is the extend annotation of memory swap
+	MemorySwapExtendAnnotation = "io.alibaba.pouch.resources.memory_swap"
 )
diff --git a/cri/v1alpha2/cri_utils.go b/cri/v1alpha2/cri_utils.go
index 6823a407c0..a393f72fe4 100644
--- a/cri/v1alpha2/cri_utils.go
+++ b/cri/v1alpha2/cri_utils.go
@@ -838,6 +838,13 @@ func (c *CriManager) updateCreateConfig(createConfig *apitypes.ContainerCreateCo
 		}
 	}
 
+	if len(config.Annotations) > 0 {
+		// Apply container config by annotation
+		if err := applyContainerConfigByAnnotation(config.Annotations, sandboxMeta.ID, &createConfig.ContainerConfig, createConfig.HostConfig); err != nil {
+			return fmt.Errorf("failed to apply container annotation for container %q: %v", config.Metadata.Name, err)
+		}
+	}
+
 	// Apply cgroupsParent derived from the sandbox config.
 	if lc := sandboxConfig.GetLinux(); lc != nil {
 		// Apply Cgroup options.
@@ -1258,3 +1265,21 @@ func toCNIPortMappings(criPortMappings []*runtime.PortMapping) []ocicni.PortMapp
 	}
 	return portMappings
 }
+
+// applyContainerConfigByAnnotation updates pouch container config according to annotation.
+func applyContainerConfigByAnnotation(annotations map[string]string, podSandboxID string, config *apitypes.ContainerConfig, hc *apitypes.HostConfig) error {
+	if len(annotations) == 0 {
+		return nil
+	}
+
+	if memorySwap, ok := annotations[anno.MemorySwapExtendAnnotation]; ok {
+		ms, err := strconv.ParseInt(memorySwap, 10, 64)
+		if err != nil {
+			return fmt.Errorf("failed to parse resources.memory_swap: %v", err)
+		}
+
+		hc.MemorySwap = ms
+	}
+
+	return nil
+}
diff --git a/cri/v1alpha2/cri_utils_test.go b/cri/v1alpha2/cri_utils_test.go
index dd6b29c2cc..5f080a2197 100644
--- a/cri/v1alpha2/cri_utils_test.go
+++ b/cri/v1alpha2/cri_utils_test.go
@@ -9,6 +9,7 @@ import (
 	"time"
 
 	apitypes "github.com/alibaba/pouch/apis/types"
+	anno "github.com/alibaba/pouch/cri/annotations"
 	runtime "github.com/alibaba/pouch/cri/apis/v1alpha2"
 	"github.com/alibaba/pouch/daemon/mgr"
 	"github.com/alibaba/pouch/pkg/utils"
@@ -1854,3 +1855,50 @@ func Test_toCNIPortMappings(t *testing.T) {
 		})
 	}
 }
+
+// CRI test: apply container config by annotation
+func Test_applyContainerConfigByAnnotation(t *testing.T) {
+	tests := []struct {
+		name          string
+		memorySwapStr string
+		memorySwap    int64
+		errMsg        string
+	}{
+		{
+			name:          "normalMemorySwapTest",
+			memorySwapStr: "200000000",
+			memorySwap:    200000000,
+			errMsg:        "",
+		},
+		{
+			name:          "errorMemorySwapTest",
+			memorySwapStr: "1g",
+			memorySwap:    0,
+			errMsg:        "failed to parse resources.memory_swap",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			annotations := map[string]string{
+				anno.MemorySwapExtendAnnotation: tt.memorySwapStr,
+			}
+
+			config := &apitypes.ContainerConfig{}
+			hc := &apitypes.HostConfig{}
+
+			err := applyContainerConfigByAnnotation(annotations, "", config, hc)
+			if tt.errMsg != "" {
+				assert.NotNil(t, err, "error should be %v", tt.errMsg)
+				if err != nil {
+					assert.Contains(t, err.Error(), tt.errMsg)
+				}
+			}
+
+			if tt.errMsg == "" {
+				assert.Nil(t, err)
+				assert.Equal(t, tt.memorySwap, hc.MemorySwap)
+			}
+		})
+	}
+}