jpictl/pkg/zones/zones.go

// Package zones contains information about the cluster
package zones

import (
	"io/fs"
	"path/filepath"
	"sort"

	"github.com/hack-pad/hackpadfs/os"

	"darvaza.org/resolver"
)

var (
	_ MachineIterator = Machines(nil)
	_ sort.Interface  = Machines(nil)

	_ MachineIterator = (*Zone)(nil)
	_ MachineIterator = (*Zones)(nil)
	_ ZoneIterator    = (*Zones)(nil)
)

// A MachineIterator is a set of Machines we can iterate on
type MachineIterator interface {
	ForEachMachine(func(*Machine) bool)
}

// A ZoneIterator is a set of Zones we can iterate on
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

	ID   int    `toml:"id"`
	Name string `toml:"name"`

	Machines `toml:"machines"`
}

func (z *Zone) String() string {
	return z.Name
}

// SetGateway configures a machine to be the zone's ring0 gateway
func (z *Zone) SetGateway(gatewayID int, enabled bool) error {
	var err error
	var found bool

	z.ForEachMachine(func(p *Machine) bool {
		if p.ID == gatewayID {
			found = true
			err = p.SetGateway(enabled)

			return true
		}
		return false
	})

	switch {
	case err != nil:
		return err
	case !found:
		return fs.ErrNotExist
	default:
		return nil
	}
}

// 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
	resolver resolver.Resolver
	domain   string

	Zones []*Zone `toml:"zones"`
}

// ForEachMachine calls a function for each Machine in the cluster
// until instructed to terminate the loop
func (m *Zones) ForEachMachine(fn func(*Machine) bool) {
	m.ForEachZone(func(z *Zone) bool {
		var term bool

		z.ForEachMachine(func(p *Machine) bool {
			term = fn(p)
			return term
		})

		return term
	})
}

// ForEachZone calls a function for each Zone in the cluster
// until instructed to terminate the loop
func (m *Zones) ForEachZone(fn func(*Zone) bool) {
	for _, p := range m.Zones {
		if fn(p) {
			// terminate
			return
		}
	}
}

// GetMachineByName looks for a machine with the specified
// name on any zone
func (m *Zones) GetMachineByName(name string) (*Machine, bool) {
	var out *Machine

	if name != "" {
		m.ForEachMachine(func(p *Machine) bool {
			if p.Name == name {
				out = p
			}

			return out != nil
		})
	}

	return out, out != nil
}

// NewFS builds a [Zones] tree using the given directory
func NewFS(dir fs.FS, domain string) (*Zones, error) {
	lockuper := resolver.NewCloudflareLookuper()

	z := &Zones{
		dir:      dir,
		resolver: resolver.NewResolver(lockuper),
		domain:   domain,
	}

	if err := z.scan(); err != nil {
		return nil, err
	}

	return z, nil
}

// New builds a [Zones] tree using the given directory
func New(dir, domain string) (*Zones, error) {
	dir, err := filepath.Abs(dir)
	if err != nil {
		return nil, err
	}

	base, err := os.NewFS().Sub(dir[1:])
	if err != nil {
		return nil, err
	}

	return NewFS(base, domain)
}