From 01341f9bb402fe7700b055d7fb83d99b65802ebf Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Sat, 7 Nov 2026 20:53:18 -0500 Subject: [PATCH] feat(natsstore): record which replica wrote each value Live sync has to tell a replica's own writes from another's when the bucket echoes them back, and an unconditional spool rewrite would clobber a local edit made after the write. Values now travel in an envelope naming the writing replica and the revision it replaced. Reads still accept the raw values the earlier backend stored, so an existing bucket keeps working and upgrades value by value. Signed-off-by: Yordis Prieto --- internal/store/natsstore/envelope.go | 40 ++++++++++++++++++ internal/store/natsstore/envelope_test.go | 51 +++++++++++++++++++++++ internal/store/natsstore/kv.go | 48 +++++++++++++++++++-- internal/store/natsstore/kv_test.go | 42 +++++++++++++++++++ internal/store/natsstore/store.go | 27 +++++++++++- internal/store/natsstore/store_test.go | 29 +++++++++++++ 6 files changed, 231 insertions(+), 6 deletions(-) create mode 100644 internal/store/natsstore/envelope.go create mode 100644 internal/store/natsstore/envelope_test.go diff --git a/internal/store/natsstore/envelope.go b/internal/store/natsstore/envelope.go new file mode 100644 index 00000000000..958fd56d4ca --- /dev/null +++ b/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 +} diff --git a/internal/store/natsstore/envelope_test.go b/internal/store/natsstore/envelope_test.go new file mode 100644 index 00000000000..ea27c5652ae --- /dev/null +++ b/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) + } + }) + } +} diff --git a/internal/store/natsstore/kv.go b/internal/store/natsstore/kv.go index aea3712fb74..bddbbfbd89e 100644 --- a/internal/store/natsstore/kv.go +++ b/internal/store/natsstore/kv.go @@ -37,6 +37,11 @@ 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 @@ -44,10 +49,36 @@ type Entry struct { 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} } diff --git a/internal/store/natsstore/kv_test.go b/internal/store/natsstore/kv_test.go index 50dd6922615..01e2b65dc2b 100644 --- a/internal/store/natsstore/kv_test.go +++ b/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) + } +} diff --git a/internal/store/natsstore/store.go b/internal/store/natsstore/store.go index 24a67fce57f..b487691bd67 100644 --- a/internal/store/natsstore/store.go +++ b/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,13 +49,19 @@ 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, @@ -60,6 +69,20 @@ func NewStore(nc *nats.Conn, boot *Bootstrapped, settings StoreSettings, spoolDi }, 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 { diff --git a/internal/store/natsstore/store_test.go b/internal/store/natsstore/store_test.go index 303c4bdf155..2c8bd7b1447 100644 --- a/internal/store/natsstore/store_test.go +++ b/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) + } +}