llb: convert envlist from slice to linked list

The semantics of envlist is that it gets many adds, and reads
but reads often only appear on a specific state of the list, usually
when all insertions have been completed.

Previous implementation did a full copy on each insert. For lookups
the options were to do a O(N) lookup per key or copy to a string
slice and then externally try to convert it to map with another copy.
In new version, the insertions are cheap and on first access optimized
map table is created that is then reused for next lookups.

Signed-off-by: Tonis Tiigi <[email protected]>

Author: tonistiigi

client/llb/meta.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"net"
 	"path"
+	"slices"
+	"sync"
 
 	"github.com/containerd/containerd/platforms"
 	"github.com/google/shlex"
@@ -111,16 +113,16 @@ func Reset(other State) StateOption {
 	}
 }
 
-func getEnv(s State) func(context.Context, *Constraints) (EnvList, error) {
-	return func(ctx context.Context, c *Constraints) (EnvList, error) {
+func getEnv(s State) func(context.Context, *Constraints) (*EnvList, error) {
+	return func(ctx context.Context, c *Constraints) (*EnvList, error) {
 		v, err := s.getValue(keyEnv)(ctx, c)
 		if err != nil {
 			return nil, err
 		}
 		if v != nil {
-			return v.(EnvList), nil
+			return v.(*EnvList), nil
 		}
-		return EnvList{}, nil
+		return &EnvList{}, nil
 	}
 }
 
@@ -346,54 +348,83 @@ func getSecurity(s State) func(context.Context, *Constraints) (pb.SecurityMode,
 	}
 }
 
-type EnvList []KeyValue
-
-type KeyValue struct {
-	key   string
-	value string
+type EnvList struct {
+	parent *EnvList
+	key    string
+	value  string
+	del    bool
+	once   sync.Once
+	l      int
+	values map[string]string
+	keys   []string
 }
 
-func (e EnvList) AddOrReplace(k, v string) EnvList {
-	e = e.Delete(k)
-	e = append(e, KeyValue{key: k, value: v})
-	return e
+func (e *EnvList) AddOrReplace(k, v string) *EnvList {
+	return &EnvList{
+		parent: e,
+		key:    k,
+		value:  v,
+		l:      e.l + 1,
+	}
 }
 
-func (e EnvList) SetDefault(k, v string) EnvList {
+func (e *EnvList) SetDefault(k, v string) *EnvList {
 	if _, ok := e.Get(k); !ok {
-		e = append(e, KeyValue{key: k, value: v})
+		return e.AddOrReplace(k, v)
 	}
 	return e
 }
 
-func (e EnvList) Delete(k string) EnvList {
-	e = append([]KeyValue(nil), e...)
-	if i, ok := e.Index(k); ok {
-		return append(e[:i], e[i+1:]...)
+func (e *EnvList) Delete(k string) EnvList {
+	return EnvList{
+		parent: e,
+		key:    k,
+		del:    true,
+		l:      e.l + 1,
 	}
-	return e
 }
 
-func (e EnvList) Get(k string) (string, bool) {
-	if index, ok := e.Index(k); ok {
-		return e[index].value, true
-	}
-	return "", false
+func (e *EnvList) makeValues() {
+	m := make(map[string]string, e.l)
+	seen := make(map[string]struct{}, e.l)
+	keys := make([]string, 0, e.l)
+	e.keys = e.addValue(keys, m, seen)
+	e.values = m
+	slices.Reverse(e.keys)
 }
 
-func (e EnvList) Index(k string) (int, bool) {
-	for i, kv := range e {
-		if kv.key == k {
-			return i, true
-		}
+func (e *EnvList) addValue(keys []string, vals map[string]string, seen map[string]struct{}) []string {
+	if e.parent == nil {
+		return keys
 	}
-	return -1, false
+	if _, ok := seen[e.key]; !e.del && !ok {
+		vals[e.key] = e.value
+		keys = append(keys, e.key)
+	}
+	seen[e.key] = struct{}{}
+	if e.parent != nil {
+		keys = e.parent.addValue(keys, vals, seen)
+	}
+	return keys
+}
+
+func (e *EnvList) Get(k string) (string, bool) {
+	e.once.Do(e.makeValues)
+	v, ok := e.values[k]
+	return v, ok
+}
+
+func (e *EnvList) Keys() []string {
+	e.once.Do(e.makeValues)
+	return e.keys
 }
 
-func (e EnvList) ToArray() []string {
-	out := make([]string, 0, len(e))
-	for _, kv := range e {
-		out = append(out, kv.key+"="+kv.value)
+func (e *EnvList) ToArray() []string {
+	keys := e.Keys()
+	out := make([]string, 0, len(keys))
+	for _, k := range keys {
+		v, _ := e.Get(k)
+		out = append(out, k+"="+v)
 	}
 	return out
 }

client/llb/state.go

@@ -351,16 +351,12 @@ func (s State) GetEnv(ctx context.Context, key string, co ...ConstraintsOpt) (st
 
 // Env returns a new [State] with the provided environment variable set.
 // See [Env]
-func (s State) Env(ctx context.Context, co ...ConstraintsOpt) ([]string, error) {
+func (s State) Env(ctx context.Context, co ...ConstraintsOpt) (*EnvList, error) {
 	c := &Constraints{}
 	for _, f := range co {
 		f.SetConstraintsOption(c)
 	}
-	env, err := getEnv(s)(ctx, c)
-	if err != nil {
-		return nil, err
-	}
-	return env.ToArray(), nil
+	return getEnv(s)(ctx, c)
 }
 
 // GetDir returns the current working directory for the state.

frontend/dockerfile/dockerfile2llb/convert.go

@@ -613,7 +613,7 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 							dfCmd(d.stage.SourceCode),
 							llb.Platform(*platform),
 							opt.ImageResolveMode,
-							llb.WithCustomName(prefixCommand(d, "FROM "+d.stage.BaseName, opt.MultiPlatformRequested, platform, nil)),
+							llb.WithCustomName(prefixCommand(d, "FROM "+d.stage.BaseName, opt.MultiPlatformRequested, platform, emptyEnvs{})),
 							location(opt.SourceMap, d.stage.Location),
 						)
 						if reachable {
@@ -884,7 +884,7 @@ func (e *envsFromState) init() {
 	if err != nil {
 		return
 	}
-	e.env = shell.EnvsFromSlice(env)
+	e.env = env
 }
 
 func (e *envsFromState) Get(key string) (string, bool) {
@@ -2352,3 +2352,13 @@ func validateBaseImagePlatform(name string, expected, actual ocispecs.Platform,
 		lint.Run(&linter.RuleInvalidBaseImagePlatform, location, msg)
 	}
 }
+
+type emptyEnvs struct{}
+
+func (emptyEnvs) Get(string) (string, bool) {
+	return "", false
+}
+
+func (emptyEnvs) Keys() []string {
+	return nil
+}