feat(natsstore): record which replica wrote each value

internal/store/natsstore/envelope.go

@@ -0,0 +1,40 @@
+package natsstore
+
+import (
+	"bytes"
+	"encoding/json"
+)
+
+// envelopeVersion marks a value written inside an envelope. A legacy value is
+// the raw payload with no marker, which no credential or config carries.
+const envelopeVersion = 1
+
+// envelope wraps a stored payload with where it came from, so a replica can
+// recognize its own writes when the bucket echoes them back.
+type envelope struct {
+	Version int `json:"natsstore_envelope"`
+	// Origin identifies the replica that wrote the value.
+	Origin string `json:"origin"`
+	// Base is the revision the write replaced, zero for a create.
+	Base uint64 `json:"base_revision"`
+	// Data is the payload byte for byte, so the file it came from is restored exactly.
+	Data []byte `json:"data"`
+}
+
+// encodeEnvelope wraps the payload for storage.
+func encodeEnvelope(origin string, base uint64, data []byte) ([]byte, error) {
+	return json.Marshal(envelope{Version: envelopeVersion, Origin: origin, Base: base, Data: data})
+}
+
+// decodeEnvelope unwraps a stored value. A value without the marker is the raw
+// format the earlier backend wrote; it is returned unchanged with no origin.
+func decodeEnvelope(raw []byte) (data []byte, origin string, base uint64) {
+	if !bytes.HasPrefix(bytes.TrimSpace(raw), []byte("{")) {
+		return raw, "", 0
+	}
+	var env envelope
+	if err := json.Unmarshal(raw, &env); err != nil || env.Version != envelopeVersion {
+		return raw, "", 0
+	}
+	return env.Data, env.Origin, env.Base
+}

internal/store/natsstore/envelope_test.go

@@ -0,0 +1,51 @@
+package natsstore
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestEnvelopeRoundTripKeepsBytesExact(t *testing.T) {
+	payloads := map[string][]byte{
+		"json":    []byte("{\"type\": \"test\",\n  \"token\": \"a b\"}\n"),
+		"yaml":    []byte("port: 8317\nhosts:\n  - \"x\"\n"),
+		"binary":  {0x00, 0xff, 0x7b, 0x22},
+		"empty":   {},
+		"looks":   []byte(`{"natsstore_envelope":1,"origin":"x","data":"AA=="}`),
+		"unicode": []byte("name: é\r\n"),
+	}
+	for name, payload := range payloads {
+		t.Run(name, func(t *testing.T) {
+			raw, err := encodeEnvelope("replica-a", 7, payload)
+			if err != nil {
+				t.Fatal(err)
+			}
+			data, origin, base := decodeEnvelope(raw)
+			if !bytes.Equal(data, payload) || origin != "replica-a" || base != 7 {
+				t.Fatalf("decoded = %q, %q, %d", data, origin, base)
+			}
+		})
+	}
+}
+
+func TestDecodeEnvelopeAcceptsLegacyValues(t *testing.T) {
+	legacy := map[string]string{
+		"credential":     `{"type":"test","access_token":"t"}`,
+		"yaml":           "port: 8317\n",
+		"flow yaml":      "{port: 8317}\n",
+		"other version":  `{"natsstore_envelope":2,"origin":"x","data":"AA=="}`,
+		"not json":       "{broken",
+		"empty":          "",
+		"json array":     `[1,2]`,
+		"spaced":         "  \n{\"type\":\"test\"}",
+		"wrong data key": `{"natsstore_envelope":"1"}`,
+	}
+	for name, raw := range legacy {
+		t.Run(name, func(t *testing.T) {
+			data, origin, base := decodeEnvelope([]byte(raw))
+			if string(data) != raw || origin != "" || base != 0 {
+				t.Fatalf("decoded = %q, %q, %d, want the raw value", data, origin, base)
+			}
+		})
+	}
+}

internal/store/natsstore/kv.go

@@ -37,17 +37,48 @@ type Entry struct {
 	Revision uint64
 	// Modified is when the revision was written.
 	Modified time.Time
+	// Origin is the replica that wrote the revision, empty for a value written
+	// before envelopes existed.
+	Origin string
+	// Base is the revision the write replaced, zero for a create or a legacy value.
+	Base uint64
 }
 
 // KV wraps a bucket so callers work with identifiers and revisions and never
 // with encoded keys. Every write is a compare-and-swap; there is no blind put.
 type KV struct {
 	bucket string
 	kv     jetstream.KeyValue
+	origin string
 }
 
