feat(network): wire Membership into node; router rebuilds on scale

node.Start() now subscribes to a Membership (DNS-backed by default, static by test injection) and rebuilds the router on every change. The transport's own reconnect loop picks up new peers naturally; we don't push dials from the node — Membership is the authoritative source of truth, transport follows. Scale flow (no pod restart): 1. kubectl scale sts foo --replicas=N (or docker compose scale) 2. Orchestrator DNS starts/stops resolving pod names 3. dnsMembership's 5s poll picks up the change 4. node.watchMembership receives updated []NodeID 5. router.setMembers([]NodeID) — O(N log N) sort 6. WriterFor / MembersFor return the new topology immediately New metric: base_network_membership_size (gauge). Drops below expected count before transport notices — useful for scale-event dashboards. SetMembership(m) is how tests (and embedding callers) inject a Membership other than the default DNS one. Tests (all with -race): - TestNodeScaleUp 1 → 3, router rebuilds, MembersFor growth - TestNodeScaleDown 3 → 2, dropped member leaves MembersFor - TestNodeScaleDownToOne 3 → 1, WriterFor flips back to local - TestNodeMembershipNotSet static BASE_PEERS path still works
2026-10-21 03:41:50 +00:00 · 2026-04-23 04:44:39 -07:00
parent 8c03b17e48
commit 785f35d408
3 changed files with 282 additions and 5 deletions
@@ -57,6 +57,11 @@ type Metrics struct {
 	ApplyErrors   prometheus.Counter
 	WALHookErrors prometheus.Counter
 	WALBytes      prometheus.Counter
+
+	// MembershipSize is the live member count as reported by the
+	// Membership watcher. Dashboards alert on unexpected drops
+	// (scale-down events show up here before the transport notices).
+	MembershipSize prometheus.Gauge
 }

 // NewMetrics constructs the collectors. Register them on a caller-owned
@@ -137,6 +142,10 @@ func NewMetrics() *Metrics {
 			Name: "base_network_wal_bytes_total",
 			Help: "Total WAL payload bytes captured by the commit hook.",
 		}),
+		MembershipSize: prometheus.NewGauge(prometheus.GaugeOpts{
+			Name: "base_network_membership_size",
+			Help: "Live member count reported by the Membership watcher.",
+		}),
 	}
 }

