LocalAI/tests/e2e/distributed/cluster_baseline_test.go

package distributed_test

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/mudler/LocalAI/pkg/httpclient"
	"github.com/mudler/LocalAI/tests/e2e/distributed/cluster"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
)

const (
	// nodeRosterTimeout bounds the wait for a worker to appear healthy in
	// /api/nodes. Registration is an HTTP call the worker retries, followed by a
	// heartbeat that has to land before the frontend calls the node healthy, so
	// the budget covers several retry intervals rather than a single round trip.
	nodeRosterTimeout = "90s"
	nodeRosterPoll    = "1s"
	// authProbeTimeout bounds the single unauthenticated request that checks the
	// admin gate is actually closed. One round trip against a ready local
	// process; anything slower is a defect, not slowness.
	authProbeTimeout = 30 * time.Second
)

// node is the subset of the /api/nodes payload these specs assert on. ID is the
// registration identity the worker minted, which is what distinguishes "the
// same node row seen from a second replica" from "a second registration that
// happens to share a name".
type node struct {
	ID     string `json:"id"`
	Name   string `json:"name"`
	Status string `json:"status"`
	// Address and HTTPAddress are what a PRE-TUNNEL worker advertised. A worker
	// running this release sends neither, which is the fact the tunnel specs
	// assert on: with nothing advertised there is no address a frontend could
	// have dialled instead of the tunnel.
	Address     string `json:"address"`
	HTTPAddress string `json:"http_address"`
	// LastHeartbeat is what separates "the worker is gone" from "the worker is
	// here and this deployment cannot reach it". A spec asserting the second
	// has to show the first is false, and the heartbeat is the only evidence
	// of that in this payload.
	LastHeartbeat time.Time `json:"last_heartbeat"`
	// keys is what the payload actually carried, which a decoded struct cannot
	// tell you. Both fields above are the zero value when a worker advertises
	// nothing AND when the key was renamed or dropped, and the whole point of
	// the change these specs cover was removing the advertisement, so a rename
	// would leave "it advertises nothing" passing for a payload that no longer
	// says anything either way.
	keys map[string]json.RawMessage `json:"-"`
}

// String is what a failing assertion prints for a node.
//
// Without it %+v renders keys, whose values are json.RawMessage, as slices of
// byte VALUES: one node becomes several hundred numbers and a roster of two
// buries the assertion that failed. The key set is still what advertisementOf
// reads; it is just not something a human ever needs to see.
func (n node) String() string {
	return fmt.Sprintf("{name:%s status:%s id:%s lastHeartbeat:%s advertised:%q/%q}",
		n.Name, n.Status, n.ID, n.LastHeartbeat.Format(time.RFC3339), n.Address, n.HTTPAddress)
}

// UnmarshalJSON decodes the fields above and keeps the raw key set beside them.
func (n *node) UnmarshalJSON(data []byte) error {
	// A distinct type, or this method calls itself.
	type decoded node
	var plain decoded
	if err := json.Unmarshal(data, &plain); err != nil {
		return err
	}
	*n = node(plain)
	return json.Unmarshal(data, &n.keys)
}

// requireBinaries reports whether a missing binary must fail the spec instead of
// skipping it. It defaults to ON under CI.
//
// Skipping is the right courtesy locally: someone who has not run `make build`
// should get a clear note, not a wall of red. In CI it is the opposite. The
// whole Cluster label partition is these two specs, so if the workflow's build
// step breaks or moves its output, a skip would leave the job reporting
// "0 Passed | 2 Skipped" and exiting 0. Ginkgo exits 0 on skips, so that job
// goes green having never started a cluster, which is precisely the silent pass
// this suite exists to make impossible.
//
// Hence the polarity: the safe behaviour is the default, keyed off CI (GitHub
// Actions always sets it), and LOCALAI_E2E_REQUIRE_BINARIES exists to be forced
// OFF rather than to be remembered ON. A future workflow author cannot reach
// the green-on-nothing state by forgetting a line, only by writing one that
// explicitly asks for it. A local developer sees no change: CI is unset in an
// ordinary shell, so a missing binary still skips.
func requireBinaries() bool {
	value := strings.TrimSpace(os.Getenv("LOCALAI_E2E_REQUIRE_BINARIES"))
	if value == "" {
		return os.Getenv("CI") != ""
	}
	// ParseBool rejects these, and the fallback below reads anything it rejects
	// as ON. Someone writing "off" plainly means off, and silently inverting
	// them would be a worse trap than the one this flag removes.
	switch strings.ToLower(value) {
	case "off", "no", "n", "disabled":
		return false
	}
	if parsed, err := strconv.ParseBool(value); err == nil {
		return parsed
	}
	// Set to something meaningless means someone meant to turn this on. Reading
	// it as false would quietly restore the silent skip the flag guards against.
	return true
}

