Merge pull request #693 from umohnani8/env

Improve performance of AddProcessEnv

generate/generate.go

@@ -29,6 +29,9 @@ var (
 type Generator struct {
 	Config       *rspec.Spec
 	HostSpecific bool
+	// This is used to keep a cache of the ENVs added to improve
+	// performance when adding a huge number of ENV variables
+	envMap map[string]int
 }
 
 // ExportOptions have toggles for exporting only certain parts of the specification
@@ -236,7 +239,12 @@ func New(os string) (generator Generator, err error) {
 		}
 	}
 
-	return Generator{Config: &config}, nil
+	envCache := map[string]int{}
+	if config.Process != nil {
+		envCache = createEnvCacheMap(config.Process.Env)
+	}
+
+	return Generator{Config: &config, envMap: envCache}, nil
 }
 
 // NewFromSpec creates a configuration Generator from a given
@@ -246,8 +254,14 @@ func New(os string) (generator Generator, err error) {
 //
 //   generator := Generator{Config: config}
 func NewFromSpec(config *rspec.Spec) Generator {
+	envCache := map[string]int{}
+	if config != nil && config.Process != nil {
+		envCache = createEnvCacheMap(config.Process.Env)
+	}
+
 	return Generator{
 		Config: config,
+		envMap: envCache,
 	}
 }
 
@@ -273,11 +287,27 @@ func NewFromTemplate(r io.Reader) (Generator, error) {
 	if err := json.NewDecoder(r).Decode(&config); err != nil {
 		return Generator{}, err
 	}
+
+	envCache := map[string]int{}
+	if config.Process != nil {
+		envCache = createEnvCacheMap(config.Process.Env)
+	}
+
 	return Generator{
 		Config: &config,
+		envMap: envCache,
 	}, nil
 }
 
+// createEnvCacheMap creates a hash map with the ENV variables given by the config
+func createEnvCacheMap(env []string) map[string]int {
+	envMap := make(map[string]int, len(env))
+	for i, val := range env {
+		envMap[val] = i
+	}
+	return envMap
+}
+
 // SetSpec sets the configuration in the Generator g.
 //
 // Deprecated: Replace with:
@@ -456,21 +486,44 @@ func (g *Generator) ClearProcessEnv() {
 		return
 	}
 	g.Config.Process.Env = []string{}
+	// Clear out the env cache map as well
+	g.envMap = map[string]int{}
 }
 
 // AddProcessEnv adds name=value into g.Config.Process.Env, or replaces an
 // existing entry with the given name.
 func (g *Generator) AddProcessEnv(name, value string) {
+	if name == "" {
+		return
+	}
+
 	g.initConfigProcess()
+	g.addEnv(fmt.Sprintf("%s=%s", name, value), name)
+}
 
-	env := fmt.Sprintf("%s=%s", name, value)
-	for idx := range g.Config.Process.Env {
-		if strings.HasPrefix(g.Config.Process.Env[idx], name+"=") {
-			g.Config.Process.Env[idx] = env
-			return
-		}
+// AddMultipleProcessEnv adds multiple name=value into g.Config.Process.Env, or replaces
+// existing entries with the given name.
+func (g *Generator) AddMultipleProcessEnv(envs []string) {
+	g.initConfigProcess()
+
+	for _, val := range envs {
+		split := strings.SplitN(val, "=", 2)
+		g.addEnv(val, split[0])
+	}
+}
+
+// addEnv looks through adds ENV to the Process and checks envMap for
+// any duplicates
+// This is called by both AddMultipleProcessEnv and AddProcessEnv
+func (g *Generator) addEnv(env, key string) {
+	if idx, ok := g.envMap[key]; ok {
+		// The ENV exists in the cache, so change its value in g.Config.Process.Env
+		g.Config.Process.Env[idx] = env
+	} else {
+		// else the env doesn't exist, so add it and add it's index to g.envMap
+		g.Config.Process.Env = append(g.Config.Process.Env, env)
+		g.envMap[key] = len(g.Config.Process.Env) - 1
 	}
-	g.Config.Process.Env = append(g.Config.Process.Env, env)
 }
 
 // AddProcessRlimits adds rlimit into g.Config.Process.Rlimits.

generate/generate_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/opencontainers/runtime-tools/generate"
 	"github.com/opencontainers/runtime-tools/specerror"
 	"github.com/opencontainers/runtime-tools/validate"
+	"github.com/stretchr/testify/assert"
 )
 
 // Smoke test to ensure that _at the very least_ our default configuration
@@ -93,3 +94,65 @@ func TestRemoveMount(t *testing.T) {
 		t.Errorf("Unable to remove /dev/shm from mounts")
 	}
 }
+
+func TestEnvCaching(t *testing.T) {
+	// Start with empty ENV and add a few
+	g, err := generate.New("windows")
+	if err != nil {
+		t.Fatal(err)
+	}
+	expected := []string{"k1=v1", "k2=v2"}
+	g.AddProcessEnv("k1", "v1")
+	g.AddProcessEnv("k2", "v2")
+	assert.Equal(t, expected, g.Config.Process.Env)
+
+	// Test override and existing ENV
+	g, err = generate.New("linux")
+	if err != nil {
+		t.Fatal(err)
+	}
+	expected = []string{"PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin", "TERM=xterm", "k1=v1", "k2=v4", "k3=v3"}
+	g.AddProcessEnv("k1", "v1")
+	g.AddProcessEnv("k2", "v2")
+	g.AddProcessEnv("k3", "v3")
+	g.AddProcessEnv("k2", "v4")
+	assert.Equal(t, expected, g.Config.Process.Env)
+
+	// Test empty ENV
+	g, err = generate.New("windows")
+	if err != nil {
+		t.Fatal(err)
+	}
+	g.AddProcessEnv("", "")
+	assert.Equal(t, []string(nil), g.Config.Process.Env)
+}
+
+func TestMultipleEnvCaching(t *testing.T) {
+	// Start with empty ENV and add a few
+	g, err := generate.New("windows")
+	if err != nil {
+		t.Fatal(err)
+	}
+	newEnvs := []string{"k1=v1", "k2=v2"}
+	expected := []string{"k1=v1", "k2=v2"}
+	g.AddMultipleProcessEnv(newEnvs)
+	assert.Equal(t, expected, g.Config.Process.Env)
+
+	// Test override and existing ENV
+	g, err = generate.New("linux")
+	if err != nil {
+		t.Fatal(err)
+	}
+	newEnvs = []string{"k1=v1", "k2=v2", "k3=v3", "k2=v4"}
+	expected = []string{"PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin", "TERM=xterm", "k1=v1", "k2=v4", "k3=v3"}
+	g.AddMultipleProcessEnv(newEnvs)
+	assert.Equal(t, expected, g.Config.Process.Env)
+
+	// Test empty ENV
+	g, err = generate.New("windows")
+	if err != nil {
+		t.Fatal(err)
+	}
+	g.AddMultipleProcessEnv([]string{})
+	assert.Equal(t, []string(nil), g.Config.Process.Env)
+}