diff --git a/pkg/zones/env.go b/pkg/zones/env.go index c51bf83..32e30a4 100644 --- a/pkg/zones/env.go +++ b/pkg/zones/env.go @@ -22,11 +22,23 @@ func (m *Zones) Env(export bool) *Env { } } +// Zones returns the list of Zone IDs +func (m *Env) Zones() []int { + var zones []int + + m.ForEachZone(func(z *Zone) bool { + zones = append(zones, z.ID) + return false + }) + + return zones +} + // WriteTo generates environment variables for shell scripts func (m *Env) WriteTo(w io.Writer) (int64, error) { var buf bytes.Buffer - m.writeEnvVarFn(&buf, genEnvZones, "ZONES") + m.writeEnvVar(&buf, genEnvInts(m.Zones()), "ZONES") m.ForEachZone(func(z *Zone) bool { m.writeEnvZone(&buf, z) return false @@ -51,16 +63,6 @@ func (m *Env) writeEnvZone(w io.Writer, z *Zone) { } } -func (m *Env) writeEnvVarFn(w io.Writer, fn func(*Env) string, name string, args ...any) { - var value string - - if fn != nil { - value = fn(m) - } - - m.writeEnvVar(w, value, name, args...) -} - func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) { var prefix string @@ -79,23 +81,27 @@ func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) { } } -func genEnvZones(m *Env) string { - var s []string - - m.ForEachZone(func(z *Zone) bool { - s = append(s, fmt.Sprintf("%v", z.ID)) - return false - }) +func genEnvInts(a []int) string { + s := make([]string, len(a)) + for i, id := range a { + s[i] = fmt.Sprintf("%v", id) + } return strings.Join(s, " ") } func genEnvZoneNodes(z *Zone) string { - s := make([]string, 0, len(z.Machines)) - for _, p := range z.Machines { - s = append(s, p.Name) + if n := z.Len(); n > 0 { + s := make([]string, 0, n) + + z.ForEachMachine(func(p *Machine) bool { + s = append(s, p.Name) + return false + }) + + return strings.Join(s, " ") } - return strings.Join(s, " ") + return "" } func getRingZeroGatewayID(z *Zone) int {