llmgw/internal/pgha/setup.go

package pgha

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/onboard"
)

// Node is a machine the setup reaches, and how.
type Node struct {
	Member
	Remote onboard.Remote
}

// Setup builds, or finds, a Patroni cluster.
type Setup struct {
	// DB are the database machines; the first is where the cluster starts.
	DB []Node
	// Witness is the third etcd vote when there are only two database
	// machines: the control plane, run locally. Ignored with three or more.
	Witness *Node
	// Clients are the addresses control planes connect from.
	Clients []string
	// Status reads the cluster from Patroni's API; nil asks over HTTP.
	Status func(ctx context.Context, hosts []string) (Cluster, error)
	// Emit reports progress: a step, what happened, and whether anything
	// was changed or found as it should be.
	Emit func(step, detail string, changed bool)
	// Wait bounds each wait for the cluster to come up.
	Wait time.Duration
}

// Local is a Remote for the machine this runs on, which init runs as root.
type Local struct{}

// Run runs script with sh here.
func (Local) Run(ctx context.Context, script string, stdin []byte) ([]byte, error) {
	cmd := exec.CommandContext(ctx, "sh", "-c", script)
	if stdin != nil {
		cmd.Stdin = bytes.NewReader(stdin)
	}
	var out, errb bytes.Buffer
	cmd.Stdout, cmd.Stderr = &out, &errb
	if err := cmd.Run(); err != nil {
		return append(out.Bytes(), errb.Bytes()...), err
	}
	return out.Bytes(), nil
}

func (s *Setup) emit(step string, changed bool, format string, a ...any) {
	if s.Emit != nil {
		s.Emit(step, fmt.Sprintf(format, a...), changed)
	}
}

// etcdMembers are the machines that vote: the database machines, and the
// witness when they are only two.
func (s *Setup) etcdMembers() []Node {
	out := append([]Node{}, s.DB...)
	if len(s.DB) == 2 && s.Witness != nil {
		out = append(out, *s.Witness)
	}
	return out
}

// Run sets the cluster up, or confirms it is, and returns the store address.
func (s *Setup) Run(ctx context.Context) (string, error) {
	if len(s.DB) < 2 {
		return "", fmt.Errorf("a replicated store needs at least two database machines, not %d", len(s.DB))
	}
	if len(s.DB) == 2 && s.Witness == nil {
		return "", fmt.Errorf("two database machines need a third etcd vote, or losing either stops failover")
	}
	if s.Wait == 0 {
		s.Wait = 3 * time.Minute
	}
	if s.Status == nil {
		s.Status = HTTPStatus
	}
	sec, err := s.secrets(ctx)
	if err != nil {
		return "", fmt.Errorf("passwords: %w", err)
	}
	if err := s.etcd(ctx); err != nil {
		return "", fmt.Errorf("etcd: %w", err)
	}
	if err := s.patroni(ctx, sec); err != nil {
		return "", fmt.Errorf("patroni: %w", err)
	}
	if err := s.database(ctx); err != nil {
		return "", fmt.Errorf("database: %w", err)
	}
	members := make([]Member, len(s.DB))
	for i, d := range s.DB {
		members[i] = d.Member
	}
	return DSN(members, sec.App), nil
}

// secretDir holds the passwords on every database machine, readable by root
// alone. Generated on the first, carried to the others.
const secretDir = "/root/llmgw-pg"

var secretNames = []string{"superuser", "replicator", "llmgw"}

