Add env cache.

Signed-off-by: Lantao Liu <lantaol@google.com>

pkg/server/container_create.go

@@ -35,7 +35,6 @@ import (
 	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/opencontainers/runc/libcontainer/devices"
 	runtimespec "github.com/opencontainers/runtime-spec/specs-go"
-	"github.com/opencontainers/runtime-tools/generate"
 	"github.com/opencontainers/runtime-tools/validate"
 	"github.com/opencontainers/selinux/go-selinux/label"
 	"github.com/pkg/errors"
@@ -518,7 +517,7 @@ func (c *criService) generateContainerMounts(sandboxID string, config *runtime.C
 
 // setOCIProcessArgs sets process args. It returns error if the final arg list
 // is empty.
-func setOCIProcessArgs(g *generate.Generator, config *runtime.ContainerConfig, imageConfig *imagespec.ImageConfig) error {
+func setOCIProcessArgs(g *generator, config *runtime.ContainerConfig, imageConfig *imagespec.ImageConfig) error {
 	command, args := config.GetCommand(), config.GetArgs()
 	// The following logic is migrated from https://github.com/moby/moby/blob/master/daemon/commit.go
 	// TODO(random-liu): Clearly define the commands overwrite behavior.
@@ -540,7 +539,7 @@ func setOCIProcessArgs(g *generate.Generator, config *runtime.ContainerConfig, i
 
 // addImageEnvs adds environment variables from image config. It returns error if
 // an invalid environment variable is encountered.
-func addImageEnvs(g *generate.Generator, imageEnvs []string) error {
+func addImageEnvs(g *generator, imageEnvs []string) error {
 	for _, e := range imageEnvs {
 		kv := strings.SplitN(e, "=", 2)
 		if len(kv) != 2 {
@@ -551,7 +550,7 @@ func addImageEnvs(g *generate.Generator, imageEnvs []string) error {
 	return nil
 }
 
-func setOCIPrivileged(g *generate.Generator, config *runtime.ContainerConfig) error {
+func setOCIPrivileged(g *generator, config *runtime.ContainerConfig) error {
 	// Add all capabilities in privileged mode.
 	g.SetupPrivileged(true)
 	setOCIBindMountsPrivileged(g)
@@ -572,7 +571,7 @@ func clearReadOnly(m *runtimespec.Mount) {
 }
 
 // addDevices set device mapping without privilege.
-func (c *criService) addOCIDevices(g *generate.Generator, devs []*runtime.Device) error {
+func (c *criService) addOCIDevices(g *generator, devs []*runtime.Device) error {
 	spec := g.Config
 	for _, device := range devs {
 		path, err := c.os.ResolveSymbolicLink(device.HostPath)
@@ -604,7 +603,7 @@ func (c *criService) addOCIDevices(g *generate.Generator, devs []*runtime.Device
 }
 
 // addDevices set device mapping with privilege.
-func setOCIDevicesPrivileged(g *generate.Generator) error {
+func setOCIDevicesPrivileged(g *generator) error {
 	spec := g.Config
 	hostDevices, err := devices.HostDevices()
 	if err != nil {
@@ -635,7 +634,7 @@ func setOCIDevicesPrivileged(g *generate.Generator) error {
 }
 
 // addOCIBindMounts adds bind mounts.
-func (c *criService) addOCIBindMounts(g *generate.Generator, mounts []*runtime.Mount, mountLabel string) error {
+func (c *criService) addOCIBindMounts(g *generator, mounts []*runtime.Mount, mountLabel string) error {
 	// Sort mounts in number of parts. This ensures that high level mounts don't
 	// shadow other mounts.
 	sort.Sort(orderedMounts(mounts))
@@ -740,7 +739,7 @@ func (c *criService) addOCIBindMounts(g *generate.Generator, mounts []*runtime.M
 	return nil
 }
 
-func setOCIBindMountsPrivileged(g *generate.Generator) {
+func setOCIBindMountsPrivileged(g *generator) {
 	spec := g.Config
 	// clear readonly for /sys and cgroup
 	for i, m := range spec.Mounts {
@@ -756,7 +755,7 @@ func setOCIBindMountsPrivileged(g *generate.Generator) {
 }
 
 // setOCILinuxResourceCgroup set container cgroup resource limit.
-func setOCILinuxResourceCgroup(g *generate.Generator, resources *runtime.LinuxContainerResources) {
+func setOCILinuxResourceCgroup(g *generator, resources *runtime.LinuxContainerResources) {
 	if resources == nil {
 		return
 	}
@@ -769,7 +768,7 @@ func setOCILinuxResourceCgroup(g *generate.Generator, resources *runtime.LinuxCo
 }
 
 // setOCILinuxResourceOOMScoreAdj set container OOMScoreAdj resource limit.
-func setOCILinuxResourceOOMScoreAdj(g *generate.Generator, resources *runtime.LinuxContainerResources, restrictOOMScoreAdjFlag bool) error {
+func setOCILinuxResourceOOMScoreAdj(g *generator, resources *runtime.LinuxContainerResources, restrictOOMScoreAdjFlag bool) error {
 	if resources == nil {
 		return nil
 	}
@@ -799,7 +798,7 @@ func getOCICapabilitiesList() []string {
 }
 
 // Adds capabilities to all sets relevant to root (bounding, permitted, effective, inheritable)
-func addProcessRootCapability(g *generate.Generator, c string) error {
+func addProcessRootCapability(g *generator, c string) error {
 	if err := g.AddProcessCapabilityBounding(c); err != nil {
 		return err
 	}
@@ -816,7 +815,7 @@ func addProcessRootCapability(g *generate.Generator, c string) error {
 }
 
 // Drops capabilities to all sets relevant to root (bounding, permitted, effective, inheritable)
-func dropProcessRootCapability(g *generate.Generator, c string) error {
+func dropProcessRootCapability(g *generator, c string) error {
 	if err := g.DropProcessCapabilityBounding(c); err != nil {
 		return err
 	}
@@ -833,7 +832,7 @@ func dropProcessRootCapability(g *generate.Generator, c string) error {
 }
 
 // setOCICapabilities adds/drops process capabilities.
-func setOCICapabilities(g *generate.Generator, capabilities *runtime.Capability) error {
+func setOCICapabilities(g *generator, capabilities *runtime.Capability) error {
 	if capabilities == nil {
 		return nil
 	}
@@ -879,7 +878,7 @@ func setOCICapabilities(g *generate.Generator, capabilities *runtime.Capability)
 }
 
 // setOCINamespaces sets namespaces.
-func setOCINamespaces(g *generate.Generator, namespaces *runtime.NamespaceOption, sandboxPid uint32) {
+func setOCINamespaces(g *generator, namespaces *runtime.NamespaceOption, sandboxPid uint32) {
 	g.AddOrReplaceLinuxNamespace(string(runtimespec.NetworkNamespace), getNetworkNamespace(sandboxPid)) // nolint: errcheck
 	g.AddOrReplaceLinuxNamespace(string(runtimespec.IPCNamespace), getIPCNamespace(sandboxPid))         // nolint: errcheck
 	g.AddOrReplaceLinuxNamespace(string(runtimespec.UTSNamespace), getUTSNamespace(sandboxPid))         // nolint: errcheck

pkg/server/container_create_test.go

@@ -450,8 +450,9 @@ func TestContainerSpecCommand(t *testing.T) {
 	} {
 
 		config, _, imageConfig, _ := getCreateContainerTestData()
-		g, err := generate.New("linux")
+		og, err := generate.New("linux")
 		assert.NoError(t, err)
+		g := newCustomGenerator(og)
 		config.Command = test.criEntrypoint
 		config.Args = test.criArgs
 		imageConfig.Entrypoint = test.imageEntrypoint
@@ -645,8 +646,9 @@ func TestPrivilegedBindMount(t *testing.T) {
 		},
 	} {
 		t.Logf("TestCase %q", desc)
-		g, err := generate.New("linux")
+		og, err := generate.New("linux")
 		assert.NoError(t, err)
+		g := newCustomGenerator(og)
 		c := newTestCRIService()
 		c.addOCIBindMounts(&g, nil, "")
 		if test.privileged {
@@ -754,8 +756,9 @@ func TestMountPropagation(t *testing.T) {
 		},
 	} {
 		t.Logf("TestCase %q", desc)
-		g, err := generate.New("linux")
+		og, err := generate.New("linux")
 		assert.NoError(t, err)
+		g := newCustomGenerator(og)
 		c := newTestCRIService()
 		c.os.(*ostesting.FakeOS).LookupMountFn = test.fakeLookupMountFn
 		err = c.addOCIBindMounts(&g, []*runtime.Mount{test.criMount}, "")

pkg/server/helpers.go

@@ -391,10 +391,44 @@ func buildLabels(configLabels map[string]string, containerType string) map[strin
 }
 
 // newSpecGenerator creates a new spec generator for the runtime spec.
-func newSpecGenerator(spec *runtimespec.Spec) generate.Generator {
+func newSpecGenerator(spec *runtimespec.Spec) generator {
 	g := generate.NewFromSpec(spec)
 	g.HostSpecific = true
-	return g
+	return newCustomGenerator(g)
+}
+
+// generator is a custom generator with some functions overridden
+// used by the cri plugin.
+// TODO(random-liu): Upstream this fix.
+type generator struct {
+	generate.Generator
+	envCache map[string]int
+}
+
+func newCustomGenerator(g generate.Generator) generator {
+	return generator{
+		Generator: g,
+		envCache:  make(map[string]int),
+	}
+}
+
+// AddProcessEnv overrides the original AddProcessEnv. It uses
+// a map to cache and override envs.
+func (g *generator) AddProcessEnv(key, value string) {
+	if len(g.envCache) == 0 {
+		// Call AddProccessEnv once to initialize the spec.
+		g.Generator.AddProcessEnv(key, value)
+		g.envCache[key] = 0
+		return
+	}
+	spec := g.Config
+	env := fmt.Sprintf("%s=%s", key, value)
+	if idx, ok := g.envCache[key]; !ok {
+		spec.Process.Env = append(spec.Process.Env, env)
+		g.envCache[key] = len(spec.Process.Env) - 1
+	} else {
+		spec.Process.Env[idx] = env
+	}
 }
 
 func getPodCNILabels(id string, config *runtime.PodSandboxConfig) map[string]string {

pkg/server/helpers_test.go

@@ -321,3 +321,58 @@ func TestRestrictOOMScoreAdj(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, got, current+1)
 }
+
+func TestCustomGenerator(t *testing.T) {
+	for desc, test := range map[string]struct {
+		kv        [][2]string
+		expected  []string
+		expectNil bool
+	}{
+		"empty": {
+			expectNil: true,
+		},
+		"single env": {
+			kv: [][2]string{
+				{"a", "b"},
+			},
+			expected: []string{"a=b"},
+		},
+		"multiple envs": {
+			kv: [][2]string{
+				{"a", "b"},
+				{"c", "d"},
+				{"e", "f"},
+			},
+			expected: []string{
+				"a=b",
+				"c=d",
+				"e=f",
+			},
+		},
+		"env override": {
+			kv: [][2]string{
+				{"k1", "v1"},
+				{"k2", "v2"},
+				{"k3", "v3"},
+				{"k3", "v4"},
+				{"k1", "v5"},
+			},
+			expected: []string{
+				"k1=v5",
+				"k2=v2",
+				"k3=v4",
+			},
+		},
+	} {
+		t.Logf("TestCase %q", desc)
+		g := newSpecGenerator(nil)
+		for _, kv := range test.kv {
+			g.AddProcessEnv(kv[0], kv[1])
+		}
+		if test.expectNil {
+			assert.Nil(t, g.Config)
+		} else {
+			assert.Equal(t, test.expected, g.Config.Process.Env)
+		}
+	}
+}