Skip to content

Commit 4c0233b

Browse files
authored
Merge pull request kubernetes#131264 from tallclair/container-util
Add ContainerIter utility for ranging over pod containers
2 parents 5bb3aa3 + 5928fc0 commit 4c0233b

File tree

4 files changed

+338
-30
lines changed

4 files changed

+338
-30
lines changed

pkg/api/pod/util.go

Lines changed: 28 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ limitations under the License.
1717
package pod
1818

1919
import (
20+
"iter"
2021
"strings"
2122

2223
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -60,28 +61,40 @@ type ContainerVisitor func(container *api.Container, containerType ContainerType
6061
// visiting is short-circuited. VisitContainers returns true if visiting completes,
6162
// false if visiting was short-circuited.
6263
func VisitContainers(podSpec *api.PodSpec, mask ContainerType, visitor ContainerVisitor) bool {
63-
if mask&InitContainers != 0 {
64-
for i := range podSpec.InitContainers {
65-
if !visitor(&podSpec.InitContainers[i], InitContainers) {
66-
return false
67-
}
64+
for c, t := range ContainerIter(podSpec, mask) {
65+
if !visitor(c, t) {
66+
return false
6867
}
6968
}
70-
if mask&Containers != 0 {
71-
for i := range podSpec.Containers {
72-
if !visitor(&podSpec.Containers[i], Containers) {
73-
return false
69+
return true
70+
}
71+
72+
// ContainerIter returns an iterator over all containers in the given pod spec with a masked type.
73+
// The iteration order is InitContainers, then main Containers, then EphemeralContainers.
74+
func ContainerIter(podSpec *api.PodSpec, mask ContainerType) iter.Seq2[*api.Container, ContainerType] {
75+
return func(yield func(*api.Container, ContainerType) bool) {
76+
if mask&InitContainers != 0 {
77+
for i := range podSpec.InitContainers {
78+
if !yield(&podSpec.InitContainers[i], InitContainers) {
79+
return
80+
}
7481
}
7582
}
76-
}
77-
if mask&EphemeralContainers != 0 {
78-
for i := range podSpec.EphemeralContainers {
79-
if !visitor((*api.Container)(&podSpec.EphemeralContainers[i].EphemeralContainerCommon), EphemeralContainers) {
80-
return false
83+
if mask&Containers != 0 {
84+
for i := range podSpec.Containers {
85+
if !yield(&podSpec.Containers[i], Containers) {
86+
return
87+
}
88+
}
89+
}
90+
if mask&EphemeralContainers != 0 {
91+
for i := range podSpec.EphemeralContainers {
92+
if !yield((*api.Container)(&podSpec.EphemeralContainers[i].EphemeralContainerCommon), EphemeralContainers) {
93+
return
94+
}
8195
}
8296
}
8397
}
84-
return true
8598
}
8699

87100
// Visitor is called with each object name, and returns true if visiting should continue

pkg/api/pod/util_test.go

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,147 @@ func TestVisitContainers(t *testing.T) {
208208
}
209209
}
210210

211+
func TestContainerIter(t *testing.T) {
212+
testCases := []struct {
213+
desc string
214+
spec *api.PodSpec
215+
wantContainers []string
216+
mask ContainerType
217+
}{
218+
{
219+
desc: "empty podspec",
220+
spec: &api.PodSpec{},
221+
wantContainers: []string{},
222+
mask: AllContainers,
223+
},
224+
{
225+
desc: "regular containers",
226+
spec: &api.PodSpec{
227+
Containers: []api.Container{
228+
{Name: "c1"},
229+
{Name: "c2"},
230+
},
231+
InitContainers: []api.Container{
232+
{Name: "i1"},
233+
{Name: "i2"},
234+
},
235+
EphemeralContainers: []api.EphemeralContainer{
236+
{EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e1"}},
237+
{EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e2"}},
238+
},
239+
},
240+
wantContainers: []string{"c1", "c2"},
241+
mask: Containers,
242+
},
243+
{
244+
desc: "init containers",
245+
spec: &api.PodSpec{
246+
Containers: []api.Container{
247+
{Name: "c1"},
248+
{Name: "c2"},
249+
},
250+
InitContainers: []api.Container{
251+
{Name: "i1"},
252+
{Name: "i2"},
253+
},
254+
EphemeralContainers: []api.EphemeralContainer{
255+
{EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e1"}},
256+
{EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e2"}},
257+
},
258+
},
259+
wantContainers: []string{"i1", "i2"},
260+
mask: InitContainers,
261+
},
262+
{
263+
desc: "init + main containers",
264+
spec: &api.PodSpec{
265+
Containers: []api.Container{
266+
{Name: "c1"},
267+
{Name: "c2"},
268+
},
269+
InitContainers: []api.Container{
270+
{Name: "i1"},
271+
{Name: "i2"},
272+
},
273+
EphemeralContainers: []api.EphemeralContainer{
274+
{EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e1"}},
275+
{EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e2"}},
276+
},
277+
},
278+
wantContainers: []string{"i1", "i2", "c1", "c2"},
279+
mask: InitContainers | Containers,
280+
},
281+
{
282+
desc: "ephemeral containers",
283+
spec: &api.PodSpec{
284+
Containers: []api.Container{
285+
{Name: "c1"},
286+
{Name: "c2"},
287+
},
288+
InitContainers: []api.Container{
289+
{Name: "i1"},
290+
{Name: "i2"},
291+
},
292+
EphemeralContainers: []api.EphemeralContainer{
293+
{EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e1"}},
294+
{EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e2"}},
295+
},
296+
},
297+
wantContainers: []string{"e1", "e2"},
298+
mask: EphemeralContainers,
299+
},
300+
{
301+
desc: "all container types",
302+
spec: &api.PodSpec{
303+
Containers: []api.Container{
304+
{Name: "c1"},
305+
{Name: "c2"},
306+
},
307+
InitContainers: []api.Container{
308+
{Name: "i1"},
309+
{Name: "i2"},
310+
},
311+
EphemeralContainers: []api.EphemeralContainer{
312+
{EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e1"}},
313+
{EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e2"}},
314+
},
315+
},
316+
wantContainers: []string{"i1", "i2", "c1", "c2", "e1", "e2"},
317+
mask: AllContainers,
318+
},
319+
}
320+
321+
for _, tc := range testCases {
322+
t.Run(tc.desc, func(t *testing.T) {
323+
gotContainers := []string{}
324+
for c, containerType := range ContainerIter(tc.spec, tc.mask) {
325+
gotContainers = append(gotContainers, c.Name)
326+
327+
switch containerType {
328+
case InitContainers:
329+
if c.Name[0] != 'i' {
330+
t.Errorf("ContainerIter() yielded container type InitContainers for container %q", c.Name)
331+
}
332+
case Containers:
333+
if c.Name[0] != 'c' {
334+
t.Errorf("ContainerIter() yielded container type Containers for container %q", c.Name)
335+
}
336+
case EphemeralContainers:
337+
if c.Name[0] != 'e' {
338+
t.Errorf("ContainerIter() yielded container type EphemeralContainers for container %q", c.Name)
339+
}
340+
default:
341+
t.Errorf("ContainerIter() yielded unknown container type %d", containerType)
342+
}
343+
}
344+
345+
if !cmp.Equal(gotContainers, tc.wantContainers) {
346+
t.Errorf("ContainerIter() = %+v, want %+v", gotContainers, tc.wantContainers)
347+
}
348+
})
349+
}
350+
}
351+
211352
func TestPodSecrets(t *testing.T) {
212353
// Stub containing all possible secret references in a pod.
213354
// The names of the referenced secrets match struct paths detected by reflection.

pkg/api/v1/pod/util.go

Lines changed: 28 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ package pod
1818

1919
import (
2020
"fmt"
21+
"iter"
2122
"time"
2223

2324
v1 "k8s.io/api/core/v1"
@@ -105,28 +106,40 @@ func skipEmptyNames(visitor Visitor) Visitor {
105106
// visiting is short-circuited. VisitContainers returns true if visiting completes,
106107
// false if visiting was short-circuited.
107108
func VisitContainers(podSpec *v1.PodSpec, mask ContainerType, visitor ContainerVisitor) bool {
108-
if mask&InitContainers != 0 {
109-
for i := range podSpec.InitContainers {
110-
if !visitor(&podSpec.InitContainers[i], InitContainers) {
111-
return false
112-
}
109+
for c, t := range ContainerIter(podSpec, mask) {
110+
if !visitor(c, t) {
111+
return false
113112
}
114113
}
115-
if mask&Containers != 0 {
116-
for i := range podSpec.Containers {
117-
if !visitor(&podSpec.Containers[i], Containers) {
118-
return false
114+
return true
115+
}
116+
117+
// ContainerIter returns an iterator over all containers in the given pod spec with a masked type.
118+
// The iteration order is InitContainers, then main Containers, then EphemeralContainers.
119+
func ContainerIter(podSpec *v1.PodSpec, mask ContainerType) iter.Seq2[*v1.Container, ContainerType] {
120+
return func(yield func(*v1.Container, ContainerType) bool) {
121+
if mask&InitContainers != 0 {
122+
for i := range podSpec.InitContainers {
123+
if !yield(&podSpec.InitContainers[i], InitContainers) {
124+
return
125+
}
119126
}
120127
}
121-
}
122-
if mask&EphemeralContainers != 0 {
123-
for i := range podSpec.EphemeralContainers {
124-
if !visitor((*v1.Container)(&podSpec.EphemeralContainers[i].EphemeralContainerCommon), EphemeralContainers) {
125-
return false
128+
if mask&Containers != 0 {
129+
for i := range podSpec.Containers {
130+
if !yield(&podSpec.Containers[i], Containers) {
131+
return
132+
}
133+
}
134+
}
135+
if mask&EphemeralContainers != 0 {
136+
for i := range podSpec.EphemeralContainers {
137+
if !yield((*v1.Container)(&podSpec.EphemeralContainers[i].EphemeralContainerCommon), EphemeralContainers) {
138+
return
139+
}
126140
}
127141
}
128142
}
129-
return true
130143
}
131144

132145
// VisitPodSecretNames invokes the visitor function with the name of every secret

0 commit comments

Comments
 (0)