Merge pull request 'env: export REGIONS and ZONE{zoneID}_REGION' (#45) from pr-amery-env-regions into main

Reviewed-on: #45

pkg/cluster/cluster_scan.go

@@ -4,6 +4,7 @@ import (
 	"io/fs"
 	"path"
 	"sort"
+	"strings"
 
 	"darvaza.org/core"
 )
@@ -12,6 +13,10 @@ const (
 	// ZoneRegionsFileName indicates the file containing
 	// region names as references
 	ZoneRegionsFileName = "regions"
+
+	// RegionClusterTokenFileName contains the kubernetes
+	// token of the cluster this region represents
+	RegionClusterTokenFileName = "k8s_token"
 )
 
 func (m *Cluster) scan(opts *ScanOptions) error {
@@ -186,6 +191,8 @@ func (z *Zone) scan() error {
 		switch {
 		case name == ZoneRegionsFileName:
 			err = z.loadRegions()
+		case name == RegionClusterTokenFileName:
+			err = z.loadClusterToken()
 		case e.IsDir():
 			err = z.scanSubdirectory(name)
 		default:
@@ -218,6 +225,32 @@ func (z *Zone) loadRegions() error {
 	return err
 }
 
+func (z *Zone) loadClusterToken() error {
+	var token string
+
+	filename := path.Join(z.Name, RegionClusterTokenFileName)
+	lines, err := z.zones.ReadLines(filename)
+	if err != nil {
+		return err
+	}
+
+	// first non-empty line
+	for _, s := range lines {
+		s = strings.TrimSpace(s)
+		if s != "" {
+			token = s
+			break
+		}
+	}
+
+	err = z.zones.setRegionClusterToken(z.Name, token)
+	if err != nil {
+		err = core.Wrap(err, filename)
+	}
+
+	return err
+}
+
 func (z *Zone) scanSubdirectory(name string) error {
 	m := &Machine{
 		zone:   z,

pkg/cluster/env.go

@@ -4,12 +4,14 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"sort"
 	"strings"
 )
 
 // Env is a shell environment factory for this cluster
 type Env struct {
 	ZoneIterator
+	RegionIterator
 
 	cephFSID string
 	export   bool
@@ -23,9 +25,10 @@ func (m *Cluster) Env(export bool) (*Env, error) {
 	}
 
 	env := &Env{
-		ZoneIterator: m,
-		cephFSID:     fsid.String(),
-		export:       export,
+		ZoneIterator:   m,
+		RegionIterator: m,
+		cephFSID:       fsid.String(),
+		export:         export,
 	}
 
 	return env, nil
@@ -43,6 +46,22 @@ func (m *Env) Zones() []int {
 	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.Cluster != nil {
+			regions = append(regions, r.Name)
+		}
+
+		return false
+	})
+
+	sort.Strings(regions)
+	return regions
+}
+
 // WriteTo generates environment variables for shell scripts
 func (m *Env) WriteTo(w io.Writer) (int64, error) {
 	var buf bytes.Buffer
@@ -51,6 +70,7 @@ func (m *Env) WriteTo(w io.Writer) (int64, error) {
 		m.writeEnvVar(&buf, m.cephFSID, "FSID")
 	}
 
+	m.writeEnvVarStrings(&buf, m.Regions(), "REGIONS")
 	m.writeEnvVarInts(&buf, m.Zones(), "ZONES")
 
 	m.ForEachZone(func(z *Zone) bool {
@@ -74,6 +94,9 @@ func (m *Env) writeEnvZone(w io.Writer, z *Zone) {
 	gateways, _ := z.GatewayIDs()
 	m.writeEnvVarInts(w, gateways, "ZONE%v_%s", zoneID, "GW")
 
+	// ZONE{zoneID}_REGION
+	m.writeEnvVar(w, genEnvZoneRegion(z), "ZONE%v_%s", zoneID, "REGION")
+
 	// Ceph
 	monitors := z.GetCephMonitors()
 	// MON{zoneID}_NAME
@@ -85,22 +108,29 @@ func (m *Env) writeEnvZone(w io.Writer, z *Zone) {
 }
 
 func (m *Env) writeEnvVarInts(w io.Writer, value []int, name string, args ...any) {
-	var s string
+	var buf bytes.Buffer
 
-	if n := len(value); n > 0 {
-		var buf bytes.Buffer
-
-		for i, v := range value {
-			if i != 0 {
-				_, _ = fmt.Fprint(&buf, " ")
-			}
-			_, _ = fmt.Fprintf(&buf, "%v", v)
+	for _, v := range value {
+		if buf.Len() > 0 {
+			_, _ = fmt.Fprint(&buf, " ")
 		}
-
-		s = buf.String()
+		_, _ = fmt.Fprintf(&buf, "%v", v)
 	}
 
-	m.writeEnvVar(w, s, name, args...)
+	m.writeEnvVar(w, buf.String(), name, args...)
+}
+
+func (m *Env) writeEnvVarStrings(w io.Writer, value []string, name string, args ...any) {
+	var buf bytes.Buffer
+
+	for _, v := range value {
+		if buf.Len() > 0 {
+			_, _ = fmt.Fprint(&buf, " ")
+		}
+		_, _ = fmt.Fprintf(&buf, "%s", v)
+	}
+
+	m.writeEnvVar(w, buf.String(), name, args...)
 }
 
 func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
@@ -117,7 +147,11 @@ func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
 	if name != "" {
 		value = strings.TrimSpace(value)
 
-		_, _ = fmt.Fprintf(w, "%s%s=%q\n", prefix, name, value)
+		if value == "" {
+			_, _ = fmt.Fprintf(w, "%s%s=\n", prefix, name)
+		} else {
+			_, _ = fmt.Fprintf(w, "%s%s=%q\n", prefix, name, value)
+		}
 	}
 }
 
@@ -135,6 +169,21 @@ func genEnvZoneNodes(z *Zone) string {
 	return ""
 }
 
+func genEnvZoneRegion(z *Zone) string {
+	var region string
+
+	z.ForEachRegion(func(r *Region) bool {
+		if r.Cluster != nil {
+			region = r.Name
+			return true
+		}
+
+		return false
+	})
+
+	return region
+}
+
 func genEnvZoneCephMonNames(m Machines) string {
 	var buf strings.Builder
 	m.ForEachMachine(func(p *Machine) bool {

pkg/cluster/regions.go

@@ -8,14 +8,23 @@ import (
 var (
 	_ MachineIterator = (*Region)(nil)
 	_ ZoneIterator    = (*Region)(nil)
+
+	_ RegionIterator = (*Zone)(nil)
+	_ RegionIterator = (*Cluster)(nil)
 )
 
+// A RegionIterator is a set of Regions we can iterate on
+type RegionIterator interface {
+	ForEachRegion(func(*Region) bool)
+}
+
 // Region represents a group of zones geographically related
 type Region struct {
 	m     *Cluster
 	zones []*Zone
 
 	Name    string
+	Cluster *string  `json:",omitempty" yaml:",omitempty"`
 	Regions []string `json:",omitempty" yaml:",omitempty"`
 }
 
@@ -106,6 +115,26 @@ func (m *Cluster) setRegionZones(name string, zones ...*Zone) {
 	})
 }
 
+func (m *Cluster) setRegionClusterToken(name string, token string) error {
+	for i := range m.Regions {
+		r := &m.Regions[i]
+
+		if r.Name == name {
+			// found
+			r.Cluster = &token
+			return nil
+		}
+	}
+
+	// new
+	m.Regions = append(m.Regions, Region{
+		m:       m,
+		Name:    name,
+		Cluster: &token,
+	})
+	return nil
+}
+
 func (m *Cluster) appendRegionRegions(name string, subs ...string) {
 	for i := range m.Regions {
 		r := &m.Regions[i]
@@ -124,6 +153,28 @@ func (m *Cluster) appendRegionRegions(name string, subs ...string) {
 	})
 }
 
+// ForEachRegion calls a function on all regions this zone belongs to.
+func (z *Zone) ForEachRegion(fn func(*Region) bool) {
+	if fn == nil {
+		return
+	}
+
+	z.zones.ForEachRegion(func(r *Region) bool {
+		var match bool
+
+		r.ForEachZone(func(z2 *Zone) bool {
+			match = (z == z2)
+			return match
+		})
+
+		if match && fn(r) {
+			return true
+		}
+
+		return false
+	})
+}
+
 func (z *Zone) appendRegions(regions ...string) error {
 	for _, s := range regions {
 		// TODO: validate
@@ -147,7 +198,7 @@ func (m *Cluster) finishRegion(r *Region) {
 	r.m = m
 	sub := []string{}
 	for _, name := range r.Regions {
-		r2, ok := m.getRegion(name)
+		r2, ok := m.getFinishRegion(name)
 		if !ok {
 			m.warn(nil).WithField("region", name).Print("unknown region")
 			continue
@@ -162,9 +213,7 @@ func (m *Cluster) finishRegion(r *Region) {
 func (m *Cluster) getRegion(name string) (*Region, bool) {
 	for i := range m.Regions {
 		r := &m.Regions[i]
-
 		if name == r.Name {
-			m.finishRegion(r)
 			return r, true
 		}
 	}
@@ -172,17 +221,22 @@ func (m *Cluster) getRegion(name string) (*Region, bool) {
 	return nil, false
 }
 
+func (m *Cluster) getFinishRegion(name string) (*Region, bool) {
+	if r, ok := m.getRegion(name); ok {
+		m.finishRegion(r)
+		return r, true
+	}
+
+	return nil, false
+}
+
 // SyncRegions writes to the file system the regions this [Zone]
 // belongs to.
 func (z *Zone) SyncRegions() error {
 	err := z.syncZoneRegions()
 	if err == nil {
 		z.ForEachMachine(func(p *Machine) bool {
-			if p.IsActive() {
-				err = p.RemoveFile("region")
-			} else {
-				err = p.WriteStringFile("none\n", "region")
-			}
+			err = z.syncMachineRegions(p)
 			return err != nil
 		})
 	}
@@ -190,8 +244,24 @@ func (z *Zone) SyncRegions() error {
 	return err
 }
 
+func (*Zone) syncMachineRegions(p *Machine) error {
+	var err error
+
+	if p.IsActive() {
+		err = p.RemoveFile("region")
+	} else {
+		err = p.WriteStringFile("none\n", "region")
+	}
+
+	if err == nil {
+		err = p.RemoveFile(RegionClusterTokenFileName)
+	}
+
+	return err
+}
+
 func (z *Zone) syncZoneRegions() error {
-	name := filepath.Join(z.Name, "regions")
+	name := filepath.Join(z.Name, ZoneRegionsFileName)
 
 	if len(z.Regions) > 0 {
 		var buf bytes.Buffer
@@ -210,9 +280,25 @@ func (z *Zone) syncZoneRegions() error {
 // SyncRegions writes to the file system the regions covered
 // by this meta-region
 func (r *Region) SyncRegions() error {
-	name := filepath.Join(r.Name, "regions")
+	if err := r.syncRegionsFile(); err != nil {
+		return err
+	}
 
-	if len(r.Regions) > 0 {
+	return r.syncClusterFile()
+}
+
+func (r *Region) mkdir() error {
+	return r.m.MkdirAll(r.Name)
+}
+
+func (r *Region) syncRegionsFile() error {
+	var err error
+
+	name := filepath.Join(r.Name, ZoneRegionsFileName)
+
+	if len(r.Regions) == 0 {
+		err = r.m.RemoveFile(name)
+	} else if err = r.mkdir(); err == nil {
 		var buf bytes.Buffer
 
 		for _, s := range r.Regions {
@@ -220,12 +306,29 @@ func (r *Region) SyncRegions() error {
 			_, _ = buf.WriteRune('\n')
 		}
 
-		if err := r.m.MkdirAll(r.Name); err != nil {
-			return err
-		}
-
-		return r.m.WriteStringFile(buf.String(), name)
+		err = r.m.WriteStringFile(buf.String(), name)
 	}
 
-	return r.m.RemoveFile(name)
+	return err
+}
+
+func (r *Region) syncClusterFile() error {
+	var err error
+
+	name := filepath.Join(r.Name, RegionClusterTokenFileName)
+
+	if r.Cluster == nil {
+		err = r.m.RemoveFile(name)
+	} else if err = r.mkdir(); err == nil {
+		var buf bytes.Buffer
+
+		_, _ = buf.WriteString(*r.Cluster)
+		if buf.Len() > 0 {
+			_, _ = buf.WriteRune('\n')
+		}
+
+		err = r.m.WriteStringFile(buf.String(), name)
+	}
+
+	return err
 }