pkg/zones/scan.go

@@ -93,12 +93,7 @@ func (m *Zones) scanSort() error {
 	})
 
 	m.ForEachZone(func(z *Zone) bool {
-		sort.SliceStable(z.Machines, func(i, j int) bool {
-			id1 := z.Machines[i].ID
-			id2 := z.Machines[j].ID
-			return id1 < id2
-		})
-
+		sort.Sort(z)
 		return false
 	})
 

pkg/zones/zones.go

@@ -4,6 +4,7 @@ package zones
 import (
 	"io/fs"
 	"path/filepath"
+	"sort"
 
 	"github.com/hack-pad/hackpadfs/os"
 
@@ -11,6 +12,9 @@ import (
 )
 
 var (
+	_ MachineIterator = Machines(nil)
+	_ sort.Interface  = Machines(nil)
+
 	_ MachineIterator = (*Zone)(nil)
 	_ MachineIterator = (*Zones)(nil)
 	_ ZoneIterator    = (*Zones)(nil)
@@ -26,6 +30,62 @@ type ZoneIterator interface {
 	ForEachZone(func(*Zone) bool)
 }
 
+// Machines is a list of Machine objects
+type Machines []*Machine
+
+// ForEachMachine calls a function for each Machine in the list
+// until instructed to terminate the loop
+func (m Machines) ForEachMachine(fn func(*Machine) bool) {
+	for _, p := range m {
+		if fn(p) {
+			return
+		}
+	}
+}
+
+// Len returns the number of machines in the list
+func (m Machines) Len() int {
+	return len(m)
+}
+
+// Less implements sort.Interface to sort the list
+func (m Machines) Less(i, j int) bool {
+	a, b := m[i], m[j]
+	za, zb := a.Zone(), b.Zone()
+
+	switch {
+	case za == zb:
+		return a.ID < b.ID
+	default:
+		return za < zb
+	}
+}
+
+// Swap implements sort.Interface to sort the list
+func (m Machines) Swap(i, j int) {
+	m[i], m[j] = m[j], m[i]
+}
+
+// FilterMachines produces a subset of the machines offered by the given
+// iterator fulfilling a condition
+func FilterMachines(m MachineIterator, cond func(*Machine) bool) (Machines, int) {
+	var out []*Machine
+
+	if cond == nil {
+		// unconditional
+		cond = func(*Machine) bool { return true }
+	}
+
+	m.ForEachMachine(func(p *Machine) bool {
+		if cond(p) {
+			out = append(out, p)
+		}
+		return false
+	})
+
+	return out, len(out)
+}
+
 // Zone represents one zone in a cluster
 type Zone struct {
 	zones *Zones
@@ -33,23 +93,13 @@ type Zone struct {
 	ID   int    `toml:"id"`
 	Name string `toml:"name"`
 
-	Machines []*Machine `toml:"machines"`
+	Machines `toml:"machines"`
 }
 
 func (z *Zone) String() string {
 	return z.Name
 }
 
-// ForEachMachine calls a function for each Machine in the zone
-// until instructed to terminate the loop
-func (z *Zone) ForEachMachine(fn func(*Machine) bool) {
-	for _, p := range z.Machines {
-		if fn(p) {
-			return
-		}
-	}
-}
-
 // SetGateway configures a machine to be the zone's ring0 gateway
 func (z *Zone) SetGateway(gatewayID int, enabled bool) error {
 	var err error
@@ -75,6 +125,19 @@ func (z *Zone) SetGateway(gatewayID int, enabled bool) error {
 	}
 }
 
+// GatewayIDs returns the list of IDs of machines that act as ring0 gateways
+func (z *Zone) GatewayIDs() ([]int, int) {
+	var out []int
+	z.ForEachMachine(func(p *Machine) bool {
+		if p.IsGateway() {
+			out = append(out, p.ID)
+		}
+		return false
+	})
+
+	return out, len(out)
+}
+
 // Zones represents all zones in a cluster
 type Zones struct {
 	dir      fs.FS