Add env cache.

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

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 {