func (s *Setup) secrets(ctx context.Context) (Secrets, error) {
	script := `set -e
sudo install -d -m 700 ` + secretDir + `
for f in superuser replicator llmgw; do
  sudo test -s ` + secretDir + `/$f.pw || sudo sh -c "umask 077; openssl rand -hex 24 > ` + secretDir + `/$f.pw"
done
for f in superuser replicator llmgw; do printf '%s=' $f; sudo cat ` + secretDir + `/$f.pw; done`
	out, err := s.DB[0].Remote.Run(ctx, script, nil)
	if err != nil {
		return Secrets{}, fmt.Errorf("on %s: %v: %s", s.DB[0].Name, err, strings.TrimSpace(string(out)))
	}
	got := parseSecrets(out)
	if got.Superuser == "" || got.Replicator == "" || got.App == "" {
		return Secrets{}, fmt.Errorf("could not read them back from %s", s.DB[0].Name)
	}
	for _, d := range s.DB[1:] {
		have, err := d.Remote.Run(ctx, `for f in superuser replicator llmgw; do
  if sudo test -s `+secretDir+`/$f.pw; then printf '%s=' $f; sudo cat `+secretDir+`/$f.pw; fi
done`, nil)
		if err != nil {
			return Secrets{}, fmt.Errorf("on %s: %v", d.Name, err)
		}
		if h := parseSecrets(have); h != (Secrets{}) {
			if h != got {
				// Not overwritten: a member holding other passwords belongs
				// to another cluster, or to an earlier one somebody meant to
				// keep, and replacing them would lock it out of itself.
				return Secrets{}, fmt.Errorf("%s holds passwords in %s that differ from %s's; "+
					"remove them there if that machine is meant to join this cluster", d.Name, secretDir, s.DB[0].Name)
			}
			continue
		}
		body := fmt.Sprintf("%s\n%s\n%s\n", got.Superuser, got.Replicator, got.App)
		if out, err := d.Remote.Run(ctx, `set -e
sudo install -d -m 700 `+secretDir+`
read su; read rep; read app
for kv in "superuser $su" "replicator $rep" "llmgw $app"; do
  set -- $kv
  printf '%s\n' "$2" | sudo sh -c "umask 077; cat > `+secretDir+`/$1.pw"
done`, []byte(body)); err != nil {
			return Secrets{}, fmt.Errorf("on %s: %v: %s", d.Name, err, strings.TrimSpace(string(out)))
		}
		s.emit("passwords", true, "carried to %s", d.Name)
	}
	s.emit("passwords", false, "kept in %s on every database machine", secretDir)
	return got, nil
}

func parseSecrets(out []byte) Secrets {
	var s Secrets
	for _, line := range strings.Split(string(out), "\n") {
		k, v, ok := strings.Cut(strings.TrimSpace(line), "=")
		if !ok {
			continue
		}
		switch k {
		case "superuser":
			s.Superuser = v
		case "replicator":
			s.Replicator = v
		case "llmgw":
			s.App = v
		}
	}
	return s
}

// etcd brings up the voting members, unless every one of them is already
// healthy in one cluster — in which case nothing is touched, since rewriting
// a member's configuration means emptying its data.
func (s *Setup) etcd(ctx context.Context) error {
	nodes := s.etcdMembers()
	var members []Member
	var endpoints []string
	for _, n := range nodes {
		members = append(members, n.Member)
		endpoints = append(endpoints, "http://"+n.IP+":2379")
	}
	if s.etcdHealthy(ctx, nodes[0], endpoints, len(nodes)) {
		s.emit("etcd", false, "%d members healthy; left as they are", len(nodes))
		return nil
	}

	changed := make([]bool, len(nodes))
	for i, n := range nodes {
		out, err := n.Remote.Run(ctx, `set -e
if ! command -v etcd >/dev/null 2>&1; then
  sudo DEBIAN_FRONTEND=noninteractive apt-get -qq update >/dev/null
  sudo DEBIAN_FRONTEND=noninteractive apt-get -qq install -y etcd-server etcd-client >/dev/null
fi
new=$(cat)
if [ "$(sudo cat /etc/default/etcd 2>/dev/null)" != "$new" ]; then
  sudo systemctl stop etcd >/dev/null 2>&1 || true
  sudo rm -rf /var/lib/etcd/default
  printf '%s\n' "$new" | sudo tee /etc/default/etcd >/dev/null
  echo changed
fi
sudo systemctl enable etcd >/dev/null 2>&1`, []byte(Etcd(n.Member, members)))
		if err != nil {
			return fmt.Errorf("on %s: %v: %s", n.Name, err, strings.TrimSpace(string(out)))
		}
		changed[i] = strings.Contains(string(out), "changed")
	}
	// Started together: a new member waits for the others before it says it
	// is up, so starting them one after another would wait for ever on the
	// first.
	var wg sync.WaitGroup
	errs := make([]error, len(nodes))
	for i, n := range nodes {
		wg.Add(1)
		go func(i int, n Node) {
			defer wg.Done()
			verb := "start"
			if changed[i] {
				verb = "restart"
			}
			if out, err := n.Remote.Run(ctx, "sudo systemctl "+verb+" --no-block etcd", nil); err != nil {
				errs[i] = fmt.Errorf("on %s: %v: %s", n.Name, err, strings.TrimSpace(string(out)))
			}
		}(i, n)
	}
	wg.Wait()
	for _, err := range errs {
		if err != nil {
			return err
		}
	}
	deadline := time.Now().Add(s.Wait)
	for !s.etcdHealthy(ctx, nodes[0], endpoints, len(nodes)) {
		if time.Now().After(deadline) {
			return fmt.Errorf("the %d members did not become healthy within %s; `journalctl -u etcd` on each says why", len(nodes), s.Wait)
		}
		if err := sleep(ctx, 2*time.Second); err != nil {
			return err
		}
	}
	s.emit("etcd", true, "%d members up", len(nodes))
	return nil
}

