fix: replicate every JetStream KV bucket across the cluster (#1117)

* fix: replicate every JetStream KV bucket across the cluster

A KV bucket created without a replica count gets one replica, and one
replica lives on one JetStream-chosen server recorded in no config.
Losing, restarting or partitioning that server makes the bucket return
"no responders available" from every node at once, so a single-node event
becomes a cluster-wide outage. The reconciler leader-lease buckets were
the worst case: acquiring a lease is a write, so a lease bucket without a
quorum cannot be acquired by anyone, and every reconciler sharing it stops
cluster-wide at exactly the moment a node has failed.

The cluster's node count is now declared once, in config.LoadConfig, and
read by the one helper every bucket creation goes through. Every process
that creates a bucket loads the cluster config first, so no service can
forget to pass it and no signature has to carry it. An undeclared size is
an error at creation rather than a silent fall back to one replica.

Existing buckets are raised rather than left behind: a service opening one
raises it, the daemon sweeps them all after formation, and
"spx admin kv replicas --repair" does it on demand. Raising reads the live
stream config and changes only the replica count, and never lowers one.

Removes the ~30 replicas/clusterSize parameters that only forwarded a
number, and JetStreamManager.UpdateReplicas, which set len(nodes) with no
clamp and would have failed outright on a cluster of six or more.

Proved by a repo-wide invariant test pinning bucket creation to the one
helper, a whole-cluster sweep in every multi-node nightly cell, and a new
kvquorum suite (cell 33) that takes away the node leading the most buckets
and requires the survivors to keep answering for all of them.

* chore: name cell 33 in the nightly dispatch input

* fix: a KV bucket that cannot be raised must still open

The first version of the replica fix treated a raise it could not place as a
hard error, and on a real three-node cluster that turned the fix into the
outage it prevents. A node whose config named a fourth server wanted four
replicas, NATS answered "no suitable peers for placement" for three, and the
daemon retried for its full five-minute budget without finishing bootstrap.
The bucket was present, quorate and serving on three nodes throughout. A
config naming a node before that node serves is the ordinary middle of
growing a cluster, so that is a state to tolerate rather than refuse.

The line is now drawn between the two cases instead of around both. Creating
a bucket below the cluster's replica count still fails, because that is the
defect and a fallback there reintroduces it where it is hardest to notice.
Failing to raise an existing bucket warns and hands the bucket back, and the
raise is retried on every open besides the daemon sweep and --repair. This
matters most in kvlease, which runs on every reconcile tick: refusing there
would stop every reconciler sharing a lease bucket that was working.

getOrCreateBucket now attaches before it creates. A create asking for more
replicas than the cluster can place is refused before the server looks at the
name, so with create-first "this bucket already exists" was an inference from
which error came back, and the wrong one comes back on a single-node server.

RaiseAllBucketReplicas attempts every bucket and joins the failures, and
--repair keeps going too. Stopping at the first would leave the rest of a
cluster short because one bucket could not be placed.

Also adds nats.ServerName/NodeFromServerName beside the template that renders
server_name, with a test tying them together. The kvquorum suite had mapped
stream peers to nodes by comparing them to node names, which they are not:
they are the node name behind a "spinifex-nats-" prefix, so nothing matched
and the node-loss test could not choose a victim.

* test: wait for the right NATS peer count after a node rejoins

/routez reports a node's routes to its peers, not its own membership, so a
fully formed three-node cluster is two peers per node. The kvquorum suite
waited for three and spent its whole five-minute rejoin budget on a cluster
that had already reformed. Every other caller passes len(Nodes)-1; the
helper's doc comment now says so.

The assertion the test exists for had already passed: node1 was leading 18 of
47 buckets when it was taken away, and node2 and node3 each still answered for
all 47.

* fix: replicate KV buckets across the largest odd count a cluster can place

A bucket was replicated across min(nodes, 5) servers, which stores a fourth
full copy on a four-node cluster and still survives exactly one loss — the
same tolerance R3 gives for three quarters of the write amplification. A Raft
majority needs an odd count, so the rule becomes the largest odd count the
cluster can place, capped at five: 1, 3, 3, 5, 5. It is one expression with two
entry points, so no caller can hold a stale copy of the policy and keep passing
after the policy changes.

Three retry loops could not tell a configuration fault from a cluster that had
yet to form. An undeclared cluster size never becomes declared by waiting, so
each of them spent its whole budget on one and then reported quorum as the
cause. They now stop on a permanent error and say what is actually wrong. This
is what the integration suite was hitting: it builds its gateway by hand and
never loads a cluster config, so the size was never declared and the IAM retry
sat in its sleep until the test timed out. Both ends are fixed.

The audit reported configuration alone. A bucket at three replicas with one
peer offline and one lagging is one failure from unavailable and read as ok.
BucketReport now carries the leader and a count of the peers JetStream itself
calls current and reachable, quorum is measured against that count, and
spx admin kv replicas exits 0, 1 or 2 so a deployment gate can tell a problem
--repair fixes from one needing a human.

Declaration is write-once: daemon.Start refuses a conflicting redeclaration
because it is about to create buckets, and it declares before JetStream init
rather than after it. config.LoadConfig warns instead, because it also reads
config files that are not this process's own cluster.

spx admin init --nodes 2 is refused — no replica count on two nodes both
survives a loss and accepts writes — and single-node replication is asserted in
the single e2e suite, so it runs in every nightly host permutation.

* fix: withhold the vLLM serving image until its artifact is published

The catalog offers ubuntu-26.04-vllm-serving-x86_64, and the qcow2 it names
404s on iso.mulgadc.com, so spx admin images import would fail for every
operator who tried it. Its GPU siblings are published; this one never was.

The entry moves to WithheldImages rather than being deleted, because the entry
is the definition — the URL, the checksum and the two tags resolveServingAMI
filters on. Publishing the artifact is then a move back, not a retype.
Bedrock already handles the image not being registered: resolveServingAMI
returns ErrServingAMINotFound, which is the honest answer while there is
nothing to launch.

The bedrock tag contract is still pinned, now against the map the definition
lives in, and it additionally asserts the entry is not offered for import. The
e2e reachability walk covers AvailableImages alone, so a withheld entry is out
of it by construction.

* fix: drop the cluster size from the call sites dev added

Two constructors on dev still pass a replica count that no longer exists as a
parameter: the cluster size is a declared process fact, so EKSServiceDeps has
no ClusterSize field and NewIAMServiceImpl takes three arguments. The merge is
textually clean and breaks the build, which is the shape of collision a
signature removal always has.

* docs: explain what an operator sees when a cluster server fails

The replication section said how many copies a bucket has and stopped there,
which leaves the question an operator actually has during an incident: one
server is gone, what still works. It now covers the majority side carrying on,
an isolated server refusing changes on purpose and why that is the behaviour
you want, guests continuing to run while anything that has to be recorded does
not, the risk of reading state from a server that is cut off, and a returning
server catching up by itself.

The surrounding prose reads from the operator's side rather than the
implementation's: what you see and what to do, not which consensus protocol is
underneath.
This commit is contained in:
Ben Duncan
2026-09-28 12:34:55 -05:00
committed by GitHub
parent ecef5e4d2f
commit 8ab96b29a1
155 changed files with 3268 additions and 871 deletions
+12
View File
@@ -178,5 +178,17 @@
"extra": "",
"suite": "lbrecovery",
"timeout_minutes": 90
},
{
"cell": 33,
"install": "binary",
"topo": "real-multi",
"net": "static",
"host": "debian-13",
"arch": "x86_64",
"runner": "ci-multi",
"extra": "",
"suite": "kvquorum",
"timeout_minutes": 55
}
]
+15 -1
View File
@@ -51,7 +51,13 @@ E2E_SUITES_SINGLE="single iam cert eks ecs storagegrowth partialblock rds quota
# the load balancer's, so taking that node away does not take the backends with
# it. It runs beside instancerecovery at the end because it also takes a node
# away and owns its environment while it does.
E2E_SUITES_MULTI="multinode lb cert quota lbrecovery instancerecovery storagefault"
#
# kvquorum is multi-only and needs three, for the same reason as those two: a
# surviving majority. It is what proves the control plane's own state is
# replicated rather than sitting on one node, so it runs before the pair that
# take a node away — a cluster whose buckets are already wrong should report
# that, not report a recovery failing.
E2E_SUITES_MULTI="multinode lb cert quota kvquorum lbrecovery instancerecovery storagefault"
# grep -xE alternation form, for narrowing a requested set to the eligible one.
E2E_SUITES_SINGLE_RE="$(printf '%s' "${E2E_SUITES_SINGLE}" | tr ' ' '|')"
@@ -75,6 +81,10 @@ E2E_SUITES_MULTI_RE="$(printf '%s' "${E2E_SUITES_MULTI}" | tr ' ' '|')"
# after a failed seal do not exist on one node, and they are the assertions
# worth a nightly.
E2E_SUITES_NIGHTLY_SINGLE="single cert iam"
#
# The multi-node nightly subset carries kvquorum's read-only half through
# `multinode`, which sweeps every KV bucket rather than the three the daemon
# owns. The node-loss half is too slow for a 35-minute cell and has its own.
E2E_SUITES_NIGHTLY_MULTI="multinode cert lb"
# How long a suite gets, for the same reason the lists above live here: the
@@ -112,6 +122,10 @@ e2e_suite_timeout() {
# the backends before it forwards anything. Then the node comes back and has
# to be shown not to answer for an address it no longer holds.
lbrecovery) echo "75m" ;;
# kvquorum audits every bucket in seconds and then spends the rest of its
# budget on a node stopping and starting: seven units down, a raft election,
# and the whole service stack back up with every stream caught up.
kvquorum) echo "40m" ;;
*) echo "30m" ;;
esac
}
+1 -1
View File
@@ -10,7 +10,7 @@ on:
workflow_dispatch:
inputs:
cells:
description: "Comma-separated cell numbers to run, no spaces (blank = all). 1-3, 8-10, 17-19 tofu permutation matrix, 20 baremetal, 21-22 ISO, 23 predastore throughput, 24 predastore fault injection, 25-27 service suites (eks, ecs, rds), 28 guest disk performance, 29 storage fault injection, 30 routed NAT on three nodes, 31 instance recovery on three nodes, 32 load balancer recovery on three nodes"
description: "Comma-separated cell numbers to run, no spaces (blank = all). 1-3, 8-10, 17-19 tofu permutation matrix, 20 baremetal, 21-22 ISO, 23 predastore throughput, 24 predastore fault injection, 25-27 service suites (eks, ecs, rds), 28 guest disk performance, 29 storage fault injection, 30 routed NAT on three nodes, 31 instance recovery on three nodes, 32 load balancer recovery on three nodes, 33 KV bucket quorum under node loss on three nodes"
required: false
permissions:
+32 -2
View File
@@ -30,6 +30,7 @@ import (
"github.com/mulgadc/bluebottle/pkg/masterkey"
"github.com/mulgadc/bluebottle/pkg/safecast"
"github.com/mulgadc/spinifex/spinifex/admin"
"github.com/mulgadc/spinifex/spinifex/clustersize"
"github.com/mulgadc/spinifex/spinifex/config"
"github.com/mulgadc/spinifex/spinifex/ebsmetadata"
"github.com/mulgadc/spinifex/spinifex/ebsprovider"
@@ -372,7 +373,7 @@ func init() {
adminInitCmd.Flags().String("region", "ap-southeast-2", "Mulga region to create")
adminInitCmd.Flags().String("az", "ap-southeast-2a", "Mulga AZ to create")
adminInitCmd.Flags().String("node", "node1", "Node name, increment for additional nodes (default, node1)")
adminInitCmd.Flags().Int("nodes", 3, "Number of nodes to expect for cluster")
adminInitCmd.Flags().Int("nodes", 3, "Number of nodes to expect for cluster (1, or 3 and up)")
adminInitCmd.Flags().String("host", "", "Leader node to join (if not specified, tries multicast discovery)")
adminInitCmd.Flags().Int("port", 4432, "Port to bind cluster services on")
adminInitCmd.Flags().String("bind", "0.0.0.0", "IP address to bind services to (e.g., 10.11.12.1 for multi-node). Default 0.0.0.0 listens on all interfaces.")
@@ -1212,6 +1213,26 @@ func runimagesListCmd(cmd *cobra.Command, args []string) {
pterm.Println("spx admin images import --name <image-name>")
}
// validateNodeCount accepts the cluster sizes Spinifex supports: one server, or
// three and up.
//
// Two is rejected rather than allowed and warned about. On two servers OVN runs
// standalone, the storage metadata quorum has no majority to lose, and the
// control plane's own KV state has no replica count that works — one replica
// disappears with its node, and two stop accepting writes when either goes. A
// cluster that cannot survive the event it was built for is worth refusing at the
// point someone asks for it.
func validateNodeCount(nodes int) error {
if nodes < 1 {
return fmt.Errorf("a cluster needs at least one node, got %d", nodes)
}
if nodes > 1 && nodes < clustersize.MinHANodes {
return fmt.Errorf("%d nodes is not a supported cluster size: use 1 for a single server, or %d and up for a cluster that can survive losing one",
nodes, clustersize.MinHANodes)
}
return nil
}
// TODO: Move all logic to a module, use minimal application logic in viper commands.
func runAdminInit(cmd *cobra.Command, args []string) {
if os.Getuid() != 0 {
@@ -1243,6 +1264,15 @@ func runAdminInit(cmd *cobra.Command, args []string) {
clusterName, _ := cmd.Flags().GetString("cluster-name")
services, _ := cmd.Flags().GetStringSlice("services")
// Rejected before anything is written, because a two-node cluster has no
// correct answer rather than a worse one: no layer of the stack has a
// majority to lose, and no replica count for its state both survives a node
// going away and still accepts writes.
if err := validateNodeCount(nodes); err != nil {
fmt.Fprintf(os.Stderr, "--nodes: %v\n", err)
os.Exit(1)
}
// Optional operator email — validated up-front so a bad address fails
// before we touch any config state. Empty is allowed here; reset / repeat
// inits on a box that already has an email in /etc/spinifex/spinifex.toml
@@ -2824,7 +2854,7 @@ func initIAMServiceFromConfig() (*handlers_iam.IAMServiceImpl, *config.ClusterCo
// Background: this runs at CLI top level, where there is no request to
// inherit a deadline from and the process exits after the one command.
svc, err := handlers_iam.NewIAMServiceImpl(context.Background(), nc, masterKey, len(cfg.Nodes))
svc, err := handlers_iam.NewIAMServiceImpl(context.Background(), nc, masterKey)
if err != nil {
nc.Close()
return nil, nil, nil, nil, fmt.Errorf("init IAM service: %w", err)
+201
View File
@@ -0,0 +1,201 @@
package cmd
import (
"context"
"encoding/json"
"fmt"
"io"
"os"
"strings"
"text/tabwriter"
"github.com/mulgadc/spinifex/spinifex/kvutil"
"github.com/nats-io/nats.go/jetstream"
"github.com/spf13/cobra"
)
var kvReplicasCmd = &cobra.Command{
Use: "replicas",
Short: "Report how many nodes each KV bucket is replicated across",
Long: `Print one line per JetStream KV bucket: how many replicas it has, how many the
cluster's size entitles it to, and which nodes currently hold it.
A bucket with fewer replicas than the cluster has nodes lives on fewer nodes
than exist, and every node depends on those. Losing one of them makes the
bucket unreadable and unwritable from every node at once, which is a
cluster-wide outage caused by a single-node event.
A bucket also has to be able to serve, which is a different question from how
many replicas it was configured for. ONLINE counts the replicas that are caught
up and reachable; a bucket whose ONLINE is not a majority of REPLICAS cannot
accept a write no matter what its replica count says.
--repair raises any under-replicated bucket to the cluster's replica count,
capped at JetStream's maximum of five. It changes nothing else about a bucket
and never lowers a replica count, so it is safe to run on a healthy cluster and
safe to run twice.
Exit status separates the two problems, so this can gate a deployment without
parsing the output:
0 every bucket is at the cluster's replica count and can serve
1 at least one bucket is under-replicated — --repair is the answer
2 at least one bucket has no quorum, or the report could not be produced —
a node needs looking at, and --repair will not help`,
Run: runKVReplicas,
}
func init() {
kvCmd.AddCommand(kvReplicasCmd)
kvReplicasCmd.Flags().Bool("repair", false, "raise every under-replicated bucket to the cluster's node count")
kvReplicasCmd.Flags().Bool("json", false, "emit the report as JSON")
}
func runKVReplicas(cmd *cobra.Command, _ []string) {
repair, _ := cmd.Flags().GetBool("repair")
asJSON, _ := cmd.Flags().GetBool("json")
// Anything that stops the report being produced exits unreachable rather
// than under-replicated: an empty answer is not a cluster with no buckets,
// and --repair is not the response to it.
_, nc, err := loadConfigAndConnect()
if err != nil {
fmt.Fprintf(os.Stderr, "connect to cluster: %v\n", err)
os.Exit(exitUnreachable)
}
defer nc.Close()
js, err := jetstream.New(nc)
if err != nil {
fmt.Fprintf(os.Stderr, "get JetStream context: %v\n", err)
os.Exit(exitUnreachable)
}
// Background: this runs at CLI top level, where there is no request to
// inherit a deadline from and the process exits after the one command.
ctx := context.Background()
reports, err := kvutil.AuditBucketReplicas(ctx, js)
if err != nil {
fmt.Fprintf(os.Stderr, "audit KV bucket replicas: %v\n", err)
os.Exit(exitUnreachable)
}
if repair {
repairReplicaReports(ctx, js, reports, os.Stderr)
}
if err := writeReplicaReport(os.Stdout, reports, asJSON); err != nil {
fmt.Fprintf(os.Stderr, "write report: %v\n", err)
os.Exit(exitUnreachable)
}
os.Exit(replicaReportExit(reports))
}
// Exit statuses, documented in the command's own help because a deployment gate
// reads them rather than the output.
const (
exitHealthy = 0
// exitUnderReplicated is the repairable problem: the buckets serve, but on
// fewer nodes than the cluster has.
exitUnderReplicated = 1
// exitUnreachable is the problem repair cannot touch — a bucket that cannot
// accept a write, or a report that could not be produced at all.
exitUnreachable = 2
)
// replicaReportExit is the exit status for a set of reports.
//
// No quorum outranks under-replicated because they need different responses: one
// is fixed by raising a count, the other by finding out why a node is not
// answering. A bucket that is both should send the operator to the node.
func replicaReportExit(reports []kvutil.BucketReport) int {
status := exitHealthy
for _, r := range reports {
if !r.HasQuorum() {
return exitUnreachable
}
if r.UnderReplicated() {
status = exitUnderReplicated
}
}
return status
}
// repairReplicaReports raises every under-replicated bucket and updates reports
// in place for the ones that were raised.
//
// One bucket that cannot be placed must not stop the others being repaired, so a
// failure is named on errOut and the sweep continues. The bucket it failed on
// stays UNDER-REPLICATED in the report, which is already what makes the exit
// status non-zero, so nothing is lost by not returning an error here.
func repairReplicaReports(ctx context.Context, js jetstream.JetStream, reports []kvutil.BucketReport, errOut io.Writer) {
for i, r := range reports {
if !r.UnderReplicated() {
continue
}
if err := kvutil.RaiseBucketReplicas(ctx, js, r.Bucket, r.Want); err != nil {
fmt.Fprintf(errOut, "repair %s: %v\n", r.Bucket, err)
continue
}
reports[i].Replicas = r.Want
}
}
// underReplicatedCount is how many buckets are below the cluster's replica count.
func underReplicatedCount(reports []kvutil.BucketReport) int {
under := 0
for _, r := range reports {
if r.UnderReplicated() {
under++
}
}
return under
}
// noQuorumCount is how many buckets cannot currently accept a write.
func noQuorumCount(reports []kvutil.BucketReport) int {
stuck := 0
for _, r := range reports {
if !r.HasQuorum() {
stuck++
}
}
return stuck
}
// replicaStatus is the one word for a bucket's condition.
//
// No quorum wins over under-replicated when both are true, matching the exit
// status: a bucket that cannot serve is the more urgent of the two and is not
// what --repair is for.
func replicaStatus(r kvutil.BucketReport) string {
switch {
case !r.HasQuorum():
return "NO-QUORUM"
case r.UnderReplicated():
return "UNDER-REPLICATED"
default:
return "ok"
}
}
// writeReplicaReport renders reports as a table, or as JSON when asJSON.
func writeReplicaReport(out io.Writer, reports []kvutil.BucketReport, asJSON bool) error {
if asJSON {
return json.NewEncoder(out).Encode(reports)
}
w := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
fmt.Fprintln(w, "BUCKET\tREPLICAS\tWANT\tONLINE\tSTATUS\tLEADER\tHELD BY")
for _, r := range reports {
fmt.Fprintf(w, "%s\t%d\t%d\t%d\t%s\t%s\t%s\n",
r.Bucket, r.Replicas, r.Want, r.Online, replicaStatus(r), r.Leader, strings.Join(r.Peers, ","))
}
if err := w.Flush(); err != nil {
return err
}
_, err := fmt.Fprintf(out, "\n%d buckets, %d under-replicated, %d without quorum\n",
len(reports), underReplicatedCount(reports), noQuorumCount(reports))
return err
}
+133
View File
@@ -0,0 +1,133 @@
//test:in-package — it tests this command's unexported rendering and repair
// helpers, which exist only to be exercised without a live cluster.
package cmd
import (
"bytes"
"encoding/json"
"strings"
"testing"
"github.com/mulgadc/spinifex/spinifex/clustersize"
"github.com/mulgadc/spinifex/spinifex/kvutil"
"github.com/mulgadc/spinifex/spinifex/testutil"
"github.com/nats-io/nats.go/jetstream"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestWriteReplicaReport_NamesEveryUnderReplicatedBucket(t *testing.T) {
reports := []kvutil.BucketReport{
{Bucket: "healthy", Replicas: 3, Want: 3, Online: 3, Leader: "a", Peers: []string{"a", "b", "c"}},
{Bucket: "short", Replicas: 1, Want: 3, Online: 1, Leader: "a", Peers: []string{"a"}},
{Bucket: "stuck", Replicas: 3, Want: 3, Online: 1, Peers: []string{"a", "b", "c"}},
}
var out bytes.Buffer
require.NoError(t, writeReplicaReport(&out, reports, false))
got := out.String()
assert.Contains(t, got, "BUCKET")
assert.Contains(t, got, "HELD BY")
assert.Regexp(t, `healthy\s+3\s+3\s+3\s+ok\s+a\s+a,b,c`, got)
assert.Regexp(t, `short\s+1\s+3\s+1\s+UNDER-REPLICATED\s+a\s+a`, got)
// A bucket that cannot serve reads as NO-QUORUM even though its configured
// count is right, because raising a count is not the response to it.
assert.Regexp(t, `stuck\s+3\s+3\s+1\s+NO-QUORUM`, got)
assert.Contains(t, got, "3 buckets, 1 under-replicated, 1 without quorum")
}
// TestReplicaReportExit pins the three exit classes a deployment gate reads. They
// are separate because the responses are: a raise fixes an under-replicated
// bucket, and nothing the command can do fixes one that will not answer.
func TestReplicaReportExit(t *testing.T) {
healthy := kvutil.BucketReport{Bucket: "a", Replicas: 3, Want: 3, Online: 3}
short := kvutil.BucketReport{Bucket: "b", Replicas: 1, Want: 3, Online: 1}
stuck := kvutil.BucketReport{Bucket: "c", Replicas: 3, Want: 3, Online: 1}
assert.Equal(t, exitHealthy, replicaReportExit(nil))
assert.Equal(t, exitHealthy, replicaReportExit([]kvutil.BucketReport{healthy}))
assert.Equal(t, exitUnderReplicated, replicaReportExit([]kvutil.BucketReport{healthy, short}))
assert.Equal(t, exitUnreachable, replicaReportExit([]kvutil.BucketReport{healthy, stuck}))
// Both problems at once sends the operator to the node rather than to --repair.
assert.Equal(t, exitUnreachable, replicaReportExit([]kvutil.BucketReport{short, stuck}))
}
// TestWriteReplicaReport_JSONIsWhatTheE2ESuiteDecodes keeps the machine-readable
// form honest: the e2e suite and anything gating a deployment parse these field
// names, so a rename here breaks them silently.
func TestWriteReplicaReport_JSONIsWhatTheE2ESuiteDecodes(t *testing.T) {
var out bytes.Buffer
require.NoError(t, writeReplicaReport(&out, []kvutil.BucketReport{
{Bucket: "short", Replicas: 1, Want: 3, Cluster: "spinifex", Online: 1, Leader: "a", Peers: []string{"a"}},
}, true))
assert.Contains(t, out.String(), `"bucket":"short"`)
var decoded []kvutil.BucketReport
require.NoError(t, json.Unmarshal(out.Bytes(), &decoded))
require.Len(t, decoded, 1)
assert.Equal(t, "short", decoded[0].Bucket)
assert.Equal(t, 1, decoded[0].Replicas)
assert.Equal(t, 3, decoded[0].Want)
assert.Equal(t, 1, decoded[0].Online)
assert.Equal(t, "a", decoded[0].Leader)
assert.True(t, decoded[0].UnderReplicated())
assert.True(t, decoded[0].HasQuorum())
}
func TestUnderReplicatedCount(t *testing.T) {
assert.Equal(t, 0, underReplicatedCount(nil))
assert.Equal(t, 2, underReplicatedCount([]kvutil.BucketReport{
{Bucket: "a", Replicas: 1, Want: 3},
{Bucket: "b", Replicas: 3, Want: 3},
{Bucket: "c", Replicas: 2, Want: 3},
}))
}
// TestRepairReplicaReports_KeepsGoingPastABucketItCannotPlace is the behaviour
// that makes --repair usable mid-growth: the buckets it can raise are raised,
// and the one it cannot is named rather than aborting the run.
func TestRepairReplicaReports_KeepsGoingPastABucketItCannotPlace(t *testing.T) {
_, nc, js := testutil.StartTestJetStream(t)
defer nc.Close()
for _, b := range []string{"repair-alpha", "repair-beta"} {
_, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{Bucket: b, History: 1})
require.NoError(t, err)
}
// One embedded server cannot hold two replicas, so every raise fails — and
// every bucket must still have been attempted and named.
reports := []kvutil.BucketReport{
{Bucket: "repair-alpha", Replicas: 1, Want: 2},
{Bucket: "repair-beta", Replicas: 1, Want: 2},
}
var errOut bytes.Buffer
repairReplicaReports(t.Context(), js, reports, &errOut)
assert.Contains(t, errOut.String(), "repair repair-alpha")
assert.Contains(t, errOut.String(), "repair repair-beta")
assert.Equal(t, 2, underReplicatedCount(reports), "a failed raise must not be reported as repaired")
}
// TestRepairReplicaReports_RaisesWhatItCan covers the success side against a real
// server, so the in-place report update is proved rather than assumed.
func TestRepairReplicaReports_RaisesWhatItCan(t *testing.T) {
_, nc, js := testutil.StartTestJetStream(t)
defer nc.Close()
clustersize.DeclareForTest(t, 1)
_, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{Bucket: "repair-ok", History: 1})
require.NoError(t, err)
// Want 1 against a bucket already at 1: the raise is a no-op that must still
// leave the report reading as healthy rather than as repaired-and-wrong.
reports := []kvutil.BucketReport{{Bucket: "repair-ok", Replicas: 1, Want: 1}}
var errOut bytes.Buffer
repairReplicaReports(t.Context(), js, reports, &errOut)
assert.Empty(t, strings.TrimSpace(errOut.String()))
assert.Equal(t, 0, underReplicatedCount(reports))
}
+39
View File
@@ -0,0 +1,39 @@
//test:in-package — validateNodeCount is an unexported guard on this command's
// flags, and exporting it would widen the package's API for a test's convenience.
package cmd
import (
"strings"
"testing"
)
// TestValidateNodeCount pins the supported cluster sizes. Two is rejected rather
// than warned about: on two servers no layer of the stack has a majority to lose,
// and the control plane's own state has no replica count that both survives a node
// going away and still accepts writes.
func TestValidateNodeCount(t *testing.T) {
for _, nodes := range []int{1, 3, 4, 5, 9} {
if err := validateNodeCount(nodes); err != nil {
t.Errorf("validateNodeCount(%d) = %v, want nil", nodes, err)
}
}
for _, nodes := range []int{0, -1} {
if err := validateNodeCount(nodes); err == nil {
t.Errorf("validateNodeCount(%d) returned no error", nodes)
}
}
err := validateNodeCount(2)
if err == nil {
t.Fatal("validateNodeCount(2) returned no error")
}
// The message has to say what to do instead, because the operator typed a
// number and needs to know which one to type next.
for _, want := range []string{"2", "1", "3"} {
if !strings.Contains(err.Error(), want) {
t.Errorf("error %q does not mention %s", err, want)
}
}
}
+8 -8
View File
@@ -854,7 +854,7 @@ func runOchreWeightsStage(cmd *cobra.Command, _ []string) {
ochreExit(1)
return
}
weightsStore := gateway_bedrock.NewWeightsStore(js, len(appConfig.Nodes))
weightsStore := gateway_bedrock.NewWeightsStore(js)
store := objectstore.NewS3ObjectStoreFromConfig(node.Predastore.Host, node.Predastore.Region, node.Predastore.AccessKey, node.Predastore.SecretKey)
provider := ebsprovider.NewNATSProvider(nc, imageImportTimeout)
@@ -930,7 +930,7 @@ func runOchreWeightsList(_ *cobra.Command, _ []string) {
ochreExit(1)
return
}
weightsStore := gateway_bedrock.NewWeightsStore(js, len(appConfig.Nodes))
weightsStore := gateway_bedrock.NewWeightsStore(js)
store := objectstore.NewS3ObjectStoreFromConfig(node.Predastore.Host, node.Predastore.Region, node.Predastore.AccessKey, node.Predastore.SecretKey)
checkSnapshotLive := newWeightsSnapshotChecker(store, node.Predastore.Bucket)
@@ -967,7 +967,7 @@ func removeWeights(ctx context.Context, weightsStore *gateway_bedrock.WeightsSto
func runOchreWeightsRemove(cmd *cobra.Command, _ []string) {
modelID, _ := cmd.Flags().GetString("model-id")
appConfig, nc, err := loadConfigAndConnectFn()
_, nc, err := loadConfigAndConnectFn()
if err != nil {
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
ochreExit(1)
@@ -981,7 +981,7 @@ func runOchreWeightsRemove(cmd *cobra.Command, _ []string) {
ochreExit(1)
return
}
weightsStore := gateway_bedrock.NewWeightsStore(js, len(appConfig.Nodes))
weightsStore := gateway_bedrock.NewWeightsStore(js)
entry, err := removeWeights(context.Background(), weightsStore, modelID)
if err != nil {
@@ -1014,7 +1014,7 @@ func resolveHFToken(ctx context.Context, cmd *cobra.Command, cfg *config.Cluster
if err != nil {
return ""
}
credStore := gateway_bedrock.NewCredentialStore(js, masterKey, len(cfg.Nodes), nil)
credStore := gateway_bedrock.NewCredentialStore(js, masterKey, nil)
token, ok, err := credStore.Resolve(ctx, utils.GlobalAccountID, vendorHuggingFace)
if err != nil || !ok {
return ""
@@ -1088,7 +1088,7 @@ func ochreCredentialsStore() (*gateway_bedrock.CredentialStore, func(), error) {
nc.Close()
return nil, nil, fmt.Errorf("jetstream context: %w", err)
}
return gateway_bedrock.NewCredentialStore(js, masterKey, len(cfg.Nodes), nil), func() { nc.Close() }, nil
return gateway_bedrock.NewCredentialStore(js, masterKey, nil), func() { nc.Close() }, nil
}
func runOchreCredentialsSet(cmd *cobra.Command, _ []string) {
@@ -1186,7 +1186,7 @@ func init() {
// ochreAccessStore connects to the cluster and returns the grant store along
// with a cleanup that closes the connection.
func ochreAccessStore() (*gateway_bedrock.ModelAccessStore, func(), error) {
cfg, nc, err := loadConfigAndConnectFn()
_, nc, err := loadConfigAndConnectFn()
if err != nil {
return nil, nil, fmt.Errorf("connect to cluster: %w", err)
}
@@ -1195,7 +1195,7 @@ func ochreAccessStore() (*gateway_bedrock.ModelAccessStore, func(), error) {
nc.Close()
return nil, nil, fmt.Errorf("jetstream context: %w", err)
}
return gateway_bedrock.NewModelAccessStore(js, len(cfg.Nodes)), func() { nc.Close() }, nil
return gateway_bedrock.NewModelAccessStore(js), func() { nc.Close() }, nil
}
// ochreTargetModels resolves the --model-id / --all-models pair into the model
+2 -2
View File
@@ -470,7 +470,7 @@ func TestResolveHFToken_FallsBackToStoredPlatformCredential(t *testing.T) {
require.NoError(t, os.MkdirAll(filepath.Join(baseDir, "config"), 0750))
require.NoError(t, handlers_iam.SaveMasterKey(filepath.Join(baseDir, "config", "master.key"), masterKey))
credStore := gateway_bedrock.NewCredentialStore(js, masterKey, 1, nil)
credStore := gateway_bedrock.NewCredentialStore(js, masterKey, nil)
require.NoError(t, credStore.PutCredential(context.Background(), utils.GlobalAccountID, vendorHuggingFace, "hf_stored_token"))
cmd := newOchreWeightsPullTestCmd(t, selfHostModelID, testHFRepo, "main", "", "")
@@ -621,7 +621,7 @@ func TestRunOchreCredentialsSet_RoundTripsAndDefaultsToPlatformAccount(t *testin
defer verifyConn.Close()
verifyJS := testutil.NewJetStream(t, verifyConn)
credStore := gateway_bedrock.NewCredentialStore(verifyJS, masterKey, 1, nil)
credStore := gateway_bedrock.NewCredentialStore(verifyJS, masterKey, nil)
token, ok, err := credStore.Resolve(context.Background(), utils.GlobalAccountID, vendorHuggingFace)
require.NoError(t, err)
require.True(t, ok)
+4 -4
View File
@@ -42,7 +42,7 @@ const (
func newWeightsStoreForTest(t *testing.T) *gateway_bedrock.WeightsStore {
t.Helper()
_, nc, _ := testutil.StartTestJetStream(t)
return gateway_bedrock.NewWeightsStore(testutil.NewJetStream(t, nc), 1)
return gateway_bedrock.NewWeightsStore(testutil.NewJetStream(t, nc))
}
// explodingObjectStore fails the test immediately if any of its methods is
@@ -775,7 +775,7 @@ func TestRunOchreWeightsList_PrintsStagedModels(t *testing.T) {
}
stubConnect(t, cfg, nc, nil)
seedStore := gateway_bedrock.NewWeightsStore(js, 1)
seedStore := gateway_bedrock.NewWeightsStore(js)
require.NoError(t, seedStore.PutWeights(context.Background(), selfHostModelID, "s3://models/llama-3.2-1b/", "snap-0001"))
out := captureStdout(t, func() { runOchreWeightsList(nil, nil) })
@@ -822,7 +822,7 @@ func TestRunOchreWeightsRemove_RoundTripPrintsSnapshotAndDropsEntry(t *testing.T
ns, nc, js := testutil.StartTestJetStream(t)
stubConnect(t, fakeClusterConfig(), nc, nil)
seedStore := gateway_bedrock.NewWeightsStore(js, 1)
seedStore := gateway_bedrock.NewWeightsStore(js)
require.NoError(t, seedStore.PutWeights(context.Background(), selfHostModelID, "s3://models/llama-3.2-1b/", "snap-0001"))
cmd := newOchreWeightsRemoveTestCmd(t, selfHostModelID)
@@ -835,7 +835,7 @@ func TestRunOchreWeightsRemove_RoundTripPrintsSnapshotAndDropsEntry(t *testing.T
verifyNC, err := nats.Connect(ns.ClientURL())
require.NoError(t, err)
t.Cleanup(func() { verifyNC.Close() })
verifyStore := gateway_bedrock.NewWeightsStore(testutil.NewJetStream(t, verifyNC), 1)
verifyStore := gateway_bedrock.NewWeightsStore(testutil.NewJetStream(t, verifyNC))
_, ok, err := verifyStore.GetWeights(context.Background(), selfHostModelID)
require.NoError(t, err)
+1 -1
View File
@@ -167,7 +167,7 @@ func runClusterShutdown(cmd *cobra.Command, args []string) {
fmt.Printf("Timeout per phase: %s\n\n", timeout)
// Write cluster shutdown marker to KV
jsm, err := daemon.NewJetStreamManager(nc, 1)
jsm, err := daemon.NewJetStreamManager(nc)
if err != nil {
fmt.Fprintf(os.Stderr, "Warning: failed to create JetStream manager: %v\n", err)
} else {
+44
View File
@@ -0,0 +1,44 @@
//test:in-package — the embedded nats.conf template it checks is unexported,
// and exporting it would widen the package's API for a test's convenience.
package cmd
import (
"strings"
"testing"
natssvc "github.com/mulgadc/spinifex/spinifex/services/nats"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// TestNATSConfServerNameMatchesTheHelper ties the rendered server_name to
// nats.ServerName. JetStream reports stream peers by server name, so anything
// mapping a replica back to the node holding it reads the prefix from that
// helper — and a template that drifted from it would break the mapping silently,
// reporting a healthy cluster as one whose buckets name no node at all.
func TestNATSConfServerNameMatchesTheHelper(t *testing.T) {
var line string
for l := range strings.SplitSeq(natsConfTemplate, "\n") {
if strings.HasPrefix(strings.TrimSpace(l), "server_name:") {
line = strings.TrimSpace(l)
break
}
}
require.NotEmpty(t, line, "nats.conf no longer sets server_name")
got := strings.TrimSpace(strings.TrimPrefix(line, "server_name:"))
assert.Equal(t, natssvc.ServerName("{{.Node}}"), got)
node, ok := natssvc.NodeFromServerName(natssvc.ServerName("node2"))
assert.True(t, ok)
assert.Equal(t, "node2", node)
_, ok = natssvc.NodeFromServerName("node2")
assert.False(t, ok, "a bare node name is not a server name")
// The bare prefix is a name we could not have generated, and mapping it to the
// empty node would put a phantom peer into anything counting nodes.
_, ok = natssvc.NodeFromServerName(natssvc.ServerNamePrefix)
assert.False(t, ok, "the prefix with no node after it is not a server name")
}
-1
View File
@@ -384,7 +384,6 @@ var viperblockStartCmd = &cobra.Command{
SecretKey: nodeConfig.Predastore.SecretKey,
BaseDir: nodeConfig.Predastore.BaseDir,
NodeName: clusterConfig.Node,
KVReplicas: len(clusterConfig.Nodes),
ShardWAL: shardWAL,
GCEnabled: gcEnabled,
WALBaseDir: nodeConfig.Viperblock.WALBaseDir,
+130 -1
View File
@@ -30,6 +30,7 @@ resources:
- [Overview](#overview)
- [Cluster sizing](#cluster-sizing)
- [How the cluster's own state is replicated](#how-the-clusters-own-state-is-replicated)
- [Hardware](#hardware)
- [Network requirements](#network-requirements)
- [Prerequisites](#prerequisites)
@@ -65,11 +66,82 @@ A Spinifex cluster distributes services across multiple servers for high availab
| **VPC networking** — OVN | Control-plane databases run clustered, surviving the loss of any one node. |
| **Object storage** — Predastore (S3) | Objects are erasure coded `RS(2,1)`, surviving the loss of any one node's shards. |
| **Block storage** — Viperblock (EBS) | Volumes are stored in Predastore, so they inherit the same durability. |
| **Cluster state** — NATS JetStream | Every key-value bucket is replicated onto all three nodes, so the control plane keeps its own state through the loss of any one. |
On one or two servers none of that holds. OVN runs standalone on the first node, and the storage metadata quorum has no majority to lose. If that node goes down, running instances keep full networking — but nothing can *change*: no new VPCs, no launches, no security group edits. See [OVN control plane on multi-node clusters](/docs/vpc-networking#ovn-control-plane-on-multi-node-clusters).
Servers beyond the third run the full set of services — storage, gateway and networking agents — and add their capacity to the pool. What they do not do is join the OVN database cluster, which stays at three members, so write latency there stays flat as the cluster grows.
### How the cluster's own state is replicated
Guest data has an obvious home: volumes live in Viperblock, objects in Predastore, both erasure coded across the cluster. The cluster's own bookkeeping has a less obvious one, and it matters just as much. Every instance record, every VPC and subnet, the server roster, IAM users and access keys, quotas and DNS records are kept in **NATS JetStream key-value buckets** — Spinifex's control-plane database, replicated across the servers you installed rather than sitting on any one of them. A cluster that has lost its state has lost the ability to describe, launch, terminate or route anything, even while every guest is still running.
**A bucket is replicated across the largest odd number of nodes the cluster has, up to five.** That number is decided once, by the `--nodes` count you give `spx admin init` in Step 4, and every service on every node reads it back out of `/etc/spinifex/spinifex.toml` before it creates anything.
| Servers | Replicas per bucket | Quorum | Survives |
|---|---|---|---|
| 1 | 1 | 1 | nothing — a single node is a single point of failure by definition |
| 2 | — | — | rejected by `spx admin init`; see below |
| 3 | 3 | 2 | any 1 node lost |
| 4 | 3 | 2 | any 1 node lost |
| 5 | 5 | 3 | any 2 nodes lost |
| 10 | 5 | 3 | any 2 of the 5 nodes holding each bucket |
**The count is always odd, and that is why four servers replicate three ways rather than four.** Agreement is by majority, and an even number of copies does not buy you one: four copies cost a fourth write every time and still survive exactly one loss, the same as three. On four servers NATS spreads which three nodes hold each bucket, so all four carry their share and none of them is idle.
Five is JetStream's own ceiling on a stream's replica count, not a choice of ours. Past five servers the buckets stay at five and NATS spreads which five nodes hold each one, so the cluster keeps growing while the cost of a write does not.
**Two servers is refused rather than allowed with a warning.** There is no replica count that works on two: one disappears with the node holding it, and two stop accepting writes when either member goes, because neither survivor is a majority on its own. OVN and the storage metadata quorum have the same problem on two servers, so `spx admin init --nodes 2` fails and tells you to use one server or three.
#### Worked example — three servers
Say you ran Step 4 with `--nodes 3` on `node1`, `node2` and `node3`. Every bucket the cluster creates is held by all three servers, with one of them responsible for coordinating changes to it:
```
KV bucket "spinifex-instance-state" KV bucket "spinifex-vpcd-reconcile"
┌──────────┬──────────┬──────────┐ ┌──────────┬──────────┬──────────┐
│ node1 │ node2 │ node3 │ │ node1 │ node2 │ node3 │
│ leader │ follower │ follower │ │ follower │ leader │ follower │
└──────────┴──────────┴──────────┘ └──────────┴──────────┴──────────┘
quorum = 2 of 3 quorum = 2 of 3
```
Three things follow from that picture, and they are the whole reason three servers is the floor:
- **A change is not confirmed until most of the cluster has it.** When you launch an instance, the API does not report success until at least two of the three servers have written the record to disk. Losing a server afterwards cannot lose that instance.
- **Losing a server costs a pause, not the data.** If `node1` goes down, `node2` and `node3` are still a majority. They take over the buckets `node1` was responsible for — seconds on a healthy cluster — and the API keeps answering throughout. Nothing becomes unreadable, because nothing lived only on `node1`.
- **The cluster keeps repairing itself.** Spinifex runs background workers that fix up VPC state, enforce quotas and publish DNS. Only one server runs each at a time, and they claim that right by writing to a bucket. If those buckets were not replicated, a server failure would stop the repair work at precisely the moment there was repair work to do.
Contrast that with a bucket on **one** server. It lives wherever it happened to be created, which is not recorded in any config file you can inspect. Stopping, rebooting or losing that one server makes the bucket unreachable from *every* node at once — a single-server event becomes a cluster-wide outage, on a server nobody chose.
#### What you see when a server fails
Majority is not an implementation detail. It is what stops two halves of a damaged cluster from both believing they are in charge, and it is the reason the minimum is three servers rather than two.
**With one server down, the cluster carries on.** `node2` and `node3` are two of three. Instances launch, VPCs change, the API answers normally, and there is nothing to do but replace the failed server. What you have lost is the *next* failure: a second server down leaves no majority and the cluster stops accepting changes until one comes back.
**A server that is cut off stops accepting changes, on purpose.** If `node1` is not down but isolated — a failed switch, a misconfigured port — it is one server out of three, which is not a majority, so it will not record anything. An operator working on `node1` who tries to launch an instance gets a failure rather than a success. That is the behaviour you want: if `node1` accepted changes while `node2` and `node3` were also accepting them, you would end up with two different versions of the same cluster and no way to merge them. Refusing the isolated side is what makes the surviving side worth trusting.
Everything else on `node1` keeps running. Guests continue to execute and their volumes keep serving. What stops is anything that has to be recorded: launching, terminating, attaching a volume, changing a security group. Storage writes fail alongside it, so a guest on the isolated server stalls rather than quietly diverging from the rest of the cluster.
**Be careful reading state from an isolated server.** A server that is cut off can still answer `describe-instances` from its own copy of the data — frozen at the moment it lost contact, with nothing in the answer to say so. **Run operator commands from a server you know is with the majority**, and if one server's answers disagree with its peers', believe the peers.
**A returning server catches up on its own.** Bring `node1` back and it synchronises from the others automatically: no command to run, no restart, no repair step. (`spx admin kv replicas --repair` exists for a different situation — a bucket configured with too few copies — not a server that is simply behind.)
#### Growing an existing cluster
Adding a fourth server does not, on its own, re-replicate the buckets that already exist — the nodes that hold them are perfectly healthy, so nothing prompts a change. Spinifex raises them for you at the next service start: each service checks every bucket it opens against the cluster's current node count and raises any that are short. It never lowers one.
**A config that names a node before that node is serving is safe.** It is the ordinary middle of growing a cluster, and services start normally through it: a bucket that cannot be raised yet keeps working at the count it has, and the node logs `Could not raise KV bucket to the cluster's replica count`. The raise is retried every time a service opens the bucket, so it completes on its own once the new server is in the cluster.
To do it immediately rather than waiting for a restart, run the audit with `--repair` from any node:
```bash
sudo spx admin kv replicas --repair
```
It changes nothing about a bucket except the replica count, never lowers one, and is safe to run twice or on a healthy cluster.
### Hardware
Per server, for the three-node minimum:
@@ -252,6 +324,8 @@ sudo spx admin init --force \
`--nodes 3` is the number of servers init waits for. Set it to your total node count if you are building a larger cluster.
It is also what decides how many nodes the cluster's JetStream state is replicated across, so it is worth getting right rather than raising later — see [How the cluster's own state is replicated](#how-the-clusters-own-state-is-replicated). Every joining server takes the count from server 1, so it is chosen once, here. `--nodes 2` is refused: use 1 for a single server, or 3 and up for a cluster that can survive losing one.
IPsec encrypts the Geneve overlay between servers and is on by default. Joining servers take the setting from server 1, so it is chosen once, on init. On servers that share a trusted private link, `--ipsec=false` leaves the overlay unencrypted in exchange for considerably higher throughput between instances.
The init output displays the join command including the token:
@@ -360,7 +434,40 @@ The top table is per-node CPU, memory and GPU usage. The bottom table is what th
If capacity looks like a single server rather than the sum of your nodes, the others have not joined.
**3. The AWS API answers.**
**3. The cluster's own state is replicated across every node.**
```bash
sudo spx admin kv replicas
```
```
spinifex@node1:~$ sudo spx admin kv replicas
BUCKET REPLICAS WANT ONLINE STATUS LEADER HELD BY
spinifex-cluster-state 3 3 3 ok spinifex-nats-node1 spinifex-nats-node1,spinifex-nats-node2,spinifex-nats-node3
spinifex-dns-reconcile 3 3 3 ok spinifex-nats-node2 spinifex-nats-node2,spinifex-nats-node1,spinifex-nats-node3
spinifex-iam 3 3 3 ok spinifex-nats-node3 spinifex-nats-node3,spinifex-nats-node1,spinifex-nats-node2
spinifex-instance-state 3 3 3 ok spinifex-nats-node1 spinifex-nats-node1,spinifex-nats-node2,spinifex-nats-node3
spinifex-quota-reconcile 3 3 3 ok spinifex-nats-node2 spinifex-nats-node2,spinifex-nats-node3,spinifex-nats-node1
spinifex-terminated-instances 3 3 3 ok spinifex-nats-node3 spinifex-nats-node3,spinifex-nats-node2,spinifex-nats-node1
spinifex-vpcd-reconcile 3 3 3 ok spinifex-nats-node2 spinifex-nats-node2,spinifex-nats-node1,spinifex-nats-node3
47 buckets, 0 under-replicated, 0 without quorum
```
Server names are the node name behind a `spinifex-nats-` prefix, which is what NATS itself calls each server.
What to check:
- **`0 under-replicated, 0 without quorum`.** These are different problems. Under-replicated means a bucket lives on fewer nodes than depend on it — see [Buckets report as under-replicated](#buckets-report-as-under-replicated), and `--repair` is the fix. Without quorum means a bucket is not able to accept a write right now, whatever its replica count says, and no repair helps until you find out why a node is not answering.
- **`WANT` is the largest odd number of nodes you have**, capped at five — 3 on three or four servers, 5 on five or more. If it says 1 on a three-server cluster, the node the command ran on is not configured for three: check `--nodes` in Step 4 and the `[nodes.*]` sections of that host's `/etc/spinifex/spinifex.toml`.
- **`ONLINE` equals `REPLICAS`.** It counts the replicas that are caught up and reachable, so a bucket at `3` replicas with `ONLINE 2` is one node away from being unavailable and needs looking at before anything else is done to the cluster.
- **`LEADER` is spread across the cluster** rather than parked on one node, and `HELD BY` names as many distinct servers as `REPLICAS`.
Exit status is what a deployment gate should read rather than the output: `0` healthy, `1` something is under-replicated, `2` something has no quorum or the report could not be produced at all.
Run it on every node, not just one. It reads cluster-global Raft metadata, so all of them should give the same answer; one that disagrees is not fully part of the cluster.
**4. The AWS API answers.**
```bash
export AWS_PROFILE=spinifex
@@ -403,6 +510,28 @@ The node has its own cluster configuration — normal for anything installed fro
Safe on a freshly installed node. On one that has been in service it orphans every volume and fragment sealed under the old key, so check before forcing.
### Buckets Report as Under-Replicated
```bash
sudo spx admin kv replicas
```
A bucket showing `UNDER-REPLICATED` lives on fewer nodes than the cluster has, so losing the node holding it takes it down for the whole cluster. Two things put a bucket in that state, and the fix is the same for both:
```bash
sudo spx admin kv replicas --repair
```
The first cause is **a cluster that grew**: buckets created when there were three servers stay on three replicas until something raises them. Services do this themselves as they restart, and `--repair` does it now.
The second is **a node that was down when the bucket was created**, so the Raft group formed short. Bring the node back first — `--repair` on an incomplete cluster raises the configured count, but the group cannot place the extra replica until the node is there to hold it. Confirm with `spx get nodes` before repairing.
If `WANT` itself is wrong — 1 on a three-server cluster — the problem is not the buckets but that host's view of the cluster. Check it has all three `[nodes.*]` sections:
```bash
sudo grep -oP '^\[nodes\.\K[^.\]]+' /etc/spinifex/spinifex.toml | sort -u
```
### OVN Database Cluster Not Forming
```bash
+21 -2
View File
@@ -7,11 +7,13 @@ Suites live in `tests/e2e/<name>/` and build to `tests/e2e/_bin/<name>.test`. Ev
| Set | Suites |
| --- | --- |
| `E2E_SUITES_SINGLE` | `single iam cert eks ecs storagegrowth partialblock rds quota storagefault` |
| `E2E_SUITES_MULTI` | `multinode lb cert quota lbrecovery instancerecovery storagefault` |
| `E2E_SUITES_MULTI` | `multinode lb cert quota kvquorum lbrecovery instancerecovery storagefault` |
| `E2E_SUITES_NIGHTLY_SINGLE` | `single cert iam` |
| `E2E_SUITES_NIGHTLY_MULTI` | `multinode cert lb` |
The nightly permutation sets are deliberately narrower than the full ones: those cells have a ~35 minute budget and exist to prove every install / network / host-OS combination boots and serves. `eks`, `ecs`, `rds`, `storagefault`, `instancerecovery` and `lbrecovery` each get a dedicated cell instead.
The nightly permutation sets are deliberately narrower than the full ones: those cells have a ~35 minute budget and exist to prove every install / network / host-OS combination boots and serves. `eks`, `ecs`, `rds`, `storagefault`, `instancerecovery`, `lbrecovery` and `kvquorum` each get a dedicated cell instead.
**`multinode` carries the KV replication sweep as well**, which is why the thing `kvquorum` proves is still checked in every multi-node permutation cell even though the suite itself is not. The read-only half costs a NATS connection and a stream listing; only the node-loss half needs a cell of its own.
## The suites
@@ -27,6 +29,7 @@ The largest suite, 25 top-level tests. Notable entry points:
- `TestInstanceMetadata` — IMDS from inside a guest.
- `TestStopStart`, `TestEarlyRebootLiveness`, `TestGuestChurnDurability`, `TestSnapshotLifecycle`, `TestSnapshotBackedLaunch`, `TestCreateImage`, `TestVolumeLifecycle`, `TestSpotInstanceLifecycle`, `TestLaunchTemplateBoot`.
- `TestNegativeErrorPaths` — the `8a`–`8h` AWS error-shape cases (malformed AMI ID, invalid instance type, volume in use, detach-root forbidden, and so on).
- **`TestKVReplication`** — the single-node half of the KV replication policy, and it is not the multi-node assertion with a smaller number in it. On one node one replica per bucket is the correct answer rather than the defect, so what it proves is the opposite: that nothing declared a replica count this deployment cannot place, that `spx admin kv replicas` exits 0 and reports `WANT 1`, and that `--repair` run twice against a cluster with nothing to repair changes no bucket's replica count and leaves every bucket serving. The multi-node suites cannot cover any of that, because on three nodes the unplaceable count is the right one. It lives here rather than in its own suite so it runs in every single-node nightly cell, across all three host and network permutations.
### `natuplink` — routed NAT (`external_mode = "nat"`), one node or three
@@ -42,6 +45,22 @@ Two nightly cells run it: `nat-single` (cell 19) and `nat-multi` (cell 30). Cell
`VPCSetup`, `SpansMultipleNodes`, `SpreadPlacement`, `EveryRunningInstanceReported`, `BastionSSH`, and its own NAT Gateway lane: `PreNATIsolation`, `NATGatewayInternet`, `NATCleanupOrdering`.
**`TestMultinodeJetStreamReplicas` is a whole-cluster sweep, not a named list, and it lives here on purpose.** It audits every KV bucket the cluster has and fails any that is replicated across fewer nodes than the cluster has, capped at JetStream's ceiling of five. Naming buckets would only ever cover the ones somebody remembered — a replica count is set where a bucket is created, and there is a bucket-creating call in most services — so the sweep is the assertion and the three daemon-owned buckets are named only so that a run enumerating nothing still fails. Because `multinode` is in `E2E_SUITES_NIGHTLY_MULTI`, this runs in every multi-node permutation cell, and no cell can be added that quietly opts out of it.
### `kvquorum` — the cluster's own state survives losing a node (multi-node, needs three)
**The control plane keeps its state in JetStream KV buckets, and this is the suite that proves those buckets are actually replicated.** A bucket on one replica lives on one JetStream-chosen server recorded in no config, so losing that node makes it return `nats: no responders available for request` from *every* node at once — a cluster-wide outage caused by a single-node event. Leader leases are the worst case: acquiring one is a write, so a lease bucket without a quorum cannot be acquired by anyone, and every reconciler sharing it stops cluster-wide at exactly the moment a node has failed and there is repair work to do.
Three entry points, in the order they run:
- **`TestKVBucketsAreReplicatedAcrossTheCluster`** is read-only and runs first, so a cluster that is already wrong is reported as wrong rather than as a failure to survive a node loss. It sweeps every bucket, then separately requires the leader-lease buckets to be among those checked — they were the worst case of the defect, and a regression reaching only them would otherwise be one line in a long list rather than the headline.
- **`TestKVReplicasCommandReportsHealthy`** runs `spx admin kv replicas --json` on **every** node, not one. The command reads cluster-global Raft metadata, so all nodes must give the same answer; a node that disagrees is reporting on a cluster it is not fully part of, which is worth knowing before a deployment is gated on its exit status.
- **`TestKVBucketsSurviveANodeLoss`** produces the fault rather than inspecting for it. It takes away **the node leading the most buckets** — the worst one to lose, and the one a test picking arbitrarily would usually miss — then requires every survivor to still answer for every bucket, brings the node back, and requires it to carry its share again.
**Configured, placed and working replicas are three different facts, and all three are checked.** A stream can carry the right replica count while its Raft group is short of peers, or hold the right peers while one of them is offline or still catching up, and either reads as healthy right up until the node it is really on goes away. The audit reports the leader and a count of peers NATS itself calls current and reachable, and the assertion is against that count rather than against the configured one. The last assertion is that no bucket came back with *fewer* replicas than it had: a cluster that "repaired" itself by lowering a count would pass every other assertion here and have quietly traded the durability this suite is about for a green run.
**Its own suite and its own cell, for lbrecovery's reasons.** It takes a node away, so anything sharing the environment reports that outage as its own failure, and the question it answers — "is our internal state actually replicated" — has to read as itself rather than as one more red subtest somewhere else. Three nodes is the floor: on two, a majority is two, so losing either leaves no quorum for any stream and there is nothing to assert beyond "Raft needs a majority", which is true of NATS and not a property of ours.
### `instancerecovery` — a guest whose node goes away (multi-node)
**Asserted from both sides, and the pair is the point.** `InstanceAutoRecovery` takes a node down and requires the guest to come back on exactly one survivor, identified by the qemu process rather than by what the API says, then brings the node back and requires it to drop the local copy instead of relaunching it. `InstanceRecoveryRefusesStorageFault` freezes predastore cluster-wide first, so the guest is paused by `werror=stop` when its node goes away, and requires that **nothing moves** — a store that refuses every node cannot be fixed by moving a guest, and the relaunch would cost it the request QEMU is holding. A reconciler that cannot tell a host failure from a storage failure passes the first and fails the second. Neither turns recovery on, because there is no setting for it: both wait for every node to report the reconciler active on a cluster nothing configured, so a build that reintroduced a switch fails here.
+2 -2
View File
@@ -20,7 +20,7 @@ func TestNewClusterEngineCoversEveryStage(t *testing.T) {
masterKey, err := handlers_iam.GenerateMasterKey()
require.NoError(t, err)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey, 1)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey)
require.NoError(t, err)
engine, err := NewClusterEngine(t.Context(), nc, 1, svc, newFakeBuckets())
@@ -48,7 +48,7 @@ func TestClusterEnginePrechecks(t *testing.T) {
masterKey, err := handlers_iam.GenerateMasterKey()
require.NoError(t, err)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey, 1)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey)
require.NoError(t, err)
engine, err := NewClusterEngine(t.Context(), nc, 1, svc, newFakeBuckets())
+1 -1
View File
@@ -398,7 +398,7 @@ func TestPolicyReaperDeletesAPolicyWithSeveralVersions(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
masterKey, err := handlers_iam.GenerateMasterKey()
require.NoError(t, err)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey, 1)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey)
require.NoError(t, err)
const accountID = "000000000042"
+157
View File
@@ -0,0 +1,157 @@
// Package clustersize holds one fact and the one number derived from it: how
// many nodes this process believes the cluster has, and therefore how many
// nodes each of its JetStream streams is replicated across.
//
// It is its own package because the fact is declared in one place,
// config.LoadConfig, and read in another, kvutil, and those sit either side of
// packages that would otherwise import each other.
//
// The number matters because a stream on one replica lives on one
// JetStream-chosen server recorded in no config. Losing that node makes the
// stream unreadable and unwritable from every node at once, so a single-node
// event becomes a cluster-wide outage.
package clustersize
import (
"errors"
"fmt"
"sync/atomic"
"testing"
)
// MaxReplicas is JetStream's own ceiling on stream replicas, so a cluster
// larger than this replicates to five nodes and no further. Asking for more is
// refused by the server, which is why the count is clamped rather than passed
// through.
const MaxReplicas = 5
// MinHANodes is the smallest cluster in which losing a node is survivable. Below
// it no layer of the stack has a majority to keep: OVN runs standalone, the
// storage metadata quorum has nothing to lose, and a KV bucket has no replica
// count that both survives a loss and accepts writes.
const MinHANodes = 3
// ErrUndeclared is returned by Replicas in a process that never loaded a
// cluster config. It is an error rather than a fallback to one replica because
// one replica is the failure this package exists to prevent, and a fallback
// would reinstate it in exactly the processes nobody checked.
var ErrUndeclared = errors.New("cluster size was never declared, so the stream replica count is unknown")
// nodes is the declared node count. Zero means undeclared.
var nodes atomic.Int64
// Declare records the cluster's node count for this process. It is called from
// config.LoadConfig, so every process that can reach a cluster's NATS has
// already declared it before any handler creates a bucket.
//
// The first declaration wins. Repeating the same count is a no-op, because two
// call sites legitimately read it from the same config; a different count is an
// error rather than an overwrite, since a process that disagrees with itself
// about how many nodes exist would silently pick whichever ran last.
func Declare(n int) error {
n = max(n, 0)
if nodes.CompareAndSwap(0, int64(n)) {
return nil
}
if have := Nodes(); have != n {
return fmt.Errorf("cluster size was already declared as %d nodes and cannot be redeclared as %d", have, n)
}
return nil
}
// ResetForTest clears the declaration so a test can declare a different size.
// It panics outside a test binary: production has exactly one cluster size for
// the life of the process, and a reset available there would undo that.
func ResetForTest() {
if !testing.Testing() {
//nolint:forbidigo // Returning an error would let a caller ignore it and carry on with the count cleared.
panic("clustersize.ResetForTest called outside a test binary")
}
nodes.Store(0)
}
// DeclareForTest declares n for the rest of the test binary, tolerating a fixture
// that already declared the same count.
//
// It is for fixtures rather than for tests. Standing up a one-node embedded server
// says something about the binary, not about one test, and several tests running
// in parallel share that one server — so restoring the previous count when any one
// of them finishes would leave its siblings undeclared mid-run.
func DeclareForTest(tb testing.TB, n int) {
tb.Helper()
if err := Declare(n); err != nil {
tb.Fatalf("declare cluster size %d: %v", n, err)
}
}
// RedeclareForTest declares n for the duration of one test, whatever was declared
// before it, and puts the previous count back afterwards.
//
// Declaring is write-once, so a test needing a different size than its fixture set
// up cannot simply call Declare again. It must not run in parallel with anything
// else that reads the count, because there is one count per process and this
// changes it.
func RedeclareForTest(tb testing.TB, n int) {
tb.Helper()
was := Nodes()
ResetForTest()
if err := Declare(n); err != nil {
tb.Fatalf("declare cluster size %d: %v", n, err)
}
tb.Cleanup(func() {
ResetForTest()
if was > 0 {
_ = Declare(was)
}
})
}
// Nodes returns the declared node count, or 0 when it was never declared.
func Nodes() int {
return int(nodes.Load())
}
// Replicas is the replica count every stream is created with: the largest odd
// number of nodes this cluster can place a replica on, capped at JetStream's
// ceiling of five.
//
// Odd, because a majority is what a raft group needs and an even count buys
// none. Four replicas store a fourth full copy of every write and still survive
// only one loss, exactly like three; two survive nothing, because losing either
// member leaves no majority to accept a write. So four nodes replicate three
// ways and five replicate five ways, which is also the most any cluster does
// however large it grows.
//
// There is no argument and no override. A replica count is a property of the
// cluster, not of the stream or of the caller, and every stream here is read or
// written by a node other than the one that created it.
func Replicas() (int, error) {
n := Nodes()
if n < 1 {
return 0, ErrUndeclared
}
return ReplicasFor(n), nil
}
// ReplicasFor is Replicas for a node count the caller already holds, so code that
// knows the cluster's size without having declared it — a test harness reading a
// cluster it did not build — gets its answer from the same rule rather than a
// second copy of it.
func ReplicasFor(nodes int) int {
r := min(max(nodes, 1), MaxReplicas)
if r%2 == 0 {
r--
}
return r
}
// Permanent reports whether err is this package's own refusal rather than a
// transient cluster condition.
//
// The retry loops that wait for JetStream quorum on a cold start must stop on
// one: no amount of waiting declares a cluster size, so a loop that cannot tell
// the two apart spends its whole budget on a misconfiguration and then reports
// a quorum problem that was never the cause.
func Permanent(err error) bool {
return errors.Is(err, ErrUndeclared)
}
+143
View File
@@ -0,0 +1,143 @@
package clustersize_test
import (
"errors"
"strings"
"testing"
"github.com/mulgadc/spinifex/spinifex/clustersize"
)
func TestDeclareAndNodes(t *testing.T) {
t.Cleanup(clustersize.ResetForTest)
clustersize.ResetForTest()
if got := clustersize.Nodes(); got != 0 {
t.Fatalf("undeclared cluster size = %d, want 0", got)
}
if err := clustersize.Declare(3); err != nil {
t.Fatalf("Declare(3): %v", err)
}
if got := clustersize.Nodes(); got != 3 {
t.Fatalf("Declare(3) = %d, want 3", got)
}
}
// TestDeclareIsWriteOnce covers the three ways a second declaration can arrive.
// Repeating the same count is how two call sites reading the same config behave
// and must be accepted; a different count is a process disagreeing with itself
// about how many nodes exist, which must not be resolved by whichever ran last.
func TestDeclareIsWriteOnce(t *testing.T) {
t.Cleanup(clustersize.ResetForTest)
clustersize.ResetForTest()
if err := clustersize.Declare(3); err != nil {
t.Fatalf("first Declare(3): %v", err)
}
if err := clustersize.Declare(3); err != nil {
t.Fatalf("repeating Declare(3) must be accepted, got %v", err)
}
err := clustersize.Declare(5)
if err == nil {
t.Fatal("Declare(5) after Declare(3) returned no error")
}
if got := clustersize.Nodes(); got != 3 {
t.Fatalf("a rejected redeclaration changed the count to %d, want 3 kept", got)
}
// The message has to name both counts, because the failure it reports is a
// disagreement and neither number alone identifies it.
for _, want := range []string{"3", "5"} {
if !strings.Contains(err.Error(), want) {
t.Errorf("error %q does not name the count %s", err, want)
}
}
}
// TestDeclareRejectsANegativeCount: a negative node count is not a smaller
// cluster, it is a broken config, and storing it as-is would read back as a
// count rather than as undeclared.
func TestDeclareRejectsANegativeCount(t *testing.T) {
t.Cleanup(clustersize.ResetForTest)
clustersize.ResetForTest()
if err := clustersize.Declare(-2); err != nil {
t.Fatalf("Declare(-2): %v", err)
}
if got := clustersize.Nodes(); got != 0 {
t.Fatalf("Declare(-2) = %d, want 0 (undeclared)", got)
}
if _, err := clustersize.Replicas(); !errors.Is(err, clustersize.ErrUndeclared) {
t.Fatalf("Replicas() after Declare(-2) = %v, want ErrUndeclared", err)
}
}
// TestReplicas covers the policy and the fail-closed answer together: a bucket is
// replicated across the largest odd number of nodes the cluster can place, capped
// at five, and a process that never declared a size gets an error rather than a
// guess at one.
//
// The even counts are the point of the table. Four replicas cost a fourth full
// copy of every write and survive exactly one loss, the same as three, and two
// survive nothing at all because losing either member leaves no majority.
func TestReplicas(t *testing.T) {
t.Cleanup(clustersize.ResetForTest)
for _, tc := range []struct{ nodes, want int }{
{1, 1}, {2, 1}, {3, 3}, {4, 3}, {5, 5}, {6, 5}, {10, 5}, {64, 5},
} {
clustersize.ResetForTest()
if err := clustersize.Declare(tc.nodes); err != nil {
t.Fatalf("Declare(%d): %v", tc.nodes, err)
}
got, err := clustersize.Replicas()
if err != nil {
t.Fatalf("Replicas() with %d nodes: %v", tc.nodes, err)
}
if got != tc.want {
t.Errorf("Replicas() with %d nodes = %d, want %d", tc.nodes, got, tc.want)
}
if got%2 == 0 {
t.Errorf("Replicas() with %d nodes = %d, which is even", tc.nodes, got)
}
if fromCount := clustersize.ReplicasFor(tc.nodes); fromCount != got {
t.Errorf("ReplicasFor(%d) = %d but Replicas() = %d; the two must be one rule",
tc.nodes, fromCount, got)
}
}
clustersize.ResetForTest()
if _, err := clustersize.Replicas(); !errors.Is(err, clustersize.ErrUndeclared) {
t.Fatalf("Replicas() with an undeclared cluster size = %v, want ErrUndeclared", err)
}
}
// TestReplicasForNeverReturnsZero: callers that hold a node count use it without
// an error return, so a count below one has to resolve to the single replica a
// lone server can hold rather than to none.
func TestReplicasForNeverReturnsZero(t *testing.T) {
for _, nodes := range []int{-1, 0} {
if got := clustersize.ReplicasFor(nodes); got != 1 {
t.Errorf("ReplicasFor(%d) = %d, want 1", nodes, got)
}
}
}
// TestPermanent separates this package's own refusal from a cluster that has yet
// to form. The retry loops waiting for JetStream quorum stop on the first and
// keep waiting through the second, so confusing them spends a whole startup
// budget on a misconfiguration.
func TestPermanent(t *testing.T) {
if !clustersize.Permanent(clustersize.ErrUndeclared) {
t.Error("Permanent(ErrUndeclared) = false")
}
wrapped := errors.Join(errors.New("create KV bucket x"), clustersize.ErrUndeclared)
if !clustersize.Permanent(wrapped) {
t.Error("Permanent() did not see through a wrapped ErrUndeclared")
}
if clustersize.Permanent(errors.New("no suitable peers for placement")) {
t.Error("Permanent() treated a placement failure as permanent")
}
if clustersize.Permanent(nil) {
t.Error("Permanent(nil) = true")
}
}
+19
View File
@@ -8,6 +8,7 @@ import (
"slices"
"strings"
"github.com/mulgadc/spinifex/spinifex/clustersize"
"github.com/spf13/viper"
)
@@ -605,6 +606,24 @@ func LoadConfig(configPath string) (*ClusterConfig, error) {
return nil, err
}
// Every process that creates a KV bucket loads the cluster config first, so
// this is the one place the node count can be declared without a service
// being able to forget to. Buckets refuse to be created until it has been.
//
// A conflicting second load is reported rather than refused, because this
// function also reads config files that are not this process's own cluster —
// an installer inspecting a peer's, a tool rendering one. The first
// declaration stands either way, so nothing that already created a bucket can
// have the count changed underneath it.
if err := clustersize.Declare(len(config.Nodes)); err != nil {
slog.Warn("Ignoring a cluster size that disagrees with the one already declared; the first one stands",
"loaded_nodes", len(config.Nodes), "declared_nodes", clustersize.Nodes(), "error", err)
}
if len(config.Nodes) > 1 && len(config.Nodes) < clustersize.MinHANodes {
slog.Warn("This cluster has too few nodes to survive losing one, so its state is replicated to a single node",
"nodes", len(config.Nodes), "ha_minimum", clustersize.MinHANodes)
}
return &config, nil
}
+1 -5
View File
@@ -22,11 +22,7 @@ func (d *Daemon) buildBedrockLaunchDeps() handlers_bedrock.LaunchDeps {
if err != nil {
slog.Warn("Bedrock: jetstream init failed; serving VMs cannot resolve staged weights", "err", err)
} else {
clusterSize := 1
if d.clusterConfig != nil {
clusterSize = len(d.clusterConfig.Nodes)
}
weights = gateway_bedrock.NewWeightsStore(js, clusterSize)
weights = gateway_bedrock.NewWeightsStore(js)
}
return handlers_bedrock.LaunchDeps{
+34 -14
View File
@@ -32,6 +32,7 @@ import (
"github.com/mulgadc/bluebottle/pkg/tlsconfig"
"github.com/mulgadc/spinifex/spinifex/admin"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/clustersize"
"github.com/mulgadc/spinifex/spinifex/config"
"github.com/mulgadc/spinifex/spinifex/ebsprovider"
"github.com/mulgadc/spinifex/spinifex/gpu"
@@ -1667,14 +1668,18 @@ func (d *Daemon) startCluster() error {
return fmt.Errorf("connect NATS: %w", err)
}
if err := d.initJetStream(); err != nil {
return fmt.Errorf("initialize JetStream: %w", err)
// Declared before JetStream, not after: initJetStream creates buckets, and
// bucket creation refuses until the count is known. A daemon assembled in a
// test may never have gone through config.LoadConfig, which is where a real
// process declares it from this same source.
if d.clusterConfig != nil {
if err := clustersize.Declare(len(d.clusterConfig.Nodes)); err != nil {
return err
}
}
// Set the default KV replica count before any handler creates a bucket, so
// lazily-created buckets are born at cluster-size replication instead of R1.
if d.clusterConfig != nil {
utils.SetDefaultKVReplicas(len(d.clusterConfig.Nodes))
if err := d.initJetStream(); err != nil {
return fmt.Errorf("initialize JetStream: %w", err)
}
// Remove the obsolete spinifex-dhcp-leases bucket (idempotent).
@@ -2397,7 +2402,7 @@ func (d *Daemon) initJetStream() error {
for {
attempt++
var err error
d.jsManager, err = NewJetStreamManager(d.natsConn, 1)
d.jsManager, err = NewJetStreamManager(d.natsConn)
if err == nil {
err = d.jsManager.InitKVBucket()
}
@@ -2412,10 +2417,20 @@ func (d *Daemon) initJetStream() error {
if err == nil {
d.jsManager.SetSyncObserver(d)
slog.Info("JetStream KV stores initialized successfully", "replicas", 1, "attempts", attempt, "elapsed_ms", otelsetup.Millis(time.Since(start)))
// Replicas cannot be undeclared here: bucket creation refuses
// before it, so reaching this line means it was declared.
replicas, _ := clustersize.Replicas()
slog.Info("JetStream KV stores initialized successfully", "replicas", replicas, "attempts", attempt, "elapsed_ms", otelsetup.Millis(time.Since(start)))
break
}
// A misdeclared cluster size is not something quorum arrives and fixes,
// so waiting for it spends the whole budget and then reports a quorum
// problem that was never the cause.
if clustersize.Permanent(err) {
return fmt.Errorf("initialize JetStream: %w", err)
}
elapsed := time.Since(start)
if elapsed >= maxWait {
return fmt.Errorf("failed to initialize JetStream after %s (%d attempts): %w", elapsed.Round(time.Second), attempt, err)
@@ -2431,15 +2446,20 @@ func (d *Daemon) initJetStream() error {
return nil
}
// upgradeJetStreamReplicas bumps KV_* stream replication to match the cluster
// size. Runs after all buckets are created and the cluster is ready.
// upgradeJetStreamReplicas raises every KV bucket to the cluster's replica
// count. Runs after all buckets are created and the cluster is ready, which is
// when a bucket created by whichever node got there first can be repaired.
func (d *Daemon) upgradeJetStreamReplicas() {
clusterSize := len(d.clusterConfig.Nodes)
if clusterSize <= 1 || d.jsManager == nil {
if d.jsManager == nil || d.jsManager.js == nil {
return
}
if err := d.jsManager.UpdateReplicas(clusterSize); err != nil {
slog.Warn("Failed to upgrade JetStream replicas", "targetReplicas", clusterSize, "error", err)
raised, err := kvutil.RaiseAllBucketReplicas(d.ctx, d.jsManager.js)
if err != nil {
slog.Warn("Failed to raise KV bucket replicas to the cluster's node count", "error", err)
return
}
if raised > 0 {
slog.Info("Raised KV buckets to the cluster's replica count", "buckets", raised)
}
}
+1 -1
View File
@@ -132,7 +132,7 @@ func createFullTestDaemonWithJetStream(t *testing.T, natsURL string) *Daemon {
daemon := createFullTestDaemon(t, natsURL)
var err error
daemon.jsManager, err = NewJetStreamManager(daemon.natsConn, 1)
daemon.jsManager, err = NewJetStreamManager(daemon.natsConn)
require.NoError(t, err)
err = daemon.jsManager.InitKVBucket()
require.NoError(t, err)
+2 -2
View File
@@ -24,7 +24,7 @@ func newTestMgmtJSM(t *testing.T) *JetStreamManager {
}
t.Cleanup(nc.Close)
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
if err != nil {
t.Fatalf("new JetStreamManager: %v", err)
}
@@ -552,7 +552,7 @@ func TestMgmtIPAllocator_Allocate_RefusedWhenKVUnhealthy(t *testing.T) {
if err != nil {
t.Fatal(err)
}
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
if err != nil {
t.Fatal(err)
}
+3 -3
View File
@@ -740,7 +740,7 @@ func TestDaemon_BootAllocation(t *testing.T) {
defer nc.Close()
daemon.natsConn = nc
daemon.jsManager, err = NewJetStreamManager(nc, 1)
daemon.jsManager, err = NewJetStreamManager(nc)
require.NoError(t, err)
err = daemon.jsManager.InitKVBucket()
require.NoError(t, err)
@@ -1692,7 +1692,7 @@ func TestTerminatedTeardownReaper_SelfHealsFailedVolumeTeardown(t *testing.T) {
// A real terminated-instance KV, so Sweep can list/update/(not yet)purge
// the record exactly as production does.
jsManager, err := NewJetStreamManager(daemon.natsConn, 1)
jsManager, err := NewJetStreamManager(daemon.natsConn)
require.NoError(t, err)
require.NoError(t, jsManager.InitTerminatedInstanceBucket())
daemon.stateStore = newStateStoreAdapter(jsManager, daemon.persistState)
@@ -4545,7 +4545,7 @@ func newEBSProviderTestDaemon(t *testing.T, provider string) *Daemon {
require.NoError(t, err)
t.Cleanup(nc.Close)
jsManager, err := NewJetStreamManager(nc, 1)
jsManager, err := NewJetStreamManager(nc)
require.NoError(t, err)
ctx := context.Background()
+2 -3
View File
@@ -45,9 +45,8 @@ func (d *Daemon) instanceStateWatchBuckets(context.Context) ([]*kvstore.Bucket,
return nil, nil
}
return []*kvstore.Bucket{kvstore.NewBucket(d.jsManager.js, kvstore.Config{
Name: InstanceStateBucket,
History: 1,
Replicas: d.jsManager.replicas,
Name: InstanceStateBucket,
History: 1,
})}, nil
}
+2 -2
View File
@@ -26,10 +26,10 @@ func dnsTestDaemon(t *testing.T) (*Daemon, func(id, publicIP, privateIP string,
t.Helper()
_, _, js := testutil.StartTestJetStream(t)
store := kvstore.New[vm.InstanceRecord](js, kvstore.Config{
Name: InstanceStateBucket, History: 1, Replicas: 1,
Name: InstanceStateBucket, History: 1,
})
d := &Daemon{
jsManager: &JetStreamManager{js: js, records: store, replicas: 1},
jsManager: &JetStreamManager{js: js, records: store},
vmMgr: vm.NewManager(),
config: &config.Config{Region: "ap-southeast-2"},
dnsBaseDomain: "spx3.net",
+2 -2
View File
@@ -62,7 +62,7 @@ func TestDNSWatchBuckets_ZeroValueDaemonIsNilSafe(t *testing.T) {
func TestInstanceStateWatchBuckets_ReturnsSharedBucket(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
d := &Daemon{jsManager: &JetStreamManager{js: js, replicas: 1}}
d := &Daemon{jsManager: &JetStreamManager{js: js}}
buckets, err := d.instanceStateWatchBuckets(t.Context())
require.NoError(t, err)
@@ -111,7 +111,7 @@ func TestEksWatchBuckets_NilGuards(t *testing.T) {
func TestEksWatchBuckets_DelegatesToAccountWatchBuckets(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
_, err := handlers_eks.GetOrCreateAccountBucket(t.Context(), js, "111111111111", 1)
_, err := handlers_eks.GetOrCreateAccountBucket(t.Context(), js, "111111111111")
require.NoError(t, err)
d := &Daemon{eksService: &handlers_eks.EKSServiceImpl{}, natsConn: nc}
+1 -5
View File
@@ -43,16 +43,12 @@ func (d *Daemon) buildECSServiceDeps() handlers_ecs.Deps {
return deps
}
clusterSize := 1
if d.clusterConfig != nil {
clusterSize = len(d.clusterConfig.Nodes)
}
// Retried, because the constructor runs a KV migration that needs JetStream
// to have responders. A daemon that started first got one attempt, failed
// it, and had capacity provisioning off for the life of the process with no
// symptom but a 500 from ProvisionCapacity.
iamSvc, iamErr := initServiceWithRetry("ECS IAM service", func() (*handlers_iam.IAMServiceImpl, error) {
return handlers_iam.NewIAMServiceImpl(d.ctx, d.natsConn, masterKey, clusterSize)
return handlers_iam.NewIAMServiceImpl(d.ctx, d.natsConn, masterKey)
})
if iamErr != nil {
// Off rather than fatal: ECS CRUD, EKS and EC2 do not need this, and a
+1 -8
View File
@@ -36,11 +36,7 @@ func (d *Daemon) systemRoleEnsurer() handlers_iam.SystemInstanceRoleEnsurer {
"err", err)
return nil
}
clusterSize := 1
if d.clusterConfig != nil {
clusterSize = len(d.clusterConfig.Nodes)
}
iamSvc, iamErr := handlers_iam.NewIAMServiceImpl(d.ctx, d.natsConn, masterKey, clusterSize)
iamSvc, iamErr := handlers_iam.NewIAMServiceImpl(d.ctx, d.natsConn, masterKey)
if iamErr != nil {
slog.Warn("System role ensurer: IAM service init failed (retried on next launch); system VMs fall back to baked static creds",
"err", iamErr)
@@ -61,10 +57,8 @@ func (d *Daemon) buildEKSServiceDeps() handlers_eks.EKSServiceDeps {
}
internalSuffix := ""
clusterSize := 1
if d.clusterConfig != nil {
internalSuffix = d.clusterConfig.AWS.ServicesDomain
clusterSize = len(d.clusterConfig.Nodes)
}
gatewayCA := ""
@@ -84,7 +78,6 @@ func (d *Daemon) buildEKSServiceDeps() handlers_eks.EKSServiceDeps {
GatewayBaseURL: d.resolveGatewayBaseURL(),
Region: d.config.Region,
HolderID: d.node,
ClusterSize: clusterSize,
InternalSuffix: internalSuffix,
SystemGatewayURL: d.resolveSystemGatewayBaseURL(),
SystemAccessKey: d.config.Predastore.AccessKey,
+1 -1
View File
@@ -103,7 +103,7 @@ func TestHeartbeatKVContract(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitClusterStateBucket())
@@ -19,7 +19,7 @@ func newRecordManagerInPackage(t *testing.T) *JetStreamManager {
t.Helper()
_, nc, _ := testutil.StartTestJetStream(t)
m, err := NewJetStreamManager(nc, 1)
m, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitKVBucket())
require.NoError(t, m.InitTerminatedInstanceBucket())
@@ -46,7 +46,7 @@ func TestRekey_MovesInstanceStateRecordsOntoTheDot(t *testing.T) {
kv := seedSlashRecord(t, nc, daemon.InstanceStateBucket, 3,
&vm.VM{ID: "i-1", InstanceType: "t3.nano", LastNode: "node-1"})
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitKVBucket())
@@ -64,7 +64,7 @@ func TestRekey_MovesTerminatedRecordsOntoTheDot(t *testing.T) {
kv := seedSlashRecord(t, nc, daemon.TerminatedInstanceBucket, 2,
&vm.VM{ID: "i-1", InstanceType: "t3.nano", LastNode: "node-1"})
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitTerminatedInstanceBucket())
+2 -2
View File
@@ -29,7 +29,7 @@ func newRecordManagerConn(t *testing.T) (*daemon.JetStreamManager, *nats.Conn) {
t.Helper()
_, nc, _ := testutil.StartTestJetStream(t)
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitKVBucket())
require.NoError(t, m.InitTerminatedInstanceBucket())
@@ -259,7 +259,7 @@ func TestInstanceRecord_CarriesDeletionTimestamp(t *testing.T) {
// store: Tier 1 boot runs before cluster KV exists.
func TestInstanceRecordAccessors_ReportAnUnopenedBucket(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
assert.Error(t, m.WriteInstanceRecord("i-1", testRecord("i-1")))
@@ -223,7 +223,7 @@ func TestInstanceStateMigration_CopiesStoppedInstancesForward(t *testing.T) {
seedLegacyBucket(t, nc, daemon.InstanceStateBucket, daemon.StoppedInstancePrefix,
&vm.VM{ID: "i-1", InstanceType: "t3.nano"})
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitKVBucket())
@@ -247,7 +247,7 @@ func TestInstanceStateMigration_RunsEveryStepFromTheOldestVersion(t *testing.T)
_, err = kv.Put(context.Background(), daemon.InstanceStatePrefix+"node-1", blob)
require.NoError(t, err)
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitKVBucket())
@@ -277,7 +277,7 @@ func TestInstanceStateMigration_DoesNotOverwriteAFresherRecord(t *testing.T) {
_, err = kv.Put(context.Background(), daemon.StoppedInstancePrefix+"i-2", untouched)
require.NoError(t, err)
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitKVBucket())
@@ -297,12 +297,12 @@ func TestInstanceStateMigration_IsSafeToRunTwice(t *testing.T) {
kv := seedLegacyBucket(t, nc, daemon.InstanceStateBucket, daemon.StoppedInstancePrefix,
&vm.VM{ID: "i-1", InstanceType: "t3.nano"})
first, err := daemon.NewJetStreamManager(nc, 1)
first, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, first.InitKVBucket())
require.NoError(t, kvutil.WriteVersion(context.Background(), kv, 1))
second, err := daemon.NewJetStreamManager(nc, 1)
second, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, second.InitKVBucket())
@@ -317,7 +317,7 @@ func TestTerminatedInstanceMigration_CopiesTerminatedInstancesForward(t *testing
seedLegacyBucket(t, nc, daemon.TerminatedInstanceBucket, daemon.TerminatedInstancePrefix,
&vm.VM{ID: "i-1", InstanceType: "t3.nano"})
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitTerminatedInstanceBucket())
+1 -1
View File
@@ -522,7 +522,7 @@ func (d *Daemon) newLiveness() *instancecache.Liveness {
if d.jsManager == nil || d.jsManager.js == nil {
return nil
}
return instancecache.NewLiveness(d.jsManager.js, clusterStateConfig(d.jsManager.replicas))
return instancecache.NewLiveness(d.jsManager.js, clusterStateConfig())
}
// forgetSupersededInstances drops any local copy of an instance the records say
+12 -12
View File
@@ -234,7 +234,7 @@ func TestWriteRunningSet_SeedsFromWhatIsAlreadyThere(t *testing.T) {
writeRunningSet(t, first, runningVM("i-1", "t3.nano"))
before := recordRevision(t, nc, "i-1")
second, err := daemon.NewJetStreamManager(nc, 1)
second, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, second.InitKVBucket())
writeRunningSet(t, second, runningVM("i-1", "t3.nano"))
@@ -293,7 +293,7 @@ func TestInstanceStateMigration_SplitsTheNodeBlob(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
seedNodeBlobBucket(t, nc, localNode, runningSet(runningVM("i-1", "t3.nano"), runningVM("i-2", "t3.micro")))
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitKVBucket())
@@ -314,7 +314,7 @@ func TestInstanceStateMigration_LeavesTheNodeBlobIntact(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
kv := seedNodeBlobBucket(t, nc, localNode, runningSet(runningVM("i-1", "t3.nano")))
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitKVBucket())
writeRunningSet(t, m, runningVM("i-1", "t3.nano"))
@@ -330,7 +330,7 @@ func TestInstanceStateMigration_SeedsAPresenceMarkerPerNode(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
kv := seedNodeBlobBucket(t, nc, localNode, runningSet(runningVM("i-1", "t3.nano")))
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitKVBucket())
@@ -359,7 +359,7 @@ func TestInstanceStateMigration_SplitsEveryNodesBlob(t *testing.T) {
_, err = kv.Put(context.Background(), daemon.InstanceStatePrefix+"node-2", data)
require.NoError(t, err)
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitKVBucket())
@@ -383,7 +383,7 @@ func TestInstanceStateMigration_CarriesOwnershipOffTheBlobKey(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
seedNodeBlobBucket(t, nc, localNode, runningSet(unownedVM("i-1", "t3.nano")))
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitKVBucket())
@@ -401,7 +401,7 @@ func TestInstanceStateMigration_ANodeRecoversAnUnownedRunningSet(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
seedNodeBlobBucket(t, nc, localNode, runningSet(unownedVM("i-1", "t3.nano"), unownedVM("i-2", "t3.micro")))
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitKVBucket())
@@ -419,7 +419,7 @@ func TestInstanceStateMigration_LeavesAnOwnedRecordAlone(t *testing.T) {
moved := &vm.VM{ID: "i-1", InstanceType: "t3.nano", Status: vm.StateRunning, LastNode: "node-2"}
seedNodeBlobBucket(t, nc, localNode, runningSet(moved))
m, err := daemon.NewJetStreamManager(nc, 1)
m, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, m.InitKVBucket())
@@ -433,13 +433,13 @@ func TestInstanceStateMigration_OwnershipStampIsSafeToRunTwice(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
kv := seedNodeBlobBucket(t, nc, localNode, runningSet(unownedVM("i-1", "t3.nano")))
first, err := daemon.NewJetStreamManager(nc, 1)
first, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, first.InitKVBucket())
before := recordRevision(t, nc, "i-1")
require.NoError(t, kvutil.WriteVersion(context.Background(), kv, 4))
second, err := daemon.NewJetStreamManager(nc, 1)
second, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, second.InitKVBucket())
@@ -450,13 +450,13 @@ func TestInstanceStateMigration_BlobSplitIsSafeToRunTwice(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
kv := seedNodeBlobBucket(t, nc, localNode, runningSet(runningVM("i-1", "t3.nano")))
first, err := daemon.NewJetStreamManager(nc, 1)
first, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, first.InitKVBucket())
before := recordRevision(t, nc, "i-1")
require.NoError(t, kvutil.WriteVersion(context.Background(), kv, 2))
second, err := daemon.NewJetStreamManager(nc, 1)
second, err := daemon.NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, second.InitKVBucket())
+1 -1
View File
@@ -21,7 +21,7 @@ func newTestIPSecBarrier(t *testing.T) (*KVIPSecBarrier, *JetStreamManager) {
require.NoError(t, err)
t.Cleanup(nc.Close)
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitClusterStateBucket())
+8 -74
View File
@@ -6,7 +6,6 @@ import (
"errors"
"fmt"
"log/slog"
"strings"
"time"
"github.com/aws/aws-sdk-go/aws"
@@ -83,7 +82,6 @@ type JetStreamManager struct {
records *kvstore.Store[vm.InstanceRecord] // i.<id>, instance-state bucket
termRecords *kvstore.Store[vm.InstanceRecord] // i.<id>, terminated bucket
clusterKV jetstream.KeyValue // spinifex-cluster-state
replicas int
obs KVSyncObserver
running runningSetState
}
@@ -117,28 +115,23 @@ func (m *JetStreamManager) KVHealthy() bool {
}
// NewJetStreamManager creates a new JetStreamManager from a NATS connection.
// replicas specifies the number of replicas for the KV bucket (typically matches cluster node count).
func NewJetStreamManager(nc *nats.Conn, replicas int) (*JetStreamManager, error) {
func NewJetStreamManager(nc *nats.Conn) (*JetStreamManager, error) {
js, err := jetstream.New(nc)
if err != nil {
return nil, err
}
return &JetStreamManager{
js: js,
replicas: replicas,
}, nil
return &JetStreamManager{js: js}, nil
}
// instanceStateConfig describes the instance-state bucket. Extracted so the
// bucket can be rebuilt against a different JetStream client without the
// description drifting from the one InitKVBucket creates.
func instanceStateConfig(replicas int) kvstore.Config {
func instanceStateConfig() kvstore.Config {
return kvstore.Config{
Name: InstanceStateBucket,
Description: "Spinifex instance state storage",
History: 1,
Replicas: replicas,
// The owning node republishes its record on its next write, so an
// emptied bucket costs a sync rather than the records themselves.
RecreateIfMissing: true,
@@ -150,12 +143,11 @@ func instanceStateConfig(replicas int) kvstore.Config {
}
// terminatedInstanceConfig describes the terminated-instances bucket.
func terminatedInstanceConfig(replicas int) kvstore.Config {
func terminatedInstanceConfig() kvstore.Config {
return kvstore.Config{
Name: TerminatedInstanceBucket,
Description: "Terminated instances (auto-expire after 1 hour)",
History: 1,
Replicas: replicas,
TTL: 1 * time.Hour,
RecreateIfMissing: true,
OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error {
@@ -176,7 +168,7 @@ func (m *JetStreamManager) setInstanceStateBucket(b *kvstore.Bucket) {
// InitKVBucket initializes the KV bucket, creating it if it doesn't exist.
func (m *JetStreamManager) InitKVBucket() error {
m.setInstanceStateBucket(kvstore.NewBucket(m.js, instanceStateConfig(m.replicas)))
m.setInstanceStateBucket(kvstore.NewBucket(m.js, instanceStateConfig()))
// Opened eagerly: a bucket that cannot be created must fail startup here,
// not on the first write, and Tier 1 boot must not reach cluster KV at all.
@@ -185,12 +177,11 @@ func (m *JetStreamManager) InitKVBucket() error {
}
// clusterStateConfig describes the cluster-state bucket.
func clusterStateConfig(replicas int) kvstore.Config {
func clusterStateConfig() kvstore.Config {
return kvstore.Config{
Name: ClusterStateBucket,
Description: "Spinifex cluster state (heartbeats, shutdown markers, service maps)",
History: 1,
Replicas: replicas,
TTL: 1 * time.Hour,
OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error {
return migrate.DefaultRegistry.RunKV(ctx, ClusterStateBucket, kv, ClusterStateBucketVersion)
@@ -204,7 +195,7 @@ func clusterStateConfig(replicas int) kvstore.Config {
// The handle is kept rather than the Bucket: this one holds several unrelated
// record types under one namespace, and its callers still read it as raw keys.
func (m *JetStreamManager) InitClusterStateBucket() error {
kv, err := kvstore.NewBucket(m.js, clusterStateConfig(m.replicas)).KV(context.Background())
kv, err := kvstore.NewBucket(m.js, clusterStateConfig()).KV(context.Background())
if err != nil {
return err
}
@@ -215,7 +206,7 @@ func (m *JetStreamManager) InitClusterStateBucket() error {
// InitTerminatedInstanceBucket initializes the terminated-instances KV bucket with a 1-hour TTL.
// JetStream automatically purges keys after 1 hour, matching AWS behavior for terminated instances.
func (m *JetStreamManager) InitTerminatedInstanceBucket() error {
m.term = kvstore.New[vm.VM](m.js, terminatedInstanceConfig(m.replicas))
m.term = kvstore.New[vm.VM](m.js, terminatedInstanceConfig())
// A second view over the same handle, for the same reason the instance-state
// bucket carries several.
m.termRecords = kvstore.On[vm.InstanceRecord](m.term.Bucket)
@@ -521,63 +512,6 @@ func (m *JetStreamManager) DeleteState(nodeID string) error {
return nil
}
// UpdateReplicas updates the replica count for ALL JetStream KV buckets.
// It iterates over every KV_* stream and bumps replicas to match the cluster size.
// This ensures service buckets (IAM, VPC, IGW, etc.) are replicated alongside daemon buckets.
// This should be called when new nodes join the cluster.
func (m *JetStreamManager) UpdateReplicas(newReplicas int) error {
if m.js == nil {
return errors.New("JetStream context not initialized")
}
m.replicas = newReplicas
ctx := context.Background()
// Iterate all streams and update any KV-backed stream (prefixed "KV_")
updated := 0
lister := m.js.StreamNames(ctx)
for name := range lister.Name() {
if !strings.HasPrefix(name, "KV_") {
continue
}
stream, err := m.js.Stream(ctx, name)
if err != nil {
slog.Warn("Failed to open stream", "stream", name, "error", err)
continue
}
info, err := stream.Info(ctx)
if err != nil {
slog.Warn("Failed to get stream info", "stream", name, "error", err)
continue
}
if info.Config.Replicas >= newReplicas {
continue
}
oldReplicas := info.Config.Replicas
info.Config.Replicas = newReplicas
if _, err := m.js.UpdateStream(ctx, info.Config); err != nil {
slog.Warn("Failed to update KV bucket replicas", "stream", name, "error", err)
continue
}
bucket := strings.TrimPrefix(name, "KV_")
slog.Info("Updated KV bucket replicas", "bucket", bucket, "oldReplicas", oldReplicas, "newReplicas", newReplicas)
updated++
}
if err := lister.Err(); err != nil {
return fmt.Errorf("list JetStream streams: %w", err)
}
if updated > 0 {
slog.Info("KV replication update complete", "bucketsUpdated", updated, "replicas", newReplicas)
}
return nil
}
// WriteStoppedInstance writes a stopped instance to the shared KV store,
// replacing any existing record for instanceID. Uses CAS internally (bounded
// retry) so a write racing a concurrent update is detected and retried
+68 -113
View File
@@ -40,7 +40,7 @@ func TestJetStreamManager_WriteAndLoadState(t *testing.T) {
defer nc.Close()
// Create JetStreamManager
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err, "Failed to create JetStreamManager")
// Initialize the KV bucket
@@ -90,7 +90,7 @@ func TestJetStreamManager_LoadState_KeyNotFound(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
@@ -112,7 +112,7 @@ func TestJetStreamManager_BucketCreation(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
// InitKVBucket should create the bucket
@@ -134,7 +134,7 @@ func TestJetStreamManager_BucketReconnection(t *testing.T) {
nc1, err := nats.Connect(natsURL)
require.NoError(t, err)
jsm1, err := NewJetStreamManager(nc1, 1)
jsm1, err := NewJetStreamManager(nc1)
require.NoError(t, err)
err = jsm1.InitKVBucket()
@@ -157,7 +157,7 @@ func TestJetStreamManager_BucketReconnection(t *testing.T) {
require.NoError(t, err)
defer nc2.Close()
jsm2, err := NewJetStreamManager(nc2, 1)
jsm2, err := NewJetStreamManager(nc2)
require.NoError(t, err)
err = jsm2.InitKVBucket()
@@ -180,7 +180,7 @@ func TestJetStreamManager_DeleteState(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
@@ -222,7 +222,7 @@ func TestJetStreamManager_DeleteState_NonExistent(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
@@ -241,7 +241,7 @@ func TestJetStreamManager_WriteState_UpdateExisting(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
@@ -289,7 +289,7 @@ func TestJetStreamManager_MultipleNodes(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
@@ -337,7 +337,7 @@ func TestJetStreamManager_KVNotInitialized(t *testing.T) {
defer nc.Close()
// Create JetStreamManager but don't call InitKVBucket
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
testInstances := make(map[string]*vm.VM)
@@ -351,51 +351,6 @@ func TestJetStreamManager_KVNotInitialized(t *testing.T) {
assert.Error(t, err, "DeleteState should error when KV not initialized")
}
// TestJetStreamManager_UpdateReplicas tests updating replica count for the KV bucket.
func TestJetStreamManager_UpdateReplicas(t *testing.T) {
natsURL := sharedJSNATSURL
nc, err := nats.Connect(natsURL)
require.NoError(t, err)
defer nc.Close()
// Create with 1 replica (typical for single node startup)
jsm, err := NewJetStreamManager(nc, 1)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
// Verify initial replica count
js, _ := jetstream.New(nc)
stream, err := js.Stream(t.Context(), "KV_"+InstanceStateBucket)
require.NoError(t, err)
streamInfo, err := stream.Info(t.Context())
require.NoError(t, err)
assert.Equal(t, 1, streamInfo.Config.Replicas, "Should start with 1 replica")
// Try to update to same replica count (should be a no-op)
err = jsm.UpdateReplicas(1)
assert.NoError(t, err, "Updating to same replica count should succeed")
// Note: Increasing replicas beyond 1 requires additional NATS servers in the cluster,
// which we don't have in the test environment. In a single-node test server,
// attempting to increase replicas will fail with "insufficient resources" error.
// This test verifies the basic functionality works.
}
// TestJetStreamManager_UpdateReplicas_NoInit tests UpdateReplicas when JS not initialized.
func TestJetStreamManager_UpdateReplicas_NoInit(t *testing.T) {
// Test with nil JetStream context
jsm := &JetStreamManager{
js: nil,
replicas: 1,
}
err := jsm.UpdateReplicas(3)
assert.Error(t, err, "UpdateReplicas should error when JetStream not initialized")
}
// --- Stopped instance KV tests ---
// TestJetStreamManager_WriteAndLoadStoppedInstance tests round-trip write and load of a stopped instance.
@@ -404,7 +359,7 @@ func TestJetStreamManager_WriteAndLoadStoppedInstance(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -439,7 +394,7 @@ func TestJetStreamManager_LoadStoppedInstance_NotFound(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -455,7 +410,7 @@ func TestJetStreamManager_DeleteStoppedInstance(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -489,7 +444,7 @@ func TestJetStreamManager_ClaimStoppedInstance_HappyPath(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -517,7 +472,7 @@ func TestJetStreamManager_ClaimStoppedInstance_NotFound(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -537,7 +492,7 @@ func TestJetStreamManager_ClaimStoppedInstance_ConcurrentClaim(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -588,7 +543,7 @@ func TestJetStreamManager_UpdateStoppedInstance_ConflictRetry(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -631,7 +586,7 @@ func TestJetStreamManager_UpdateStoppedInstance_NotFound(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -652,7 +607,7 @@ func TestJetStreamManager_UpdateStoppedInstance_NoResurrectAfterClaim(t *testing
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -696,7 +651,7 @@ func TestJetStreamManager_ListStoppedInstances(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -738,7 +693,7 @@ func TestJetStreamManager_StoppedInstances_NoInterference(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -778,7 +733,7 @@ func TestJetStreamManager_StoppedInstance_KVNotInitialized(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
// Don't call InitKVBucket
@@ -803,7 +758,7 @@ func TestJetStreamManager_WriteServiceManifest(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitClusterStateBucket()
require.NoError(t, err)
@@ -840,7 +795,7 @@ func TestJetStreamManager_WriteServiceManifest_EmptyServices(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitClusterStateBucket()
require.NoError(t, err)
@@ -867,7 +822,7 @@ func TestJetStreamManager_WriteServiceManifest_ClusterKVNotInitialized(t *testin
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
// Don't call InitClusterStateBucket
@@ -895,7 +850,7 @@ func TestJetStreamManager_WriteState_RecoverAfterStreamLost(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -922,7 +877,7 @@ func TestJetStreamManager_LoadState_RecoverAfterStreamLost(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -941,7 +896,7 @@ func TestJetStreamManager_DeleteState_RecoverAfterStreamLost(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -958,7 +913,7 @@ func TestJetStreamManager_WriteStoppedInstance_RecoverAfterStreamLost(t *testing
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -984,7 +939,7 @@ func TestJetStreamManager_LoadStoppedInstance_RecoverAfterStreamLost(t *testing.
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -1002,7 +957,7 @@ func TestJetStreamManager_DeleteStoppedInstance_RecoverAfterStreamLost(t *testin
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -1018,7 +973,7 @@ func TestJetStreamManager_ListStoppedInstances_RecoverAfterStreamLost(t *testing
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -1055,12 +1010,12 @@ func swapToNonJSContext(t *testing.T, jsm *JetStreamManager) {
if jsm.stateB != nil {
kv, err := jsm.stateB.KV(t.Context())
require.NoError(t, err)
jsm.setInstanceStateBucket(kvstore.NewOpenBucket(js, kv, instanceStateConfig(jsm.replicas)))
jsm.setInstanceStateBucket(kvstore.NewOpenBucket(js, kv, instanceStateConfig()))
}
if jsm.term != nil {
kv, err := jsm.term.KV(t.Context())
require.NoError(t, err)
jsm.term = kvstore.Over[vm.VM](js, kv, terminatedInstanceConfig(jsm.replicas))
jsm.term = kvstore.Over[vm.VM](js, kv, terminatedInstanceConfig())
}
}
@@ -1069,7 +1024,7 @@ func TestJetStreamManager_WriteState_RecoveryFailure(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -1087,7 +1042,7 @@ func TestJetStreamManager_LoadState_RecoveryFailure(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -1105,7 +1060,7 @@ func TestJetStreamManager_DeleteState_RecoveryFailure(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -1122,7 +1077,7 @@ func TestJetStreamManager_WriteStoppedInstance_RecoveryFailure(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -1140,7 +1095,7 @@ func TestJetStreamManager_LoadStoppedInstance_RecoveryFailure(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -1158,7 +1113,7 @@ func TestJetStreamManager_DeleteStoppedInstance_RecoveryFailure(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -1175,7 +1130,7 @@ func TestJetStreamManager_ListStoppedInstances_RecoveryFailure(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitKVBucket()
require.NoError(t, err)
@@ -1195,7 +1150,7 @@ func TestJetStreamManager_WriteAndLoadTerminatedInstance(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
@@ -1225,7 +1180,7 @@ func TestJetStreamManager_LoadTerminatedInstance_NotFound(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
@@ -1239,7 +1194,7 @@ func TestJetStreamManager_DeleteTerminatedInstance(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
@@ -1268,7 +1223,7 @@ func TestJetStreamManager_ListTerminatedInstances(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
@@ -1304,7 +1259,7 @@ func TestJetStreamManager_WriteLoadDeleteTerminatedInstance_RoundTrip(t *testing
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
@@ -1336,7 +1291,7 @@ func TestJetStreamManager_TerminatedInstance_KVNotInitialized(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
// Don't call InitTerminatedInstanceBucket
@@ -1365,7 +1320,7 @@ func TestJetStreamManager_UpdateTerminatedInstance_ConflictRetry(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
@@ -1414,7 +1369,7 @@ func TestJetStreamManager_UpdateTerminatedInstance_NotFound(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
@@ -1432,7 +1387,7 @@ func TestJetStreamManager_WriteStoppedInstance_OverwritesConcurrentValue(t *test
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -1470,7 +1425,7 @@ func TestJetStreamManager_WriteTerminatedInstance_RecoverAfterStreamLost(t *test
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
@@ -1493,7 +1448,7 @@ func TestJetStreamManager_LoadTerminatedInstance_RecoverAfterStreamLost(t *testi
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
@@ -1509,7 +1464,7 @@ func TestJetStreamManager_DeleteTerminatedInstance_RecoverAfterStreamLost(t *tes
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
@@ -1524,7 +1479,7 @@ func TestJetStreamManager_ListTerminatedInstances_RecoverAfterStreamLost(t *test
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
@@ -1540,7 +1495,7 @@ func TestJetStreamManager_InitBuckets_WritesVersion(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -1570,7 +1525,7 @@ func TestJetStreamManager_ListStoppedInstances_SkipsVersionKey(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -1585,7 +1540,7 @@ func TestJetStreamManager_ListTerminatedInstances_SkipsVersionKey(t *testing.T)
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
@@ -1632,7 +1587,7 @@ func TestJetStreamManager_BestEffort_Success_NotifiesObserver(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -1651,7 +1606,7 @@ func TestJetStreamManager_BestEffort_PutError_NotifiesObserver(t *testing.T) {
nc, err := nats.Connect(sharedJSNATSURL)
require.NoError(t, err)
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -1675,7 +1630,7 @@ func TestJetStreamManager_BestEffort_NilObserver_NoPanic(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -1696,7 +1651,7 @@ func TestJetStreamManager_UpdateMgmtIPAM_ConflictRetry(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitClusterStateBucket())
@@ -1746,7 +1701,7 @@ func TestJetStreamManager_UpdateMgmtIPAM_NotFound(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitClusterStateBucket())
@@ -1761,7 +1716,7 @@ func TestJetStreamManager_UpdateMgmtIPAM_ClusterKVNotInitialized(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
// Don't call InitClusterStateBucket
@@ -1781,7 +1736,7 @@ func TestJetStreamManager_UpdateMgmtIPAM_ClusterKVNotInitialized(t *testing.T) {
func TestJetStreamManager_ListStoppedInstances_FailsOnAnUndecodableRecord(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -1806,7 +1761,7 @@ func TestJetStreamManager_UpdateStoppedInstance_AbsentKeyIsNotFound(t *testing.T
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
require.NoError(t, jsm.InitTerminatedInstanceBucket())
@@ -1842,7 +1797,7 @@ func TestJetStreamManager_LoadState_UnversionedRecordStillReads(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -1867,7 +1822,7 @@ func TestJetStreamManager_LoadState_FutureRecordIsRefused(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
@@ -1887,7 +1842,7 @@ func TestJetStreamManager_WriteState_StampsSchemaVersion(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
+1 -1
View File
@@ -30,7 +30,7 @@ func TestClusterShutdownStateKVRoundTrip(t *testing.T) {
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
jsm, err := NewJetStreamManager(nc)
require.NoError(t, err)
err = jsm.InitClusterStateBucket()
require.NoError(t, err)
+1 -1
View File
@@ -41,7 +41,7 @@ func createDaemonWithJetStream(t *testing.T) *Daemon {
daemon.config = &config.Config{BaseDir: tmpDir, DataDir: tmpDir}
daemon.natsConn = nc
daemon.jsManager, err = NewJetStreamManager(nc, 1)
daemon.jsManager, err = NewJetStreamManager(nc)
require.NoError(t, err)
require.NoError(t, daemon.jsManager.InitKVBucket())
require.NoError(t, daemon.jsManager.InitClusterStateBucket())
+1 -1
View File
@@ -241,7 +241,7 @@ func newCreateAccountGateway(t *testing.T) *GatewayConfig {
masterKey, err := handlers_iam.GenerateMasterKey()
require.NoError(t, err)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey, 1)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey)
require.NoError(t, err)
sub, err := nc.Subscribe(utils.SubjectEnsureDefaultVpc, func(msg *nats.Msg) {
@@ -40,7 +40,7 @@ func newDeleteAccountGateway(t *testing.T) *GatewayConfig {
masterKey, err := handlers_iam.GenerateMasterKey()
require.NoError(t, err)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey, 1)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey)
require.NoError(t, err)
for _, subject := range teardownSubjects {
+3 -3
View File
@@ -235,7 +235,7 @@ func (gw *GatewayConfig) bedrockLoggingConfigStore() *gateway_bedrock.LoggingCon
if gw.BedrockLoggingConfig != nil {
return gw.BedrockLoggingConfig
}
return gateway_bedrock.NewLoggingConfigStore(nil, 1)
return gateway_bedrock.NewLoggingConfigStore(nil)
}
// bedrockRecorder returns gw.BedrockRecorder, or the no-op fallback when no
@@ -269,7 +269,7 @@ func (gw *GatewayConfig) bedrockProvisionedStore() *gateway_bedrock.ProvisionedS
if gw.BedrockProvisioned != nil {
return gw.BedrockProvisioned
}
return gateway_bedrock.NewProvisionedStore(nil, 1, gw.Region, nil)
return gateway_bedrock.NewProvisionedStore(nil, gw.Region, nil)
}
// bedrockGuardrailStore returns gw.BedrockGuardrails, or a store backed by no
@@ -281,7 +281,7 @@ func (gw *GatewayConfig) bedrockGuardrailStore() *gateway_bedrock.GuardrailStore
if gw.BedrockGuardrails != nil {
return gw.BedrockGuardrails
}
return gateway_bedrock.NewGuardrailStore(nil, 1, gw.Region)
return gateway_bedrock.NewGuardrailStore(nil, gw.Region)
}
// bedrockEndpointResolver returns the registry-backed resolver when one is
+5 -6
View File
@@ -65,13 +65,12 @@ type ModelAccessStore struct {
var _ AccessResolver = (*ModelAccessStore)(nil)
// NewModelAccessStore constructs a ModelAccessStore over the cluster's
// JetStream client, replicated across replicas nodes.
func NewModelAccessStore(js jetstream.JetStream, replicas int) *ModelAccessStore {
// JetStream client.
func NewModelAccessStore(js jetstream.JetStream) *ModelAccessStore {
return &ModelAccessStore{bucket: kvstore.NewBucket(js, kvstore.Config{
Name: modelAccessBucket,
History: modelAccessHistory,
Replicas: replicas,
Missing: "bedrock: model access store has no JetStream client configured",
Name: modelAccessBucket,
History: modelAccessHistory,
Missing: "bedrock: model access store has no JetStream client configured",
})}
}
+10 -10
View File
@@ -72,7 +72,7 @@ func TestAccessKey_EncodesModelID(t *testing.T) {
// Revoke.
func TestModelAccessStore_GrantResolveRevoke(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewModelAccessStore(js, 1)
store := NewModelAccessStore(js)
ctx := context.Background()
granted, err := store.Granted(ctx, "000000000001", selfHostTestModel)
@@ -100,7 +100,7 @@ func TestModelAccessStore_GrantResolveRevoke(t *testing.T) {
func TestModelAccessStore_GrantIsIdempotent(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewModelAccessStore(js, 1)
store := NewModelAccessStore(js)
ctx := context.Background()
require.NoError(t, store.Grant(ctx, "000000000001", selfHostTestModel))
@@ -116,7 +116,7 @@ func TestModelAccessStore_GrantIsIdempotent(t *testing.T) {
// error.
func TestModelAccessStore_RevokeMissingGrantSucceeds(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewModelAccessStore(js, 1)
store := NewModelAccessStore(js)
assert.NoError(t, store.Revoke(context.Background(), "000000000001", selfHostTestModel))
}
@@ -125,7 +125,7 @@ func TestModelAccessStore_RevokeMissingGrantSucceeds(t *testing.T) {
// reversible, including the ':' that forced it.
func TestModelAccessStore_List_RoundTripsModelIDs(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewModelAccessStore(js, 1)
store := NewModelAccessStore(js)
ctx := context.Background()
require.NoError(t, store.Grant(ctx, "000000000001", selfHostTestModel))
@@ -143,7 +143,7 @@ func TestModelAccessStore_List_RoundTripsModelIDs(t *testing.T) {
func TestModelAccessStore_ListEmptyAccount(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewModelAccessStore(js, 1)
store := NewModelAccessStore(js)
models, err := store.List(context.Background(), "000000000009")
require.NoError(t, err)
@@ -154,7 +154,7 @@ func TestModelAccessStore_ListEmptyAccount(t *testing.T) {
// handlers_quota exempts the system account from every quota dimension.
func TestModelAccessStore_SystemAccountBypassesGrants(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewModelAccessStore(js, 1)
store := NewModelAccessStore(js)
granted, err := store.Granted(context.Background(), utils.GlobalAccountID, selfHostTestModel)
require.NoError(t, err)
@@ -189,7 +189,7 @@ func TestGrantedCatalogEntry_ErrorClasses(t *testing.T) {
// who then revokes a grant does not have it restored by the next restart.
func TestSeedAccountGrants_SeedsOnceThenRespectsRevoke(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewModelAccessStore(js, 1)
store := NewModelAccessStore(js)
ctx := context.Background()
const account = "000000000001"
@@ -218,7 +218,7 @@ func TestSeedAccountGrants_SeedsOnceThenRespectsRevoke(t *testing.T) {
// account to end up with the catalog.
func TestSeedAccountGrants_LeavesOtherAccountsDenied(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewModelAccessStore(js, 1)
store := NewModelAccessStore(js)
ctx := context.Background()
_, err := store.SeedAccountGrants(ctx, "000000000001", CatalogModelIDs())
@@ -233,7 +233,7 @@ func TestSeedAccountGrants_LeavesOtherAccountsDenied(t *testing.T) {
// read back as a grant would show up as a bogus model ID in List.
func TestSeedAccountGrants_MarkerIsNotAGrant(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewModelAccessStore(js, 1)
store := NewModelAccessStore(js)
ctx := context.Background()
_, err := store.SeedAccountGrants(ctx, "000000000001", nil)
@@ -248,7 +248,7 @@ func TestSeedAccountGrants_MarkerIsNotAGrant(t *testing.T) {
// JetStream client: every accessor must report the misconfiguration rather
// than panic on a nil handle.
func TestModelAccessStore_RequiresJetStream(t *testing.T) {
store := NewModelAccessStore(nil, 1)
store := NewModelAccessStore(nil)
ctx := context.Background()
_, err := store.Granted(ctx, "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0")
+2 -2
View File
@@ -210,7 +210,7 @@ func TestAdmitSelfHost_OnDemandAdmitsThenThrottles(t *testing.T) {
// resolving committed ModelUnits surfaces as that error, never masquerading as
// a throttle or admitting the request.
func TestAdmitSelfHost_CommittedUnitsErrorPropagates(t *testing.T) {
store := NewProvisionedStore(nil, 1, ptTestRegion, newStubEndpointProvisioner())
store := NewProvisionedStore(nil, ptTestRegion, newStubEndpointProvisioner())
release, err := admitSelfHost(context.Background(), store, ptCallerAccount, selfHostTestModel,
catalogEntry{MaxConcurrency: 1})
@@ -282,7 +282,7 @@ func TestCommittedModelUnits_EmptyAccountReturnsZero(t *testing.T) {
// TestCommittedModelUnits_BucketErrorPropagates asserts an unusable store
// surfaces its error rather than silently reading zero committed units.
func TestCommittedModelUnits_BucketErrorPropagates(t *testing.T) {
store := NewProvisionedStore(nil, 1, ptTestRegion, newStubEndpointProvisioner())
store := NewProvisionedStore(nil, ptTestRegion, newStubEndpointProvisioner())
_, err := committedModelUnits(context.Background(), store, ptCallerAccount, selfHostTestModel)
require.Error(t, err)
}
+4 -5
View File
@@ -38,13 +38,12 @@ var _ CredentialResolver = (*CredentialStore)(nil)
// NewCredentialStore constructs a CredentialStore. platformDefaults maps
// vendor to a platform-wide API key used when an account has no key of its
// own; pass an empty map to disable platform defaults entirely.
func NewCredentialStore(js jetstream.JetStream, masterKey []byte, replicas int, platformDefaults map[string]string) *CredentialStore {
func NewCredentialStore(js jetstream.JetStream, masterKey []byte, platformDefaults map[string]string) *CredentialStore {
return &CredentialStore{
bucket: kvstore.NewBucket(js, kvstore.Config{
Name: bedrockCredentialsBucket,
History: bedrockCredentialsHistory,
Replicas: replicas,
Missing: "bedrock: credential store has no JetStream client configured",
Name: bedrockCredentialsBucket,
History: bedrockCredentialsHistory,
Missing: "bedrock: credential store has no JetStream client configured",
}),
masterKey: masterKey,
platformDefaults: platformDefaults,
+6 -6
View File
@@ -21,7 +21,7 @@ func fixedMasterKey() []byte {
// live NATS server is out of scope here). Resolve must consult
// platformDefaults without dereferencing js.
func TestCredentialStore_Resolve_NilJS(t *testing.T) {
store := NewCredentialStore(nil, fixedMasterKey(), 1, map[string]string{"anthropic": "sk-platform-default"})
store := NewCredentialStore(nil, fixedMasterKey(), map[string]string{"anthropic": "sk-platform-default"})
cases := []struct {
name string
@@ -43,7 +43,7 @@ func TestCredentialStore_Resolve_NilJS(t *testing.T) {
}
func TestCredentialStore_Resolve_NoPlatformDefaults(t *testing.T) {
store := NewCredentialStore(nil, fixedMasterKey(), 1, nil)
store := NewCredentialStore(nil, fixedMasterKey(), nil)
key, ok, err := store.Resolve(context.Background(), "000000000001", "anthropic")
require.NoError(t, err)
@@ -63,7 +63,7 @@ func TestNoopCredentialResolver_ResolvesNothing(t *testing.T) {
// KV-hit branch of Resolve (Get+decrypt).
func TestCredentialStore_PutAndResolve_KV(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewCredentialStore(testutil.NewJetStream(t, nc), fixedMasterKey(), 1, nil)
store := NewCredentialStore(testutil.NewJetStream(t, nc), fixedMasterKey(), nil)
ctx := context.Background()
require.NoError(t, store.PutCredential(ctx, "000000000001", "anthropic", "sk-test"))
@@ -78,7 +78,7 @@ func TestCredentialStore_PutAndResolve_KV(t *testing.T) {
// miss (jetstream.ErrKeyNotFound) that falls through to the platform default.
func TestCredentialStore_Resolve_KVMissFallsBackToPlatformDefault(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewCredentialStore(testutil.NewJetStream(t, nc), fixedMasterKey(), 1, map[string]string{"anthropic": "sk-platform-default"})
store := NewCredentialStore(testutil.NewJetStream(t, nc), fixedMasterKey(), map[string]string{"anthropic": "sk-platform-default"})
key, ok, err := store.Resolve(context.Background(), "999999999999", "anthropic")
require.NoError(t, err)
@@ -90,7 +90,7 @@ func TestCredentialStore_Resolve_KVMissFallsBackToPlatformDefault(t *testing.T)
// unknown account with no platform defaults configured: ("", false, nil).
func TestCredentialStore_Resolve_KVMissNoPlatformDefault(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewCredentialStore(testutil.NewJetStream(t, nc), fixedMasterKey(), 1, nil)
store := NewCredentialStore(testutil.NewJetStream(t, nc), fixedMasterKey(), nil)
key, ok, err := store.Resolve(context.Background(), "999999999999", "anthropic")
require.NoError(t, err)
@@ -119,7 +119,7 @@ func TestEncryptDecryptSecret_RoundTrip(t *testing.T) {
// Resolve's tolerated nil: a write has no platform-default fallback to fall
// back to, so it must report the misconfiguration.
func TestCredentialStore_PutCredential_RequiresJetStream(t *testing.T) {
store := NewCredentialStore(nil, fixedMasterKey(), 1, map[string]string{"anthropic": "sk-default"})
store := NewCredentialStore(nil, fixedMasterKey(), map[string]string{"anthropic": "sk-default"})
err := store.PutCredential(context.Background(), "000000000001", "anthropic", "sk-account")
require.Error(t, err)
+6 -7
View File
@@ -89,15 +89,14 @@ type GuardrailStore struct {
region string
}
// NewGuardrailStore constructs a GuardrailStore over js, replicated across
// replicas nodes, minting ARNs for region.
func NewGuardrailStore(js jetstream.JetStream, replicas int, region string) *GuardrailStore {
// NewGuardrailStore constructs a GuardrailStore over js, minting ARNs for
// region.
func NewGuardrailStore(js jetstream.JetStream, region string) *GuardrailStore {
return &GuardrailStore{
store: kvstore.New[GuardrailRecord](js, kvstore.Config{
Name: bedrockGuardrailBucket,
History: bedrockGuardrailHistory,
Replicas: replicas,
Missing: "bedrock: guardrail store has no JetStream client configured",
Name: bedrockGuardrailBucket,
History: bedrockGuardrailHistory,
Missing: "bedrock: guardrail store has no JetStream client configured",
}),
region: region,
}
+1 -1
View File
@@ -20,7 +20,7 @@ const (
func newGuardrailTestStore(t *testing.T) *GuardrailStore {
t.Helper()
_, _, js := testutil.StartTestJetStream(t)
return NewGuardrailStore(js, 1, ptTestRegion)
return NewGuardrailStore(js, ptTestRegion)
}
// createGuardrailInput builds a minimal valid CreateGuardrailInput carrying
+14 -10
View File
@@ -16,6 +16,7 @@ import (
"github.com/aws/aws-sdk-go/service/bedrock"
"github.com/aws/aws-sdk-go/service/s3"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/clustersize"
"github.com/mulgadc/spinifex/spinifex/kvstore"
"github.com/mulgadc/spinifex/spinifex/objectstore"
"github.com/nats-io/nats.go/jetstream"
@@ -51,11 +52,15 @@ const (
)
// EnsureInvocationStream idempotently creates (or updates) the invocation
// stream. Safe to call from every gateway node at boot. replicas must match
// the cluster's node count: at replicas=1, losing the one node holding this
// stream loses billing/audit records even though the control plane survives
// on the others, defeating the point of using JetStream over a lossy buffer.
func EnsureInvocationStream(ctx context.Context, js jetstream.JetStream, replicas int) (jetstream.Stream, error) {
// stream, replicated across the cluster. Safe to call from every gateway node
// at boot. On one replica, losing the one node holding this stream loses
// billing and audit records even though the control plane survives on the
// others, defeating the point of using JetStream over a lossy buffer.
func EnsureInvocationStream(ctx context.Context, js jetstream.JetStream) (jetstream.Stream, error) {
replicas, err := clustersize.Replicas()
if err != nil {
return nil, fmt.Errorf("ensure invocation stream: %w", err)
}
stream, err := js.CreateOrUpdateStream(ctx, jetstream.StreamConfig{
Name: InvocationStreamName,
Subjects: []string{InvocationStreamSubject},
@@ -310,12 +315,11 @@ type LoggingConfigStore struct {
var _ LoggingConfigReader = (*LoggingConfigStore)(nil)
// NewLoggingConfigStore constructs a LoggingConfigStore.
func NewLoggingConfigStore(js jetstream.JetStream, replicas int) *LoggingConfigStore {
func NewLoggingConfigStore(js jetstream.JetStream) *LoggingConfigStore {
return &LoggingConfigStore{store: kvstore.New[LoggingConfig](js, kvstore.Config{
Name: bedrockLoggingConfigBucket,
History: bedrockLoggingConfigHistory,
Replicas: replicas,
Missing: "bedrock: logging config store has no JetStream client configured",
Name: bedrockLoggingConfigBucket,
History: bedrockLoggingConfigHistory,
Missing: "bedrock: logging config store has no JetStream client configured",
})}
}
+9 -9
View File
@@ -33,7 +33,7 @@ func (s stubLoggingConfigReader) Get(context.Context, string) (LoggingConfig, bo
func TestLoggingConfigStore_PutGetDelete(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewLoggingConfigStore(js, 1)
store := NewLoggingConfigStore(js)
ctx := context.Background()
_, ok, err := store.Get(ctx, "000000000001")
@@ -59,7 +59,7 @@ func TestLoggingConfigStore_PutGetDelete(t *testing.T) {
func TestPutModelInvocationLoggingConfiguration_RejectsCloudWatch(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewLoggingConfigStore(js, 1)
store := NewLoggingConfigStore(js)
input := &bedrock.PutModelInvocationLoggingConfigurationInput{
LoggingConfig: &bedrock.LoggingConfig{
@@ -74,7 +74,7 @@ func TestPutModelInvocationLoggingConfiguration_RejectsCloudWatch(t *testing.T)
func TestPutModelInvocationLoggingConfiguration_RequiresBucketWhenDeliveryEnabled(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewLoggingConfigStore(js, 1)
store := NewLoggingConfigStore(js)
input := &bedrock.PutModelInvocationLoggingConfigurationInput{
LoggingConfig: &bedrock.LoggingConfig{TextDataDeliveryEnabled: aws.Bool(true)},
@@ -89,7 +89,7 @@ func TestPutModelInvocationLoggingConfiguration_RequiresBucketWhenDeliveryEnable
// undeliverable as text or image, so it must not slip past the same gate.
func TestPutModelInvocationLoggingConfiguration_RequiresBucketForEmbeddingDelivery(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewLoggingConfigStore(js, 1)
store := NewLoggingConfigStore(js)
input := &bedrock.PutModelInvocationLoggingConfigurationInput{
LoggingConfig: &bedrock.LoggingConfig{EmbeddingDataDeliveryEnabled: aws.Bool(true)},
@@ -105,7 +105,7 @@ func TestPutModelInvocationLoggingConfiguration_RequiresBucketForEmbeddingDelive
// internal LoggingConfig.
func TestModelInvocationLoggingConfiguration_RoundTripsViaSDKStructs(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewLoggingConfigStore(js, 1)
store := NewLoggingConfigStore(js)
ctx := context.Background()
putInput := &bedrock.PutModelInvocationLoggingConfigurationInput{
@@ -159,7 +159,7 @@ func TestStreamRecorder_DropsAndCountsWhenStreamMissing(t *testing.T) {
func TestStreamRecorder_Record_BodyOnlyWhenLoggingEnabled(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
ctx := context.Background()
_, err := EnsureInvocationStream(ctx, js, 1)
_, err := EnsureInvocationStream(ctx, js)
require.NoError(t, err)
consumer, err := EnsureDeliveryConsumer(ctx, js)
require.NoError(t, err)
@@ -198,7 +198,7 @@ func TestDeliveryConsumer_Run_WritesS3AndNeverLogsBodyText(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
_, err := EnsureInvocationStream(ctx, js, 1)
_, err := EnsureInvocationStream(ctx, js)
require.NoError(t, err)
consumer, err := EnsureDeliveryConsumer(ctx, js)
require.NoError(t, err)
@@ -247,7 +247,7 @@ func TestDeliveryConsumer_Run_DeliversMetadataOnlyRecordToS3(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
_, err := EnsureInvocationStream(ctx, js, 1)
_, err := EnsureInvocationStream(ctx, js)
require.NoError(t, err)
consumer, err := EnsureDeliveryConsumer(ctx, js)
require.NoError(t, err)
@@ -294,7 +294,7 @@ func TestDeliveryConsumer_Run_DeliversMetadataOnlyRecordToS3(t *testing.T) {
// JetStream client: every accessor must report the misconfiguration rather
// than panic on a nil handle.
func TestLoggingConfigStore_RequiresJetStream(t *testing.T) {
store := NewLoggingConfigStore(nil, 1)
store := NewLoggingConfigStore(nil)
ctx := context.Background()
_, _, err := store.Get(ctx, "000000000001")
+5 -6
View File
@@ -76,13 +76,12 @@ type PriceStore struct {
var _ PriceResolver = (*PriceStore)(nil)
// NewPriceStore constructs a PriceStore over the cluster's JetStream client,
// replicated across replicas nodes.
func NewPriceStore(js jetstream.JetStream, replicas int) *PriceStore {
func NewPriceStore(js jetstream.JetStream) *PriceStore {
return &PriceStore{store: kvstore.New[Price](js, kvstore.Config{
Name: bedrockPricesBucket,
History: bedrockPricesHistory,
Replicas: replicas,
Missing: "bedrock: price store has no JetStream client configured",
Name: bedrockPricesBucket,
History: bedrockPricesHistory,
Missing: "bedrock: price store has no JetStream client configured",
})}
}
+2 -2
View File
@@ -94,7 +94,7 @@ func TestResolvePrice_ResolverErrorPropagates(t *testing.T) {
// TestWeightsStore_PutAndResolve_KV.
func TestPriceStore_PutGetDelete(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewPriceStore(testutil.NewJetStream(t, nc), 1)
store := NewPriceStore(testutil.NewJetStream(t, nc))
ctx := context.Background()
_, ok, err := store.Resolve(ctx, priceTestProviderModelID)
@@ -123,7 +123,7 @@ func TestPriceStore_PutGetDelete(t *testing.T) {
// JetStream client: every accessor must report the misconfiguration rather
// than panic on a nil handle.
func TestPriceStore_RequiresJetStream(t *testing.T) {
store := NewPriceStore(nil, 1)
store := NewPriceStore(nil)
ctx := context.Background()
_, _, err := store.Resolve(ctx, priceTestProviderModelID)
+6 -8
View File
@@ -81,16 +81,14 @@ type ProvisionedStore struct {
endpoint EndpointProvisioner
}
// NewProvisionedStore constructs a ProvisionedStore over js, replicated
// across replicas nodes, driving endpoint launches/teardowns through
// endpoint.
func NewProvisionedStore(js jetstream.JetStream, replicas int, region string, endpoint EndpointProvisioner) *ProvisionedStore {
// NewProvisionedStore constructs a ProvisionedStore over js, driving endpoint
// launches and teardowns through endpoint.
func NewProvisionedStore(js jetstream.JetStream, region string, endpoint EndpointProvisioner) *ProvisionedStore {
return &ProvisionedStore{
store: kvstore.New[ProvisionedModelRecord](js, kvstore.Config{
Name: bedrockProvisionedBucket,
History: bedrockProvisionedHistory,
Replicas: replicas,
Missing: "bedrock: provisioned store has no JetStream client configured",
Name: bedrockProvisionedBucket,
History: bedrockProvisionedHistory,
Missing: "bedrock: provisioned store has no JetStream client configured",
}),
region: region,
endpoint: endpoint,
+1 -1
View File
@@ -103,7 +103,7 @@ const (
func newProvisionedTestStore(t *testing.T, endpoint EndpointProvisioner) *ProvisionedStore {
t.Helper()
_, _, js := testutil.StartTestJetStream(t)
return NewProvisionedStore(js, 1, ptTestRegion, endpoint)
return NewProvisionedStore(js, ptTestRegion, endpoint)
}
func createInput(modelID, name string, units int64) *bedrock.CreateProvisionedModelThroughputInput {
+2 -3
View File
@@ -98,13 +98,12 @@ type UsageStore struct {
var _ UsageReader = (*UsageStore)(nil)
// NewUsageStore constructs a UsageStore over the cluster's JetStream client,
// replicated across replicas nodes.
func NewUsageStore(js jetstream.JetStream, replicas int) *UsageStore {
func NewUsageStore(js jetstream.JetStream) *UsageStore {
return &UsageStore{
store: kvstore.New[UsageCounters](js, kvstore.Config{
Name: bedrockUsageBucket,
History: bedrockUsageHistory,
Replicas: replicas,
Missing: "bedrock: usage store has no JetStream client configured",
Attempts: usageCASRetries,
Exhausted: usageCASExhausted,
+6 -6
View File
@@ -21,7 +21,7 @@ const (
func newTestUsageStore(t *testing.T) *UsageStore {
t.Helper()
_, nc, _ := testutil.StartTestJetStream(t)
return NewUsageStore(testutil.NewJetStream(t, nc), 1)
return NewUsageStore(testutil.NewJetStream(t, nc))
}
// TestMicroUSDCost_Arithmetic covers D12's integer accrual, including
@@ -184,7 +184,7 @@ func TestUsageConsumer_Run_RedeliveredRecordDoesNotDoubleCount(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
_, err := EnsureInvocationStream(ctx, js, 1)
_, err := EnsureInvocationStream(ctx, js)
require.NoError(t, err)
consumer, err := EnsureUsageConsumer(ctx, js)
require.NoError(t, err)
@@ -206,7 +206,7 @@ func TestUsageConsumer_Run_RedeliveredRecordDoesNotDoubleCount(t *testing.T) {
_, err = js.Publish(ctx, InvocationStreamSubject, payload)
require.NoError(t, err)
usage := NewUsageStore(js, 1)
usage := NewUsageStore(js)
uc := NewUsageConsumer(usage, nil)
done := make(chan struct{})
go func() {
@@ -243,7 +243,7 @@ func TestUsageConsumer_Run_UnknownModelStillRecordsTokensCostUnknown(t *testing.
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
_, err := EnsureInvocationStream(ctx, js, 1)
_, err := EnsureInvocationStream(ctx, js)
require.NoError(t, err)
consumer, err := EnsureUsageConsumer(ctx, js)
require.NoError(t, err)
@@ -258,7 +258,7 @@ func TestUsageConsumer_Run_UnknownModelStillRecordsTokensCostUnknown(t *testing.
_, err = js.Publish(ctx, InvocationStreamSubject, payload)
require.NoError(t, err)
usage := NewUsageStore(js, 1)
usage := NewUsageStore(js)
uc := NewUsageConsumer(usage, nil)
done := make(chan struct{})
go func() {
@@ -293,7 +293,7 @@ func TestBillingPeriod_MonthlyFormat(t *testing.T) {
// JetStream client: both the counter bucket and the dedupe bucket must report
// the misconfiguration rather than panic on a nil handle.
func TestUsageStore_RequiresJetStream(t *testing.T) {
store := NewUsageStore(nil, 1)
store := NewUsageStore(nil)
ctx := context.Background()
_, _, err := store.Get(ctx, "000000000001", "m", "2026-08")
+5 -6
View File
@@ -57,13 +57,12 @@ type WeightsStore struct {
var _ WeightsResolver = (*WeightsStore)(nil)
// NewWeightsStore constructs a WeightsStore over the cluster's JetStream
// client, replicated across replicas nodes.
func NewWeightsStore(js jetstream.JetStream, replicas int) *WeightsStore {
// client.
func NewWeightsStore(js jetstream.JetStream) *WeightsStore {
return &WeightsStore{store: kvstore.New[weightsRecord](js, kvstore.Config{
Name: bedrockWeightsBucket,
History: bedrockWeightsHistory,
Replicas: replicas,
Missing: "bedrock: weights store has no JetStream client configured",
Name: bedrockWeightsBucket,
History: bedrockWeightsHistory,
Missing: "bedrock: weights store has no JetStream client configured",
})}
}
+12 -12
View File
@@ -20,7 +20,7 @@ func TestNoopWeightsResolver_ResolvesNothing(t *testing.T) {
// snapshot: ("", false, nil), not an error.
func TestWeightsStore_ResolveMiss(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
store := NewWeightsStore(testutil.NewJetStream(t, nc))
snapshotID, ok, err := store.Resolve(context.Background(), "meta.llama3-2-1b-instruct-v1:0")
require.NoError(t, err)
@@ -33,7 +33,7 @@ func TestWeightsStore_ResolveMiss(t *testing.T) {
// Resolve.
func TestWeightsStore_PutAndResolve_KV(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
store := NewWeightsStore(testutil.NewJetStream(t, nc))
ctx := context.Background()
require.NoError(t, store.PutWeights(ctx, "meta.llama3-2-1b-instruct-v1:0", "s3://models/llama-3.2-1b/", "snap-0001"))
@@ -49,7 +49,7 @@ func TestWeightsStore_PutAndResolve_KV(t *testing.T) {
// history.
func TestWeightsStore_PutOverwritesPrevious(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
store := NewWeightsStore(testutil.NewJetStream(t, nc))
ctx := context.Background()
require.NoError(t, store.PutWeights(ctx, "meta.llama3-2-1b-instruct-v1:0", "s3://models/llama-3.2-1b/", "snap-0001"))
@@ -66,7 +66,7 @@ func TestWeightsStore_PutOverwritesPrevious(t *testing.T) {
// reports ok=false, not an error.
func TestWeightsStore_GetWeights_Miss(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
store := NewWeightsStore(testutil.NewJetStream(t, nc))
entry, ok, err := store.GetWeights(context.Background(), "meta.llama3-2-1b-instruct-v1:0")
require.NoError(t, err)
@@ -80,7 +80,7 @@ func TestWeightsStore_GetWeights_Miss(t *testing.T) {
// snapshot it is about to replace.
func TestWeightsStore_GetWeights_ReturnsSourceURIAndSnapshot(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
store := NewWeightsStore(testutil.NewJetStream(t, nc))
ctx := context.Background()
require.NoError(t, store.PutWeights(ctx, "meta.llama3-2-1b-instruct-v1:0", "s3://models/llama-3.2-1b/", "snap-0001"))
@@ -100,7 +100,7 @@ func TestWeightsStore_GetWeights_ReturnsSourceURIAndSnapshot(t *testing.T) {
// both GetWeights and ListWeights.
func TestWeightsStore_PutWeightsWithRevision_RoundTrips(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
store := NewWeightsStore(testutil.NewJetStream(t, nc))
ctx := context.Background()
require.NoError(t, store.PutWeightsWithRevision(ctx, "meta.llama3-2-1b-instruct-v1:0", "s3://ochre-weights/meta-llama/Llama-3.2-1B-Instruct/abc123/", "snap-0001", "abc123def456"))
@@ -121,7 +121,7 @@ func TestWeightsStore_PutWeightsWithRevision_RoundTrips(t *testing.T) {
// SourceRevision, not fail or fabricate one.
func TestWeightsStore_PutWeights_LeavesSourceRevisionEmpty(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
store := NewWeightsStore(testutil.NewJetStream(t, nc))
ctx := context.Background()
require.NoError(t, store.PutWeights(ctx, "meta.llama3-2-1b-instruct-v1:0", "s3://models/llama-3.2-1b/", "snap-0001"))
@@ -136,7 +136,7 @@ func TestWeightsStore_PutWeights_LeavesSourceRevisionEmpty(t *testing.T) {
// bucket must report an empty list, not an error (jetstream.ErrNoKeysFound).
func TestWeightsStore_ListWeights_Empty(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
store := NewWeightsStore(testutil.NewJetStream(t, nc))
entries, err := store.ListWeights(context.Background())
require.NoError(t, err)
@@ -148,7 +148,7 @@ func TestWeightsStore_ListWeights_Empty(t *testing.T) {
// base64url KV key), its source URI and its snapshot ID, sorted by model ID.
func TestWeightsStore_ListWeights_ReturnsAllStaged(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
store := NewWeightsStore(testutil.NewJetStream(t, nc))
ctx := context.Background()
require.NoError(t, store.PutWeights(ctx, "meta.llama3-2-1b-instruct-v1:0", "s3://models/llama-3.2-1b/", "snap-0001"))
@@ -167,7 +167,7 @@ func TestWeightsStore_ListWeights_ReturnsAllStaged(t *testing.T) {
// DeleteWeights drops the record and Resolve reports not-found afterwards.
func TestWeightsStore_DeleteWeights_RemovesEntry(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
store := NewWeightsStore(testutil.NewJetStream(t, nc))
ctx := context.Background()
require.NoError(t, store.PutWeights(ctx, "meta.llama3-2-1b-instruct-v1:0", "s3://models/llama-3.2-1b/", "snap-0001"))
@@ -184,7 +184,7 @@ func TestWeightsStore_DeleteWeights_RemovesEntry(t *testing.T) {
// must not error on a missing key either.
func TestWeightsStore_DeleteWeights_NeverStagedIsNotAnError(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
store := NewWeightsStore(testutil.NewJetStream(t, nc))
require.NoError(t, store.DeleteWeights(context.Background(), "meta.llama3-2-1b-instruct-v1:0"))
}
@@ -208,7 +208,7 @@ func TestSetWeightsResolver_NilRestoresNoop(t *testing.T) {
// JetStream client: every accessor must report the misconfiguration rather
// than panic on a nil handle.
func TestWeightsStore_RequiresJetStream(t *testing.T) {
store := NewWeightsStore(nil, 1)
store := NewWeightsStore(nil)
ctx := context.Background()
_, _, err := store.Resolve(ctx, "meta.llama3-2-1b-instruct-v1:0")
+18 -18
View File
@@ -41,13 +41,13 @@ func TestLoadOrCreateSigningKey_CreatesThenReloads(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
key1, verify1, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key1, verify1, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
require.NotEmpty(t, key1.Kid)
require.Contains(t, verify1, key1.Kid)
// Second call must reload the same persisted key, not mint a new one.
key2, verify2, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key2, verify2, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
assert.Equal(t, key1.Kid, key2.Kid, "persisted signing key must be reused")
assert.Len(t, verify2, 1)
@@ -60,7 +60,7 @@ func TestLoadOrCreateSigningKey_SkipsAnUnreadableKey(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
good, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
good, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
kv, err := js.KeyValue(t.Context(), SigningBucket)
@@ -68,7 +68,7 @@ func TestLoadOrCreateSigningKey_SkipsAnUnreadableKey(t *testing.T) {
_, err = kv.Put(t.Context(), signingKeyName("stale-kid"), []byte("not decryptable with this master key"))
require.NoError(t, err)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
assert.Equal(t, good.Kid, key.Kid, "the readable key must still be active")
assert.Len(t, verify, 1, "an unreadable key contributes no verify key")
@@ -80,14 +80,14 @@ func TestLoadOrCreateSigningKey_SkipsAKeyDeletedWhileReading(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
good, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
good, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
kv, err := js.KeyValue(t.Context(), SigningBucket)
require.NoError(t, err)
require.NoError(t, kv.Delete(t.Context(), signingKeyName("never-existed")))
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
assert.Equal(t, good.Kid, key.Kid)
assert.Len(t, verify, 1)
@@ -100,14 +100,14 @@ func TestLoadOrCreateSigningKey_FailsWhenNoKeyDecrypts(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
_, err := openSigningBucket(t.Context(), js, testMasterKey, 1)
_, err := openSigningBucket(t.Context(), js, testMasterKey)
require.NoError(t, err)
kv, err := js.KeyValue(t.Context(), SigningBucket)
require.NoError(t, err)
_, err = kv.Put(t.Context(), signingKeyName("stale-kid"), []byte("not decryptable with this master key"))
require.NoError(t, err)
_, _, err = LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
_, _, err = LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.Error(t, err)
assert.Contains(t, err.Error(), "master key does not match")
@@ -119,7 +119,7 @@ func TestLoadOrCreateSigningKey_FailsWhenNoKeyDecrypts(t *testing.T) {
func TestLoadOrCreateSigningKey_StoresEncrypted(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
key, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
kv, err := js.KeyValue(t.Context(), SigningBucket)
@@ -133,14 +133,14 @@ func TestLoadOrCreateSigningKey_StoresEncrypted(t *testing.T) {
func TestLoadOrCreateSigningKey_EmptyMasterKeyRejected(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
_, _, err := LoadOrCreateSigningKey(t.Context(), js, nil, 1)
_, _, err := LoadOrCreateSigningKey(t.Context(), js, nil)
require.Error(t, err)
}
func TestIssuerVerifier_RoundTrip(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
iss := NewIssuer(key, testAudience)
@@ -158,7 +158,7 @@ func TestIssuerVerifier_RoundTrip(t *testing.T) {
func TestIssuer_MintRequiresAccountID(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
key, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
p := samplePrincipal()
@@ -174,7 +174,7 @@ func TestIssuer_MintRequiresAccountID(t *testing.T) {
func TestMint_RequiresCompleteIdentityPointer(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
key, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
iss := NewIssuer(key, testAudience)
@@ -208,7 +208,7 @@ func TestSupportedPrincipalType(t *testing.T) {
func TestVerifier_RejectsWrongAudience(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
tok, _, err := NewIssuer(key, testAudience).Mint(samplePrincipal())
@@ -221,7 +221,7 @@ func TestVerifier_RejectsWrongAudience(t *testing.T) {
func TestVerifier_RejectsUnknownKid(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
key, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key, _, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
tok, _, err := NewIssuer(key, testAudience).Mint(samplePrincipal())
@@ -235,7 +235,7 @@ func TestVerifier_RejectsUnknownKid(t *testing.T) {
func TestVerifier_RejectsExpiredToken(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
// Hand-mint an already-expired token with the same signing key.
@@ -263,7 +263,7 @@ func TestVerifier_RejectsExpiredToken(t *testing.T) {
func TestVerifier_RejectsIncompleteIdentityPointer(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
base := func() Claims {
@@ -310,7 +310,7 @@ func TestVerifier_RejectsIncompleteIdentityPointer(t *testing.T) {
func TestVerifier_RejectsNonES256(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
_, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
_, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
// HS256 token must be refused regardless of signature.
+2 -2
View File
@@ -78,8 +78,8 @@ type Rotator struct {
// NewRotator opens the signing-key bucket and builds a rotator that keeps issuer
// and verifier current. replicas matches the cluster size for first-run bucket
// creation.
func NewRotator(ctx context.Context, js jetstream.JetStream, masterKey []byte, replicas int, issuer *Issuer, verifier *Verifier) (*Rotator, error) {
kv, err := openSigningBucket(ctx, js, masterKey, replicas)
func NewRotator(ctx context.Context, js jetstream.JetStream, masterKey []byte, issuer *Issuer, verifier *Verifier) (*Rotator, error) {
kv, err := openSigningBucket(ctx, js, masterKey)
if err != nil {
return nil, err
}
+8 -8
View File
@@ -79,12 +79,12 @@ func TestPlanRotation(t *testing.T) {
func TestRotateOnce_MintsAndSwaps(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
issuer := NewIssuer(key, testAudience)
verifier := NewVerifier(verify, testAudience)
kv, err := openSigningBucket(t.Context(), js, testMasterKey, 1)
kv, err := openSigningBucket(t.Context(), js, testMasterKey)
require.NoError(t, err)
// Token minted under the original key before rotation.
@@ -120,9 +120,9 @@ func TestRotateOnce_MintsAndSwaps(t *testing.T) {
func TestRotateOnce_PrunesExpiredKey(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
first, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
first, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
kv, err := openSigningBucket(t.Context(), js, testMasterKey, 1)
kv, err := openSigningBucket(t.Context(), js, testMasterKey)
require.NoError(t, err)
second, err := generateSigningKey(t.Context(), kv, testMasterKey)
require.NoError(t, err)
@@ -151,11 +151,11 @@ func TestRotateOnce_PrunesExpiredKey(t *testing.T) {
func TestRotator_RunRotatesThenStops(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
key, verify, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
issuer := NewIssuer(key, testAudience)
verifier := NewVerifier(verify, testAudience)
kv, err := openSigningBucket(t.Context(), js, testMasterKey, 1)
kv, err := openSigningBucket(t.Context(), js, testMasterKey)
require.NoError(t, err)
r := &Rotator{
@@ -185,9 +185,9 @@ func TestRotator_RunRotatesThenStops(t *testing.T) {
func TestIssuerVerifier_ConcurrentSwap(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
k1, v1, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey, 1)
k1, v1, err := LoadOrCreateSigningKey(t.Context(), js, testMasterKey)
require.NoError(t, err)
kv, err := openSigningBucket(t.Context(), js, testMasterKey, 1)
kv, err := openSigningBucket(t.Context(), js, testMasterKey)
require.NoError(t, err)
k2, err := generateSigningKey(t.Context(), kv, testMasterKey)
require.NoError(t, err)
+4 -4
View File
@@ -76,8 +76,8 @@ func kidFor(pub *ecdsa.PublicKey) (string, error) {
// returns the active signing key. On first run it generates and persists one.
// The active key is the newest by KV creation time (ties broken by kid), so all
// nodes converge on the same signer and a freshly rotated key takes over.
func LoadOrCreateSigningKey(ctx context.Context, js jetstream.JetStream, masterKey []byte, replicas int) (*SigningKey, map[string]*ecdsa.PublicKey, error) {
kv, err := openSigningBucket(ctx, js, masterKey, replicas)
func LoadOrCreateSigningKey(ctx context.Context, js jetstream.JetStream, masterKey []byte) (*SigningKey, map[string]*ecdsa.PublicKey, error) {
kv, err := openSigningBucket(ctx, js, masterKey)
if err != nil {
return nil, nil, err
}
@@ -104,14 +104,14 @@ func LoadOrCreateSigningKey(ctx context.Context, js jetstream.JetStream, masterK
// openSigningBucket validates inputs and returns the awsgw-keys KV handle,
// creating the cluster-replicated bucket on first use.
func openSigningBucket(ctx context.Context, js jetstream.JetStream, masterKey []byte, replicas int) (jetstream.KeyValue, error) {
func openSigningBucket(ctx context.Context, js jetstream.JetStream, masterKey []byte) (jetstream.KeyValue, error) {
if js == nil {
return nil, errors.New("ecrauth: nil JetStream context")
}
if len(masterKey) == 0 {
return nil, errors.New("ecrauth: empty master key")
}
return kvutil.GetOrCreateBucketWithReplicas(ctx, js, SigningBucket, signingKeyHistory, replicas)
return kvutil.GetOrCreateBucket(ctx, js, SigningBucket, signingKeyHistory)
}
// reloadKeys reads every stored signing key, returning the verify set, each
+1 -1
View File
@@ -34,7 +34,7 @@ func newECRAuth(t *testing.T) (*gateway_ecrauth.Issuer, *gateway_ecrauth.Verifie
t.Helper()
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
key, verify, err := gateway_ecrauth.LoadOrCreateSigningKey(t.Context(), js, ecrTestMasterKey(t), 1)
key, verify, err := gateway_ecrauth.LoadOrCreateSigningKey(t.Context(), js, ecrTestMasterKey(t))
require.NoError(t, err)
return gateway_ecrauth.NewIssuer(key, ecrTestAudience), gateway_ecrauth.NewVerifier(verify, ecrTestAudience)
}
+1 -1
View File
@@ -23,7 +23,7 @@ func seedDiscoveryCluster(t *testing.T) http.Handler {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
kv, err := handlers_eks.GetOrCreateAccountBucket(t.Context(), js, testDiscAccount, 1)
kv, err := handlers_eks.GetOrCreateAccountBucket(t.Context(), js, testDiscAccount)
if err != nil {
t.Fatalf("account bucket: %v", err)
}
@@ -48,7 +48,7 @@ func accessGateway(t *testing.T) *GatewayConfig {
_, _, js := testutil.StartTestJetStream(t)
return &GatewayConfig{
DisableLogging: true,
BedrockAccessAdmin: gateway_bedrock.NewModelAccessStore(js, 1),
BedrockAccessAdmin: gateway_bedrock.NewModelAccessStore(js),
IAMService: allowAllIAMService(),
}
}
+1 -1
View File
@@ -90,7 +90,7 @@ func NewService(nc *nats.Conn, deps ServiceDeps) *Service {
return &Service{
nc: nc,
deps: deps,
store: newEndpointStore(js, deps.Replicas),
store: newEndpointStore(js),
leader: newLeaderBucket(js),
}
}
+6 -8
View File
@@ -66,15 +66,13 @@ type endpointStore struct {
*kvstore.Store[EndpointRecord]
}
// newEndpointStore returns the endpoints store over js, replicated across
// replicas nodes. A nil js is permitted: every accessor then reports
// missingJetStream rather than panicking.
func newEndpointStore(js jetstream.JetStream, replicas int) *endpointStore {
// newEndpointStore returns the endpoints store over js. A nil js is permitted:
// every accessor then reports missingJetStream rather than panicking.
func newEndpointStore(js jetstream.JetStream) *endpointStore {
return &endpointStore{Store: kvstore.New[EndpointRecord](js, kvstore.Config{
Name: KVBucketEndpoints,
History: KVBucketEndpointsHistory,
Replicas: replicas,
Missing: missingJetStream,
Name: KVBucketEndpoints,
History: KVBucketEndpointsHistory,
Missing: missingJetStream,
})}
}
+1 -1
View File
@@ -17,7 +17,7 @@ func TestEndpointKey_Base64EncodesColonBearingModelID(t *testing.T) {
func newTestBucket(t *testing.T) *endpointStore {
t.Helper()
_, _, js := testutil.StartTestJetStream(t)
return newEndpointStore(js, 1)
return newEndpointStore(js)
}
func TestStore_CreateGetUpdateDelete(t *testing.T) {
+4 -11
View File
@@ -84,23 +84,16 @@ func NewNatGatewayServiceImplWithNATS(ctx context.Context, natsConn *nats.Conn)
}, nil
}
// getOrCreateDeletedBucket creates the expiring bucket without changing an
// existing bucket's replica or placement configuration.
// getOrCreateDeletedBucket creates the expiring bucket that records a deletion
// for as long as a client may still ask about it.
func getOrCreateDeletedBucket(ctx context.Context, js jetstream.KeyValueManager) (jetstream.KeyValue, error) {
// Terraform polls DescribeNatGateways after delete and expects state=deleted.
kv, err := js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
Bucket: KVBucketDeletedNatGateways,
return kvutil.GetOrCreateBucketWithOptions(ctx, js, kvutil.BucketOptions{
Name: KVBucketDeletedNatGateways,
Description: "Deleted NAT Gateways (auto-expire after 1 hour)",
History: 1,
TTL: time.Hour,
})
if err == nil {
return kv, nil
}
if !errors.Is(err, jetstream.ErrBucketExists) {
return nil, err
}
return js.KeyValue(ctx, KVBucketDeletedNatGateways)
}
// natGatewayEvent is published on vpc.add-nat-gateway / vpc.delete-nat-gateway topics.
@@ -71,24 +71,12 @@ func NewSpotInstanceServiceImplWithNATS(ctx context.Context, cfg *config.Config,
// getOrCreateTerminalBucket opens the terminal bucket without changing its
// existing replica or placement configuration, creating it only when absent.
func getOrCreateTerminalBucket(ctx context.Context, js jetstream.KeyValueManager) (jetstream.KeyValue, error) {
kv, err := js.KeyValue(ctx, KVBucketSpotRequestsTerminal)
if err == nil {
return kv, nil
}
if !errors.Is(err, jetstream.ErrBucketNotFound) {
return nil, err
}
kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
Bucket: KVBucketSpotRequestsTerminal,
return kvutil.GetOrCreateBucketWithOptions(ctx, js, kvutil.BucketOptions{
Name: KVBucketSpotRequestsTerminal,
Description: "Terminal Spot Instance Requests (auto-expire after 1 hour)",
History: 1,
TTL: spotTerminalTTL,
})
if errors.Is(err, jetstream.ErrBucketExists) {
return js.KeyValue(ctx, KVBucketSpotRequestsTerminal)
}
return kv, err
}
// PutSpotInstanceRequests stores each request in the active bucket.
+5 -6
View File
@@ -37,16 +37,15 @@ type StagedAddonManifest struct {
// The VM-side delivery slice applies it via the K3s auto-deploy dir; until then
// the add-on sits CREATING.
type stagingInstaller struct {
nc *nats.Conn
clusterSize int
nc *nats.Conn
}
var _ AddonInstaller = (*stagingInstaller)(nil)
// newStagingInstaller returns a stagingInstaller bound to the daemon NATS
// connection, using clusterSize as the per-account bucket's replica count.
func newStagingInstaller(nc *nats.Conn, clusterSize int) *stagingInstaller {
return &stagingInstaller{nc: nc, clusterSize: clusterSize}
// connection.
func newStagingInstaller(nc *nats.Conn) *stagingInstaller {
return &stagingInstaller{nc: nc}
}
func (i *stagingInstaller) acctKV(ctx context.Context, accountID string) (jetstream.KeyValue, error) {
@@ -57,7 +56,7 @@ func (i *stagingInstaller) acctKV(ctx context.Context, accountID string) (jetstr
if err != nil {
return nil, fmt.Errorf("jetstream: %w", err)
}
return GetOrCreateAccountBucket(ctx, js, accountID, max(i.clusterSize, 1))
return GetOrCreateAccountBucket(ctx, js, accountID)
}
func (i *stagingInstaller) Install(ctx context.Context, accountID, cluster string, rec *AddonRecord) error {
@@ -15,7 +15,7 @@ import (
func acctKVForTest(t *testing.T, svc *EKSServiceImpl) jetstream.KeyValue {
t.Helper()
js := testutil.NewJetStream(t, svc.deps.NATSConn)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
return kv
}
+1 -1
View File
@@ -331,7 +331,7 @@ func (s *EKSServiceImpl) addonInstaller() AddonInstaller {
if s.deps.AddonInstaller != nil {
return s.deps.AddonInstaller
}
return newStagingInstaller(s.deps.NATSConn, s.deps.ClusterSize)
return newStagingInstaller(s.deps.NATSConn)
}
// markAddonFailed best-effort flips a record to CREATE_FAILED with the error reason.
+1 -1
View File
@@ -243,7 +243,7 @@ func TestStagingInstaller_StagesManifest(t *testing.T) {
require.NoError(t, err)
js := testutil.NewJetStream(t, svc.deps.NATSConn)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
entry, err := kv.Get(t.Context(), AddonManifestKey("c1", albController))
require.NoError(t, err, "installer must stage a manifest for VM-side delivery")
@@ -21,9 +21,9 @@ func newStateReconcilerHarness(t *testing.T, opts ...ReconcilerOption) (*Cluster
t.Helper()
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
leaderKV, err := InitLeaderBucket(t.Context(), js, 1)
leaderKV, err := InitLeaderBucket(t.Context(), js)
require.NoError(t, err)
acctKV, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
acctKV, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
require.NoError(t, PutClusterMeta(t.Context(), acctKV, sampleClusterMeta("alpha")))
@@ -58,9 +58,9 @@ func newReconcilerHarness(t *testing.T, healthURL string, opts ...ReconcilerOpti
t.Helper()
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
leaderKV, err := InitLeaderBucket(t.Context(), js, 1)
leaderKV, err := InitLeaderBucket(t.Context(), js)
require.NoError(t, err)
acctKV, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
acctKV, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
require.NoError(t, PutClusterMeta(t.Context(), acctKV, sampleClusterMeta("alpha")))
@@ -72,9 +72,9 @@ func newReconcilerHarness(t *testing.T, healthURL string, opts ...ReconcilerOpti
func TestNewClusterReconciler_EmptyInputsRejected(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
leaderKV, err := InitLeaderBucket(t.Context(), js, 1)
leaderKV, err := InitLeaderBucket(t.Context(), js)
require.NoError(t, err)
acctKV, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
acctKV, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
_, err = NewClusterReconciler(nil, acctKV, testAccountID, "alpha", "h", "")
+1 -1
View File
@@ -15,7 +15,7 @@ func newClusterStateTestKV(t *testing.T) jetstream.KeyValue {
t.Helper()
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
return kv
}
+1 -1
View File
@@ -60,7 +60,7 @@ func newEKSServiceFixture(t *testing.T) *eksServiceFixture {
t.Helper()
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
nlb := newFakeNLBProvisioner()
+1 -1
View File
@@ -62,7 +62,7 @@ func TestDesiredDNSChanges_MetadataReadFailureIsNotAuthoritative(t *testing.T) {
require.NoError(t, PutClusterMeta(t.Context(), fixture.kv, active))
js := testutil.NewJetStream(t, fixture.svc.deps.NATSConn)
corruptKV, err := GetOrCreateAccountBucket(t.Context(), js, "444455556666", 1)
corruptKV, err := GetOrCreateAccountBucket(t.Context(), js, "444455556666")
require.NoError(t, err)
_, err = corruptKV.Put(t.Context(), ClusterMetaKey("unreadable"), []byte("{not json"))
require.NoError(t, err)
+1 -1
View File
@@ -69,7 +69,7 @@ func (r *EKSBillableReaper) Sweep(ctx context.Context) (int, error) {
continue // ENI gone or untagged: cannot confirm orphan-hood, never reap
}
acctKV, err := GetOrCreateAccountBucket(ctx, js, clusterAccount, max(r.svc.deps.ClusterSize, 1))
acctKV, err := GetOrCreateAccountBucket(ctx, js, clusterAccount)
if err != nil {
slog.WarnContext(ctx, "eks-billable: account bucket lookup failed", "account", clusterAccount, "err", err)
continue
+2 -2
View File
@@ -39,7 +39,7 @@ func newBootstrapHarness(t *testing.T) *natsBootstrapHarness {
t.Helper()
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
require.NoError(t, PutClusterMeta(t.Context(), kv, sampleClusterMeta(bootstrapTestCluster)))
@@ -85,7 +85,7 @@ func (h *natsBootstrapHarness) waitSubsBound(base, want int) {
func TestNewNATSBootstrap_RejectsBadInputs(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
_, err = NewNATSBootstrap(nil, kv, bootstrapTestMasterKey, testAccountID, "alpha")
+1 -1
View File
@@ -195,7 +195,7 @@ func (s *EKSServiceImpl) nodegroupAcctKV(ctx context.Context, accountID string)
if err != nil {
return nil, fmt.Errorf("jetstream: %w", err)
}
acctKV, err := GetOrCreateAccountBucket(ctx, js, accountID, max(s.deps.ClusterSize, 1))
acctKV, err := GetOrCreateAccountBucket(ctx, js, accountID)
if err != nil {
return nil, fmt.Errorf("get account bucket: %w", err)
}
+1 -1
View File
@@ -242,7 +242,7 @@ func TestReclaimOrphanedNodegroups_TerminatesStrandedWorkers(t *testing.T) {
func TestNodegroupRecord_CRUDRoundTrip(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
rec := &NodegroupRecord{
@@ -13,7 +13,7 @@ import (
func TestLoadRecoveryDirective_AbsentIsZeroNone(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
acctKV, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
acctKV, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
d, err := LoadRecoveryDirective(t.Context(), acctKV, "alpha", "i-missing")
@@ -25,7 +25,7 @@ func TestLoadRecoveryDirective_AbsentIsZeroNone(t *testing.T) {
func TestStoreRecoveryDirective_EpochIncrements(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
acctKV, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
acctKV, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
first, err := StoreRecoveryDirective(t.Context(), acctKV, "alpha", "i-0", RecoveryActionClusterReset, "", false)
@@ -46,7 +46,7 @@ func TestStoreRecoveryDirective_EpochIncrements(t *testing.T) {
func TestStoreRecoveryDirective_SnapshotRequiredRoundTrips(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
acctKV, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
acctKV, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
stored, err := StoreRecoveryDirective(t.Context(), acctKV, "alpha", "i-seed", RecoveryActionClusterReset, "etcd-frequent-20260709T010000Z.snap", true)
+8 -13
View File
@@ -50,11 +50,6 @@ type EKSServiceDeps struct {
Region string
HolderID string
// ClusterSize is the daemon's node count, used as the JetStream replica
// count for the lazily-created per-account and leader KV buckets so they
// match the cluster's other R3 streams instead of staying stuck at R1.
ClusterSize int
// InternalSuffix is the AWS-parity internal DNS suffix (e.g. spinifex.internal)
// used to compose the worker's ECR registry host.
InternalSuffix string
@@ -305,7 +300,7 @@ func NewEKSServiceImpl(deps EKSServiceDeps) (*EKSServiceImpl, error) {
// The leader bucket outlives every request, so its open is bounded by the
// service's own background context rather than a caller's.
ctx, cancel := context.WithCancel(context.Background())
leaderKV, err := InitLeaderBucket(ctx, js, max(deps.ClusterSize, 1))
leaderKV, err := InitLeaderBucket(ctx, js)
if err != nil {
cancel()
return nil, err
@@ -519,7 +514,7 @@ func (s *EKSServiceImpl) CreateCluster(ctx context.Context, input *eks.CreateClu
if err != nil {
return nil, logCreateErr(name, accountID, "jetstream", err)
}
acctKV, err := GetOrCreateAccountBucket(ctx, js, accountID, max(s.deps.ClusterSize, 1))
acctKV, err := GetOrCreateAccountBucket(ctx, js, accountID)
if err != nil {
return nil, logCreateErr(name, accountID, "get account bucket", err)
}
@@ -1051,7 +1046,7 @@ func (s *EKSServiceImpl) DescribeCluster(ctx context.Context, input *eks.Describ
if err != nil {
return nil, fmt.Errorf("jetstream: %w", err)
}
acctKV, err := GetOrCreateAccountBucket(ctx, js, accountID, max(s.deps.ClusterSize, 1))
acctKV, err := GetOrCreateAccountBucket(ctx, js, accountID)
if err != nil {
return nil, fmt.Errorf("get account bucket: %w", err)
}
@@ -1070,7 +1065,7 @@ func (s *EKSServiceImpl) ListClusters(ctx context.Context, input *eks.ListCluste
if err != nil {
return nil, eksReadUnavailableOr(err, "jetstream")
}
acctKV, err := GetOrCreateAccountBucket(ctx, js, accountID, max(s.deps.ClusterSize, 1))
acctKV, err := GetOrCreateAccountBucket(ctx, js, accountID)
if err != nil {
return nil, eksReadUnavailableOr(err, "get account bucket")
}
@@ -1117,7 +1112,7 @@ func (s *EKSServiceImpl) DeleteCluster(ctx context.Context, input *eks.DeleteClu
if err != nil {
return nil, fmt.Errorf("jetstream: %w", err)
}
acctKV, err := GetOrCreateAccountBucket(ctx, js, accountID, max(s.deps.ClusterSize, 1))
acctKV, err := GetOrCreateAccountBucket(ctx, js, accountID)
if err != nil {
return nil, fmt.Errorf("get account bucket: %w", err)
}
@@ -1536,7 +1531,7 @@ func (s *EKSServiceImpl) acctKVForCluster(ctx context.Context, accountID, cluste
if err != nil {
return nil, fmt.Errorf("jetstream: %w", err)
}
acctKV, err := GetOrCreateAccountBucket(ctx, js, accountID, max(s.deps.ClusterSize, 1))
acctKV, err := GetOrCreateAccountBucket(ctx, js, accountID)
if err != nil {
return nil, fmt.Errorf("get account bucket: %w", err)
}
@@ -2002,7 +1997,7 @@ func (s *EKSServiceImpl) accountBucket(ctx context.Context, accountID string) (j
if err != nil {
return nil, fmt.Errorf("jetstream: %w", err)
}
return GetOrCreateAccountBucket(ctx, js, accountID, max(s.deps.ClusterSize, 1))
return GetOrCreateAccountBucket(ctx, js, accountID)
}
// clusterNameFromARN extracts the cluster name from an EKS cluster ARN
@@ -2232,7 +2227,7 @@ func (s *EKSServiceImpl) spawnReconciler(accountID, clusterName string, _ *Clust
slog.Error("spawnReconciler: jetstream", "err", err)
return
}
acctKV, err := GetOrCreateAccountBucket(s.bgCtx, js, accountID, max(s.deps.ClusterSize, 1))
acctKV, err := GetOrCreateAccountBucket(s.bgCtx, js, accountID)
if err != nil {
slog.Error("spawnReconciler: account bucket", "err", err)
return
+4 -4
View File
@@ -223,7 +223,7 @@ func TestEKSServiceImpl_NodegroupMethodsShimMode(t *testing.T) {
func seedTestCluster(t *testing.T, svc *EKSServiceImpl, cluster string) {
t.Helper()
js := testutil.NewJetStream(t, svc.deps.NATSConn)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
require.NoError(t, PutClusterMeta(t.Context(), kv, &ClusterMeta{Name: cluster, Status: ClusterStatusActive}))
}
@@ -414,7 +414,7 @@ func TestAccessPolicy_DisassociateRemovesProjectedGroup(t *testing.T) {
require.NoError(t, err)
js := testutil.NewJetStream(t, svc.deps.NATSConn)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
rec, err := GetAccessEntryRecord(t.Context(), kv, "c1", testPrincipalARN)
require.NoError(t, err)
@@ -461,7 +461,7 @@ func TestEKSServiceImpl_OIDCMethodsReturnNotImplemented(t *testing.T) {
func TestEKSServiceImpl_ClusterTagRoundTrip(t *testing.T) {
svc := setupTestService(t)
js := testutil.NewJetStream(t, svc.deps.NATSConn)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
const arn = "arn:aws:eks:us-east-1:111122223333:cluster/c1"
@@ -503,7 +503,7 @@ func TestEKSServiceImpl_ClusterTagRoundTrip(t *testing.T) {
func TestEKSServiceImpl_NodegroupTagRoundTrip(t *testing.T) {
svc := setupTestService(t)
js := testutil.NewJetStream(t, svc.deps.NATSConn)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
const arn = "arn:aws:eks:us-east-1:111122223333:nodegroup/c1/ng1/abc123"
+7 -9
View File
@@ -232,9 +232,9 @@ func AccountWatchBuckets(ctx context.Context, nc *nats.Conn) ([]*kvstore.Bucket,
// creating it on first use at the given replica count (clamped to a minimum
// of 1). Idempotent: subsequent calls with the same accountID return the
// existing handle.
func GetOrCreateAccountBucket(ctx context.Context, js jetstream.JetStream, accountID string, replicas int) (jetstream.KeyValue, error) {
func GetOrCreateAccountBucket(ctx context.Context, js jetstream.JetStream, accountID string) (jetstream.KeyValue, error) {
bucket := AccountBucketName(accountID)
kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, js, bucket, KVBucketEKSAccountHistory, replicas)
kv, err := kvutil.GetOrCreateBucket(ctx, js, bucket, KVBucketEKSAccountHistory)
if err != nil {
return nil, fmt.Errorf("failed to create EKS per-account KV bucket %s: %w", bucket, err)
}
@@ -245,13 +245,11 @@ func GetOrCreateAccountBucket(ctx context.Context, js jetstream.JetStream, accou
}
// InitLeaderBucket creates (or attaches to) the shared spinifex-eks-leader
// bucket used for per-cluster reconciler leader-lease CAS locks, at the given
// replica count (clamped to a minimum of 1).
func InitLeaderBucket(ctx context.Context, js jetstream.JetStream, replicas int) (jetstream.KeyValue, error) {
// bucket used for per-cluster reconciler leader-lease CAS locks.
func InitLeaderBucket(ctx context.Context, js jetstream.JetStream) (jetstream.KeyValue, error) {
return kvlease.OpenBucket(ctx, js, kvlease.BucketConfig{
Name: KVBucketEKSLeader,
TTL: KVBucketEKSLeaderTTL,
Replicas: max(replicas, 1),
Version: KVBucketEKSLeaderVersion,
Name: KVBucketEKSLeader,
TTL: KVBucketEKSLeaderTTL,
Version: KVBucketEKSLeaderVersion,
})
}
+6 -6
View File
@@ -14,11 +14,11 @@ func TestGetOrCreateAccountBucket_Idempotent(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
kv1, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
kv1, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
require.NotNil(t, kv1)
kv2, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 1)
kv2, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
require.NotNil(t, kv2)
@@ -36,7 +36,7 @@ func TestGetOrCreateAccountBucket_ReplicasClamped(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID, 0)
kv, err := GetOrCreateAccountBucket(t.Context(), js, testAccountID)
require.NoError(t, err)
require.NotNil(t, kv)
@@ -51,12 +51,12 @@ func TestInitLeaderBucket_Idempotent(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
kv1, err := InitLeaderBucket(t.Context(), js, 1)
kv1, err := InitLeaderBucket(t.Context(), js)
require.NoError(t, err)
require.NotNil(t, kv1)
assert.Equal(t, KVBucketEKSLeader, kv1.Bucket())
kv2, err := InitLeaderBucket(t.Context(), js, 1)
kv2, err := InitLeaderBucket(t.Context(), js)
require.NoError(t, err)
require.NotNil(t, kv2)
assert.Equal(t, KVBucketEKSLeader, kv2.Bucket())
@@ -68,7 +68,7 @@ func TestInitLeaderBucket_ReplicasClamped(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
kv, err := InitLeaderBucket(t.Context(), js, -1)
kv, err := InitLeaderBucket(t.Context(), js)
require.NoError(t, err)
require.NotNil(t, kv)
+3 -3
View File
@@ -18,9 +18,9 @@ func TestAccountWatchBuckets_ReturnsOnePerAccountBucket(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
_, err := handlers_eks.GetOrCreateAccountBucket(t.Context(), js, "111111111111", 1)
_, err := handlers_eks.GetOrCreateAccountBucket(t.Context(), js, "111111111111")
require.NoError(t, err)
_, err = handlers_eks.GetOrCreateAccountBucket(t.Context(), js, "222222222222", 1)
_, err = handlers_eks.GetOrCreateAccountBucket(t.Context(), js, "222222222222")
require.NoError(t, err)
_, err = kvutil.GetOrCreateBucket(t.Context(), js, "something-else", 1)
require.NoError(t, err)
@@ -51,7 +51,7 @@ func TestAccountWatchBuckets_NoAccountsIsNotAnError(t *testing.T) {
func TestAccountWatchBuckets_BucketsAreWatchable(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
js := testutil.NewJetStream(t, nc)
_, err := handlers_eks.GetOrCreateAccountBucket(t.Context(), js, "111111111111", 1)
_, err := handlers_eks.GetOrCreateAccountBucket(t.Context(), js, "111111111111")
require.NoError(t, err)
buckets, err := handlers_eks.AccountWatchBuckets(t.Context(), nc)
@@ -16,7 +16,7 @@ func newLifecycleService(t *testing.T) handlers_iam.IAMService {
masterKey, err := handlers_iam.GenerateMasterKey()
require.NoError(t, err)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey, 1)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey)
require.NoError(t, err)
return svc
}
+1 -1
View File
@@ -407,7 +407,7 @@ func TestCreateAccount_PublishesEvent(t *testing.T) {
masterKey, err := GenerateMasterKey()
require.NoError(t, err)
svc, err := NewIAMServiceImpl(t.Context(), nc, masterKey, 1)
svc, err := NewIAMServiceImpl(t.Context(), nc, masterKey)
require.NoError(t, err)
eventCh := make(chan *nats.Msg, 1)
@@ -67,7 +67,7 @@ func TestNewIAMServiceImpl_PropagatesBuiltinManagedPolicyParseError(t *testing.T
builtinManagedPolicyParseErr = parseErr
t.Cleanup(func() { builtinManagedPolicyParseErr = previousErr })
_, err := NewIAMServiceImpl(t.Context(), nil, make([]byte, 32), 1)
_, err := NewIAMServiceImpl(t.Context(), nil, make([]byte, 32))
require.ErrorIs(t, err, parseErr)
assert.ErrorContains(t, err, "init builtin managed policies")
}
+3 -3
View File
@@ -50,8 +50,8 @@ func OIDCProviderARN(accountID, issuerHostPath string) string {
}
// GetOrCreateIAMAccountBucket opens the per-account IAM bucket, creating it on first use.
func GetOrCreateIAMAccountBucket(ctx context.Context, js jetstream.JetStream, accountID string, replicas int) (jetstream.KeyValue, error) {
return kvutil.GetOrCreateBucketWithReplicas(ctx, js, IAMAccountBucketName(accountID), 1, max(replicas, 1))
func GetOrCreateIAMAccountBucket(ctx context.Context, js jetstream.JetStream, accountID string) (jetstream.KeyValue, error) {
return kvutil.GetOrCreateBucket(ctx, js, IAMAccountBucketName(accountID), 1)
}
// OIDCProviderRecord is the stored shape of a registered OIDC identity provider.
@@ -141,7 +141,7 @@ func (s *IAMServiceImpl) CreateOpenIDConnectProvider(accountID string, input *ia
return nil, fmt.Errorf("marshal OIDC provider: %w", err)
}
kv, err := GetOrCreateIAMAccountBucket(ctx, s.js, accountID, s.replicas)
kv, err := GetOrCreateIAMAccountBucket(ctx, s.js, accountID)
if err != nil {
return nil, fmt.Errorf("open IAM account bucket: %w", err)
}
+1 -1
View File
@@ -20,7 +20,7 @@ func newIAMService(t *testing.T) *handlers_iam.IAMServiceImpl {
masterKey, err := handlers_iam.GenerateMasterKey()
require.NoError(t, err)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey, 1)
svc, err := handlers_iam.NewIAMServiceImpl(t.Context(), nc, masterKey)
require.NoError(t, err)
return svc
}
+12 -3
View File
@@ -3,10 +3,12 @@ package handlers_iam
import (
"context"
"fmt"
"github.com/mulgadc/spinifex/spinifex/otelsetup"
"log/slog"
"time"
"github.com/mulgadc/spinifex/spinifex/clustersize"
"github.com/mulgadc/spinifex/spinifex/otelsetup"
"github.com/nats-io/nats.go"
)
@@ -16,7 +18,7 @@ import (
// leave the service nil for the process lifetime. Blocks up to maxWait, then
// returns the last error. Callers that legitimately run without a master key
// must guard the call themselves.
func NewIAMServiceWithRetry(ctx context.Context, natsConn *nats.Conn, masterKey []byte, clusterSize int) (*IAMServiceImpl, error) {
func NewIAMServiceWithRetry(ctx context.Context, natsConn *nats.Conn, masterKey []byte) (*IAMServiceImpl, error) {
const maxWait = 5 * time.Minute
retryDelay := 500 * time.Millisecond
start := time.Now()
@@ -24,7 +26,7 @@ func NewIAMServiceWithRetry(ctx context.Context, natsConn *nats.Conn, masterKey
for {
attempt++
svc, err := NewIAMServiceImpl(ctx, natsConn, masterKey, clusterSize)
svc, err := NewIAMServiceImpl(ctx, natsConn, masterKey)
if err == nil {
if attempt > 1 {
slog.Info("IAM service initialized after retry", "attempts", attempt, "elapsed_ms", otelsetup.Millis(time.Since(start)))
@@ -32,6 +34,13 @@ func NewIAMServiceWithRetry(ctx context.Context, natsConn *nats.Conn, masterKey
return svc, nil
}
// An undeclared cluster size is a configuration fault, not a cluster that
// has yet to form, so retrying it turns a clear error into five minutes
// of waiting followed by a misleading one.
if clustersize.Permanent(err) {
return nil, err
}
elapsed := time.Since(start)
if elapsed >= maxWait {
return nil, fmt.Errorf("IAM service unavailable after %s (%d attempts): %w", elapsed.Round(time.Second), attempt, err)
+38
View File
@@ -0,0 +1,38 @@
package handlers_iam_test
import (
"testing"
"time"
"github.com/mulgadc/spinifex/spinifex/clustersize"
handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
"github.com/mulgadc/spinifex/spinifex/testutil"
"github.com/stretchr/testify/require"
)
// TestNewIAMServiceWithRetry_StopsOnAConfigFault covers the difference between a
// cluster that has yet to form and one that was misconfigured.
//
// The retry loop exists for the first: on a cold multi-node boot the KV store
// needs a quorum that does not exist yet, so waiting is right. An undeclared
// cluster size is the second, and waiting five minutes for it produces a timeout
// that names quorum as the cause when the cause was a config the process never
// loaded. This test bounds that: it must come back fast, and it must come back
// saying what is actually wrong.
func TestNewIAMServiceWithRetry_StopsOnAConfigFault(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
clustersize.RedeclareForTest(t, 0)
masterKey, err := handlers_iam.GenerateMasterKey()
require.NoError(t, err)
start := time.Now()
_, err = handlers_iam.NewIAMServiceWithRetry(t.Context(), nc, masterKey)
elapsed := time.Since(start)
require.Error(t, err, "an undeclared cluster size must not look like a working IAM service")
require.ErrorIs(t, err, clustersize.ErrUndeclared,
"the error has to name the config fault, not the quorum the loop was waiting for")
require.Lessf(t, elapsed, 10*time.Second,
"returned after %s; a permanent fault must not consume the retry budget", elapsed)
}

Some files were not shown because too many files have changed in this diff Show More