feat(natsstore): bootstrap the store buckets idempotently

internal/store/natsstore/bootstrap.go

@@ -0,0 +1,173 @@
+package natsstore
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	// SchemaVersion is the layout version this binary reads and writes.
+	SchemaVersion = 1
+	// SchemaKey holds the schema marker in the meta bucket.
+	SchemaKey = "schema"
+)
+
+// ErrSchemaTooNew is returned when the stored layout was written by a newer binary.
+var ErrSchemaTooNew = errors.New("stored schema is newer than this binary supports")
+
+// ErrBucketMissing is returned in CreateNever mode when a bucket has not been provisioned.
+var ErrBucketMissing = errors.New("bucket does not exist")
+
+// DriftError reports a bucket whose configuration differs from the expected shape.
+type DriftError struct {
+	Bucket string
+	// Fields lists each mismatch as `name: have X, want Y`.
+	Fields []string
+}
+
+func (e *DriftError) Error() string {
+	return fmt.Sprintf("bucket %s differs from the expected configuration (%s); update it to match or set %s=%s to let startup reconcile it",
+		e.Bucket, strings.Join(e.Fields, "; "), EnvStoreCreate, CreateAuto)
+}
+
+// Bootstrapped is a connected JetStream context with every bucket verified.
+type Bootstrapped struct {
+	JS       jetstream.JetStream
+	Features ServerFeatures
+	// Buckets maps full bucket name to its handle.
+	Buckets map[string]jetstream.KeyValue
+}
+
+// Bucket returns the handle for a purpose such as BucketAuth.
+func (b *Bootstrapped) Bucket(settings StoreSettings, purpose string) jetstream.KeyValue {
+	return b.Buckets[settings.BucketName(purpose)]
+}
+
+// Bootstrap probes the server, ensures every bucket and records the schema
+// marker. It is idempotent and safe for many replicas to run concurrently.
+func Bootstrap(ctx context.Context, nc *nats.Conn, settings StoreSettings) (*Bootstrapped, error) {
+	if err := settings.Validate(); err != nil {
+		return nil, err
+	}
+	js, err := jetstream.New(nc)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create JetStream context: %w", err)
+	}
+	features, err := Probe(ctx, nc, js)
+	if err != nil {
+		return nil, err
+	}
+	if !features.KeyTTL {
+		return nil, fmt.Errorf("NATS server %s does not support per-key TTL (JetStream API level %d); use a server of version %d.%d or newer",
+			features.Version, features.APILevel, 2, 11)
+	}
+	if !features.AtLeast(TargetMajor, TargetMinor) {
+		log.Warnf("NATS server %s is older than the tested %d.%d line", features.Version, TargetMajor, TargetMinor)
+	}
+
+	buckets := make(map[string]jetstream.KeyValue)
+	for _, spec := range settings.Buckets() {
+		kv, err := EnsureBucket(ctx, js, spec, settings.Create)
+		if err != nil {
+			return nil, err
+		}
+		buckets[spec.Name] = kv
+	}
+	meta := buckets[settings.BucketName(BucketMeta)]
+	if err := EnsureSchema(ctx, meta); err != nil {
+		return nil, err
+	}
+	return &Bootstrapped{JS: js, Features: features, Buckets: buckets}, nil
+}
+
+// EnsureBucket makes the bucket exist with the expected shape. With CreateAuto
+// it creates or reconciles; with CreateNever it only looks. Either way the
+// final shape is verified, so a drifted bucket is reported instead of used.
+func EnsureBucket(ctx context.Context, js jetstream.JetStream, spec BucketSpec, mode CreateMode) (jetstream.KeyValue, error) {
+	var (
+		kv  jetstream.KeyValue
+		err error
+	)
+	switch mode {
+	case CreateAuto:
+		kv, err = js.CreateOrUpdateKeyValue(ctx, spec.config())
+	case CreateNever:
+		kv, err = js.KeyValue(ctx, spec.Name)
+		if errors.Is(err, jetstream.ErrBucketNotFound) {
+			return nil, fmt.Errorf("%w: %s; provision it or set %s=%s", ErrBucketMissing, spec.Name, EnvStoreCreate, CreateAuto)
+		}
+	default:
+		return nil, fmt.Errorf("unknown create mode %q", mode)
+	}
+	if err != nil {
+		return nil, fmt.Errorf("failed to ensure bucket %s: %w", spec.Name, err)
+	}
+	if err := verifyBucket(ctx, kv, spec); err != nil {
+		return nil, err
+	}
+	return kv, nil
+}
+
+func verifyBucket(ctx context.Context, kv jetstream.KeyValue, spec BucketSpec) error {
+	status, err := kv.Status(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to read bucket %s: %w", spec.Name, err)
+	}
+	cfg := status.(*jetstream.KeyValueBucketStatus).StreamInfo().Config
+
+	var drift []string
+	check := func(name string, have, want any) {
+		if have != want {
+			drift = append(drift, fmt.Sprintf("%s: have %v, want %v", name, have, want))
+		}
+	}
+	check("storage", cfg.Storage.String(), spec.Storage.String())
+	check("replicas", cfg.Replicas, spec.Replicas)
+	check("history", cfg.MaxMsgsPerSubject, int64(spec.History))
+	check("ttl", cfg.MaxAge, time.Duration(0))
+	check("limit marker ttl", cfg.SubjectDeleteMarkerTTL, spec.MarkerTTL)
+	if len(drift) > 0 {
+		return &DriftError{Bucket: spec.Name, Fields: drift}
+	}
+	return nil
+}
+
+type schemaMarker struct {
+	Version int `json:"v"`
+}
+
+// EnsureSchema records the schema version once and refuses a newer stored one,
+// so an old binary in a rolling upgrade never writes into a layout it does not know.
+func EnsureSchema(ctx context.Context, meta jetstream.KeyValue) error {
+	payload, err := json.Marshal(schemaMarker{Version: SchemaVersion})
+	if err != nil {
+		return err
+	}
+	_, err = meta.Create(ctx, SchemaKey, payload)
+	if err == nil {
+		return nil
+	}
+	if !errors.Is(err, jetstream.ErrKeyExists) {
+		return fmt.Errorf("failed to record schema marker: %w", err)
+	}
+	entry, err := meta.Get(ctx, SchemaKey)
+	if err != nil {
+		return fmt.Errorf("failed to read schema marker: %w", err)
+	}
+	var stored schemaMarker
+	if err := json.Unmarshal(entry.Value(), &stored); err != nil {
+		return fmt.Errorf("schema marker is not valid: %w", err)
+	}
+	if stored.Version > SchemaVersion {
+		return fmt.Errorf("%w: stored v%d, supported v%d", ErrSchemaTooNew, stored.Version, SchemaVersion)
+	}
+	return nil
+}