// missingBinary skips or fails, naming the path and how to produce it.
func missingBinary(what, path, remedy string) {
	GinkgoHelper()
	message := fmt.Sprintf("%s not found at %s; %s", what, path, remedy)
	if requireBinaries() {
		Fail(message + " (binaries are required here, either under CI or via " +
			"LOCALAI_E2E_REQUIRE_BINARIES, so this fails rather than skips: a skipped " +
			"cluster spec is indistinguishable from a passing one)")
	}
	Skip(message)
}

// localAIBinary resolves the built binary and refuses a stale one.
func localAIBinary() string {
	GinkgoHelper()
	path := os.Getenv("LOCALAI_E2E_BINARY")
	if path == "" {
		wd, err := os.Getwd()
		Expect(err).ToNot(HaveOccurred())
		path = filepath.Join(wd, "..", "..", "..", "local-ai")
	}
	info, err := os.Stat(path)
	if err != nil {
		missingBinary("local-ai binary", path, "run `make test-e2e-cluster`, which builds it, or set LOCALAI_E2E_BINARY")
	}
	requireFreshBinary(path, info.ModTime())
	return path
}

// requireFreshBinary fails when the binary about to be exec'd predates the
// sources it is supposed to contain.
//
// A missing binary is a loud failure; a STALE one is a silent lie. The suite
// spawns this file as its frontends and workers, so an out-of-date build makes
// every spec a statement about code nobody can name, and the pass or fail it
// produces belongs to a commit that is not the one under test. That has already
// happened here once: a spec reported caller line numbers that matched no line
// in any source file, which was the only reason anyone noticed.
//
// It FAILS rather than skipping, and it does so on a developer box as well as
// under CI, which is the opposite polarity to missingBinary above. The
// distinction is what a wrong answer costs: a skipped spec reports nothing and
// a stale spec reports something false, so there is no environment in which
// continuing is the kinder choice.
//
// _test.go files are excluded on purpose and the exclusion is load-bearing.
// They are compiled into the ginkgo suite, never into local-ai, so counting
// them would fire on every edit to the spec being run and would train everyone
// to route around the check within a day.
func requireFreshBinary(path string, built time.Time) {
	GinkgoHelper()
	newest, at, err := newestGoSource()
	Expect(err).ToNot(HaveOccurred(), "walking the tree for source timestamps")
	if at == "" || !newest.After(built) {
		return
	}
	Fail(fmt.Sprintf(
		"%s was built at %s but %s changed at %s, so this suite would exec a binary that does not contain the tree it is about to report on; run `make test-e2e-cluster`, which builds it",
		path, built.Format(time.RFC3339), at, newest.Format(time.RFC3339)))
}

// sourceScanSkip names the directories the freshness walk does not descend
// into. Each holds either no Go the binary links, or enough files to make the
// walk cost more than the check is worth.
var sourceScanSkip = map[string]bool{
	".git":           true,
	"node_modules":   true,
	"dist":           true,
	"models":         true,
	"backends":       true,
	"local-backends": true,
	"vendor":         true,
}

// newestGoSource returns the newest modification time among the Go sources that
// end up in the binary, and the path that carries it.
//
// go.mod and go.sum count: a dependency bump changes neither a .go file nor the
// working tree's own code, and produces a different binary all the same.
func newestGoSource() (time.Time, string, error) {
	root, err := os.Getwd()
	if err != nil {
		return time.Time{}, "", err
	}
	root = filepath.Join(root, "..", "..", "..")
	var newest time.Time
	var at string
	err = filepath.WalkDir(root, func(p string, d os.DirEntry, walkErr error) error {
		if walkErr != nil {
			// A file that vanished mid-walk is not this check's business, and
			// aborting on it would turn a race in someone's editor into a
			// failed suite.
			return nil //nolint:nilerr // see above
		}
		if d.IsDir() {
			if sourceScanSkip[d.Name()] {
				return filepath.SkipDir
			}
			return nil
		}
		name := d.Name()
		if strings.HasSuffix(name, "_test.go") {
			return nil
		}
		if !strings.HasSuffix(name, ".go") && name != "go.mod" && name != "go.sum" {
			return nil
		}
		info, err := d.Info()
		if err != nil {
			return nil //nolint:nilerr // same reasoning as the walk error above
		}
		if info.ModTime().After(newest) {
			newest = info.ModTime()
			at = p
		}
		return nil
	})
	return newest, at, err
}

func mockBackendBinary() string {
	GinkgoHelper()
	wd, err := os.Getwd()
	Expect(err).ToNot(HaveOccurred())
	path := filepath.Join(wd, "..", "mock-backend", "mock-backend")
	if _, err := os.Stat(path); err != nil {
		missingBinary("mock-backend", path, "run `make build-mock-backend`")
	}
	return path
}