@@ -155,6 +164,7 @@ func (m *Metrics) Register(r prometheus.Registerer) error {
 		m.FramesDuplicate, m.FramesInvalid,
 		m.FramesRejectedShardMismatch, m.FramesRejectedSeqGap,
 		m.ApplyErrors, m.WALHookErrors, m.WALBytes,
+		m.MembershipSize,
 	}
 	for _, c := range collectors {
 		if err := r.Register(c); err != nil {
@@ -19,6 +19,11 @@ type node struct {
 	// transport; production will inject a QUIC/gRPC transport. Always set.
 	transport Transport

+	// membership is the live view of reachable members. When set, Start
+	// subscribes to change events and rebuilds the router on every
+	// update. Dynamic reconnects are delegated to the transport.
+	membership Membership
+
 	mu     sync.RWMutex
 	shards map[string]*Shard
 	walSrc walSource
@@ -26,10 +31,19 @@ type node struct {
 	cancel context.CancelFunc
 }

-// newNode constructs the live node with the production ZAP transport.
-// Tests that wire peers together in-process use NewWithTransport.
+// newNode constructs the live node with the production ZAP transport and
+// a DNS-based Membership built from BASE_PEERS. Tests that wire peers
+// together in-process use NewWithTransport.
 func newNode(cfg Config) (*node, error) {
-	return newNodeWithTransport(cfg, newZapTransport(cfg))
+	n, err := newNodeWithTransport(cfg, newZapTransport(cfg))
+	if err != nil {
+		return nil, err
+	}
+	// Production Membership tracks BASE_PEERS through DNS: a seed that
+	// resolves to N addresses yields N members; scale events propagate
+	// within refreshInterval (5s by default).
+	n.membership = NewDNSMembership(context.Background(), cfg.NodeID, cfg.Peers)
+	return n, nil
 }

 func newNodeWithTransport(cfg Config, t Transport) (*node, error) {
@@ -63,6 +77,16 @@ func newNodeWithTransport(cfg Config, t Transport) (*node, error) {
 	}, nil
 }

+// SetMembership injects a Membership source. The node subscribes on
+// Start and rebuilds the router on every change. Tests use this to
+// drive scale events; production wires a dnsMembership automatically
+// from newNode.
+func (n *node) SetMembership(m Membership) {
+	n.mu.Lock()
+	n.membership = m
+	n.mu.Unlock()
+}
+
 // NewWithTransport constructs a Network with a caller-supplied transport.
 // In-process tests inject a memory transport here; production callers use
 // the default wired by New.
@@ -80,22 +104,53 @@ func (n *node) Enabled() bool { return true }

 func (n *node) Start(ctx context.Context) error {
 	n.mu.Lock()
-	defer n.mu.Unlock()
 	if n.ctx != nil {
+		n.mu.Unlock()
 		return fmt.Errorf("network: already started")
 	}
 	n.ctx, n.cancel = context.WithCancel(ctx)
+	ourCtx := n.ctx
+	membership := n.membership
+	n.mu.Unlock()

 	// Transport fan-in: peer frames arrive here and are submitted to the
 	// local shard engine so every member converges on the same DAG.
-	if err := n.transport.Start(n.ctx, n.onPeerFrame); err != nil {
+	if err := n.transport.Start(ourCtx, n.onPeerFrame); err != nil {
 		n.cancel()
+		n.mu.Lock()
 		n.ctx = nil
+		n.mu.Unlock()
 		return fmt.Errorf("network: transport start: %w", err)
 	}
+
+	// Dynamic-membership watch: every Membership change rebuilds the
+	// router's member set. Transport learns about new peers via its own
+	// reconnect loop — we don't have to push dials from here. Dropped
+	// peers naturally time out on the transport side.
+	if membership != nil {
+		go n.watchMembership(ourCtx, membership)
+	}
 	return nil
 }

+// watchMembership blocks on the Membership change stream and forwards
+// updates to the router. Exits when ctx is cancelled (Stop).
+func (n *node) watchMembership(ctx context.Context, m Membership) {
+	updates := m.Watch(ctx)
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case members, ok := <-updates:
+			if !ok {
+				return
+			}
+			n.router.setMembers(members)
+			n.metrics.MembershipSize.Set(float64(len(members)))
+		}
+	}
+}
+
 func (n *node) Stop(ctx context.Context) error {
 	n.mu.Lock()
 	cancel := n.cancel
@@ -0,0 +1,212 @@
+// Copyright (c) 2025, Hanzo Industries Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+//
+// End-to-end: node + Membership + router. Covers the scale-1→N→1
+// contract that every Base service needs:
+//
+//   - singleton start → WriterFor returns local
+//   - Membership change → router rebuilds → WriterFor may flip
+//   - scale-down → member drops out → router recomputes
+//
+// This is the integration story for consumers (BD, ATS, TA, IAM).
+// If these tests pass, `kubectl scale` (or `docker compose scale`)
+// works without pod restart.
+
+package network
+
+import (
+	"context"
+	"testing"
+	"time"
+)
+
+// TestNodeScaleUp: node starts singleton, operator scales to 3,
+// router picks up new members, WriterFor flips appropriately.
+func TestNodeScaleUp(t *testing.T) {
+	cfg := Config{
+		Enabled:     true,
+		ShardKey:    "user_id",
+		Replication: 3,
+		NodeID:      "a",
+		Peers:       []string{"b", "c"},
+		Role:        RoleValidator,
+		Archive:     "off",
+		ListenHTTP:  ":8090",
+		ListenP2P:   ":9651",
+	}
+	n, err := newNodeWithTransport(cfg, &nopTransport{})
+	if err != nil {
+		t.Fatalf("newNodeWithTransport: %v", err)
+	}
+
+	mem := NewStaticMembership("a", []string{"a"})
+	n.SetMembership(mem)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	if err := n.Start(ctx); err != nil {
+		t.Fatalf("Start: %v", err)
+	}
+	defer n.Stop(context.Background())
+
+	// Singleton: only self is member → WriterFor(x) is local.
+	waitForMemberCount(t, n, 1)
+	_, local := n.WriterFor("user-7")
+	if !local {
+		t.Errorf("singleton WriterFor(user-7): expected local=true")
+	}
+
+	// Scale to 3: self + b + c.
+	mem.Set([]string{"b", "c"})
+
+	waitForMemberCount(t, n, 3)
+
+	// Router now knows 3 members. WriterFor is deterministic via
+	// rendezvous hash — some shard IDs land local, others don't.
+	// We just assert the set is stable and of size 3.
+	ms := n.MembersFor("user-7")
+	if len(ms) != 3 {
+		t.Errorf("MembersFor after scale-up: got %d, want 3", len(ms))
+	}
+}
+
+// TestNodeScaleDown: node starts with 3 peers, one drops out,
+// router rebuilds without that member.
+func TestNodeScaleDown(t *testing.T) {
+	cfg := Config{
+		Enabled:     true,
+		ShardKey:    "user_id",
+		Replication: 3,
+		NodeID:      "a",
+		Peers:       []string{"b", "c"},
+		Role:        RoleValidator,
+		Archive:     "off",
+		ListenHTTP:  ":8090",
+		ListenP2P:   ":9651",
+	}
+	n, err := newNodeWithTransport(cfg, &nopTransport{})
+	if err != nil {
+		t.Fatalf("newNodeWithTransport: %v", err)
+	}
+
+	mem := NewStaticMembership("a", []string{"b", "c"})
+	n.SetMembership(mem)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	if err := n.Start(ctx); err != nil {
+		t.Fatalf("Start: %v", err)
+	}
+	defer n.Stop(context.Background())
+
+	waitForMemberCount(t, n, 3)
+
+	// c drops out (pod terminated).
+	mem.Set([]string{"b"})
+
+	waitForMemberCount(t, n, 2)
+	ms := n.MembersFor("user-42")
+	if len(ms) > 2 {
+		t.Errorf("MembersFor after scale-down: got %d, want <= 2", len(ms))
+	}
+	for _, m := range ms {
+		if m == "c" {
+			t.Errorf("scaled-down member %q still in MembersFor", m)
+		}
+	}
+}
+
+// TestNodeScaleDownToOne: collapse all the way back to singleton.
+// Router must have exactly self and WriterFor returns local for
+// every shard.
+func TestNodeScaleDownToOne(t *testing.T) {
+	cfg := Config{
+		Enabled:     true,
+		ShardKey:    "user_id",
+		Replication: 3,
+		NodeID:      "a",
+		Peers:       []string{"b", "c"},
+		Role:        RoleValidator,
+		Archive:     "off",
+		ListenHTTP:  ":8090",
+		ListenP2P:   ":9651",
+	}
+	n, err := newNodeWithTransport(cfg, &nopTransport{})
+	if err != nil {
+		t.Fatalf("newNodeWithTransport: %v", err)
+	}
+
+	mem := NewStaticMembership("a", []string{"b", "c"})
+	n.SetMembership(mem)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	if err := n.Start(ctx); err != nil {
+		t.Fatalf("Start: %v", err)
+	}
+	defer n.Stop(context.Background())
+
+	waitForMemberCount(t, n, 3)
+
+	// All peers drop out (cluster scaled to 1 during low traffic).
+	mem.Set(nil)
+
+	waitForMemberCount(t, n, 1)
+
+	for _, shard := range []string{"user-1", "user-2", "user-99", "org-7"} {
+		_, local := n.WriterFor(shard)
+		if !local {
+			t.Errorf("WriterFor(%q) after scale-to-1: expected local=true", shard)
+		}
+	}
+}
+
+// TestNodeMembershipNotSet: node without a Membership still works
+// (static Peers list from Config). Back-compat for any call site
+// that doesn't wire a Membership.
+func TestNodeMembershipNotSet(t *testing.T) {
+	cfg := Config{
+		Enabled:     true,
+		ShardKey:    "user_id",
+		Replication: 3,
+		NodeID:      "a",
+		Peers:       []string{"b", "c"},
+		Role:        RoleValidator,
+		Archive:     "off",
+		ListenHTTP:  ":8090",
+		ListenP2P:   ":9651",
+	}
+	n, err := newNodeWithTransport(cfg, &nopTransport{})
+	if err != nil {
+		t.Fatalf("newNodeWithTransport: %v", err)
+	}
+	// Deliberately NOT calling SetMembership.
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	if err := n.Start(ctx); err != nil {
+		t.Fatalf("Start: %v", err)
+	}
+	defer n.Stop(context.Background())
+
+	// Static Peers gave us self + b + c = 3.
+	ms := n.MembersFor("x")
+	if len(ms) != 3 {
+		t.Errorf("MembersFor without Membership: got %d, want 3", len(ms))
+	}
+}
+
+// waitForMemberCount polls n.MembersFor until the count matches,
+// bounded by a 500ms deadline. Membership updates propagate through
+// a channel; this keeps tests non-flaky without hard-coded sleeps.
+func waitForMemberCount(t *testing.T, n *node, want int) {
+	t.Helper()
+	deadline := time.Now().Add(500 * time.Millisecond)
+	for time.Now().Before(deadline) {
+		if len(n.MembersFor("probe")) == want {
+			return
+		}
+		time.Sleep(5 * time.Millisecond)
+	}
+	t.Fatalf("waitForMemberCount: got %d, want %d", len(n.MembersFor("probe")), want)
+}