jpictl/pkg/cluster/env.go

package cluster

import (
	"bytes"
	"fmt"
	"io"
	"sort"
	"strings"

	"darvaza.org/core"

	"git.jpi.io/amery/jpictl/pkg/rings"
)

// Env is a shell environment factory for this cluster
type Env struct {
	ZoneIterator
	RegionIterator

	cephFSID string
	export   bool
}

// Env returns a shell environment factory
func (m *Cluster) Env(export bool) (*Env, error) {
	fsid, err := m.GetCephFSID()
	if err != nil {
		return nil, err
	}

	env := &Env{
		ZoneIterator:   m,
		RegionIterator: m,
		cephFSID:       fsid.String(),
		export:         export,
	}

	return env, nil
}

// Zones returns the list of Zone IDs of a region,
// or from all if none is specified.
func (m *Env) Zones(r *Region) []rings.ZoneID {
	var zones []rings.ZoneID

	iter := core.IIf[ZoneIterator](r != nil, r, m)

	iter.ForEachZone(func(z *Zone) bool {
		zones = append(zones, z.ID)
		return false
	})

	core.SliceSort(zones, cmpOrdered[rings.ZoneID])
	return zones
}

// Regions returns the list of primary regions
func (m *Env) Regions() []string {
	var regions []string

	m.ForEachRegion(func(r *Region) bool {
		if r.IsPrimary() {
			regions = append(regions, r.Name)
		}

		return false
	})

	sort.Strings(regions)
	return regions
}

func (m *Env) RegionsID() (regions []rings.RegionID) {
	m.ForEachRegion(func(r *Region) bool {
		if r.IsPrimary() {
			regions = append(regions, r.ID)
		}

		return false
	})

	core.SliceSort(regions, cmpOrdered[rings.RegionID])
	return regions
}

// WriteTo generates environment variables for shell scripts
func (m *Env) WriteTo(w io.Writer) (int64, error) {
	var buf bytes.Buffer

	if m.cephFSID != "" {
		m.writeEnvVar(&buf, m.cephFSID, "FSID")
	}

	m.writeEnvVar(&buf, genEnvStrings(m.Regions()), "REGIONS")
	m.writeEnvVar(&buf, genEnvInts(m.RegionsID()), "REGIONS_ID")

	m.ForEachRegion(func(r *Region) bool {
		if r.IsPrimary() {
			m.writeEnvRegion(&buf, r)
		}
		return false
	})

	return buf.WriteTo(w)
}

func (m *Env) writeEnvRegion(w io.Writer, r *Region) {
	regionID := r.ID

	// REGION{regionID}_NAME
	m.writeEnvVar(w, r.Name, "REGION%v_%s", regionID, "NAME")

	// REGION{regionID}_ZONES
	m.writeEnvVar(w, genEnvInts(m.Zones(r)), "REGION%v_%s", regionID, "ZONES")

	r.ForEachZone(func(z *Zone) bool {
		m.writeEnvZone(w, r, z)
		return false
	})
}

func (m *Env) writeEnvZone(w io.Writer, r *Region, z *Zone) {
	zonePrefix := fmt.Sprintf("REGION%v_ZONE%v", r.ID, z.ID)
	monPrefix := zonePrefix + "_MON"

	// ZONE{zoneID}
	m.writeEnvVar(w, genEnvZoneNodes(z), zonePrefix)

	// ZONE{zoneID}_NAME
	m.writeEnvVar(w, z.Name, zonePrefix+"_NAME")

	// ZONE{zoneID}_GW
	gateways, _ := z.GatewayIDs()
	m.writeEnvVar(w, genEnvInts(gateways), zonePrefix+"_GW")

	// Ceph
	monitors := z.GetCephMonitors()
	// MON{zoneID}_NAME
	m.writeEnvVar(w, genEnvZoneCephMonNames(monitors), monPrefix)
	// MON{zoneID}_IP
	m.writeEnvVar(w, genEnvZoneCephMonIPs(monitors), monPrefix+"_IP")
	// MON{zoneID}_ID
	m.writeEnvVar(w, genEnvZoneCephMonIDs(monitors), monPrefix+"_ID")
}

func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
	var prefix string

	if m.export {
		prefix = "export "
	}

	if len(args) > 0 {
		name = fmt.Sprintf(name, args...)
	}

	if name != "" {
		value = strings.TrimSpace(value)

		if value == "" {
			_, _ = fmt.Fprintf(w, "%s%s=\n", prefix, name)
		} else {
			_, _ = fmt.Fprintf(w, "%s%s=%q\n", prefix, name, value)
		}
	}
}

func genEnvInts[T core.Signed](values []T) string {
	var buf bytes.Buffer

	for _, v := range values {
		if buf.Len() > 0 {
			_, _ = buf.WriteRune(' ')
		}

		_, _ = buf.WriteString(fmt.Sprintf("%v", int64(v)))
	}

	return buf.String()
}

func genEnvStrings(values []string) string {
	return strings.Join(values, " ")
}

func genEnvZoneNodes(z *Zone) string {
	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 genEnvStrings(s)
	}
	return ""
}

func genEnvZoneRegion(z *Zone) string {
	if z != nil && z.region != nil {
		return z.region.Name
	}
	return ""
}

func genEnvZoneCephMonNames(m Machines) string {
	var buf strings.Builder
	m.ForEachMachine(func(p *Machine) bool {
		if buf.Len() > 0 {
			_, _ = buf.WriteRune(' ')
		}
		_, _ = buf.WriteString(p.Name)

		return false
	})
	return buf.String()
}

func genEnvZoneCephMonIPs(m Machines) string {
	var buf strings.Builder
	m.ForEachMachine(func(p *Machine) bool {
		addr := p.RingOneAddress()

		if buf.Len() > 0 {
			_, _ = buf.WriteRune(' ')
		}
		_, _ = buf.WriteString(addr.String())

		return false
	})
	return buf.String()
}

func genEnvZoneCephMonIDs(m Machines) string {
	var buf strings.Builder
	m.ForEachMachine(func(p *Machine) bool {
		if buf.Len() > 0 {
			_, _ = buf.WriteRune(' ')
		}
		_, _ = fmt.Fprintf(&buf, "%v", p.ID)

		return false
	})
	return buf.String()
}

func cmpOrdered[T core.Ordered](a, b T) int {
	switch {
	case a == b:
		return 0
	case a < b:
		return -1
	default:
		return 1
	}
}