func (s *Setup) etcdHealthy(ctx context.Context, n Node, endpoints []string, want int) bool {
	out, err := n.Remote.Run(ctx, "ETCDCTL_API=3 etcdctl --dial-timeout=3s --endpoints="+
		strings.Join(endpoints, ",")+" endpoint health 2>&1", nil)
	return err == nil && strings.Count(string(out), "is healthy") == want
}

// patroni installs PostgreSQL and Patroni on every database machine, unless
// the cluster already has a leader and every member streaming.
func (s *Setup) patroni(ctx context.Context, sec Secrets) error {
	hosts := make([]string, len(s.DB))
	for i, d := range s.DB {
		hosts[i] = d.IP
	}
	if c, err := s.Status(ctx, hosts); err == nil && c.Complete(len(s.DB)) {
		s.emit("patroni", false, "%s leads, %d streaming; left as it is", c.Leader(), len(s.DB)-1)
		return nil
	}

	var version string
	for _, d := range s.DB {
		v, err := s.preparePostgres(ctx, d)
		if err != nil {
			return err
		}
		if version != "" && v != version {
			// Streaming replication is between one major version.
			return fmt.Errorf("%s has PostgreSQL %s and %s has %s; they must match", s.DB[0].Name, version, d.Name, v)
		}
		version = v
	}

	var etcd, db []Member
	for _, n := range s.etcdMembers() {
		etcd = append(etcd, n.Member)
	}
	for _, d := range s.DB {
		db = append(db, d.Member)
	}
	for i, d := range s.DB {
		cfg := Patroni(d.Member, db, etcd, s.Clients, version, sec)
		if out, err := d.Remote.Run(ctx, `set -e
sudo install -d -m 755 /etc/patroni
sudo install -d -o postgres -g postgres -m 700 /var/lib/postgresql/`+version+`/llmgw
sudo sh -c 'umask 027; cat > /etc/patroni/config.yml.new'
sudo chown root:postgres /etc/patroni/config.yml.new
sudo chmod 640 /etc/patroni/config.yml.new
sudo mv -f /etc/patroni/config.yml.new /etc/patroni/config.yml
sudo systemctl enable patroni >/dev/null 2>&1
sudo systemctl restart patroni`, []byte(cfg)); err != nil {
			return fmt.Errorf("on %s: %v: %s", d.Name, err, strings.TrimSpace(string(out)))
		}
		// The first starts the cluster; the rest clone it, which they can
		// only do once it has a leader.
		want := func(c Cluster) bool { return c.Leader() != "" }
		if i > 0 {
			want = func(c Cluster) bool { return c.Complete(i + 1) }
		}
		if err := s.waitFor(ctx, hosts, want); err != nil {
			return fmt.Errorf("after starting %s: %w", d.Name, err)
		}
		s.emit("patroni", true, "%s up", d.Name)
	}
	return nil
}

