llmgw/internal/pgha/memory.go

package pgha

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"slices"
	"strconv"
	"strings"
	"time"
)

// PostgreSQL's own defaults — 128MB of cache, 4MB to sort in — are for a
// machine it barely uses. A knowledge base's vector index is searched well
// only from memory: one that does not fit is read from disk on every search.
// So each member is sized to the memory it has, read when setup runs, and
// again on every later run, so a machine given more memory gets the use of it.

// MemorySettings are PostgreSQL's memory parameters for a machine with total
// bytes of memory. shared is a machine that also runs a control plane or an
// agent: the database is given less of it, since a gateway swapped out to
// give the database cache is worse than a smaller cache.
func MemorySettings(total int64, shared bool) map[string]string {
	const mb = 1 << 20
	ram := total / mb
	// Shares of memory, in MB: the cache PostgreSQL keeps itself, what the
	// planner may count on the kernel caching besides, what building or
	// vacuuming an index may take, and what one sort may.
	buffers, cache, maint, sort := ram/4, ram*3/4, ram/16, ram/256
	if shared {
		buffers, cache, maint, sort = ram*3/20, ram/2, ram/32, ram/512
	}
	clamp := func(v, lo, hi int64) int64 { return min(max(v, lo), hi) }
	return map[string]string{
		"shared_buffers":       strconv.FormatInt(clamp(buffers, 96, 32<<10), 10) + "MB",
		"effective_cache_size": strconv.FormatInt(clamp(cache, 256, 1<<20), 10) + "MB",
		"maintenance_work_mem": strconv.FormatInt(clamp(maint, 64, 2<<10), 10) + "MB",
		"work_mem":             strconv.FormatInt(clamp(sort, 4, 64), 10) + "MB",
	}
}

// MemTotal reads the machine's memory, in bytes, from /proc/meminfo.
func MemTotal(meminfo string) (int64, bool) {
	for _, line := range strings.Split(meminfo, "\n") {
		f := strings.Fields(line)
		if len(f) >= 2 && f[0] == "MemTotal:" {
			kb, err := strconv.ParseInt(f[1], 10, 64)
			return kb << 10, err == nil && kb > 0
		}
	}
	return 0, false
}

// memory sizes each member to its machine, in its own configuration — the
// machines need not be alike. What takes effect on reload is applied so;
// shared_buffers needs a restart, given one member at a time, the leader
// last, each waited on before the next.
func (s *Setup) memory(ctx context.Context, sec Secrets) error {
	var restart []Node
	for _, d := range s.DB {
		out, err := d.Remote.Run(ctx, `cat /proc/meminfo
if [ -f /etc/llmgw/server.yaml ] || [ -f /etc/llmgw/agent.yaml ]; then echo "llmgw-shared"; fi`, nil)
		total, ok := MemTotal(string(out))
		if err != nil || !ok {
			s.emit("memory", false, "%s: could not read its memory; PostgreSQL left as it is", d.Name)
			continue
		}
		shared := strings.Contains(string(out), "llmgw-shared")
		params := MemorySettings(total, shared)
		body, _ := json.Marshal(params)
		out, err = d.Remote.Run(ctx, `set -e
test -f /etc/patroni/config.yml || { echo missing; exit 0; }
want=$(cat)
r=$(sudo env WANT="$want" python3 - <<'PY'
import json, os, yaml
p = "/etc/patroni/config.yml"
want = json.loads(os.environ["WANT"])
c = yaml.safe_load(open(p))
pg = c.setdefault("postgresql", {})
params = pg.get("parameters") or {}
pg["parameters"] = params
if all(str(params.get(k)) == v for k, v in want.items()):
    print("same")
else:
    restart = str(params.get("shared_buffers")) != want["shared_buffers"]
    params.update(want)
    open(p + ".new", "w").write("# llmgw: the control planes' store, kept by Patroni, its leader lock in etcd.\n" +
        yaml.safe_dump(c, default_flow_style=False, sort_keys=False))
    print("restart" if restart else "changed")
PY
)
if [ "$r" != same ]; then
  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 reload patroni
fi
echo "$r"`, body)
		if err != nil {
			return fmt.Errorf("on %s: %v: %s", d.Name, err, strings.TrimSpace(string(out)))
		}
		what := fmt.Sprintf("%s: %d MB of memory%s; shared_buffers %s, effective_cache_size %s, maintenance_work_mem %s, work_mem %s",
			d.Name, total>>20, map[bool]string{true: ", shared with llmgw"}[shared],
			params["shared_buffers"], params["effective_cache_size"], params["maintenance_work_mem"], params["work_mem"])
		switch r := strings.TrimSpace(lastLine(out)); r {
		case "same", "missing":
			s.emit("memory", false, "%s", what)
		case "restart":
			s.emit("memory", true, "%s", what)
			restart = append(restart, d)
		default:
			s.emit("memory", true, "%s", what)
		}
	}
	// Patroni takes a reload in its own time: restarted before it has written
	// the new configuration, PostgreSQL comes back on the old one. So each is
	// waited on until Patroni says a restart is pending — and a member found
	// pending, by a run that stopped before restarting it, is restarted too.
	pending := func(c Cluster, name string) bool {
		for _, m := range c.Members {
			if m.Name == name {
				return m.PendingRestart
			}
		}
		return false
	}
	var c Cluster
	err := s.waitFor(ctx, s.dbHosts(), func(got Cluster) bool {
		c = got
		for _, d := range restart {
			if !pending(got, d.Name) {
				return false
			}
		}
		return true
	})
	if err != nil {
		return fmt.Errorf("waiting for Patroni to take the new configuration: %w", err)
	}
	for _, d := range s.DB {
		if pending(c, d.Name) && !slices.ContainsFunc(restart, func(n Node) bool { return n.Name == d.Name }) {
			restart = append(restart, d)
		}
	}
	if len(restart) == 0 {
		return nil
	}
	// The leader last: its restart is the only one the control planes see,
	// a few seconds without writes, and by then the rest are back.
	slices.SortStableFunc(restart, func(a, b Node) int {
		al, bl := a.Name == c.Leader(), b.Name == c.Leader()
		switch {
		case al && !bl:
			return 1
		case bl && !al:
			return -1
		}
		return 0
	})
	for _, d := range restart {
		if _, err := restJSON(ctx, d.IP, http.MethodPost, "/restart", []byte(`{"restart_pending": true}`), sec.RestAPI, 2*time.Minute); err != nil {
			return fmt.Errorf("restarting PostgreSQL on %s for its new shared_buffers: %w", d.Name, err)
		}
		if err := s.waitFor(ctx, s.dbHosts(), func(c Cluster) bool { return c.Complete(len(s.DB)) }); err != nil {
			return fmt.Errorf("after restarting %s: %w", d.Name, err)
		}
		s.emit("memory", true, "%s: PostgreSQL restarted with the new shared_buffers", d.Name)
	}
	return nil
}