Skip to content
This repository was archived by the owner on Mar 9, 2022. It is now read-only.

Commit 03cd5a3

Browse files
Random-LiuRandom-Liu
committed
Add env cache.
Signed-off-by: Lantao Liu <[email protected]>
1 parent b33f16e commit 03cd5a3

4 files changed

Lines changed: 105 additions & 17 deletions

File tree

‎pkg/server/container_create.go‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,6 @@ import (
3636
imagespec "github.com/opencontainers/image-spec/specs-go/v1"
3737
"github.com/opencontainers/runc/libcontainer/devices"
3838
runtimespec "github.com/opencontainers/runtime-spec/specs-go"
39-
"github.com/opencontainers/runtime-tools/generate"
4039
"github.com/opencontainers/runtime-tools/validate"
4140
"github.com/opencontainers/selinux/go-selinux/label"
4241
"github.com/pkg/errors"
@@ -501,7 +500,7 @@ func (c *criService) generateContainerMounts(sandboxID string, config *runtime.C
501500

502501
// setOCIProcessArgs sets process args. It returns error if the final arg list
503502
// is empty.
504-
func setOCIProcessArgs(g *generate.Generator, config *runtime.ContainerConfig, imageConfig *imagespec.ImageConfig) error {
503+
func setOCIProcessArgs(g *generator, config *runtime.ContainerConfig, imageConfig *imagespec.ImageConfig) error {
505504
command, args := config.GetCommand(), config.GetArgs()
506505
// The following logic is migrated from https://github.com/moby/moby/blob/master/daemon/commit.go
507506
// TODO(random-liu): Clearly define the commands overwrite behavior.
@@ -523,7 +522,7 @@ func setOCIProcessArgs(g *generate.Generator, config *runtime.ContainerConfig, i
523522

524523
// addImageEnvs adds environment variables from image config. It returns error if
525524
// an invalid environment variable is encountered.
526-
func addImageEnvs(g *generate.Generator, imageEnvs []string) error {
525+
func addImageEnvs(g *generator, imageEnvs []string) error {
527526
for _, e := range imageEnvs {
528527
kv := strings.SplitN(e, "=", 2)
529528
if len(kv) != 2 {
@@ -534,7 +533,7 @@ func addImageEnvs(g *generate.Generator, imageEnvs []string) error {
534533
return nil
535534
}
536535

537-
func setOCIPrivileged(g *generate.Generator, config *runtime.ContainerConfig) error {
536+
func setOCIPrivileged(g *generator, config *runtime.ContainerConfig) error {
538537
// Add all capabilities in privileged mode.
539538
g.SetupPrivileged(true)
540539
setOCIBindMountsPrivileged(g)
@@ -555,7 +554,7 @@ func clearReadOnly(m *runtimespec.Mount) {
555554
}
556555

557556
// addDevices set device mapping without privilege.
558-
func (c *criService) addOCIDevices(g *generate.Generator, devs []*runtime.Device) error {
557+
func (c *criService) addOCIDevices(g *generator, devs []*runtime.Device) error {
559558
spec := g.Spec()
560559
for _, device := range devs {
561560
path, err := c.os.ResolveSymbolicLink(device.HostPath)
@@ -587,7 +586,7 @@ func (c *criService) addOCIDevices(g *generate.Generator, devs []*runtime.Device
587586
}
588587

589588
// addDevices set device mapping with privilege.
590-
func setOCIDevicesPrivileged(g *generate.Generator) error {
589+
func setOCIDevicesPrivileged(g *generator) error {
591590
spec := g.Spec()
592591
hostDevices, err := devices.HostDevices()
593592
if err != nil {
@@ -618,7 +617,7 @@ func setOCIDevicesPrivileged(g *generate.Generator) error {
618617
}
619618

620619
// addOCIBindMounts adds bind mounts.
621-
func (c *criService) addOCIBindMounts(g *generate.Generator, mounts []*runtime.Mount, mountLabel string) error {
620+
func (c *criService) addOCIBindMounts(g *generator, mounts []*runtime.Mount, mountLabel string) error {
622621
// Sort mounts in number of parts. This ensures that high level mounts don't
623622
// shadow other mounts.
624623
sort.Sort(orderedMounts(mounts))
@@ -713,7 +712,7 @@ func (c *criService) addOCIBindMounts(g *generate.Generator, mounts []*runtime.M
713712
return nil
714713
}
715714

716-
func setOCIBindMountsPrivileged(g *generate.Generator) {
715+
func setOCIBindMountsPrivileged(g *generator) {
717716
spec := g.Spec()
718717
// clear readonly for /sys and cgroup
719718
for i, m := range spec.Mounts {
@@ -728,8 +727,8 @@ func setOCIBindMountsPrivileged(g *generate.Generator) {
728727
spec.Linux.MaskedPaths = nil
729728
}
730729

731-
// setOCILinuxResource set container resource limit.
732-
func setOCILinuxResource(g *generate.Generator, resources *runtime.LinuxContainerResources) {
730+
// setOCILinuxResource set container cgroup resource limit.
731+
func setOCILinuxResource(g *generator, resources *runtime.LinuxContainerResources) {
733732
if resources == nil {
734733
return
735734
}
@@ -755,7 +754,7 @@ func getOCICapabilitiesList() []string {
755754
}
756755

757756
// setOCICapabilities adds/drops process capabilities.
758-
func setOCICapabilities(g *generate.Generator, capabilities *runtime.Capability) error {
757+
func setOCICapabilities(g *generator, capabilities *runtime.Capability) error {
759758
if capabilities == nil {
760759
return nil
761760
}
@@ -801,7 +800,7 @@ func setOCICapabilities(g *generate.Generator, capabilities *runtime.Capability)
801800
}
802801

803802
// setOCINamespaces sets namespaces.
804-
func setOCINamespaces(g *generate.Generator, namespaces *runtime.NamespaceOption, sandboxPid uint32) {
803+
func setOCINamespaces(g *generator, namespaces *runtime.NamespaceOption, sandboxPid uint32) {
805804
g.AddOrReplaceLinuxNamespace(string(runtimespec.NetworkNamespace), getNetworkNamespace(sandboxPid)) // nolint: errcheck
806805
g.AddOrReplaceLinuxNamespace(string(runtimespec.IPCNamespace), getIPCNamespace(sandboxPid)) // nolint: errcheck
807806
g.AddOrReplaceLinuxNamespace(string(runtimespec.UTSNamespace), getUTSNamespace(sandboxPid)) // nolint: errcheck

‎pkg/server/container_create_test.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -451,7 +451,7 @@ func TestContainerSpecCommand(t *testing.T) {
451451
} {
452452

453453
config, _, imageConfig, _ := getCreateContainerTestData()
454-
g := generate.New()
454+
g := newCustomGenerator(generate.New())
455455
config.Command = test.criEntrypoint
456456
config.Args = test.criArgs
457457
imageConfig.Entrypoint = test.imageEntrypoint
@@ -652,7 +652,7 @@ func TestPrivilegedBindMount(t *testing.T) {
652652
},
653653
} {
654654
t.Logf("TestCase %q", desc)
655-
g := generate.New()
655+
g := newCustomGenerator(generate.New())
656656
g.SetRootReadonly(test.readonlyRootFS)
657657
c := newTestCRIService()
658658
c.addOCIBindMounts(&g, nil, "")
@@ -761,7 +761,7 @@ func TestMountPropagation(t *testing.T) {
761761
},
762762
} {
763763
t.Logf("TestCase %q", desc)
764-
g := generate.New()
764+
g := newCustomGenerator(generate.New())
765765
c := newTestCRIService()
766766
c.os.(*ostesting.FakeOS).LookupMountFn = test.fakeLookupMountFn
767767
err := c.addOCIBindMounts(&g, []*runtime.Mount{test.criMount}, "")

‎pkg/server/helpers.go‎

Lines changed: 36 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -431,10 +431,44 @@ func buildLabels(configLabels map[string]string, containerType string) map[strin
431431
}
432432

433433
// newSpecGenerator creates a new spec generator for the runtime spec.
434-
func newSpecGenerator(spec *runtimespec.Spec) generate.Generator {
434+
func newSpecGenerator(spec *runtimespec.Spec) generator {
435435
g := generate.NewFromSpec(spec)
436436
g.HostSpecific = true
437-
return g
437+
return newCustomGenerator(g)
438+
}
439+
440+
// generator is a custom generator with some functions overridden
441+
// used by the cri plugin.
442+
// TODO(random-liu): Upstream this fix.
443+
type generator struct {
444+
generate.Generator
445+
envCache map[string]int
446+
}
447+
448+
func newCustomGenerator(g generate.Generator) generator {
449+
return generator{
450+
Generator: g,
451+
envCache: make(map[string]int),
452+
}
453+
}
454+
455+
// AddProcessEnv overrides the original AddProcessEnv. It uses
456+
// a map to cache and override envs.
457+
func (g *generator) AddProcessEnv(key, value string) {
458+
if len(g.envCache) == 0 {
459+
// Call AddProccessEnv once to initialize the spec.
460+
g.Generator.AddProcessEnv(key, value)
461+
g.envCache[key] = 0
462+
return
463+
}
464+
spec := g.Spec()
465+
env := fmt.Sprintf("%s=%s", key, value)
466+
if idx, ok := g.envCache[key]; !ok {
467+
spec.Process.Env = append(spec.Process.Env, env)
468+
g.envCache[key] = len(spec.Process.Env) - 1
469+
} else {
470+
spec.Process.Env[idx] = env
471+
}
438472
}
439473

440474
func getPodCNILabels(id string, config *runtime.PodSandboxConfig) map[string]string {

‎pkg/server/helpers_test.go‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -213,3 +213,58 @@ func TestOrderedMounts(t *testing.T) {
213213
sort.Stable(orderedMounts(mounts))
214214
assert.Equal(t, expected, mounts)
215215
}
216+
217+
func TestCustomGenerator(t *testing.T) {
218+
for desc, test := range map[string]struct {
219+
kv [][2]string
220+
expected []string
221+
expectNil bool
222+
}{
223+
"empty": {
224+
expectNil: true,
225+
},
226+
"single env": {
227+
kv: [][2]string{
228+
{"a", "b"},
229+
},
230+
expected: []string{"a=b"},
231+
},
232+
"multiple envs": {
233+
kv: [][2]string{
234+
{"a", "b"},
235+
{"c", "d"},
236+
{"e", "f"},
237+
},
238+
expected: []string{
239+
"a=b",
240+
"c=d",
241+
"e=f",
242+
},
243+
},
244+
"env override": {
245+
kv: [][2]string{
246+
{"k1", "v1"},
247+
{"k2", "v2"},
248+
{"k3", "v3"},
249+
{"k3", "v4"},
250+
{"k1", "v5"},
251+
},
252+
expected: []string{
253+
"k1=v5",
254+
"k2=v2",
255+
"k3=v4",
256+
},
257+
},
258+
} {
259+
t.Logf("TestCase %q", desc)
260+
g := newSpecGenerator(nil)
261+
for _, kv := range test.kv {
262+
g.AddProcessEnv(kv[0], kv[1])
263+
}
264+
if test.expectNil {
265+
assert.Nil(t, g.Spec())
266+
} else {
267+
assert.Equal(t, test.expected, g.Spec().Process.Env)
268+
}
269+
}
270+
}

0 commit comments

Comments
 (0)