// startCluster brings up a cluster against a freshly provisioned database and
// registers cleanup, including a log dump on failure.
//
// customise runs against the assembled Options immediately before Start, for
// the one spec that needs a non-default topology. It is variadic so every
// existing caller keeps the plain two-argument form and the default shape.
func startCluster(frontends, workers int, customise ...func(*cluster.Options)) *cluster.Cluster {
	GinkgoHelper()
	c, _ := startClusterOnFreshDB(frontends, workers, customise...)
	return c
}

// startClusterOnFreshDB is startCluster plus the DSN of the database it was
// given, for a spec that has to read a table no endpoint exposes.
func startClusterOnFreshDB(frontends, workers int, customise ...func(*cluster.Options)) (*cluster.Cluster, string) {
	GinkgoHelper()

	// Resolved before SetupInfra so a missing binary skips without having paid
	// for a database that the skip would then leave to DeferCleanup.
	binary := localAIBinary()
	mockBackend := mockBackendBinary()

	infra := SetupInfra("cluster")

	// The log directory must be predictable so CI can upload it as an artifact.
	// GinkgoT().TempDir() lands under TMPDIR, which on a GitHub runner is not
	// /tmp, so an artifact glob would silently match nothing.
	logDir := os.Getenv("LOCALAI_E2E_LOG_DIR")
	if logDir == "" {
		logDir = GinkgoT().TempDir()
	} else {
		logDir = filepath.Join(logDir, sanitizeDBName(CurrentSpecReport().LeafNodeText))
		Expect(os.MkdirAll(logDir, 0o755)).To(Succeed())
	}

	options := cluster.Options{
		Binary:      binary,
		MockBackend: mockBackend,
		PGDSN:       infra.PGURL,
		// There is no bus URL to pass. Frontends and agent workers are still
		// handed a dead LOCALAI_NATS_URL, so this suite keeps covering the
		// promise that an operator's existing command line starts unchanged
		// after the broker is shut down, but the value is the harness's own
		// cluster.StaleBusURL rather than a caller's choice: no process reads
		// it, so there was nothing left for a caller to choose.
		LogDir:    logDir,
		Frontends: frontends,
		Workers:   workers,
	}
	for _, apply := range customise {
		apply(&options)
	}

	c, err := cluster.Start(options)
	Expect(err).ToNot(HaveOccurred())

	DeferCleanup(func() {
		if CurrentSpecReport().Failed() {
			c.DumpLogs()
		}
		c.Stop()
	})
	return c, infra.PGURL
}

// rosterProbe polls one frontend's node roster.
//
// It keeps the last error and the last roster it saw so a failing Eventually can
// name the cause. Returning a bare nil on error makes a 401 at the second
// replica, a JSON decode failure and "the worker never registered" all present
// identically as an empty list, which is the least useful thing a failover
// suite can say when it goes red.
type rosterProbe struct {
	cluster  *cluster.Cluster
	client   *http.Client
	frontend int

	lastErr  error
	lastSeen []node
}

func newRosterProbe(c *cluster.Cluster, client *http.Client, frontend int) *rosterProbe {
	return &rosterProbe{cluster: c, client: client, frontend: frontend}
}

// healthyNames returns nil on any error so Eventually keeps retrying: the roster
// is unreachable for the first moments of a replica's life, and failing hard
// there would only re-report a startup race.
func (p *rosterProbe) healthyNames() []string {
	var roster []node
	if err := p.cluster.GetJSON(p.client, p.frontend, "/api/nodes", &roster); err != nil {
		p.lastErr = err
		return nil
	}
	p.lastErr = nil
	p.lastSeen = roster
	names := []string{}
	for _, n := range roster {
		if n.Status == "healthy" {
			names = append(names, n.Name)
		}
	}
	return names
}

// idOf returns the registration ID the roster last reported for a node name.
func (p *rosterProbe) idOf(name string) string {
	for _, n := range p.lastSeen {
		if n.Name == name {
			return n.ID
		}
	}
	return ""
}

// advertisementOf returns whatever endpoints the roster last reported a node
// advertising, joined for a failure message, and whether the payload carried
// both advertisement keys at all.
//
// The second result is the assertion, not a detail. Removing the advertisement
// is what the change under test did, so "the node advertises nothing" and "the
// keys that would have carried it are gone from the payload" are the two
// outcomes a spec has to keep apart: the first is the feature working, the
// second is the spec having lost its subject and reporting the feature working
// for any node at all, including one that advertises plenty.
func (p *rosterProbe) advertisementOf(name string) (string, bool) {
	for _, n := range p.lastSeen {
		if n.Name != name {
			continue
		}
		_, hasAddress := n.keys["address"]
		_, hasHTTP := n.keys["http_address"]
		return strings.TrimSpace(strings.Join([]string{n.Address, n.HTTPAddress}, " ")), hasAddress && hasHTTP
	}
	return "", false
}