internal/store/natsstore/bootstrap_test.go

@@ -0,0 +1,185 @@
+package natsstore
+
+import (
+	"errors"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/nats-io/nats.go/jetstream"
+
+	"github.com/router-for-me/CLIProxyAPI/v8/internal/store/natsstore/natstest"
+)
+
+func autoSettings(prefix string) StoreSettings {
+	settings := DefaultStoreSettings()
+	settings.Prefix = prefix
+	settings.Create = CreateAuto
+	return settings
+}
+
+func TestBootstrapIsIdempotentAndConcurrencySafe(t *testing.T) {
+	srv := natstest.Start(t)
+	settings := autoSettings("idem")
+	ctx := testContext(t)
+
+	var wg sync.WaitGroup
+	errs := make(chan error, 6)
+	for range 6 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			nc, err := Connect(ctx, SettingsFromEnv(mapEnv(map[string]string{EnvURL: srv.URL})).Resolve(), OSLookup, 5*time.Second)
+			if err != nil {
+				errs <- err
+				return
+			}
+			defer nc.Close()
+			_, err = Bootstrap(ctx, nc, settings)
+			errs <- err
+		}()
+	}
+	wg.Wait()
+	close(errs)
+	for err := range errs {
+		if err != nil {
+			t.Fatalf("concurrent bootstrap: %v", err)
+		}
+	}
+
+	nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	boot, err := Bootstrap(ctx, nc, settings)
+	if err != nil {
+		t.Fatalf("repeat bootstrap: %v", err)
+	}
+	if len(boot.Buckets) != len(settings.Buckets()) {
+		t.Fatalf("buckets = %d", len(boot.Buckets))
+	}
+	entry, err := boot.Bucket(settings, BucketMeta).Get(ctx, SchemaKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if string(entry.Value()) != `{"v":1}` {
+		t.Fatalf("schema marker = %s", entry.Value())
+	}
+	if entry.Revision() != 1 {
+		t.Fatalf("schema marker was rewritten, revision %d", entry.Revision())
+	}
+}
+
+func TestBootstrapNeverModeVerifiesOnly(t *testing.T) {
+	srv := natstest.Start(t)
+	nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx := testContext(t)
+	verify := DefaultStoreSettings()
+	verify.Prefix = "verify"
+
+	_, err = Bootstrap(ctx, nc, verify)
+	if !errors.Is(err, ErrBucketMissing) {
+		t.Fatalf("err = %v, want ErrBucketMissing", err)
+	}
+	js, _ := jetstream.New(nc)
+	if _, lookupErr := js.KeyValue(ctx, "verify-config"); !errors.Is(lookupErr, jetstream.ErrBucketNotFound) {
+		t.Fatalf("never mode created a bucket: %v", lookupErr)
+	}
+
+	if _, err := Bootstrap(ctx, nc, autoSettings("verify")); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := Bootstrap(ctx, nc, verify); err != nil {
+		t.Fatalf("verify after provisioning: %v", err)
+	}
+}
+
+func TestBootstrapReportsDrift(t *testing.T) {
+	srv := natstest.Start(t)
+	nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx := testContext(t)
+	js, _ := jetstream.New(nc)
+	if _, err := js.CreateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: "drift-config", History: 3}); err != nil {
+		t.Fatal(err)
+	}
+
+	settings := DefaultStoreSettings()
+	settings.Prefix = "drift"
+	_, err = Bootstrap(ctx, nc, settings)
+	var drift *DriftError
+	if !errors.As(err, &drift) {
+		t.Fatalf("err = %v, want DriftError", err)
+	}
+	if drift.Bucket != "drift-config" || len(drift.Fields) != 1 || !strings.HasPrefix(drift.Fields[0], "history") {
+		t.Fatalf("drift = %+v", drift)
+	}
+}
+
+func TestBootstrapRefusesNewerSchema(t *testing.T) {
+	srv := natstest.Start(t)
+	nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx := testContext(t)
+	settings := autoSettings("schema")
+	boot, err := Bootstrap(ctx, nc, settings)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := boot.Bucket(settings, BucketMeta).Put(ctx, SchemaKey, []byte(`{"v":2}`)); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := Bootstrap(ctx, nc, settings); !errors.Is(err, ErrSchemaTooNew) {
+		t.Fatalf("err = %v, want ErrSchemaTooNew", err)
+	}
+}
+
+func TestPerKeyTTLExpiresInBootstrappedBuckets(t *testing.T) {
+	srv := natstest.Start(t)
+	nc, err := connect(t, map[string]string{EnvURL: srv.URL}, 5*time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx := testContext(t)
+	settings := autoSettings("ttl")
+	boot, err := Bootstrap(ctx, nc, settings)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	lease := boot.Bucket(settings, BucketLease)
+	if _, err := lease.Create(ctx, "short", []byte("x"), jetstream.KeyTTL(time.Second)); err != nil {
+		t.Fatalf("create with key TTL: %v", err)
+	}
+	if _, err := lease.Create(ctx, "short", []byte("y")); !errors.Is(err, jetstream.ErrKeyExists) {
+		t.Fatalf("second create = %v, want ErrKeyExists", err)
+	}
+	deadline := time.Now().Add(15 * time.Second)
+	for {
+		_, err := lease.Get(ctx, "short")
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
+			break
+		}
+		if time.Now().After(deadline) {
+			t.Fatalf("key still readable after its TTL: %v", err)
+		}
+		time.Sleep(250 * time.Millisecond)
+	}
+	if _, err := lease.Create(ctx, "short", []byte("again")); err != nil {
+		t.Fatalf("create after expiry: %v", err)
+	}
+
+	// Keys in the refresh journal bucket never expire: the bucket has no limit marker TTL.
+	journal := boot.Bucket(settings, BucketRefresh)
+	if _, err := journal.Create(ctx, "claim", []byte("x"), jetstream.KeyTTL(time.Second)); err == nil {
+		t.Fatal("refresh bucket accepted a per-key TTL")
+	}
+}