diff --git a/.agents/building-and-testing.md b/.agents/building-and-testing.md index 021d555ec993..eee5123ac82c 100644 --- a/.agents/building-and-testing.md +++ b/.agents/building-and-testing.md @@ -45,3 +45,23 @@ Rules (both gates): - **Don't weaken the gate:** never hand-lower a baseline or widen a tolerance to turn a red gate green. The ratchet only moves up. - If a change drops coverage, **add tests** (sort `coverage-summary.json` by line% ascending to find untested code) rather than editing the baseline. When coverage legitimately rises, commit the regenerated baseline (`make test-coverage-baseline` / `test-ui-coverage-baseline`). - The Go gate is **strict — no tolerance**; `covermode=atomic` keeps it deterministic. The UI gate keeps a small tolerance only because its e2e coverage isn't. + +## Distributed-mode test suites + +Two suites cover distributed mode (frontend replicas, worker nodes, PostgreSQL, NATS), split by a Ginkgo label: + +- `make test-e2e-distributed` runs `Distributed && !VLLMMultinode && !Cluster` over `./tests/e2e/distributed` recursively. Services are wired directly into the test binary. ~240 specs, ~75s. +- `make test-e2e-cluster` runs `Cluster` and spawns real `local-ai` child processes through the `tests/e2e/distributed/cluster` helper package. 6 specs, about 8m30s measured over three consecutive runs (509.1s / 509.8s / 512.3s, so 8m29s to 8m32s). + +Both jobs live in `.github/workflows/tests-e2e-distributed.yml`, with `timeout-minutes: 45` each. They trigger on pull requests *and* on every push to `master`; the `paths-ignore` filter (see [.agents/ci-caching.md](ci-caching.md)) sits on the pull-request trigger only, so a master push always runs both. They are advisory only because `master` carries no branch protection, which is a repository setting and not a YAML key: `continue-on-error: true` would flip the run's *conclusion* to success and hide the failure, so it is not used. + +- **Containers are suite-scoped, not spec-scoped.** `SetupInfra` used to start a PostgreSQL (~10s) and a NATS (~3.5s) per spec. Across the 213 specs behind it that was roughly **48 minutes of pure container startup per run**, which is why this suite was never in CI. (213 rather than the ~240 above: the larger number is everything the label filter selects, the smaller one is just the specs that call `SetupInfra`.) Containers now start once in `BeforeSuite` and each spec gets its own database via `CREATE DATABASE` (~67ms), which is what the `dbName` argument was always describing. Adding a spec needs no change: call `SetupInfra("some-name")` as before, the name is a prefix and a counter keeps it unique. +- **Consequence for new specs:** the NATS bus is now *shared* within a Ginkgo process, so a wildcard subscriber can observe another spec's traffic. Filter assertions on an identifier your spec owns (a node ID, a job ID) instead of counting everything on `jobs.*.progress`, and verify the spec with `--randomize-all`. +- **`BeforeSuite`, not `SynchronizedBeforeSuite`.** Under `ginkgo -p` each process then gets its own container pair, keeping NATS subjects isolated per process. A single shared NATS across parallel processes would let specs on different processes see each other's messages on the same subject. +- **The label split.** The 8 argument-validation specs under `tests/e2e/distributed/cluster/` carry `Label("Distributed")` only, on purpose: they need no binary, no PostgreSQL and no NATS, so they belong in the fast job. That is why `test-e2e-distributed` keeps `-r` (it must reach the subpackage) and `test-e2e-cluster` deliberately does **not** (the subpackage is out of its scope). +- **`--fail-on-empty` is load-bearing on both targets.** Ginkgo exits 0 when a label filter selects nothing, so without it a refactor that renames or drops `Label("Cluster")` leaves the target reporting "Test Suite Passed" having started no cluster at all. `LOCALAI_E2E_REQUIRE_BINARIES` does not cover this case: it only fires inside a spec that is actually running. +- **The binary gate.** `localAIBinary()` and `mockBackendBinary()` **fail** rather than skip when `CI` is set, or when `LOCALAI_E2E_REQUIRE_BINARIES` is truthy; `LOCALAI_E2E_REQUIRE_BINARIES=0` (also `off`, `no`, `n`, `disabled`, and anything `strconv.ParseBool` reads as false) forces skipping even under CI. **Any value that parses as neither reads as ON**, not off: setting the variable to something meaningless means someone meant to turn the gate on, and reading it as false would quietly restore the silent skip the flag exists to remove. The whole polarity is deliberate, because in CI a skipped cluster spec is indistinguishable from a passing one: Ginkgo exits 0 on skips. Locally a missing binary still just skips, since `CI` is unset in an ordinary shell. +- **Flake budget: no retries at all.** `--flake-attempts` is *total attempts*, not retries (ginkgo v2.29.0 `internal/group.go` sets `maxAttempts = FlakeAttempts` and loops `attempt < maxAttempts`; the flag's own usage string reads "0 - failed tests are not retried"). `DISTRIBUTED_TEST_FLAKES` defaults to **1**, so each spec runs once and a failure is a failure, and `test-e2e-cluster` pins `--flake-attempts 1` outright rather than reading the variable. The repo-wide `TEST_FLAKES=5` means up to five attempts, so up to four retries. These suites exist to surface nondeterminism, and a retry converts exactly that signal into a green run. Raise it locally when bisecting something unrelated, not in the Makefile. +- **Coverage:** `tests/e2e/distributed` is excluded from the coverage roots (`COVERAGE_E2E_ROOTS = ./tests/e2e`, run non-recursively), and so is the `cluster` helper package beneath it. Neither suite moves the baseline, so production code that these suites are the only cover for reads as **uncovered**. Unit tests for such code belong under `./core/...` with `testutil.SetupTestDB()`. +- **The cluster job builds against a stubbed React UI.** `core/http/react-ui/dist` is gitignored and built by Node, so the workflow writes a one-line `index.html` there to satisfy the `//go:embed react-ui/dist/*` in `core/http/app.go` and skips a full Node and Vite install. That holds only while the suite drives the HTTP API and never the UI, which has its own e2e suite. A spec that ever asserts on a UI asset would pass locally, where a real `dist/` exists, and be served the stub in CI: if you write one, the stub step has to go and the real build come back. +- **Do not shorten the cluster suite's waits.** Three of its six specs sit at ~167s each because they wait out a 60s staleness threshold plus a 15s health-check tick. That wait is what stops the assertions from passing before the system could have reacted, which was a real false green earlier on. If the job has to get faster, the levers are CI concurrency or making the thresholds configurable, not shorter waits. diff --git a/.agents/ci-caching.md b/.agents/ci-caching.md index 6742049e68ff..8dc243747384 100644 --- a/.agents/ci-caching.md +++ b/.agents/ci-caching.md @@ -153,7 +153,7 @@ This is worth more than it looks. Measured over the week to 2026-07-30, **97% of The volume is real: 13 gallery-only PRs merged that week with 10 open at once, and 78 of the 137 PRs opened were bot-generated. -`paths-ignore` on the PR trigger of `image-pr.yml` (7 jobs), `build-test.yaml` (3), `lint.yml` (2) and `tests-e2e.yml` (1) drops 13 of those 20. The excluded set: +`paths-ignore` on the PR trigger of `image-pr.yml` (7 jobs), `build-test.yaml` (3), `lint.yml` (2) and `tests-e2e.yml` (1) drops 13 of those 20, measured before `tests-e2e-distributed.yml` (2 jobs) landed. That workflow carries the same exclusion set for the same reason: its dependency graph is 99 packages, so an allowlist of paths would silently stop guarding the moment code moved, while a diff confined to the paths below provably cannot reach it. The excluded set: | Path | Why no image or Go build can see it | |---|---| @@ -192,7 +192,7 @@ What still runs, and why it has to: Two properties this relies on: - `paths-ignore` skips a run only when **every** changed file matches, so a PR touching the gallery *and* Go code still runs everything. That is what makes the exclusion safe rather than a hole. -- `master` carries no branch protection and no rulesets, so a skipped workflow reports no status and nothing waits on it. If required status checks are ever introduced, these four entries must be excluded from the required set or PRs will hang on "Expected — Waiting for status to be reported". +- `master` carries no branch protection and no rulesets, so a skipped workflow reports no status and nothing waits on it. If required status checks are ever introduced, these five entries must be excluded from the required set or PRs will hang on "Expected — Waiting for status to be reported". ### `image.yml` on master push is gated too, by a job rather than a path filter diff --git a/.github/workflows/tests-e2e-distributed.yml b/.github/workflows/tests-e2e-distributed.yml new file mode 100644 index 000000000000..8f17e72f169a --- /dev/null +++ b/.github/workflows/tests-e2e-distributed.yml @@ -0,0 +1,204 @@ +--- +name: 'E2E Distributed Tests' + +on: + pull_request: + # The suite's dependency graph is 99 packages, so an allowlist of paths + # silently stops guarding the moment code moves. At ~75s the job is cheap + # enough to run unless the diff is confined to paths it provably cannot + # reach. See .agents/ci-caching.md. + paths-ignore: + - 'gallery/**' + - 'docs/**' + - 'examples/**' + - '**/*.md' + push: + branches: + - master + +concurrency: + group: ci-tests-e2e-distributed-${{ github.event.pull_request.number || github.sha }}-${{ github.repository }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + tests-e2e-distributed: + runs-on: ubuntu-latest + # Advisory because it is deliberately not in branch protection, so a failure + # is a visible red X rather than a blocked merge. Promoting it to a required + # check is a repository-settings change, to be made once it has a track + # record; a heavy suite made required on day one gets disabled instead of + # fixed. + timeout-minutes: 45 + steps: + - name: Clone + uses: actions/checkout@v7 + with: + submodules: true + - name: Configure apt mirror on runner + uses: ./.github/actions/configure-apt-mirror + - name: Setup Go + uses: actions/setup-go@v5 + with: + go-version: '1.26.0' + cache: false + - name: Dependencies + run: | + sudo apt-get update + sudo apt-get install -y build-essential libopus-dev + - name: Proto Dependencies + run: | + curl -L -s https://github.com/protocolbuffers/protobuf/releases/download/v26.1/protoc-26.1-linux-x86_64.zip -o protoc.zip && \ + unzip -j -d /usr/local/bin protoc.zip bin/protoc && \ + rm protoc.zip + go install google.golang.org/protobuf/cmd/protoc-gen-go@v1.34.2 + go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@1958fcbe2ca8bd93af633f11e97d44e567e945af + PATH="$PATH:$HOME/go/bin" make protogen-go + - name: Pre-pull test images + # Pulling here rather than inside the suite keeps container-start timing + # out of the spec timeouts and makes a registry outage read as a + # setup failure instead of a test failure. These two are the only images + # the suite needs once the testcontainers reaper is disabled below. + run: | + docker pull postgres:16-alpine + docker pull nats:2-alpine + - name: Distributed E2E + # TESTCONTAINERS_RYUK_DISABLED keeps the pre-pull above meaningful. The + # reaper exists to clean up leaked containers on a long-lived host, but + # this runner is ephemeral and every container dies with the VM. Leaving + # it enabled would pull a third, unpinned image (testcontainers/ryuk) + # from Docker Hub mid-suite: exactly the registry dependency the + # pre-pull step exists to remove. + env: + TESTCONTAINERS_RYUK_DISABLED: "true" + run: | + PATH="$PATH:$HOME/go/bin" make test-e2e-distributed + - name: Setup tmate session if tests fail + if: ${{ failure() }} + uses: mxschmitt/action-tmate@v3.23 + with: + detached: true + connect-timeout-seconds: 180 + limit-access-to-actor: true + + tests-e2e-cluster: + runs-on: ubuntu-latest + # Advisory for the same reason as the job above: master has no branch + # protection, so a failure here is a visible red X rather than a blocked + # merge. That is a repository-settings property, not a YAML key. The key + # that looks like it says "advisory" instead flips the run's conclusion to + # success, which hides the failure rather than flagging it, so it appears in + # none of this repo's workflows and must not be added here. + # + # Separate job from tests-e2e-distributed so the fast in-process suite is + # not held behind a Go build of local-ai. Serial on purpose: each Ginkgo + # process would get its own PostgreSQL and NATS container and each spec + # spawns two or three local-ai children, so --procs on an unmeasured runner + # is a change to make with numbers, not by default. + # + # The two timeouts bound different things and are not alternatives. Ginkgo's + # --timeout=20m bounds the SUITE only; this job timeout must additionally + # cover setup, which here is the larger and more variable half: submodule + # checkout, apt, protoc plus two go installs plus protogen-go, a cold-cache + # module download (cache: false), a full go build of ./cmd/local-ai, and a + # separate ginkgo test compile. That build alone is ~316s of CPU, so on a + # 4-vCPU runner setup is realistically 8-12 minutes. + # + # 45 minutes therefore, matching the sibling job. A tighter number does not + # make a hang fail faster, it just moves the kill from Ginkgo, which prints + # which spec hung, to the runner, which prints nothing: a red job with no + # evidence, which is how a suite gets disabled rather than fixed. + # + # The suite itself is about 8m30s over three consecutive runs (509.1s / + # 509.8s / 512.3s, so 8m29s to 8m32s) on a developer box, and will be slower + # here. Three specs sit at ~167s each because they wait out a 60s staleness + # threshold plus a 15s health-check tick (HealthCheckInterval, in + # core/config/distributed_config.go; core/services/nodes/health.go runs the + # ticker on the unexported checkInterval, not one of the reconcilers). Do + # not shorten those windows to make this job faster: the wait is what stops + # the assertions from passing before the system could have reacted, which + # was a real false green earlier on. + timeout-minutes: 45 + steps: + - name: Clone + uses: actions/checkout@v7 + with: + submodules: true + - name: Configure apt mirror on runner + uses: ./.github/actions/configure-apt-mirror + - name: Setup Go + uses: actions/setup-go@v5 + with: + go-version: '1.26.0' + cache: false + - name: Dependencies + run: | + sudo apt-get update + sudo apt-get install -y build-essential libopus-dev + - name: Proto Dependencies + run: | + curl -L -s https://github.com/protocolbuffers/protobuf/releases/download/v26.1/protoc-26.1-linux-x86_64.zip -o protoc.zip && \ + unzip -j -d /usr/local/bin protoc.zip bin/protoc && \ + rm protoc.zip + go install google.golang.org/protobuf/cmd/protoc-gen-go@v1.34.2 + go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@1958fcbe2ca8bd93af633f11e97d44e567e945af + PATH="$PATH:$HOME/go/bin" make protogen-go + - name: Stub the embedded React UI + # core/http/react-ui/dist is gitignored and built by Node, but this + # suite drives the HTTP API and never the UI, which has its own e2e + # suite. A single index.html satisfies the //go:embed react-ui/dist/* + # in core/http/app.go, so the job skips a full Node and Vite install. + # If a cluster spec ever asserts on a UI asset, this step must go and + # the real build come back: a developer box has a real dist/, so such a + # spec would pass locally and fail only here, or worse be served the + # stub and pass in both places. + run: | + mkdir -p core/http/react-ui/dist + printf 'stub\n' > core/http/react-ui/dist/index.html + - name: Build local-ai + # Not `make build`: that target pulls in the React UI build. The specs + # exec this binary directly via LOCALAI_E2E_BINARY. + run: | + PATH="$PATH:$HOME/go/bin" go build -o local-ai ./cmd/local-ai + - name: Pre-pull test images + # Same reasoning as the job above: pulling here keeps container-start + # timing out of the spec timeouts and makes a registry outage read as a + # setup failure rather than a test failure. + run: | + docker pull postgres:16-alpine + docker pull nats:2-alpine + - name: Cluster E2E + env: + LOCALAI_E2E_BINARY: ${{ github.workspace }}/local-ai + # Must live under the workspace so the upload step below can reach it. + # The harness defaults to GinkgoT().TempDir(), which lands under + # TMPDIR and would leave the artifact glob matching nothing. + LOCALAI_E2E_LOG_DIR: ${{ github.workspace }}/cluster-logs + # Belt and braces: the harness already fails rather than skips when CI + # is set, and GitHub Actions always sets CI. Stating it here means a + # future edit to that default cannot silently turn this job into one + # that passes without ever starting a cluster, since a skipped cluster + # spec is indistinguishable from a passing one. + LOCALAI_E2E_REQUIRE_BINARIES: "true" + # See the job above: the runner is ephemeral, so the reaper buys + # nothing and would pull a third, unpinned Docker Hub image mid-suite. + TESTCONTAINERS_RYUK_DISABLED: "true" + run: | + PATH="$PATH:$HOME/go/bin" make test-e2e-cluster + - name: Upload process logs + # The per-process logs are the only way to read a cluster failure: the + # Ginkgo output says which assertion failed, not what the four child + # processes were doing. Without this a red job is undebuggable. + if: ${{ failure() }} + uses: actions/upload-artifact@v7 + with: + name: cluster-process-logs + path: cluster-logs/**/*.log + if-no-files-found: ignore + retention-days: 7 + - name: Setup tmate session if tests fail + if: ${{ failure() }} + uses: mxschmitt/action-tmate@v3.23 + with: + detached: true + connect-timeout-seconds: 180 + limit-access-to-actor: true diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d87db37eae63..f48c7c4a337c 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -265,6 +265,37 @@ The e2e tests run LocalAI in a Docker container and exercise the API: make test-e2e ``` +### Running distributed-mode tests + +Distributed mode (several frontend replicas, worker nodes, PostgreSQL and NATS) has two suites. Both bring up their PostgreSQL and NATS with testcontainers, so Docker has to be available: + +```bash +make test-e2e-distributed # in-process: services wired directly into the test binary +make test-e2e-cluster # process-level: real local-ai child processes +``` + +`make test-e2e-distributed` is the fast one (around 240 specs in roughly 75 seconds). It starts one PostgreSQL and one NATS for the whole run and gives each spec its own database. It runs each spec exactly once, with no retry: `DISTRIBUTED_TEST_FLAKES` defaults to 1 and feeds ginkgo's `--flake-attempts`, which counts *total attempts*, not retries. That is deliberately below the repo-wide `TEST_FLAKES=5`, because this suite exists to catch nondeterministic cluster behaviour and a retry hides exactly the failure it is meant to catch. Raise it locally when bisecting something unrelated. + +`make test-e2e-cluster` runs `local-ai` as real child processes, one per frontend replica and one per worker, so a spec can kill a replica and assert what the survivors do. Budget about 8m30s (measured 509.1s / 509.8s / 512.3s over three consecutive runs): three of its six specs wait out real staleness and health-check windows. It needs a built binary and the mock backend: + +```bash +make build build-mock-backend +make test-e2e-cluster +``` + +Two environment variables steer it: + +| Variable | Purpose | +|---|---| +| `LOCALAI_E2E_BINARY` | path to the `local-ai` binary (default: `local-ai` in the repository root) | +| `LOCALAI_E2E_LOG_DIR` | directory for the per-process logs (default: a Ginkgo temp dir) | + +Set `LOCALAI_E2E_LOG_DIR` when debugging. A cluster failure is unreadable without the individual frontend and worker logs, and Ginkgo only tells you which assertion failed. + +A missing binary skips the cluster specs locally but fails them whenever `CI` is set, so a build problem cannot turn the CI job green without ever starting a cluster. `LOCALAI_E2E_REQUIRE_BINARIES=1` forces that failing behaviour anywhere; `LOCALAI_E2E_REQUIRE_BINARIES=0` forces the skip back on even under CI. + +Both suites run in `.github/workflows/tests-e2e-distributed.yml`, on pull requests and on every push to `master`. The `paths-ignore` filter is on the pull-request trigger only, so a master push always runs them. + ### React UI tests and coverage The React UI (`core/http/react-ui/`) is covered by Playwright e2e specs, gated by a **monotonic line-coverage ratchet** (`make test-ui-coverage-check`, run in CI). The metric is non-deterministic — a fast local box reads higher than a slow CI runner for the same code — so a small tolerance is unavoidable. diff --git a/Makefile b/Makefile index ebedb2c98248..2cdca27fc44c 100644 --- a/Makefile +++ b/Makefile @@ -340,12 +340,47 @@ run-e2e-aio: protogen-go @echo 'Running e2e AIO tests' $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts $(TEST_FLAKES) -v -r ./tests/e2e-aio +# Total ginkgo attempts per spec for the distributed suite: --flake-attempts counts +# attempts, not retries. Defaults to 1, so each spec runs once and is never retried, +# unlike TEST_FLAKES=5. This suite exists to catch nondeterministic cluster behaviour, +# and a retry hides exactly the failures it is meant to surface. Raise it locally if +# you are bisecting something unrelated. +DISTRIBUTED_TEST_FLAKES?=1 + # Distributed architecture e2e (PostgreSQL + NATS via testcontainers). # Includes NatsJWT specs (JWT-enabled NATS). Requires Docker. # VLLMMultinode is excluded here; use test-e2e-vllm-multinode for that. +# Cluster is excluded too and runs in test-e2e-cluster below, which needs a +# built binary. The argument-validation specs under tests/e2e/distributed/cluster +# carry Label("Distributed") only, so they run here and not there, on purpose. +# -r stays because of those: they are in a subpackage this target must reach. +# --fail-on-empty because ginkgo exits 0 when a label filter matches nothing, so +# without it a rename of the label would turn this target into a silent no-op +# that still reports "Test Suite Passed". test-e2e-distributed: protogen-go @echo 'Running distributed e2e tests (label Distributed, incl. NatsJWT)' - $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --label-filter='Distributed && !VLLMMultinode' --flake-attempts $(TEST_FLAKES) -v -r ./tests/e2e/distributed + $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --label-filter='Distributed && !VLLMMultinode && !Cluster' --fail-on-empty --flake-attempts $(DISTRIBUTED_TEST_FLAKES) --timeout=40m -v -r ./tests/e2e/distributed + +# Cluster e2e: runs local-ai as real child processes (frontend replicas + +# workers) against PostgreSQL and NATS, and kills them to assert failover. +# Needs a built ./local-ai (or LOCALAI_E2E_BINARY) plus the mock backend. +# +# The argument-validation specs in tests/e2e/distributed/cluster deliberately +# stay in test-e2e-distributed above: they need no binary, no PostgreSQL and no +# NATS, so no -r here and that package is simply out of scope. +# +# --fail-on-empty is load-bearing, not tidiness. Ginkgo exits 0 when a label +# filter selects nothing, so without it a refactor that renames or drops +# Label("Cluster") leaves this target reporting "Test Suite Passed" having +# started no cluster at all. LOCALAI_E2E_REQUIRE_BINARIES does not cover this: +# it only fires inside a spec that is actually running. +# +# --flake-attempts is pinned to 1 rather than $(DISTRIBUTED_TEST_FLAKES), and +# should stay there: this suite exists to catch nondeterministic cluster +# behaviour, and a retry turns exactly that signal into a green run. +test-e2e-cluster: protogen-go build-mock-backend + @echo 'Running cluster e2e tests (label Cluster, real local-ai processes)' + $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --label-filter='Cluster' --fail-on-empty --flake-attempts 1 --timeout=20m -v ./tests/e2e/distributed # vLLM multi-node DP smoke (CPU). Builds local-ai:tests and the # cpu-vllm backend from the current working tree, then drives a diff --git a/core/http/endpoints/localai/backend_logs.go b/core/http/endpoints/localai/backend_logs.go index 6072b8483708..f2112502d253 100644 --- a/core/http/endpoints/localai/backend_logs.go +++ b/core/http/endpoints/localai/backend_logs.go @@ -121,6 +121,17 @@ func BackendLogsWebSocketEndpoint(ml *model.ModelLoader) echo.HandlerFunc { conn := &backendLogsConn{Conn: ws} + // KNOWN RACE: the snapshot is sent before the subscription is registered, so + // a line appended in that window is never streamed. A viewer attaching while + // a model loads (when a backend is at its noisiest) can silently miss lines; + // they stay in the buffer, so a reload shows them. Fixing it needs an atomic + // snapshot-plus-subscribe held under the buffer's own lock (buf.mu in + // pkg/model/backend_log_store.go), because that is the lock AppendLine takes + // while it enqueues and fans out to subscribers. The store-level s.mu guards + // only the buffers map and excludes nothing an appender does, so taking it + // leaves this race exactly where it is. Reordering these two calls is not a + // fix either: it would duplicate instead of drop. + // Send existing lines as initial batch existingLines := ml.BackendLogs().GetLines(modelID) initialMsg := map[string]any{ diff --git a/core/services/nodes/file_transfer_server.go b/core/services/nodes/file_transfer_server.go index 0fc5ac6343f4..9a3f0ca2f5bf 100644 --- a/core/services/nodes/file_transfer_server.go +++ b/core/services/nodes/file_transfer_server.go @@ -839,6 +839,17 @@ func handleBackendLogsWS(w http.ResponseWriter, r *http.Request, logStore *model conn := &backendLogsWSConn{Conn: ws} + // KNOWN RACE: the snapshot is sent before the subscription is registered, so + // a line appended in that window is never streamed. A viewer attaching while + // a model loads (when a backend is at its noisiest) can silently miss lines; + // they stay in the buffer, so a reload shows them. Fixing it needs an atomic + // snapshot-plus-subscribe held under the buffer's own lock (buf.mu in + // pkg/model/backend_log_store.go), because that is the lock AppendLine takes + // while it enqueues and fans out to subscribers. The store-level s.mu guards + // only the buffers map and excludes nothing an appender does, so taking it + // leaves this race exactly where it is. Reordering these two calls is not a + // fix either: it would duplicate instead of drop. + // Send existing lines as initial batch existingLines := logStore.GetLines(modelID) initialMsg := map[string]any{ diff --git a/pkg/model/backend_log_store.go b/pkg/model/backend_log_store.go index c5b5253ddc40..3c60f34a3736 100644 --- a/pkg/model/backend_log_store.go +++ b/pkg/model/backend_log_store.go @@ -344,3 +344,45 @@ func (s *BackendLogStore) Subscribe(modelID string) (chan BackendLogLine, func() return ch, unsubscribe } + +// SubscriberCount reports how many live subscriptions exist for modelID, +// resolving the ID with the same exact-key / replica-prefix rules as Subscribe. +// +// Streaming handlers send a GetLines snapshot before they call Subscribe, so a +// line appended between those two calls reaches the buffer but no channel. A +// caller that has to observe a line it appends itself must therefore wait for +// the subscription to exist rather than assume the handler got there first. +func (s *BackendLogStore) SubscriberCount(modelID string) int { + s.mu.RLock() + exact, exactOK := s.buffers[modelID] + var replicas []*backendLogBuffer + if !strings.Contains(modelID, replicaSeparator) { + prefix := modelID + replicaSeparator + for k, b := range s.buffers { + if strings.HasPrefix(k, prefix) { + replicas = append(replicas, b) + } + } + } + s.mu.RUnlock() + + // Lock order in this type is always s.mu before any buffer lock — Subscribe + // holds s.mu.RLock across its replica registrations, which take buf.mu — so + // counting after releasing s.mu keeps that order rather than inverting it. + // The total is therefore a sample, not a snapshot: a concurrent Subscribe + // can register a further buffer while this loop runs. + count := func(buf *backendLogBuffer) int { + buf.mu.Lock() + defer buf.mu.Unlock() + return len(buf.subscribers) + } + + total := 0 + if exactOK { + total += count(exact) + } + for _, b := range replicas { + total += count(b) + } + return total +} diff --git a/pkg/model/backend_log_store_test.go b/pkg/model/backend_log_store_test.go index 775e07cdb561..593bcf5a23fa 100644 --- a/pkg/model/backend_log_store_test.go +++ b/pkg/model/backend_log_store_test.go @@ -76,6 +76,38 @@ var _ = Describe("BackendLogStore", func() { }) }) + Describe("SubscriberCount", func() { + It("reports zero before anyone subscribes and drops back after unsubscribe", func() { + s.AppendLine("model-a", "stderr", "preload") + Expect(s.SubscriberCount("model-a")).To(Equal(0)) + + _, unsubscribe := s.Subscribe("model-a") + Expect(s.SubscriberCount("model-a")).To(Equal(1)) + + unsubscribe() + Expect(s.SubscriberCount("model-a")).To(Equal(0)) + }) + + // Subscribe resolves a bare model ID across every replica buffer, so the + // count has to follow the same rule or a caller waiting on it would give + // up while a perfectly good subscription was in place. + It("sums the replica buffers a bare model ID resolves to", func() { + s.AppendLine("model-a#0", "stderr", "preload-r0") + s.AppendLine("model-a#1", "stderr", "preload-r1") + + _, unsubscribe := s.Subscribe("model-a") + defer unsubscribe() + + Expect(s.SubscriberCount("model-a")).To(Equal(2)) + Expect(s.SubscriberCount("model-a#0")).To(Equal(1)) + Expect(s.SubscriberCount("model-b")).To(Equal(0)) + }) + + It("returns zero for a model that has no buffer at all", func() { + Expect(s.SubscriberCount("never-seen")).To(Equal(0)) + }) + }) + Describe("Subscribe", func() { // Confirms the WebSocket streaming path (the live tail UI) receives // lines from every replica when the caller subscribes by bare modelID. diff --git a/tests/e2e/distributed/backend_logs_test.go b/tests/e2e/distributed/backend_logs_test.go index 79dea3902d01..82e8ac156401 100644 --- a/tests/e2e/distributed/backend_logs_test.go +++ b/tests/e2e/distributed/backend_logs_test.go @@ -25,6 +25,31 @@ import ( "gorm.io/gorm/logger" ) +// waitForSingleLogSubscriber blocks until the worker's WebSocket log handler has +// registered its subscription on the store. +// +// The handler writes the "initial" batch first and subscribes only afterwards, +// so a line appended the instant that batch lands is buffered but never +// streamed, and the spec then waits out its full read deadline. Measured at +// roughly one run in seventeen with `--repeat`, which is far too often for CI. +// Waiting on the subscription removes the race from the spec; the handler's own +// snapshot/subscribe window is a separate production question, marked at both +// production sites. +// +// Only valid where BackendLogStore.Subscribe resolves modelID to exactly ONE +// buffer: a bare model ID with no "#N" replica buffers in the store, or +// a full process key. Subscribe registers the exact-key buffer and each replica +// buffer one at a time, so for a model that does have replicas the count goes +// positive while later replicas are still unattached and the race survives. +// Hence the assertion is on exactly 1 rather than "at least 1": a spec that +// misapplies this to a replicated model fails loudly on the count instead of +// going quietly back to being flaky. +func waitForSingleLogSubscriber(logStore *model.BackendLogStore, modelID string) { + GinkgoHelper() + Eventually(func() int { return logStore.SubscriberCount(modelID) }, "10s", "5ms"). + Should(Equal(1), "the WebSocket handler never subscribed to %q exactly once", modelID) +} + var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func() { Context("Worker HTTP log endpoints", func() { @@ -212,6 +237,7 @@ var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func Expect(initialLines[1].Text).To(Equal("line-2")) // Now append a new line and verify it arrives via WebSocket + waitForSingleLogSubscriber(logStore, "ws-model") logStore.AppendLine("ws-model", "stdout", "line-3-realtime") conn.SetReadDeadline(time.Now().Add(5 * time.Second)) @@ -280,6 +306,7 @@ var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func Expect(conn.ReadJSON(&initialMsg)).To(Succeed()) // Append line to a different model + waitForSingleLogSubscriber(logStore, "ws-model") logStore.AppendLine("other-model", "stdout", "should not appear") // Append line to our model logStore.AppendLine("ws-model", "stdout", "should appear") @@ -475,6 +502,7 @@ var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func Expect(initialLines[0].Text).To(Equal("initial line from worker")) // Append a new line on the worker's log store + waitForSingleLogSubscriber(logStore, "proxy-model") logStore.AppendLine("proxy-model", "stderr", "realtime via proxy") // Read the streamed line through the proxy diff --git a/tests/e2e/distributed/cluster/admin.go b/tests/e2e/distributed/cluster/admin.go new file mode 100644 index 000000000000..617bcef99086 --- /dev/null +++ b/tests/e2e/distributed/cluster/admin.go @@ -0,0 +1,176 @@ +package cluster + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/cookiejar" + "net/url" + "time" + + "github.com/mudler/LocalAI/pkg/httpclient" +) + +const ( + // adminPassword is sent with "acknowledge_weak_password": true, which sets + // PasswordPolicy{AllowWeak: true} and skips the length floor and the zxcvbn + // score entirely (core/http/auth/password.go). Only the technical + // invariants still apply: non-empty, at most 72 bytes, no NUL. The + // acknowledgement is deliberate rather than incidental, so a future + // tightening of the policy cannot break every failover spec at setup time. + adminPassword = "e2e-admin-password" + // sessionCookieName mirrors the unexported constant in core/http/auth. + // The register handler returns 201 both for "user created, here is your + // session" and for "this email already exists" (a deliberate account + // enumeration defence), so the status code alone cannot tell the two + // apart: the presence of this cookie is the only reliable signal. + sessionCookieName = "session" + // authRequestTimeout bounds one register/login round trip. + authRequestTimeout = 30 * time.Second + // bodyExcerptLimit caps how much of an error response is quoted back. + bodyExcerptLimit = 512 +) + +// ForTestingEmpty returns a Cluster with no processes. It exists so the package's +// own argument-validation specs do not need to start anything. +func ForTestingEmpty() *Cluster { + return &Cluster{} +} + +// AdminSession registers the admin user on frontend i and returns a client +// carrying the resulting session cookie. The email matches LOCALAI_ADMIN_EMAIL, +// which core/http/auth exempts from the approval gate and assigns the admin +// role, so registration alone yields an active admin session. +// +// Call this ONCE per cluster and share the client. Two reasons: +// +// One, a single rate limiter of 5 requests per minute per client IP guards +// POST /api/auth/token-login, POST /api/auth/register, POST /api/auth/login AND +// PUT /api/auth/password (core/http/routes/auth.go:190). They share one budget, +// and every e2e request arrives from 127.0.0.1, so a spec that changes a +// password spends from the same five. +// +// Two, the returned client is already good for every frontend: sessions live in +// the shared Postgres auth DB, the harness pins one HMAC secret across replicas +// so the session row resolves at any of them, and Go's cookie jar keys cookies +// by host without the port. +func (c *Cluster) AdminSession(i int) (*http.Client, error) { + base, err := c.frontendBaseURL(i) + if err != nil { + return nil, err + } + + jar, err := cookiejar.New(nil) + if err != nil { + return nil, fmt.Errorf("creating cookie jar: %w", err) + } + // httpclient hardens the transport and refuses redirects; the jar is the one + // thing it does not configure, and a session cookie is the whole point here. + client := httpclient.NewWithTimeout(authRequestTimeout) + client.Jar = jar + + credentials := map[string]any{ + "email": c.opts.AdminEmail, + "password": adminPassword, + } + registration := map[string]any{ + "email": c.opts.AdminEmail, + "password": adminPassword, + "name": "E2E Admin", + "acknowledge_weak_password": true, + } + + registerStatus, registerBody, err := postJSON(client, base+"/api/auth/register", registration) + if err != nil { + return nil, fmt.Errorf("registering admin on frontend %d: %w", i, err) + } + if hasSessionCookie(jar, base) { + return client, nil + } + + // No cookie means the user already existed (a repeat call against the same + // Postgres), or registration was rejected. Log in; on failure the + // registration response is the diagnosis, so carry it into the error. + loginStatus, loginBody, err := postJSON(client, base+"/api/auth/login", credentials) + if err != nil { + return nil, fmt.Errorf("logging in admin on frontend %d: %w", i, err) + } + if loginStatus != http.StatusOK { + return nil, fmt.Errorf( + "admin login on frontend %d returned %d (%s); registration had returned %d (%s)", + i, loginStatus, loginBody, registerStatus, registerBody) + } + if !hasSessionCookie(jar, base) { + return nil, fmt.Errorf("admin login on frontend %d returned 200 but set no %q cookie: %s", i, sessionCookieName, loginBody) + } + return client, nil +} + +// GetJSON performs an authenticated GET against a frontend and decodes the body. +func (c *Cluster) GetJSON(client *http.Client, frontend int, path string, out any) error { + base, err := c.frontendBaseURL(frontend) + if err != nil { + return err + } + resp, err := client.Get(base + path) + if err != nil { + return fmt.Errorf("GET %s on frontend %d: %w", path, frontend, err) + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("GET %s on frontend %d returned %d: %s", path, frontend, resp.StatusCode, excerpt(resp.Body)) + } + if err := json.NewDecoder(resp.Body).Decode(out); err != nil { + return fmt.Errorf("decoding %s from frontend %d: %w", path, frontend, err) + } + return nil +} + +// frontendBaseURL validates the index before FrontendURL indexes the slice: a +// bare index panic in a helper every failover spec calls is far harder to read +// than a named error. +func (c *Cluster) frontendBaseURL(i int) (string, error) { + if i < 0 || i >= len(c.frontends) { + return "", fmt.Errorf("frontend %d out of range (cluster has %d)", i, len(c.frontends)) + } + return c.FrontendURL(i), nil +} + +// postJSON sends body as JSON and returns the status plus an excerpt of the +// response, closing the body in every path. +func postJSON(client *http.Client, endpoint string, body any) (int, string, error) { + encoded, err := json.Marshal(body) + if err != nil { + return 0, "", fmt.Errorf("marshalling request body: %w", err) + } + resp, err := client.Post(endpoint, "application/json", bytes.NewReader(encoded)) + if err != nil { + return 0, "", err + } + defer func() { _ = resp.Body.Close() }() + return resp.StatusCode, excerpt(resp.Body), nil +} + +// hasSessionCookie reports whether the jar holds a usable session for base. +func hasSessionCookie(jar *cookiejar.Jar, base string) bool { + u, err := url.Parse(base) + if err != nil { + return false + } + for _, cookie := range jar.Cookies(u) { + if cookie.Name == sessionCookieName && cookie.Value != "" { + return true + } + } + return false +} + +func excerpt(r io.Reader) string { + data, err := io.ReadAll(io.LimitReader(r, bodyExcerptLimit)) + if err != nil { + return fmt.Sprintf("", err) + } + return string(bytes.TrimSpace(data)) +} diff --git a/tests/e2e/distributed/cluster/cluster.go b/tests/e2e/distributed/cluster/cluster.go new file mode 100644 index 000000000000..1783d71e16f0 --- /dev/null +++ b/tests/e2e/distributed/cluster/cluster.go @@ -0,0 +1,438 @@ +// Package cluster runs LocalAI as real child processes for end-to-end tests. +// +// The in-process suites cannot express frontend-replica failure: there is no +// process to kill, no second replica to race, and no real HTTP boundary between +// a worker and the frontend it registered with. This package starts the same +// binary an operator runs, one process per frontend replica and one per worker, +// against containerised Postgres and NATS. +package cluster + +import ( + "fmt" + "net/http" + "os" + "os/exec" + "path/filepath" + "time" + + "github.com/mudler/LocalAI/pkg/httpclient" + + "github.com/phayes/freeport" +) + +// Options configures a cluster. Every field without a default is required. +type Options struct { + // Binary is the path to a built local-ai. + Binary string + // MockBackend is the path to the mock-backend binary. When set it is copied + // into each worker's backends directory as "mock-backend", which is the name + // model YAML refers to (see tests/e2e/e2e_suite_test.go:75). + MockBackend string + // PGDSN and NatsURL point at infrastructure the caller already started. + PGDSN string + NatsURL string + // LogDir receives one file per process. Never empty: a cluster failure is + // unreadable without them. + LogDir string + + RegistrationToken string // default "e2e-token" + AdminEmail string // default "admin@e2e.local" + + Frontends int + Workers int + + // SpreadWorkerRegistrations sends worker i to frontend i%Frontends instead + // of sending every worker to frontend 0. + // + // Off by default, and deliberately so: the cross-replica session specs read + // a node at frontend 1 that only frontend 0 was ever told about, and + // spreading registrations would leave them passing while proving nothing. + // It exists for the racing-replicas spec, which is about two replicas + // writing to one roster concurrently and cannot express that at all while + // every worker registers through the same process. + SpreadWorkerRegistrations bool +} + +// Process is one running local-ai. +type Process struct { + Name string + // Cmd is exposed for signalling only. Never call Cmd.Wait on it: the reaper + // goroutine started in spawn owns it, a second Wait races the first, and + // waitErr is only safe to read after <-p.exited. + Cmd *exec.Cmd + Port int + LogPath string + + logFile *os.File + // exited closes once the reaper has collected the process. Only the reaper + // calls Cmd.Wait, so nothing else may: a second Wait on the same Cmd races + // the first and corrupts ProcessState. + // + // A closed exited proves the child is gone. An open one proves nothing: it + // is still open for the whole interval between the child exiting and waitid + // collecting it, during which the child is a zombie that signal 0 reports as + // alive. Anything asserting on a process being dead must poll, not sample. + exited chan struct{} + waitErr error +} + +// Cluster is a running set of frontend and worker processes. +type Cluster struct { + opts Options + frontends []*Process + workers []*Process + baseDir string +} + +const ( + defaultRegistrationToken = "e2e-token" + defaultAdminEmail = "admin@e2e.local" + // testHMACSecret is shared by every frontend so a session minted at one + // replica validates at all of them. See the note in startFrontend. + testHMACSecret = "e2e-cluster-hmac-secret" + readinessTimeout = 90 * time.Second + readinessPoll = 200 * time.Millisecond + // processExitTimeout bounds the post-SIGKILL wait in terminate. An unbounded + // wait turns one stuck child (D state, or a Wait that never returns) into a + // suite-wide Ginkgo timeout that names nothing. + processExitTimeout = 10 * time.Second +) + +func (o *Options) applyDefaults() { + if o.RegistrationToken == "" { + o.RegistrationToken = defaultRegistrationToken + } + if o.AdminEmail == "" { + o.AdminEmail = defaultAdminEmail + } +} + +func (o Options) validate() error { + if o.Frontends < 1 { + return fmt.Errorf("cluster needs at least one frontend, got %d", o.Frontends) + } + if o.LogDir == "" { + return fmt.Errorf("cluster needs a LogDir: process logs are the only way to read a cluster failure") + } + if st, err := os.Stat(o.Binary); err != nil || st.IsDir() { + return fmt.Errorf("local-ai binary not found at %q (run: make build)", o.Binary) + } + return nil +} + +// Start brings up the cluster. It blocks until every frontend answers /readyz +// and every worker process has been spawned. It does NOT wait for workers to +// register: that needs an authenticated admin session, so a caller that depends +// on registration must poll /api/nodes itself. +func Start(opts Options) (*Cluster, error) { + opts.applyDefaults() + if err := opts.validate(); err != nil { + return nil, err + } + + baseDir, err := os.MkdirTemp("", "localai-cluster-*") + if err != nil { + return nil, fmt.Errorf("creating cluster work dir: %w", err) + } + + c := &Cluster{opts: opts, baseDir: baseDir} + + for i := 0; i < opts.Frontends; i++ { + p, err := c.startFrontend(i, 0) + if err != nil { + c.Stop() + return nil, err + } + c.frontends = append(c.frontends, p) + } + for i := 0; i < opts.Workers; i++ { + p, err := c.startWorker(i) + if err != nil { + c.Stop() + return nil, err + } + c.workers = append(c.workers, p) + } + return c, nil +} + +// startFrontend starts frontend i. A port <= 0 allocates a fresh one; a pinned +// port exists for restart: workers take LOCALAI_REGISTER_TO once at boot and +// never re-resolve it, so a replica that comes back on a new port is +// unreachable by exactly the workers that registered with it. +func (c *Cluster) startFrontend(i int, port int) (*Process, error) { + if port <= 0 { + allocated, err := freeport.GetFreePort() + if err != nil { + return nil, fmt.Errorf("allocating frontend port: %w", err) + } + port = allocated + } + name := frontendName(i) + dir := c.frontendDir(i) + if err := os.MkdirAll(filepath.Join(dir, "models"), 0o755); err != nil { + return nil, fmt.Errorf("creating %s dirs: %w", name, err) + } + if err := os.MkdirAll(filepath.Join(dir, "backends"), 0o755); err != nil { + return nil, fmt.Errorf("creating %s dirs: %w", name, err) + } + // Without an explicit LOCALAI_DATA_PATH every child resolves DataPath to + // ${cwd}/data (core/cli/run.go:48), which under `go test` is inside the + // source tree and shared by every replica: one collectiondb, one task and + // job store for processes that are meant to be independent. + dataPath := c.frontendDataDir(i) + if err := os.MkdirAll(dataPath, 0o750); err != nil { + return nil, fmt.Errorf("creating %s dirs: %w", name, err) + } + + cmd := exec.Command(c.opts.Binary, "run", + "--address", fmt.Sprintf("127.0.0.1:%d", port), + "--models-path", filepath.Join(dir, "models"), + "--backends-path", filepath.Join(dir, "backends"), + ) + // Cmd.Environ() is the parent environment this Cmd would already run with; + // the children need PATH, HOME and the Go/CI environment intact. + cmd.Env = append(cmd.Environ(), + "LOCALAI_DISTRIBUTED=true", + "LOCALAI_NATS_URL="+c.opts.NatsURL, + "LOCALAI_AUTH=true", + "LOCALAI_AUTH_DATABASE_URL="+c.opts.PGDSN, + "LOCALAI_ADMIN_EMAIL="+c.opts.AdminEmail, + "LOCALAI_DATA_PATH="+dataPath, + // Session rows are keyed by HMAC-SHA256(token, APIKeyHMACSecret), and + // the secret is generated per instance into {DataPath}/.hmac_secret + // unless pinned (core/application/startup.go:141-148). Now that each + // replica owns its data directory, an unpinned secret would differ per + // replica, so the cookie minted at frontend 0 would hash to a session + // row that does not exist at frontend 1 and every post-failover + // /api/nodes call would 401 with nothing in the logs to explain it. + // Pinning makes the cross-replica session a property of the harness. + "LOCALAI_AUTH_HMAC_SECRET="+testHMACSecret, + "LOCALAI_REGISTRATION_TOKEN="+c.opts.RegistrationToken, + "LOCALAI_AUTO_APPROVE_NODES=true", + "DEBUG=true", + ) + + p, err := c.spawn(name, cmd, port) + if err != nil { + return nil, err + } + if err := waitReady(p, fmt.Sprintf("http://127.0.0.1:%d/readyz", port)); err != nil { + // The caller never sees this process, so Stop() will never reach it: + // reap it here or it outlives the suite holding a port and a log handle. + p.terminate() + return nil, fmt.Errorf("%s never became ready (see %s): %w", name, p.LogPath, err) + } + return p, nil +} + +func (c *Cluster) startWorker(i int) (*Process, error) { + // Two independent free ports: the worker's file-transfer server defaults to + // basePort-1, which freeport never reserved and which is basePort of another + // worker whenever two allocations land adjacent. + ports, err := freeport.GetFreePorts(2) + if err != nil { + return nil, fmt.Errorf("allocating worker ports: %w", err) + } + grpcPort, httpPort := ports[0], ports[1] + name := fmt.Sprintf("worker-%d", i) + dir := filepath.Join(c.baseDir, name) + backends := filepath.Join(dir, "backends") + if err := os.MkdirAll(filepath.Join(dir, "models"), 0o755); err != nil { + return nil, fmt.Errorf("creating %s dirs: %w", name, err) + } + if err := os.MkdirAll(backends, 0o755); err != nil { + return nil, fmt.Errorf("creating %s dirs: %w", name, err) + } + if c.opts.MockBackend != "" { + if err := copyExecutable(c.opts.MockBackend, filepath.Join(backends, "mock-backend")); err != nil { + return nil, fmt.Errorf("installing mock backend for %s: %w", name, err) + } + } + + cmd := exec.Command(c.opts.Binary, "worker", + "--models-path", filepath.Join(dir, "models"), + "--backends-path", backends, + ) + cmd.Env = append(cmd.Environ(), + fmt.Sprintf("LOCALAI_SERVE_ADDR=127.0.0.1:%d", grpcPort), + fmt.Sprintf("LOCALAI_ADVERTISE_ADDR=127.0.0.1:%d", grpcPort), + fmt.Sprintf("LOCALAI_HTTP_ADDR=127.0.0.1:%d", httpPort), + fmt.Sprintf("LOCALAI_ADVERTISE_HTTP_ADDR=127.0.0.1:%d", httpPort), + // Workers register with frontend 0 ONLY unless the caller opts into + // SpreadWorkerRegistrations, and the cross-replica session specs depend + // on that default. They prove a session minted at frontend 0 resolves at + // frontend 1 by reading a node that only frontend 0 was ever told about; + // register the worker everywhere and they still pass while proving + // nothing. + // + // Nothing in those specs can detect the change. The registry keys nodes + // by name and preserves ids across the shared Postgres + // (core/services/nodes/registry.go:522-527), so a roster read at + // frontend 1 looks identical either way. Anyone changing the default + // here must revisit tests/e2e/distributed/cluster_baseline_test.go by + // hand. + // + // The registrar also fixes where this worker's heartbeats go for the + // rest of its life: the loop posts to the URL it was given at boot and + // never re-resolves it (core/cli/workerregistry/client.go), so killing a + // worker's registrar orphans that worker rather than failing it over. + "LOCALAI_REGISTER_TO="+c.FrontendURL(c.registrarFor(i)), + "LOCALAI_NODE_NAME="+name, + "LOCALAI_REGISTRATION_TOKEN="+c.opts.RegistrationToken, + "LOCALAI_NATS_URL="+c.opts.NatsURL, + "DEBUG=true", + ) + + return c.spawn(name, cmd, grpcPort) +} + +func (c *Cluster) spawn(name string, cmd *exec.Cmd, port int) (*Process, error) { + logPath := filepath.Join(c.opts.LogDir, name+".log") + // Append rather than truncate: a restarted process reopens the same path, and + // the log of the instance that died is the one a failover post-mortem needs. + f, err := os.OpenFile(logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644) + if err != nil { + return nil, fmt.Errorf("creating log file for %s: %w", name, err) + } + cmd.Stdout = f + cmd.Stderr = f + if err := cmd.Start(); err != nil { + _ = f.Close() + return nil, fmt.Errorf("starting %s: %w", name, err) + } + p := &Process{Name: name, Cmd: cmd, Port: port, LogPath: logPath, logFile: f, exited: make(chan struct{})} + // One reaper per process, joined by terminate(): a child that dies on its own + // is collected immediately, so waiters learn about it instead of polling a + // dead port until the readiness timeout. + go func() { + p.waitErr = cmd.Wait() + close(p.exited) + }() + return p, nil +} + +// terminate kills the process, waits for the reaper, and releases the log +// handle. Safe to call more than once and on a process that already exited. +func (p *Process) terminate() { + if p == nil || p.Cmd == nil || p.Cmd.Process == nil { + return + } + _ = p.Cmd.Process.Kill() + select { + case <-p.exited: + case <-time.After(processExitTimeout): + fmt.Printf("warning: %s did not exit within %s after SIGKILL; continuing teardown\n", p.Name, processExitTimeout) + } + if p.logFile != nil { + _ = p.logFile.Close() + } +} + +// FrontendURL is the base URL of frontend i. +func (c *Cluster) FrontendURL(i int) string { + return fmt.Sprintf("http://127.0.0.1:%d", c.frontends[i].Port) +} + +// WorkerName is the node name worker i registered under. +func (c *Cluster) WorkerName(i int) string { + return c.workers[i].Name +} + +// registrarFor is the frontend index worker i registers and heartbeats with. +// +// It is read at spawn time and baked into the worker's environment, so it is +// also the answer to "which replica's death orphans this worker". +func (c *Cluster) registrarFor(worker int) int { + if !c.opts.SpreadWorkerRegistrations || c.opts.Frontends < 1 { + return 0 + } + return worker % c.opts.Frontends +} + +// WorkerRegistrar is registrarFor, exported so a spec can say which replica it +// is about to kill relative to a worker instead of re-deriving the rule. +// +// It returns an error rather than indexing blindly, like every other exported +// method here that takes an index. Gomega treats the trailing error as one that +// must be nil, so Expect(c.WorkerRegistrar(0)).To(...) reads unchanged at the +// call site while an out-of-range index fails the spec by name instead of +// silently answering 0, which is a real frontend index and would send a spec +// off to kill the wrong replica. +func (c *Cluster) WorkerRegistrar(worker int) (int, error) { + if err := c.checkWorkerIndex(worker); err != nil { + return 0, err + } + return c.registrarFor(worker), nil +} + +// Stop terminates every process and removes the work directory. Logs survive in +// LogDir, which the caller owns. +func (c *Cluster) Stop() { + // Start returns (nil, err) after stopping itself, so a spec that defers + // c.Stop before asserting the error would otherwise nil-deref. + if c == nil { + return + } + for _, p := range append(append([]*Process{}, c.workers...), c.frontends...) { + p.terminate() + } + if c.baseDir != "" { + _ = os.RemoveAll(c.baseDir) + } +} + +// DumpLogs writes every process log to stdout. Call from an AfterEach guarded by +// CurrentSpecReport().Failed(). +func (c *Cluster) DumpLogs() { + for _, p := range append(append([]*Process{}, c.frontends...), c.workers...) { + if p == nil { + continue + } + data, err := os.ReadFile(p.LogPath) + if err != nil { + fmt.Printf("=== %s: log unreadable: %v\n", p.Name, err) + continue + } + fmt.Printf("=== %s (%s) ===\n%s\n", p.Name, p.LogPath, string(data)) + } +} + +func waitReady(p *Process, url string) error { + deadline := time.Now().Add(readinessTimeout) + client := httpclient.NewWithTimeout(2 * time.Second) + var last error + for time.Now().Before(deadline) { + select { + case <-p.exited: + if p.waitErr == nil { + return fmt.Errorf("process exited cleanly before becoming ready") + } + return fmt.Errorf("process exited before becoming ready: %w", p.waitErr) + default: + } + resp, err := client.Get(url) + if err == nil { + _ = resp.Body.Close() + if resp.StatusCode == http.StatusOK { + return nil + } + last = fmt.Errorf("status %d", resp.StatusCode) + } else { + last = err + } + time.Sleep(readinessPoll) + } + return fmt.Errorf("not ready within %s: %w", readinessTimeout, last) +} + +func copyExecutable(src, dst string) error { + data, err := os.ReadFile(src) + if err != nil { + return fmt.Errorf("reading %s: %w", src, err) + } + if err := os.WriteFile(dst, data, 0o755); err != nil { + return fmt.Errorf("writing %s: %w", dst, err) + } + return nil +} diff --git a/tests/e2e/distributed/cluster/cluster_suite_test.go b/tests/e2e/distributed/cluster/cluster_suite_test.go new file mode 100644 index 000000000000..69f785ba4cad --- /dev/null +++ b/tests/e2e/distributed/cluster/cluster_suite_test.go @@ -0,0 +1,13 @@ +package cluster_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestCluster(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Cluster Harness Suite") +} diff --git a/tests/e2e/distributed/cluster/cluster_test.go b/tests/e2e/distributed/cluster/cluster_test.go new file mode 100644 index 000000000000..2f63d0a28e76 --- /dev/null +++ b/tests/e2e/distributed/cluster/cluster_test.go @@ -0,0 +1,87 @@ +package cluster_test + +import ( + "os" + "path/filepath" + "time" + + "github.com/mudler/LocalAI/pkg/httpclient" + "github.com/mudler/LocalAI/tests/e2e/distributed/cluster" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Cluster options", Label("Distributed"), func() { + It("rejects a binary path that does not exist", func() { + _, err := cluster.Start(cluster.Options{ + Binary: filepath.Join(os.TempDir(), "definitely-not-local-ai"), + PGDSN: "postgres://test:test@127.0.0.1:5432/x?sslmode=disable", + NatsURL: "nats://127.0.0.1:4222", + LogDir: GinkgoT().TempDir(), + Frontends: 1, + }) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("local-ai binary")) + }) + + It("rejects a cluster with no frontends", func() { + _, err := cluster.Start(cluster.Options{ + Binary: "/bin/true", + PGDSN: "postgres://test:test@127.0.0.1:5432/x?sslmode=disable", + NatsURL: "nats://127.0.0.1:4222", + LogDir: GinkgoT().TempDir(), + Frontends: 0, + }) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("at least one frontend")) + }) +}) + +// The HTTP flow inside AdminSession and GetJSON cannot run here: it needs a +// built local-ai plus real Postgres and NATS, which arrive with the failover +// suites. These specs cover the argument validation that would otherwise panic +// on an out-of-range slice index inside a helper every later spec calls. +var _ = Describe("Admin session", Label("Distributed"), func() { + It("reports a clear error when the frontend index is out of range", func() { + c := cluster.ForTestingEmpty() + _, err := c.AdminSession(3) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("frontend 3")) + }) + + It("reports a clear error when GetJSON names a frontend that does not exist", func() { + c := cluster.ForTestingEmpty() + err := c.GetJSON(httpclient.NewWithTimeout(time.Second), 1, "/api/nodes", &struct{}{}) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("frontend 1")) + }) +}) + +// Like the admin specs above, these cover argument validation only. Killing, +// stopping and restarting a real replica needs a built local-ai plus Postgres +// and NATS, so those paths stay unexecuted until the failover suites land. +var _ = Describe("Failure primitives", Label("Distributed"), func() { + It("rejects an out-of-range frontend index rather than panicking", func() { + c := cluster.ForTestingEmpty() + Expect(c.KillFrontend(0)).To(MatchError(ContainSubstring("frontend 0 out of range"))) + Expect(c.StopFrontendGracefully(2)).To(MatchError(ContainSubstring("frontend 2 out of range"))) + Expect(c.KillWorker(1)).To(MatchError(ContainSubstring("worker 1 out of range"))) + }) + + It("rejects a restart of a frontend index that does not exist", func() { + Expect(cluster.ForTestingEmpty().RestartFrontend(0)). + To(MatchError(ContainSubstring("frontend 0 out of range"))) + }) + + It("rejects a negative index without treating it as an offset from the end", func() { + c := cluster.ForTestingEmpty() + Expect(c.KillFrontend(-1)).To(MatchError(ContainSubstring("frontend -1 out of range"))) + Expect(c.KillWorker(-1)).To(MatchError(ContainSubstring("worker -1 out of range"))) + Expect(c.FrontendAlive(-1)).To(BeFalse()) + }) + + It("reports a frontend that was never started as not alive", func() { + Expect(cluster.ForTestingEmpty().FrontendAlive(0)).To(BeFalse()) + }) +}) diff --git a/tests/e2e/distributed/cluster/failure.go b/tests/e2e/distributed/cluster/failure.go new file mode 100644 index 000000000000..fe808401cd56 --- /dev/null +++ b/tests/e2e/distributed/cluster/failure.go @@ -0,0 +1,195 @@ +package cluster + +import ( + "fmt" + "os" + "path/filepath" + "syscall" +) + +// KillFrontend SIGKILLs frontend i. This is the "replica died" case: no drain, +// no graceful deregistration, sockets drop without a FIN from the application. +// +// The signal is delivered but not waited on, because a spec that asserts on the +// cluster's reaction wants to observe the window between the death and the +// survivors noticing it. Poll FrontendAlive with Eventually to join the exit. +func (c *Cluster) KillFrontend(i int) error { + if err := c.checkFrontendIndex(i); err != nil { + return err + } + return signalProcess(c.frontends[i], syscall.SIGKILL) +} + +// StopFrontendGracefully SIGTERMs frontend i. This is the rolling-update case: +// the process gets a chance to drain and deregister. Like KillFrontend it does +// not wait; the point of the distinction between the two is what the process +// does with the time between the signal and its exit. +func (c *Cluster) StopFrontendGracefully(i int) error { + if err := c.checkFrontendIndex(i); err != nil { + return err + } + return signalProcess(c.frontends[i], syscall.SIGTERM) +} + +// KillWorker SIGKILLs worker i. +func (c *Cluster) KillWorker(i int) error { + if err := c.checkWorkerIndex(i); err != nil { + return err + } + return signalProcess(c.workers[i], syscall.SIGKILL) +} + +// RestartFrontend brings frontend i back on its original port with an empty +// data directory, modelling a replaced pod rather than a resumed one. +// +// The port is pinned rather than reallocated: workers read LOCALAI_REGISTER_TO +// once at boot and never re-resolve it, so a replica that returns on a new port +// is unreachable by exactly the workers that registered with it, and the +// failover the spec means to observe never happens. Rebinding is safe because +// the previous listener is fully closed before the new process starts (see the +// terminate below) and Go's listeners set SO_REUSEADDR, so a lingering +// TIME_WAIT on an accepted connection does not block the bind. +// +// The data directory is wiped so the replica must rehydrate node, session and +// job state from the shared Postgres and NATS. Keeping it would model a pod +// with a persistent volume and would hide the very class of bug these tests +// exist to find. This is only safe because startFrontend pins +// LOCALAI_AUTH_HMAC_SECRET: the secret otherwise lives at +// {DataPath}/.hmac_secret, and wiping it would make every session minted before +// the restart hash to a row the restarted replica cannot find, turning a +// failover assertion into an unexplained 401. +// +// The wipe also destroys state that nothing can rebuild. This harness sets no +// LOCALAI_STORAGE_URL, so the distributed object store is a directory under +// {DataPath} (core/application/distributed.go:146), and quantization and +// fine-tune jobs write their outputs to {DataPath}/quantization and +// {DataPath}/fine-tune (core/services/{quantization,finetune}/service.go:95); +// agent state, router-corpus, the voiceprofile store and {DataPath}/traces go +// the same way. Postgres keeps the job row, the artifact it points at is gone. +// So a spec that finishes a quantization or fine-tune on a replica, restarts +// it, and then asserts the artifact is retrievable fails for a storage reason +// dressed up as a failover one. No spec does that today; this note is here so +// the first one that tries does not spend a day on it. +// +// After StopFrontendGracefully, wait for the process to actually go before +// restarting: +// +// Eventually(func() bool { return c.FrontendAlive(i) }, "20s", "500ms"). +// Should(BeFalse()) +// +// FrontendAlive takes an index, so it has to be wrapped in a closure; handing +// Gomega the method value directly fails immediately: Eventually reports that +// the function it was given takes one argument and none were provided, and +// points at Eventually().WithArguments(). Restart +// terminates whatever is still running with SIGKILL, so restarting straight +// after a SIGTERM cuts the drain short and quietly turns the rolling-update +// case into the crash case, which is the opposite of what pairing those two +// calls is meant to express. +func (c *Cluster) RestartFrontend(i int) error { + if err := c.checkFrontendIndex(i); err != nil { + return err + } + old := c.frontends[i] + if old == nil { + return fmt.Errorf("frontend %d was never started, nothing to restart", i) + } + // frontendDataDir is relative when baseDir is empty, and this deletes it: + // a Cluster assembled by a future test helper without a work dir would have + // RemoveAll walking "frontend-N/data" under the package source directory. + if c.baseDir == "" { + return fmt.Errorf("refusing to wipe the data dir of frontend %d: cluster has no work dir", i) + } + // The old process may still be running (a restart with no preceding kill) or + // already dead but unreaped. terminate is idempotent, bounds its wait, and + // releases the log handle the replacement is about to reopen; without it the + // replacement races the old listener for the port and leaks a file + // descriptor per restart. + old.terminate() + if err := os.RemoveAll(c.frontendDataDir(i)); err != nil { + return fmt.Errorf("wiping data dir of frontend %d: %w", i, err) + } + + p, err := c.startFrontend(i, old.Port) + if err != nil { + return fmt.Errorf("restarting frontend %d: %w", i, err) + } + c.frontends[i] = p + return nil +} + +// FrontendAlive reports whether frontend i's process is still running. +func (c *Cluster) FrontendAlive(i int) bool { + if i < 0 || i >= len(c.frontends) { + return false + } + return c.frontends[i].alive() +} + +// alive reports whether the process is still running. +// +// The exited check is cheap hygiene, not a fix for the zombie window. The +// reaper closes exited only after Cmd.Wait returns, and Wait marks the +// os.Process done before it returns (runtime/os pidfd path), so by the time +// exited is closed signal 0 already errors: this branch cannot fire earlier +// than the one it precedes. The window that stays open is the other one, +// between the child exiting and waitid collecting it: there the child is a +// zombie, signal 0 to a zombie succeeds, and alive reports true for a process +// that is already dead. There is no local fix; the caller's is to poll rather +// than assert once, wrapping the index-taking FrontendAlive in a closure: +// +// Eventually(func() bool { return c.FrontendAlive(i) }, "20s", "500ms"). +// Should(BeFalse()) +func (p *Process) alive() bool { + if p == nil || p.Cmd == nil || p.Cmd.Process == nil { + return false + } + select { + case <-p.exited: + return false + default: + } + // Signal 0 tests for existence without delivering anything. + return p.Cmd.Process.Signal(syscall.Signal(0)) == nil +} + +func signalProcess(p *Process, sig syscall.Signal) error { + if p == nil || p.Cmd == nil || p.Cmd.Process == nil { + return fmt.Errorf("process is not running") + } + if err := p.Cmd.Process.Signal(sig); err != nil { + return fmt.Errorf("signalling %s with %v: %w", p.Name, sig, err) + } + return nil +} + +// checkFrontendIndex keeps the out-of-range wording identical across every +// primitive, so a failing spec reads the same whichever one tripped. +func (c *Cluster) checkFrontendIndex(i int) error { + if i < 0 || i >= len(c.frontends) { + return fmt.Errorf("frontend %d out of range (cluster has %d)", i, len(c.frontends)) + } + return nil +} + +// checkWorkerIndex is checkFrontendIndex for workers, and exists for the same +// reason: one wording, whichever primitive tripped. +func (c *Cluster) checkWorkerIndex(i int) error { + if i < 0 || i >= len(c.workers) { + return fmt.Errorf("worker %d out of range (cluster has %d)", i, len(c.workers)) + } + return nil +} + +func frontendName(i int) string { + return fmt.Sprintf("frontend-%d", i) +} + +func (c *Cluster) frontendDir(i int) string { + return filepath.Join(c.baseDir, frontendName(i)) +} + +// frontendDataDir is LOCALAI_DATA_PATH for frontend i. RestartFrontend wipes it, +// so it must be the exact path startFrontend hands the child. +func (c *Cluster) frontendDataDir(i int) string { + return filepath.Join(c.frontendDir(i), "data") +} diff --git a/tests/e2e/distributed/cluster_baseline_test.go b/tests/e2e/distributed/cluster_baseline_test.go new file mode 100644 index 000000000000..6cbaaa0ea915 --- /dev/null +++ b/tests/e2e/distributed/cluster_baseline_test.go @@ -0,0 +1,291 @@ +package distributed_test + +import ( + "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"` +} + +// 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. +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") + } + if _, err := os.Stat(path); err != nil { + missingBinary("local-ai binary", path, "run `make build` or set LOCALAI_E2E_BINARY") + } + return path +} + +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() + + // 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, + NatsURL: infra.NatsURL, + 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 +} + +// 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 "" +} + +// 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") + }) +}) diff --git a/tests/e2e/distributed/cluster_failover_test.go b/tests/e2e/distributed/cluster_failover_test.go new file mode 100644 index 000000000000..d007a8ae6ac8 --- /dev/null +++ b/tests/e2e/distributed/cluster_failover_test.go @@ -0,0 +1,401 @@ +package distributed_test + +import ( + "fmt" + + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/tests/e2e/distributed/cluster" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// The two sentinels below are returned by statusOf in place of a node status. +// Neither can ever equal one of the nodes.Status* constants, which is the whole +// point: every assertion in this file compares against a real status with +// Equal, so an unreachable frontend or a vanished row fails the assertion and +// names itself rather than quietly satisfying it. +// +// This is not hypothetical. The obvious way to write "the dead worker is gone" +// is ShouldNot(ContainElement(name)) over a list of healthy names, and the +// probe returns an empty list on any error, so an expired session, a 401 at the +// second replica or a decode failure all satisfy that matcher. The spec would +// go green having observed nothing at all. +const ( + statusUnreachable = "" + statusAbsent = "" +) + +const ( + // orphanEvictionWindow is how long a spec watches a roster to be sure the + // system had a real chance to evict a node and chose not to. + // + // It is sized from the only eviction path there is. Node liveness is + // heartbeat freshness: the health monitor wakes every HealthCheckInterval + // (15s) and marks any node whose last heartbeat is older than + // StaleNodeThreshold (60s) offline (core/services/nodes/health.go, defaults + // in core/config/distributed_config.go). Neither is settable from the CLI, + // so 60s + one 15s tick = 75s is the worst case and cannot be shortened. + // + // Measured rather than assumed: a worker whose registrar was killed goes + // offline at both surviving replicas at t=74.2s. A window shorter than that + // would be the classic false green, a Consistently that passes because + // nothing has had time to happen yet. 100s clears the measured latency by a + // third. + orphanEvictionWindow = "100s" + rosterPollInterval = "2s" + + // workerDeathTimeout bounds the wait for a killed worker to be marked + // offline. Roughly twice the measured 74.3s, which absorbs a health tick + // landing just before the kill plus a slow CI runner. + workerDeathTimeout = "150s" + + // settledStatusWindow is how long an observed status has to hold before the + // spec believes it. A killed worker does not go straight to offline: it + // flaps to unhealthy at ~8s and back to healthy at ~14s (see the note in + // the worker-death spec), so a status has to outlast that transient and two + // further 15s health ticks to count as the settled state. + settledStatusWindow = "45s" + + // restartRehydrationTimeout bounds the wait for a cold-restarted replica to + // answer with the roster. The restart itself measured 1.0s; the budget is + // for a loaded CI runner, not for a slow code path. + restartRehydrationTimeout = "60s" + + // frontendExitTimeout bounds the wait for a signalled frontend to be + // collected. SIGTERM measured 0.2s. It is polled rather than sampled + // because FrontendAlive reports true for the zombie window between the + // child exiting and the reaper calling waitid. + frontendExitTimeout = "30s" + frontendExitPoll = "200ms" +) + +// statusOf refreshes the roster at the probe's frontend and returns the status +// that frontend reports for name. +// +// It shares rosterProbe's lastErr/lastSeen so describe() still explains a +// failure, but unlike healthyNames it never collapses an error into an empty +// result: the caller is comparing against an exact status, so an error has to +// be a value that no assertion can accept. +func (p *rosterProbe) statusOf(name string) string { + var roster []node + if err := p.cluster.GetJSON(p.client, p.frontend, "/api/nodes", &roster); err != nil { + p.lastErr = err + return statusUnreachable + } + p.lastErr = nil + p.lastSeen = roster + for _, n := range roster { + if n.Name == name { + return n.Status + } + } + return statusAbsent +} + +// explain builds a lazy failure description. +// +// Gomega formats a (string, args...) description as soon as the assertion is +// constructed, which for an Eventually or a Consistently is before anything has +// gone wrong; the roster it quoted would be the one from before the wait. A +// func() string is called only on failure, so describe() reports the last +// observation the assertion actually made. +func (p *rosterProbe) explain(format string, args ...any) func() string { + return func() string { + return fmt.Sprintf(format, args...) + ": " + p.describe() + } +} + +// explainStuckOffline is explain with one extra diagnosis attached. +// +// An assertion waiting for offline has a failure mode that looks like a harness +// bug and is not one, so the message names it rather than leaving the reader to +// find it. See proveHealthCheckingIsAlive and the comment on the dead-worker +// spec for the mechanism. +func (p *rosterProbe) explainStuckOffline(worker, format string, args ...any) func() string { + return func() string { + message := fmt.Sprintf(format, args...) + ": " + p.describe() + for _, n := range p.lastSeen { + if n.Name != worker || n.Status != nodes.StatusUnhealthy { + continue + } + message += "\n\nThe node is stuck at unhealthy, which is a LocalAI defect rather than " + + "a harness one: core/services/nodes/health.go:153-155 skips MarkOffline for a node " + + "already marked unhealthy, so a node whose unhealthy mark lands after its heartbeat " + + "has gone stale never reaches offline at all. Start there, not here." + } + return message + } +} + +// proveHealthCheckingIsAlive kills a worker and waits for the roster to settle +// it to offline. +// +// It is the terminating positive control for the two specs that assert a +// healthy worker STAYS healthy. On their own those are pure negative +// assertions: a cluster whose health checking had wedged entirely, say by +// leaking the Postgres advisory lock the monitor takes +// (core/services/nodes/health.go:112), would freeze the roster and satisfy them +// while observing a corpse. +// +// WHAT IT ACTUALLY PROVES, which is less than it looks like. Killing a worker +// afterwards and requiring the roster to react proves the monitor was alive at +// the END of the preceding window. It does not observe the window itself. The +// inference back across it holds only if a wedge would have been sticky, i.e. +// still present when this helper ran. +// +// THE RESIDUAL GAP, and it is not hypothetical in the peer-replica-death spec. +// Health checks are single-flighted across replicas by a session-scoped +// pg_try_advisory_lock (advisorylock.TryWithLockCtx, non-blocking: a replica +// that does not get the lock returns immediately and checks nothing, silently, +// because checkAll discards the acquired flag). That spec SIGKILLs frontend 1, +// which may have been holding the lock at the moment it died. Postgres releases +// a session-level advisory lock only when it reaps the dead backend, so until +// then frontend 0's ticks acquire nothing and no check runs. The roster freezes, +// Consistently(healthy) passes BECAUSE NOTHING WAS CHECKING, and this helper +// still succeeds afterwards once the session is reaped and the lock comes free. +// That wedge is transient rather than permanent, which is exactly the shape the +// backwards inference cannot see. Low probability, real, and bounded by how +// fast Postgres reaps the dead backend, usually immediate on a local socket +// close. +// +// So treat this as a floor and not a proof: it rules out a health monitor that +// is permanently dead, which is the failure that would otherwise make the +// preceding Consistently a statement about a stopped clock, and it does not rule +// out a monitor that was idle for part of the window. Closing the gap needs a +// positive observation from inside the window (a log or metric assertion that a +// check ran), not a stronger assertion here. +// +// It costs a full detection cycle, which is why it is a shared helper: the +// wall-clock price should be paid once per spec and explained once. +func proveHealthCheckingIsAlive(c *cluster.Cluster, probe *rosterProbe, workerIndex int) { + GinkgoHelper() + worker := c.WorkerName(workerIndex) + Expect(c.KillWorker(workerIndex)).To(Succeed()) + Eventually(probe.statusOf, workerDeathTimeout, rosterPollInterval). + WithArguments(worker). + Should(Equal(nodes.StatusOffline), + probe.explainStuckOffline(worker, + "frontend %d never reacted to a killed worker, so health checking was not running during the window above and the assertion before this one proved nothing", + probe.frontend)) +} + +var _ = Describe("Cluster failover", Label("Distributed"), Label("Cluster"), func() { + It("keeps a healthy worker in the roster when a peer replica dies", func() { + // Two replicas, one worker. The worker registers and heartbeats with + // frontend 0 only (the harness default), so frontend 1 is a replica it + // has never spoken to. + c := startCluster(2, 1) + worker := c.WorkerName(0) + + // One session for the whole cluster: register/login/token-login/password + // share a five-per-minute-per-IP budget at every frontend + // (core/http/routes/auth.go:190) and everything here comes from + // 127.0.0.1. The cookie is valid at both replicas because sessions live + // in the shared Postgres and the harness pins one HMAC secret. + client, err := c.AdminSession(0) + Expect(err).ToNot(HaveOccurred()) + + survivor := newRosterProbe(c, client, 0) + Eventually(survivor.healthyNames, nodeRosterTimeout, nodeRosterPoll). + Should(ContainElement(worker), survivor.describe) + + // Kill the replica this worker never registered with. That choice is + // what makes the assertion below mean anything. + // + // Killing frontend 0 instead would sever the worker's only heartbeat + // path, because the heartbeat loop posts to the URL it was handed at + // boot and never re-resolves it; the worker is then genuinely orphaned + // and IS evicted, at a measured 74.2s. A spec written that way can only + // pass by watching for less time than the eviction takes. + Expect(c.WorkerRegistrar(0)).ToNot(Equal(1), + "this spec kills frontend 1 precisely because worker 0 does not depend on it") + Expect(c.KillFrontend(1)).To(Succeed()) + Eventually(func() bool { return c.FrontendAlive(1) }, frontendExitTimeout, frontendExitPoll). + Should(BeFalse(), "frontend 1 did not die, so nothing below is a failover assertion") + + // The survivor must keep answering, and must keep the worker healthy. + // + // A GET that returns 200 with a decodable roster is the "keeps serving" + // half; statusUnreachable would fail this matcher. The window outlasts + // the full 75s stale-plus-one-tick eviction path, so an implementation + // that reacted to a dead peer by sweeping its nodes, by resetting + // heartbeats, or by marking the whole roster stale would be caught + // whether it reacted immediately or on a health tick. + Consistently(survivor.statusOf, orphanEvictionWindow, rosterPollInterval). + WithArguments(worker). + Should(Equal(nodes.StatusHealthy), + survivor.explain("killing a peer replica must not disturb a worker that never depended on it")) + + // Everything above is a negative: nothing happened. Prove that the + // survivor was capable of making something happen the whole time. + proveHealthCheckingIsAlive(c, survivor, 0) + }) + + It("rediscovers the worker from shared state after a cold rolling restart", func() { + c := startCluster(2, 1) + worker := c.WorkerName(0) + + client, err := c.AdminSession(0) + Expect(err).ToNot(HaveOccurred()) + + probe := newRosterProbe(c, client, 0) + Eventually(probe.healthyNames, nodeRosterTimeout, nodeRosterPoll). + Should(ContainElement(worker), probe.describe) + registeredID := probe.idOf(worker) + Expect(registeredID).ToNot(BeEmpty(), "frontend 0 reported the worker without a registration ID") + + // The rolling-update shape: drain, wait for the process to actually go, + // then bring the replacement up. Restarting without that wait would + // SIGKILL the replica mid-drain and silently turn this into the crash + // case. + Expect(c.StopFrontendGracefully(0)).To(Succeed()) + Eventually(func() bool { return c.FrontendAlive(0) }, frontendExitTimeout, frontendExitPoll). + Should(BeFalse(), "frontend 0 ignored SIGTERM, so the restart below would be a SIGKILL mid-drain") + + // RestartFrontend wipes the replica's data directory, so the process + // that comes back has no local memory of the cluster. Everything the + // assertions below observe has to come out of the shared Postgres. + Expect(c.RestartFrontend(0)).To(Succeed()) + + restarted := newRosterProbe(c, client, 0) + Eventually(restarted.statusOf, restartRehydrationTimeout, nodeRosterPoll). + WithArguments(worker). + Should(Equal(nodes.StatusHealthy), + restarted.explain("a replica with an empty data directory must rehydrate the roster from shared state")) + Expect(restarted.idOf(worker)).To(Equal(registeredID), + "the restarted replica invented a new row for the worker instead of resolving the shared one") + + // Rehydration alone is a weak claim: the row was written before the + // restart and would still read healthy for up to 75s even if the + // replacement never accepted another heartbeat. Holding it past that + // window is what proves the worker's heartbeats are landing again, + // which is the part a restart can plausibly break (a replacement on a + // different port, or one that rejects the node id it did not issue). + Consistently(restarted.statusOf, orphanEvictionWindow, rosterPollInterval). + WithArguments(worker). + Should(Equal(nodes.StatusHealthy), + restarted.explain("the worker went stale after the restart, so its heartbeats are not reaching the replacement")) + + // Same hole as the peer-death spec, and it is worse here: a cold + // restart is exactly the event that could leave a replacement unable to + // run health checks at all, and a frozen roster reads identically to a + // healthy one. This is the assertion that tells the two apart. + proveHealthCheckingIsAlive(c, restarted, 0) + }) + + It("settles a dead worker to offline and both replicas report it offline", func() { + c := startCluster(2, 1) + worker := c.WorkerName(0) + + client, err := c.AdminSession(0) + Expect(err).ToNot(HaveOccurred()) + + at0 := newRosterProbe(c, client, 0) + at1 := newRosterProbe(c, client, 1) + Eventually(at0.healthyNames, nodeRosterTimeout, nodeRosterPoll). + Should(ContainElement(worker), at0.describe) + + Expect(c.KillWorker(0)).To(Succeed()) + + // Assert the settled status, not the absence of a healthy name. + // + // A killed worker does not move monotonically. Measured: healthy until + // ~8s, unhealthy at 8s, healthy again at 14s, offline from 74s. The + // unhealthy blip comes from a liveness probe; the health monitor's + // "heartbeat is still fresh" branch then marks it healthy again + // (core/services/nodes/health.go), and only the stale-heartbeat branch + // reaches MarkOffline. So requiring exactly offline is what pins this to + // the stale-detection path rather than to the transient, which any + // not-healthy or not-present matcher would accept at t=8s. + // + // KNOWN HAZARD, read this before blaming the harness for a timeout here. + // The staleness branch skips a node that is already unhealthy + // (core/services/nodes/health.go:153-155, `if node.Status == + // StatusOffline || node.Status == StatusUnhealthy { continue }`). The + // skip exists to stop the monitor re-logging nodes an operator took + // down, but it applies to the flap too: if the transient unhealthy mark + // lands AFTER the heartbeat has already gone stale, rather than at the + // ~8s observed here, MarkOffline is never called and this node stays + // unhealthy forever. This spec would then hang to workerDeathTimeout + // and fail with a roster that looks perfectly ordinary. The ordering + // that triggers it did not occur in any run so far, but nothing + // prevents it, so explainStuckOffline says so in the failure message + // when it sees a node stuck at unhealthy. Fixing it is LocalAI work, + // not test work. + // Reading the same verdict at both replicas proves shared-verdict + // propagation, NOT two independent detectors. Health checks are + // single-flighted by the advisory lock (see proveHealthCheckingIsAlive), + // so exactly one replica ran the check that wrote the offline status, and + // both probes then read that one Postgres row back. What this rules out + // is a replica that keeps a private roster, or one that reads the shared + // row and reports something else. A spec claiming both replicas can + // detect death on their own would have to isolate them from each other, + // which the shared database makes impossible by design. + for _, probe := range []*rosterProbe{at0, at1} { + Eventually(probe.statusOf, workerDeathTimeout, rosterPollInterval). + WithArguments(worker). + Should(Equal(nodes.StatusOffline), + probe.explainStuckOffline(worker, "frontend %d never settled the dead worker to offline", probe.frontend)) + } + + // And it has to stay offline. Nothing may resurrect a row for a process + // that no longer exists, and this window covers three health ticks. + for _, probe := range []*rosterProbe{at0, at1} { + Consistently(probe.statusOf, settledStatusWindow, rosterPollInterval). + WithArguments(worker). + Should(Equal(nodes.StatusOffline), + probe.explain("frontend %d flipped the dead worker away from offline", probe.frontend)) + } + }) + + It("converges on one roster when two replicas register a worker each", func() { + // SpreadWorkerRegistrations sends worker 0 to frontend 0 and worker 1 to + // frontend 1, so the roster is written through two different replicas. + // + // This is a shared-roster identity test, NOT a concurrency test, and the + // distinction matters because the obvious reading of the spec name is + // the wrong one. Start spawns workers one after another and waits for + // neither, and the registrations land about a second apart in practice; + // there is no synchronisation point and nothing here is tuned to make + // the two writes collide. What it does establish is that a roster + // written through two replicas is one roster and not two: same rows, + // same identities, read back from either process. A genuine concurrent + // registration test would need workers released together against a + // shared barrier, and does not exist yet. + c := startCluster(2, 2, func(o *cluster.Options) { + o.SpreadWorkerRegistrations = true + }) + registrar0, err := c.WorkerRegistrar(0) + Expect(err).ToNot(HaveOccurred()) + registrar1, err := c.WorkerRegistrar(1) + Expect(err).ToNot(HaveOccurred()) + Expect(registrar0).ToNot(Equal(registrar1), + "both workers registered through the same replica, so nothing below says anything about two replicas sharing a roster") + + client, err := c.AdminSession(0) + Expect(err).ToNot(HaveOccurred()) + + at0 := newRosterProbe(c, client, 0) + at1 := newRosterProbe(c, client, 1) + expected := []string{c.WorkerName(0), c.WorkerName(1)} + + // ConsistOf, not ContainElements: it fails on a third entry, which is + // how a duplicated row for one worker would show up. + Eventually(at0.healthyNames, nodeRosterTimeout, nodeRosterPoll). + Should(ConsistOf(expected), at0.describe) + Eventually(at1.healthyNames, nodeRosterTimeout, nodeRosterPoll). + Should(ConsistOf(expected), at1.describe) + + // Same names is not the same roster. Compare the registration ids, which + // is the only way to tell "both replicas read one set of rows" from + // "each replica has its own row per worker that happens to share a + // name". + for _, name := range expected { + id := at0.idOf(name) + Expect(id).ToNot(BeEmpty(), fmt.Sprintf("frontend 0 reported %s without a registration ID", name)) + Expect(at1.idOf(name)).To(Equal(id), + fmt.Sprintf("the replicas disagree on the identity of %s, so they are not sharing one roster", name)) + } + }) +}) diff --git a/tests/e2e/distributed/dbname_test.go b/tests/e2e/distributed/dbname_test.go new file mode 100644 index 000000000000..33305533d3da --- /dev/null +++ b/tests/e2e/distributed/dbname_test.go @@ -0,0 +1,34 @@ +package distributed_test + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Test database naming", Label("Distributed"), func() { + Describe("sanitizeDBName", func() { + It("lowercases and replaces characters Postgres will not accept unquoted", func() { + Expect(sanitizeDBName("LocalAI-Test.Suite")).To(Equal("localai_test_suite")) + }) + + It("truncates to fit the 63-byte identifier limit with room for a suffix", func() { + long := "" + for i := 0; i < 100; i++ { + long += "a" + } + Expect(len(sanitizeDBName(long))).To(Equal(50), + "an over-long name must be truncated to exactly the 50-byte budget; <= 50 would also accept an empty name") + }) + + It("never produces an empty name", func() { + Expect(sanitizeDBName("---")).ToNot(BeEmpty()) + }) + }) + + Describe("replaceDBName", func() { + It("swaps the database in a testcontainers DSN and keeps the query string", func() { + dsn := "postgres://test:test@127.0.0.1:32768/localai_suite?sslmode=disable" + Expect(replaceDBName(dsn, "spec_7")).To(Equal("postgres://test:test@127.0.0.1:32768/spec_7?sslmode=disable")) + }) + }) +}) diff --git a/tests/e2e/distributed/testhelpers_test.go b/tests/e2e/distributed/testhelpers_test.go index 68cf537e30bd..17ee72ce2e89 100644 --- a/tests/e2e/distributed/testhelpers_test.go +++ b/tests/e2e/distributed/testhelpers_test.go @@ -2,6 +2,10 @@ package distributed_test import ( "context" + "fmt" + "net/url" + "strings" + "sync/atomic" "time" "github.com/mudler/LocalAI/core/services/messaging" @@ -13,9 +17,17 @@ import ( tcnats "github.com/testcontainers/testcontainers-go/modules/nats" tcpostgres "github.com/testcontainers/testcontainers-go/modules/postgres" "github.com/testcontainers/testcontainers-go/wait" + "gorm.io/driver/postgres" + "gorm.io/gorm" + gormlogger "gorm.io/gorm/logger" ) // TestInfra holds shared test containers and connection strings. +// +// PGContainer and NATSContainer are the SUITE-WIDE containers, shared by every +// spec. Never call Terminate or Stop on them from a spec: it ends the run for +// everything after it. They are exposed only because nats_jwt_helpers_test.go +// builds its own TestInfra around a dedicated NATS container. type TestInfra struct { Ctx context.Context PGContainer *tcpostgres.PostgresContainer @@ -25,71 +37,188 @@ type TestInfra struct { NC *messaging.Client } -// SetupInfra starts PostgreSQL and NATS containers and connects a messaging client. -// Call in BeforeEach. Use DeferCleanup or call Teardown in AfterEach. -func SetupInfra(dbName string) *TestInfra { - GinkgoHelper() +// Containers are suite-scoped, not spec-scoped. Starting a Postgres (~10s) and a +// NATS (~3.5s) per spec cost roughly 48 minutes of pure startup across the 213 +// specs behind SetupInfra, which is why this suite was never wired into CI. +// Isolation now comes from a database per spec (~67ms), which is what the dbName +// argument was always describing. +// +// Plain BeforeSuite rather than SynchronizedBeforeSuite is deliberate: under +// `ginkgo -p` each process gets its own container pair, which keeps NATS subjects +// isolated per process. A single shared NATS across parallel processes would let +// specs on different processes see each other's messages on the same subject. +var ( + suitePG *tcpostgres.PostgresContainer + suiteNATS *tcnats.NATSContainer + suitePGDSN string + suiteNatsURL string + dbCounter atomic.Int64 +) - infra := &TestInfra{Ctx: context.Background()} +var _ = BeforeSuite(func() { + ctx := context.Background() var err error - // Start PostgreSQL container - infra.PGContainer, err = tcpostgres.Run(infra.Ctx, "postgres:16-alpine", - tcpostgres.WithDatabase(dbName), + suitePG, err = tcpostgres.Run(ctx, "postgres:16-alpine", + tcpostgres.WithDatabase("localai_suite"), tcpostgres.WithUsername("test"), tcpostgres.WithPassword("test"), testcontainers.WithWaitStrategy( wait.ForLog("database system is ready to accept connections"). WithOccurrence(2). - WithStartupTimeout(30*time.Second), + WithStartupTimeout(90*time.Second), ), ) Expect(err).ToNot(HaveOccurred()) - infra.PGURL, err = infra.PGContainer.ConnectionString(infra.Ctx, "sslmode=disable") + suitePGDSN, err = suitePG.ConnectionString(ctx, "sslmode=disable") Expect(err).ToNot(HaveOccurred()) - // Start NATS container - infra.NATSContainer, err = tcnats.Run(infra.Ctx, "nats:2-alpine") + suiteNATS, err = tcnats.Run(ctx, "nats:2-alpine") Expect(err).ToNot(HaveOccurred()) - infra.NatsURL, err = infra.NATSContainer.ConnectionString(infra.Ctx) + suiteNatsURL, err = suiteNATS.ConnectionString(ctx) Expect(err).ToNot(HaveOccurred()) +}) + +var _ = AfterSuite(func() { + ctx := context.Background() + if suitePG != nil { + _ = suitePG.Terminate(ctx) + } + if suiteNATS != nil { + _ = suiteNATS.Terminate(ctx) + } +}) + +// sanitizeDBName maps a spec-supplied label onto a legal unquoted Postgres +// identifier, leaving headroom for the uniqueness suffix appended by SetupInfra. +func sanitizeDBName(name string) string { + var b strings.Builder + for _, r := range strings.ToLower(name) { + switch { + case r >= 'a' && r <= 'z', r >= '0' && r <= '9', r == '_': + b.WriteRune(r) + default: + b.WriteRune('_') + } + } + out := strings.Trim(b.String(), "_") + if out == "" { + out = "spec" + } + // Postgres identifiers cap at 63 bytes; reserve the rest for "_". + if len(out) > 50 { + out = out[:50] + } + return out +} - // Connect messaging client - infra.NC, err = messaging.New(infra.NatsURL) +// replaceDBName swaps the database component of a DSN, preserving credentials, +// host, port and query parameters. +func replaceDBName(dsn, name string) string { + GinkgoHelper() + u, err := url.Parse(dsn) Expect(err).ToNot(HaveOccurred()) + u.Path = "/" + name + return u.String() +} + +// tryAdminDB opens a short-lived connection to the suite's maintenance +// database. Cleanup paths use this rather than adminDB: once connections are +// scarce, a fatal assertion here would convert one Postgres hiccup into a +// suite-wide cascade that buries the original failure. +// +// CREATE/DROP DATABASE cannot run inside a transaction or against the target +// database itself, so every call gets its own connection and closes it. +func tryAdminDB() (*gorm.DB, error) { + db, err := gorm.Open(postgres.Open(suitePGDSN), &gorm.Config{Logger: gormlogger.Discard}) + if err != nil { + return nil, fmt.Errorf("connecting to the suite maintenance database: %w", err) + } + return db, nil +} + +func adminDB() *gorm.DB { + GinkgoHelper() + db, err := tryAdminDB() + Expect(err).ToNot(HaveOccurred()) + return db +} + +func closeDB(db *gorm.DB) { + if db == nil { + return + } + if sqlDB, err := db.DB(); err == nil { + _ = sqlDB.Close() + } +} - // Register cleanup in LIFO order +// SetupInfra provisions a dedicated database on the suite-scoped Postgres and +// returns a client connected to the suite-scoped NATS. Call in BeforeEach; +// cleanup is registered with DeferCleanup. +func SetupInfra(dbName string) *TestInfra { + GinkgoHelper() + Expect(suitePG).ToNot(BeNil(), "SetupInfra called before BeforeSuite started the shared containers") + + infra := &TestInfra{ + Ctx: context.Background(), + PGContainer: suitePG, + NATSContainer: suiteNATS, + NatsURL: suiteNatsURL, + } + + db := fmt.Sprintf("%s_%d", sanitizeDBName(dbName), dbCounter.Add(1)) + + // Scoped so a failed CREATE cannot leak the pool: the assertion panics, and a + // leaked pgx pool per failing spec exhausts the server's connection limit. + func() { + admin := adminDB() + defer closeDB(admin) + Expect(admin.Exec(fmt.Sprintf("CREATE DATABASE %q", db)).Error).To(Succeed()) + }() + + // Registered before anything else can fail: a NATS connect error below would + // otherwise leave the database behind for the rest of the suite. DeferCleanup(func() { if infra.NC != nil { infra.NC.Close() } - if infra.PGContainer != nil { - infra.PGContainer.Terminate(context.Background()) + drop, err := tryAdminDB() + if err != nil { + AddReportEntry("drop database skipped", fmt.Sprintf("%s: %v", db, err)) + return } - if infra.NATSContainer != nil { - infra.NATSContainer.Terminate(context.Background()) + defer closeDB(drop) + // FORCE terminates any connection the spec left open (Postgres 13+). + if err := drop.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %q WITH (FORCE)", db)).Error; err != nil { + AddReportEntry("drop database failed", fmt.Sprintf("%s: %v", db, err)) } }) + infra.PGURL = replaceDBName(suitePGDSN, db) + + var err error + infra.NC, err = messaging.New(infra.NatsURL) + Expect(err).ToNot(HaveOccurred()) + return infra } -// SetupNATSOnly starts only a NATS container and connects a messaging client. -// Useful for tests that don't need PostgreSQL. +// SetupNATSOnly returns a client on the suite-scoped NATS for specs that need no +// database. func SetupNATSOnly() *TestInfra { GinkgoHelper() + Expect(suiteNATS).ToNot(BeNil(), "SetupNATSOnly called before BeforeSuite started the shared containers") - infra := &TestInfra{Ctx: context.Background()} - var err error - - infra.NATSContainer, err = tcnats.Run(infra.Ctx, "nats:2-alpine") - Expect(err).ToNot(HaveOccurred()) - - infra.NatsURL, err = infra.NATSContainer.ConnectionString(infra.Ctx) - Expect(err).ToNot(HaveOccurred()) + infra := &TestInfra{ + Ctx: context.Background(), + NATSContainer: suiteNATS, + NatsURL: suiteNatsURL, + } + var err error infra.NC, err = messaging.New(infra.NatsURL) Expect(err).ToNot(HaveOccurred()) @@ -97,16 +226,12 @@ func SetupNATSOnly() *TestInfra { if infra.NC != nil { infra.NC.Close() } - if infra.NATSContainer != nil { - infra.NATSContainer.Terminate(context.Background()) - } }) return infra } // FlushNATS ensures all subscriptions are registered server-side before publishing. -// Replaces time.Sleep(100ms) after Subscribe calls. func FlushNATS(nc *messaging.Client) { GinkgoHelper() Expect(nc.Conn().Flush()).To(Succeed())