+// KVOption customizes a KV.
+type KVOption func(*KV)
+
+// WithOrigin makes the KV wrap every value it writes with the replica that
+// wrote it. Without it values are stored raw, as the earlier backend did.
+// Reads accept both formats either way.
+func WithOrigin(origin string) KVOption { return func(k *KV) { k.origin = origin } }
+
 // NewKV wraps the bucket handle.
-func NewKV(bucket string, kv jetstream.KeyValue) *KV { return &KV{bucket: bucket, kv: kv} }
+func NewKV(bucket string, kv jetstream.KeyValue, opts ...KVOption) *KV {
+	k := &KV{bucket: bucket, kv: kv}
+	for _, opt := range opts {
+		opt(k)
+	}
+	return k
+}
+
+// Origin returns the replica identifier written with values, empty when values are stored raw.
+func (k *KV) Origin() string { return k.origin }
+
+// wrap prepares a value for storage.
+func (k *KV) wrap(value []byte, base uint64) ([]byte, error) {
+	if k.origin == "" {
+		return value, nil
+	}
+	return encodeEnvelope(k.origin, base, value)
+}
 
 // Get returns the live value for the identifier, or ErrNotFound.
 func (k *KV) Get(ctx context.Context, id string) (Entry, error) {
@@ -71,7 +102,11 @@ func (k *KV) Create(ctx context.Context, id string, value []byte) (uint64, error
 	if err != nil {
 		return 0, err
 	}
-	rev, err := k.kv.Create(ctx, key, value)
+	stored, err := k.wrap(value, 0)
+	if err != nil {
+		return 0, fmt.Errorf("failed to encode %s for %s: %w", id, k.bucket, err)
+	}
+	rev, err := k.kv.Create(ctx, key, stored)
 	if err != nil {
 		return 0, k.writeError("create", id, err)
 	}
@@ -84,7 +119,11 @@ func (k *KV) Update(ctx context.Context, id string, value []byte, expected uint6
 	if err != nil {
 		return 0, err
 	}
-	rev, err := k.kv.Update(ctx, key, value, expected)
+	stored, err := k.wrap(value, expected)
+	if err != nil {
+		return 0, fmt.Errorf("failed to encode %s for %s: %w", id, k.bucket, err)
+	}
+	rev, err := k.kv.Update(ctx, key, stored, expected)
 	if err != nil {
 		return 0, k.writeError("update", id, err)
 	}
@@ -138,5 +177,6 @@ func (k *KV) writeError(op, id string, err error) error {
 }
 
 func entryOf(id string, entry jetstream.KeyValueEntry) Entry {
-	return Entry{ID: id, Value: entry.Value(), Revision: entry.Revision(), Modified: entry.Created()}
+	value, origin, base := decodeEnvelope(entry.Value())
+	return Entry{ID: id, Value: value, Revision: entry.Revision(), Modified: entry.Created(), Origin: origin, Base: base}
 }

internal/store/natsstore/kv_test.go

@@ -103,3 +103,45 @@ func TestKVListSkipsDeletedAndForeignKeys(t *testing.T) {
 		t.Fatalf("List = %v", got)
 	}
 }
+
+func TestKVEnvelopeOriginAndLegacyReads(t *testing.T) {
+	raw := openAuthKV(t)
+	wrapped := NewKV(raw.bucket, raw.kv, WithOrigin("replica-a"))
+	ctx := testContext(t)
+
+	// A value written raw, as the earlier backend did, still reads back.
+	legacyRev, err := raw.Create(ctx, "legacy.json", []byte(`{"v":1}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	got, err := wrapped.Get(ctx, "legacy.json")
+	if err != nil || string(got.Value) != `{"v":1}` || got.Origin != "" || got.Base != 0 {
+		t.Fatalf("legacy Get = %+v, %v", got, err)
+	}
+
+	// Updating it wraps the value and records the revision it replaced.
+	rev, err := wrapped.Update(ctx, "legacy.json", []byte(`{"v":2}`), legacyRev)
+	if err != nil {
+		t.Fatal(err)
+	}
+	got, err = raw.Get(ctx, "legacy.json")
+	if err != nil || string(got.Value) != `{"v":2}` || got.Origin != "replica-a" || got.Base != legacyRev || got.Revision != rev {
+		t.Fatalf("wrapped Get = %+v, %v", got, err)
+	}
+
+	created, err := wrapped.Create(ctx, "new.json", []byte(`{"v":3}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	entries, err := raw.List(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	byID := map[string]Entry{}
+	for _, entry := range entries {
+		byID[entry.ID] = entry
+	}
+	if e := byID["new.json"]; string(e.Value) != `{"v":3}` || e.Origin != "replica-a" || e.Base != 0 || e.Revision != created {
+		t.Fatalf("listed = %+v", e)
+	}
+}

internal/store/natsstore/store.go

@@ -3,6 +3,8 @@ package natsstore
 import (
 	"bytes"
 	"context"
+	"crypto/rand"
+	"encoding/hex"
 	"errors"
 	"fmt"
 	"io/fs"
@@ -24,6 +26,7 @@ import (
 // config and credential files the rest of the process reads keep working.
 type Store struct {
 	nc         *nats.Conn
+	replicaID  string
 	auth       *AuthStore
 	config     *ConfigStore
 	spoolRoot  string
@@ -46,20 +49,40 @@ func NewStore(nc *nats.Conn, boot *Bootstrapped, settings StoreSettings, spoolDi
 	if err := os.MkdirAll(configDir, 0o700); err != nil {
 		return nil, fmt.Errorf("nats store: create config directory: %w", err)
 	}
-	auth, err := NewAuthStore(NewKV(settings.BucketName(BucketAuth), boot.Bucket(settings, BucketAuth)), filepath.Join(root, "auths"))
+	replicaID, err := newReplicaID()
 	if err != nil {
 		return nil, err
 	}
-	cfg := NewConfigStore(NewKV(settings.BucketName(BucketConfig), boot.Bucket(settings, BucketConfig)))
+	origin := WithOrigin(replicaID)
+	auth, err := NewAuthStore(NewKV(settings.BucketName(BucketAuth), boot.Bucket(settings, BucketAuth), origin), filepath.Join(root, "auths"))
+	if err != nil {
+		return nil, err
+	}
+	cfg := NewConfigStore(NewKV(settings.BucketName(BucketConfig), boot.Bucket(settings, BucketConfig), origin))
 	return &Store{
 		nc:         nc,
+		replicaID:  replicaID,
 		auth:       auth,
 		config:     cfg,
 		spoolRoot:  root,
 		configPath: filepath.Join(configDir, "config.yaml"),
 	}, nil
 }
 
+// newReplicaID returns an identifier unique to this process. It is never
+// reused across restarts, so values a previous run wrote are not mistaken for
+// this run's own.
+func newReplicaID() (string, error) {
+	raw := make([]byte, 8)
+	if _, err := rand.Read(raw); err != nil {
+		return "", fmt.Errorf("nats store: generate replica id: %w", err)
+	}
+	return hex.EncodeToString(raw), nil
+}
+
+// ReplicaID identifies this process in the values it writes.
+func (s *Store) ReplicaID() string { return s.replicaID }
+
 // Close closes the NATS connection.
 func (s *Store) Close() error {
 	if s != nil && s.nc != nil {

internal/store/natsstore/store_test.go

@@ -298,3 +298,32 @@ func TestStoreSeedRefusesPlaintextManagementSecret(t *testing.T) {
 		t.Fatalf("refused config reached the bucket: %v", err)
 	}
 }
+
+func TestStoreWritesCarryTheReplicaOrigin(t *testing.T) {
+	f := newStoreFixture(t)
+	ctx := testContext(t)
+	first, second := f.open(), f.open()
+	if first.ReplicaID() == "" || first.ReplicaID() == second.ReplicaID() {
+		t.Fatalf("replica ids = %q and %q, want distinct and set", first.ReplicaID(), second.ReplicaID())
+	}
+	if err := first.Bootstrap(ctx, f.example); err != nil {
+		t.Fatal(err)
+	}
+	auth := testAuth("origin.json", map[string]any{"type": "test", "access_token": "tok"})
+	if _, err := first.Save(ctx, auth); err != nil {
+		t.Fatal(err)
+	}
+
+	cfgEntry, err := first.config.kv.Get(ctx, configID)
+	if err != nil || cfgEntry.Origin != first.ReplicaID() {
+		t.Fatalf("config entry = %+v, %v", cfgEntry, err)
+	}
+	authEntry, err := first.auth.kv.Get(ctx, "origin.json")
+	if err != nil || authEntry.Origin != first.ReplicaID() || authEntry.Base != 0 {
+		t.Fatalf("auth entry = %+v, %v", authEntry, err)
+	}
+	// The payload is the file, untouched by the wrapping.
+	if got := string(authEntry.Value); got != readFile(t, filepath.Join(first.AuthDir(), "origin.json")) {
+		t.Fatalf("stored payload = %q", got)
+	}
+}