// heartbeatOf is the last heartbeat the roster reported for a node, refreshed
// on every call.
//
// A zero time is returned for a node the roster does not carry, which no
// freshness assertion can accept: the spec that reads this is proving the
// worker is still alive, and a missing node must not read as a recent
// heartbeat.
func (p *rosterProbe) heartbeatOf(name string) time.Time {
	var roster []node
	if err := p.cluster.GetJSON(p.client, p.frontend, "/api/nodes", &roster); err != nil {
		p.lastErr = err
		return time.Time{}
	}
	p.lastErr = nil
	p.lastSeen = roster
	for _, n := range roster {
		if n.Name == name {
			return n.LastHeartbeat
		}
	}
	return time.Time{}
}

// describe is handed to Should as the failure message. Gomega calls a
// func() string description lazily, so this runs only on failure and reports
// whichever of the two distinct causes actually occurred.
func (p *rosterProbe) describe() string {
	if p.lastErr != nil {
		return fmt.Sprintf("frontend %d: the last GET /api/nodes failed: %v", p.frontend, p.lastErr)
	}
	return fmt.Sprintf("frontend %d: GET /api/nodes succeeded but the roster held %d node(s): %+v",
		p.frontend, len(p.lastSeen), p.lastSeen)
}

var _ = Describe("Cluster baseline", Label("Distributed"), Label("Cluster"), func() {
	It("brings up a frontend and a worker, and the worker appears in the roster", func() {
		c := startCluster(1, 1)

		client, err := c.AdminSession(0)
		Expect(err).ToNot(HaveOccurred())

		probe := newRosterProbe(c, client, 0)
		Eventually(probe.healthyNames, nodeRosterTimeout, nodeRosterPoll).
			Should(ContainElement(c.WorkerName(0)), probe.describe)
	})

	It("runs two frontends against one database and both see the same worker", func() {
		c := startCluster(2, 1)

		// Observe that frontend 1 really is gated before proving a session opens
		// it. Without this, "the cookie minted at frontend 0 works here" is
		// indistinguishable from "this endpoint needs no auth at all". The probe
		// is free: it touches no auth route, so it spends nothing from the
		// five-per-minute-per-IP budget those routes share.
		anonymous := httpclient.NewWithTimeout(authProbeTimeout)
		refused, err := anonymous.Get(c.FrontendURL(1) + "/api/nodes")
		Expect(err).ToNot(HaveOccurred())
		defer func() { _ = refused.Body.Close() }()
		Expect(refused.StatusCode).To(Equal(http.StatusUnauthorized),
			"an unauthenticated GET /api/nodes must be refused, otherwise this spec proves nothing about sessions")

		// One session for the whole cluster, minted at frontend 0. Registering or
		// logging in again per frontend would spend from the same five-per-minute
		// budget, and every request here comes from 127.0.0.1. The single client
		// is valid at both replicas: sessions live in the shared Postgres, the
		// harness pins one HMAC secret so the row resolves anywhere, and Go's
		// cookie jar keys by host without port.
		client, err := c.AdminSession(0)
		Expect(err).ToNot(HaveOccurred())

		// The worker is pointed at frontend 0 alone (the harness sets
		// LOCALAI_REGISTER_TO to frontend 0), so read its identity there first.
		at0 := newRosterProbe(c, client, 0)
		Eventually(at0.healthyNames, nodeRosterTimeout, nodeRosterPoll).
			Should(ContainElement(c.WorkerName(0)), at0.describe)
		registeredID := at0.idOf(c.WorkerName(0))
		Expect(registeredID).ToNot(BeEmpty(), "frontend 0 reported the worker without a registration ID")

		// Then assert frontend 1 serves the same row, by id and not merely by name.
		//
		// Be precise about what this proves. It does NOT pin the topology:
		// NodeRegistry.Register looks a node up by name and preserves the
		// existing id (core/services/nodes/registry.go:522-527), and both
		// replicas read one Postgres, so a harness that registered the worker
		// with every frontend would yield identical ids here too. What it does
		// catch is a frontend answering from its own registry or its own
		// database rather than the shared one, which is a different regression
		// and just as silent. The topology fact is not asserted anywhere; it is
		// recorded next to LOCALAI_REGISTER_TO in cluster.go.
		at1 := newRosterProbe(c, client, 1)
		Eventually(at1.healthyNames, nodeRosterTimeout, nodeRosterPoll).
			Should(ContainElement(c.WorkerName(0)), at1.describe)
		Expect(at1.idOf(c.WorkerName(0))).To(Equal(registeredID),
			"frontend 1 must resolve the same node row as frontend 0; a differing id means it is not reading the shared state")
	})
})