From 5928fc0e600bd9a846e8b4e325bfb983a0ccb0d8 Mon Sep 17 00:00:00 2001 From: Tim Allclair Date: Thu, 27 Mar 2025 11:11:42 -0700 Subject: [PATCH] Add ContainerIter utility for ranging over pod containers --- pkg/api/pod/util.go | 51 ++++++++----- pkg/api/pod/util_test.go | 141 ++++++++++++++++++++++++++++++++++++ pkg/api/v1/pod/util.go | 51 ++++++++----- pkg/api/v1/pod/util_test.go | 141 ++++++++++++++++++++++++++++++++++++ 4 files changed, 346 insertions(+), 38 deletions(-) diff --git a/pkg/api/pod/util.go b/pkg/api/pod/util.go index c23714d0c95..c671cbe526c 100644 --- a/pkg/api/pod/util.go +++ b/pkg/api/pod/util.go @@ -17,6 +17,7 @@ limitations under the License. package pod import ( + "iter" "strings" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" @@ -60,30 +61,42 @@ type ContainerVisitor func(container *api.Container, containerType ContainerType // visiting is short-circuited. VisitContainers returns true if visiting completes, // false if visiting was short-circuited. func VisitContainers(podSpec *api.PodSpec, mask ContainerType, visitor ContainerVisitor) bool { - if mask&InitContainers != 0 { - for i := range podSpec.InitContainers { - if !visitor(&podSpec.InitContainers[i], InitContainers) { - return false - } - } - } - if mask&Containers != 0 { - for i := range podSpec.Containers { - if !visitor(&podSpec.Containers[i], Containers) { - return false - } - } - } - if mask&EphemeralContainers != 0 { - for i := range podSpec.EphemeralContainers { - if !visitor((*api.Container)(&podSpec.EphemeralContainers[i].EphemeralContainerCommon), EphemeralContainers) { - return false - } + for c, t := range ContainerIter(podSpec, mask) { + if !visitor(c, t) { + return false } } return true } +// ContainerIter returns an iterator over all containers in the given pod spec with a masked type. +// The iteration order is InitContainers, then main Containers, then EphemeralContainers. +func ContainerIter(podSpec *api.PodSpec, mask ContainerType) iter.Seq2[*api.Container, ContainerType] { + return func(yield func(*api.Container, ContainerType) bool) { + if mask&InitContainers != 0 { + for i := range podSpec.InitContainers { + if !yield(&podSpec.InitContainers[i], InitContainers) { + return + } + } + } + if mask&Containers != 0 { + for i := range podSpec.Containers { + if !yield(&podSpec.Containers[i], Containers) { + return + } + } + } + if mask&EphemeralContainers != 0 { + for i := range podSpec.EphemeralContainers { + if !yield((*api.Container)(&podSpec.EphemeralContainers[i].EphemeralContainerCommon), EphemeralContainers) { + return + } + } + } + } +} + // Visitor is called with each object name, and returns true if visiting should continue type Visitor func(name string) (shouldContinue bool) diff --git a/pkg/api/pod/util_test.go b/pkg/api/pod/util_test.go index 8f8de2fc8b3..1d2f24e2a88 100644 --- a/pkg/api/pod/util_test.go +++ b/pkg/api/pod/util_test.go @@ -208,6 +208,147 @@ func TestVisitContainers(t *testing.T) { } } +func TestContainerIter(t *testing.T) { + testCases := []struct { + desc string + spec *api.PodSpec + wantContainers []string + mask ContainerType + }{ + { + desc: "empty podspec", + spec: &api.PodSpec{}, + wantContainers: []string{}, + mask: AllContainers, + }, + { + desc: "regular containers", + spec: &api.PodSpec{ + Containers: []api.Container{ + {Name: "c1"}, + {Name: "c2"}, + }, + InitContainers: []api.Container{ + {Name: "i1"}, + {Name: "i2"}, + }, + EphemeralContainers: []api.EphemeralContainer{ + {EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e1"}}, + {EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e2"}}, + }, + }, + wantContainers: []string{"c1", "c2"}, + mask: Containers, + }, + { + desc: "init containers", + spec: &api.PodSpec{ + Containers: []api.Container{ + {Name: "c1"}, + {Name: "c2"}, + }, + InitContainers: []api.Container{ + {Name: "i1"}, + {Name: "i2"}, + }, + EphemeralContainers: []api.EphemeralContainer{ + {EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e1"}}, + {EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e2"}}, + }, + }, + wantContainers: []string{"i1", "i2"}, + mask: InitContainers, + }, + { + desc: "init + main containers", + spec: &api.PodSpec{ + Containers: []api.Container{ + {Name: "c1"}, + {Name: "c2"}, + }, + InitContainers: []api.Container{ + {Name: "i1"}, + {Name: "i2"}, + }, + EphemeralContainers: []api.EphemeralContainer{ + {EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e1"}}, + {EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e2"}}, + }, + }, + wantContainers: []string{"i1", "i2", "c1", "c2"}, + mask: InitContainers | Containers, + }, + { + desc: "ephemeral containers", + spec: &api.PodSpec{ + Containers: []api.Container{ + {Name: "c1"}, + {Name: "c2"}, + }, + InitContainers: []api.Container{ + {Name: "i1"}, + {Name: "i2"}, + }, + EphemeralContainers: []api.EphemeralContainer{ + {EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e1"}}, + {EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e2"}}, + }, + }, + wantContainers: []string{"e1", "e2"}, + mask: EphemeralContainers, + }, + { + desc: "all container types", + spec: &api.PodSpec{ + Containers: []api.Container{ + {Name: "c1"}, + {Name: "c2"}, + }, + InitContainers: []api.Container{ + {Name: "i1"}, + {Name: "i2"}, + }, + EphemeralContainers: []api.EphemeralContainer{ + {EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e1"}}, + {EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e2"}}, + }, + }, + wantContainers: []string{"i1", "i2", "c1", "c2", "e1", "e2"}, + mask: AllContainers, + }, + } + + for _, tc := range testCases { + t.Run(tc.desc, func(t *testing.T) { + gotContainers := []string{} + for c, containerType := range ContainerIter(tc.spec, tc.mask) { + gotContainers = append(gotContainers, c.Name) + + switch containerType { + case InitContainers: + if c.Name[0] != 'i' { + t.Errorf("ContainerIter() yielded container type InitContainers for container %q", c.Name) + } + case Containers: + if c.Name[0] != 'c' { + t.Errorf("ContainerIter() yielded container type Containers for container %q", c.Name) + } + case EphemeralContainers: + if c.Name[0] != 'e' { + t.Errorf("ContainerIter() yielded container type EphemeralContainers for container %q", c.Name) + } + default: + t.Errorf("ContainerIter() yielded unknown container type %d", containerType) + } + } + + if !cmp.Equal(gotContainers, tc.wantContainers) { + t.Errorf("ContainerIter() = %+v, want %+v", gotContainers, tc.wantContainers) + } + }) + } +} + func TestPodSecrets(t *testing.T) { // Stub containing all possible secret references in a pod. // The names of the referenced secrets match struct paths detected by reflection. diff --git a/pkg/api/v1/pod/util.go b/pkg/api/v1/pod/util.go index 7ba94b4219b..7bf5563124e 100644 --- a/pkg/api/v1/pod/util.go +++ b/pkg/api/v1/pod/util.go @@ -18,6 +18,7 @@ package pod import ( "fmt" + "iter" "time" v1 "k8s.io/api/core/v1" @@ -105,30 +106,42 @@ func skipEmptyNames(visitor Visitor) Visitor { // visiting is short-circuited. VisitContainers returns true if visiting completes, // false if visiting was short-circuited. func VisitContainers(podSpec *v1.PodSpec, mask ContainerType, visitor ContainerVisitor) bool { - if mask&InitContainers != 0 { - for i := range podSpec.InitContainers { - if !visitor(&podSpec.InitContainers[i], InitContainers) { - return false - } - } - } - if mask&Containers != 0 { - for i := range podSpec.Containers { - if !visitor(&podSpec.Containers[i], Containers) { - return false - } - } - } - if mask&EphemeralContainers != 0 { - for i := range podSpec.EphemeralContainers { - if !visitor((*v1.Container)(&podSpec.EphemeralContainers[i].EphemeralContainerCommon), EphemeralContainers) { - return false - } + for c, t := range ContainerIter(podSpec, mask) { + if !visitor(c, t) { + return false } } return true } +// ContainerIter returns an iterator over all containers in the given pod spec with a masked type. +// The iteration order is InitContainers, then main Containers, then EphemeralContainers. +func ContainerIter(podSpec *v1.PodSpec, mask ContainerType) iter.Seq2[*v1.Container, ContainerType] { + return func(yield func(*v1.Container, ContainerType) bool) { + if mask&InitContainers != 0 { + for i := range podSpec.InitContainers { + if !yield(&podSpec.InitContainers[i], InitContainers) { + return + } + } + } + if mask&Containers != 0 { + for i := range podSpec.Containers { + if !yield(&podSpec.Containers[i], Containers) { + return + } + } + } + if mask&EphemeralContainers != 0 { + for i := range podSpec.EphemeralContainers { + if !yield((*v1.Container)(&podSpec.EphemeralContainers[i].EphemeralContainerCommon), EphemeralContainers) { + return + } + } + } + } +} + // VisitPodSecretNames invokes the visitor function with the name of every secret // referenced by the pod spec. If visitor returns false, visiting is short-circuited. // Transitive references (e.g. pod -> pvc -> pv -> secret) are not visited. diff --git a/pkg/api/v1/pod/util_test.go b/pkg/api/v1/pod/util_test.go index a172f802742..002f217d3b0 100644 --- a/pkg/api/v1/pod/util_test.go +++ b/pkg/api/v1/pod/util_test.go @@ -426,6 +426,147 @@ func TestVisitContainers(t *testing.T) { } } +func TestContainerIter(t *testing.T) { + testCases := []struct { + desc string + spec *v1.PodSpec + wantContainers []string + mask ContainerType + }{ + { + desc: "empty podspec", + spec: &v1.PodSpec{}, + wantContainers: []string{}, + mask: AllContainers, + }, + { + desc: "regular containers", + spec: &v1.PodSpec{ + Containers: []v1.Container{ + {Name: "c1"}, + {Name: "c2"}, + }, + InitContainers: []v1.Container{ + {Name: "i1"}, + {Name: "i2"}, + }, + EphemeralContainers: []v1.EphemeralContainer{ + {EphemeralContainerCommon: v1.EphemeralContainerCommon{Name: "e1"}}, + {EphemeralContainerCommon: v1.EphemeralContainerCommon{Name: "e2"}}, + }, + }, + wantContainers: []string{"c1", "c2"}, + mask: Containers, + }, + { + desc: "init containers", + spec: &v1.PodSpec{ + Containers: []v1.Container{ + {Name: "c1"}, + {Name: "c2"}, + }, + InitContainers: []v1.Container{ + {Name: "i1"}, + {Name: "i2"}, + }, + EphemeralContainers: []v1.EphemeralContainer{ + {EphemeralContainerCommon: v1.EphemeralContainerCommon{Name: "e1"}}, + {EphemeralContainerCommon: v1.EphemeralContainerCommon{Name: "e2"}}, + }, + }, + wantContainers: []string{"i1", "i2"}, + mask: InitContainers, + }, + { + desc: "init + main containers", + spec: &v1.PodSpec{ + Containers: []v1.Container{ + {Name: "c1"}, + {Name: "c2"}, + }, + InitContainers: []v1.Container{ + {Name: "i1"}, + {Name: "i2"}, + }, + EphemeralContainers: []v1.EphemeralContainer{ + {EphemeralContainerCommon: v1.EphemeralContainerCommon{Name: "e1"}}, + {EphemeralContainerCommon: v1.EphemeralContainerCommon{Name: "e2"}}, + }, + }, + wantContainers: []string{"i1", "i2", "c1", "c2"}, + mask: InitContainers | Containers, + }, + { + desc: "ephemeral containers", + spec: &v1.PodSpec{ + Containers: []v1.Container{ + {Name: "c1"}, + {Name: "c2"}, + }, + InitContainers: []v1.Container{ + {Name: "i1"}, + {Name: "i2"}, + }, + EphemeralContainers: []v1.EphemeralContainer{ + {EphemeralContainerCommon: v1.EphemeralContainerCommon{Name: "e1"}}, + {EphemeralContainerCommon: v1.EphemeralContainerCommon{Name: "e2"}}, + }, + }, + wantContainers: []string{"e1", "e2"}, + mask: EphemeralContainers, + }, + { + desc: "all container types", + spec: &v1.PodSpec{ + Containers: []v1.Container{ + {Name: "c1"}, + {Name: "c2"}, + }, + InitContainers: []v1.Container{ + {Name: "i1"}, + {Name: "i2"}, + }, + EphemeralContainers: []v1.EphemeralContainer{ + {EphemeralContainerCommon: v1.EphemeralContainerCommon{Name: "e1"}}, + {EphemeralContainerCommon: v1.EphemeralContainerCommon{Name: "e2"}}, + }, + }, + wantContainers: []string{"i1", "i2", "c1", "c2", "e1", "e2"}, + mask: AllContainers, + }, + } + + for _, tc := range testCases { + t.Run(tc.desc, func(t *testing.T) { + gotContainers := []string{} + for c, containerType := range ContainerIter(tc.spec, tc.mask) { + gotContainers = append(gotContainers, c.Name) + + switch containerType { + case InitContainers: + if c.Name[0] != 'i' { + t.Errorf("ContainerIter() yielded container type InitContainers for container %q", c.Name) + } + case Containers: + if c.Name[0] != 'c' { + t.Errorf("ContainerIter() yielded container type Containers for container %q", c.Name) + } + case EphemeralContainers: + if c.Name[0] != 'e' { + t.Errorf("ContainerIter() yielded container type EphemeralContainers for container %q", c.Name) + } + default: + t.Errorf("ContainerIter() yielded unknown container type %d", containerType) + } + } + + if !cmp.Equal(gotContainers, tc.wantContainers) { + t.Errorf("ContainerIter() = %+v, want %+v", gotContainers, tc.wantContainers) + } + }) + } +} + func TestPodSecrets(t *testing.T) { // Stub containing all possible secret references in a pod. // The names of the referenced secrets match struct paths detected by reflection.