// preparePostgres installs the packages and clears the distribution's own
// cluster, which would hold the port Patroni's needs — unless it holds a
// database somebody made, which is not this setup's to delete.
func (s *Setup) preparePostgres(ctx context.Context, d Node) (string, error) {
	out, err := d.Remote.Run(ctx, `set -e
need=""
command -v patroni >/dev/null 2>&1 || need="$need patroni"
ls /usr/lib/postgresql/*/bin/postgres >/dev/null 2>&1 || need="$need postgresql"
if [ -n "$need" ]; then
  sudo DEBIAN_FRONTEND=noninteractive apt-get -qq update >/dev/null
  sudo DEBIAN_FRONTEND=noninteractive apt-get -qq install -y $need >/dev/null
fi
v=$(ls /usr/lib/postgresql | sort -n | tail -1)
if [ -d /var/lib/postgresql/$v/main ]; then
  if sudo systemctl is-active --quiet postgresql@$v-main; then
    n=$(sudo -u postgres psql -Atc "select count(*) from pg_database where datname not in ('postgres','template0','template1')")
    if [ "$n" != "0" ]; then echo "holds-data"; exit 0; fi
  fi
  sudo systemctl disable --now postgresql@$v-main postgresql >/dev/null 2>&1 || true
  sudo pg_dropcluster --stop $v main >/dev/null 2>&1 || sudo rm -rf /var/lib/postgresql/$v/main
fi
sudo systemctl disable --now postgresql >/dev/null 2>&1 || true
echo "version=$v"`, nil)
	if err != nil {
		return "", fmt.Errorf("on %s: %v: %s", d.Name, err, strings.TrimSpace(string(out)))
	}
	if strings.Contains(string(out), "holds-data") {
		return "", fmt.Errorf("%s runs a PostgreSQL with databases in it; move or drop them first — "+
			"this setup replaces that cluster", d.Name)
	}
	for _, line := range strings.Split(string(out), "\n") {
		if v, ok := strings.CutPrefix(strings.TrimSpace(line), "version="); ok && v != "" {
			return v, nil
		}
	}
	return "", fmt.Errorf("no PostgreSQL found on %s after installing it", d.Name)
}

// database makes the role and database on whichever member leads.
func (s *Setup) database(ctx context.Context) error {
	hosts := make([]string, len(s.DB))
	for i, d := range s.DB {
		hosts[i] = d.IP
	}
	c, err := s.Status(ctx, hosts)
	if err != nil {
		return err
	}
	var leader *Node
	for i := range s.DB {
		if s.DB[i].Name == c.Leader() {
			leader = &s.DB[i]
		}
	}
	if leader == nil {
		return fmt.Errorf("no member leads")
	}
	out, err := leader.Remote.Run(ctx, `set -e
pw=$(sudo cat `+secretDir+`/llmgw.pw)
sudo -u postgres psql -q -v ON_ERROR_STOP=1 -v pw="$pw" <<'SQL'
SELECT format('CREATE ROLE %I LOGIN', 'llmgw') WHERE NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'llmgw') \gexec
ALTER ROLE llmgw PASSWORD :'pw';
SELECT format('CREATE DATABASE %I OWNER %I', 'llmgw', 'llmgw') WHERE NOT EXISTS (SELECT 1 FROM pg_database WHERE datname = 'llmgw') \gexec
SQL`, nil)
	if err != nil {
		return fmt.Errorf("on %s: %v: %s", leader.Name, err, strings.TrimSpace(string(out)))
	}
	s.emit("database", true, "role and database %s on %s", Database, leader.Name)
	return nil
}

func (s *Setup) waitFor(ctx context.Context, hosts []string, ok func(Cluster) bool) error {
	deadline := time.Now().Add(s.Wait)
	var last error
	for {
		c, err := s.Status(ctx, hosts)
		if err == nil && ok(c) {
			return nil
		}
		if err != nil {
			last = err
		} else {
			last = fmt.Errorf("the cluster is %s", c)
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("not up within %s: %v; `journalctl -u patroni` says why", s.Wait, last)
		}
		if err := sleep(ctx, 3*time.Second); err != nil {
			return err
		}
	}
}

func sleep(ctx context.Context, d time.Duration) error {
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-time.After(d):
		return nil
	}
}