[proxy] Add middleware snapshot revisions (#8114)

Version middleware snapshots so callers can reject stale policies. Existing Rebuild callers remain supported. Validation: middleware race tests and scoped lint passed.
2026-10-13 08:59:08 +02:00 · 2026-10-12 12:09:09 +02:00
parent 5c9c64e3e0
commit a7a2575bad
2 changed files with 280 additions and 13 deletions
@@ -29,6 +29,10 @@ type PathTargetBinding struct {
 	Specs     []Spec
 }

+// Revision identifies a service's middleware binding snapshot. Zero is
+// reserved for routes that are not coupled to a manager snapshot.
+type Revision uint64
+
 // LiveServiceCheck reports whether the given service ID is still
 // present in the proxy's live mapping cache. The Manager calls it
 // during InvalidateMiddleware so a chain whose service has been
@@ -42,14 +46,16 @@ type LiveServiceCheck func(serviceID string) bool
 // InvalidateMiddleware find the chain keys that reference a given
 // middleware without scanning the whole table.
 type chainTable struct {
-	byTarget     map[string]*Chain
-	byMiddleware map[string]map[string]struct{}
+	byTarget        map[string]*Chain
+	byMiddleware    map[string]map[string]struct{}
+	serviceRevision map[string]Revision
 }

 func newChainTable() *chainTable {
 	return &chainTable{
-		byTarget:     make(map[string]*Chain),
-		byMiddleware: make(map[string]map[string]struct{}),
+		byTarget:        make(map[string]*Chain),
+		byMiddleware:    make(map[string]map[string]struct{}),
+		serviceRevision: make(map[string]Revision),
 	}
 }

@@ -65,6 +71,9 @@ func (c *chainTable) clone() *chainTable {
 		}
 		out.byMiddleware[id] = set
 	}
+	for serviceID, revision := range c.serviceRevision {
+		out.serviceRevision[serviceID] = revision
+	}
 	return out
 }

@@ -121,6 +130,7 @@ type Manager struct {
 	resolver         *Resolver
 	lastBindings     map[string]PathTargetBinding
 	liveServiceCheck atomic.Pointer[LiveServiceCheck]
+	nextRevision     Revision
 }

 // NewManager constructs a Manager with the given capture budget size.
@@ -179,14 +189,29 @@ func (m *Manager) Dispatcher() *Dispatcher {
 	return m.dispatcher
 }

-// Rebuild replaces every chain keyed by serviceID with the provided
-// bindings. Entries for other services are preserved. Replaced chains
-// are closed asynchronously after the atomic swap so in-flight
-// requests against the previous chain finish before middleware
-// resources are released.
+// Rebuild replaces every chain keyed by serviceID with the provided bindings.
+// It is the compatibility wrapper for callers that do not consume revisions.
 func (m *Manager) Rebuild(serviceID string, bindings []PathTargetBinding) error {
+	_, err := m.RebuildSnapshot(serviceID, bindings)
+	return err
+}
+
+// RebuildSnapshot atomically replaces a service's chains and assigns a nonzero
+// revision, including for empty bindings. Retired chains close asynchronously.
+func (m *Manager) RebuildSnapshot(serviceID string, bindings []PathTargetBinding) (Revision, error) {
+	for _, b := range bindings {
+		if b.ServiceID != serviceID {
+			return 0, fmt.Errorf("binding service %q does not match rebuild service %q", b.ServiceID, serviceID)
+		}
+	}
+
 	m.writeMu.Lock()
 	defer m.writeMu.Unlock()
+	if m.nextRevision == ^Revision(0) {
+		return 0, fmt.Errorf("middleware snapshot revision exhausted")
+	}
+	m.nextRevision++
+	revision := m.nextRevision

 	cur := m.chains.Load()
 	next := cur.clone()
@@ -205,9 +230,6 @@ func (m *Manager) Rebuild(serviceID string, bindings []PathTargetBinding) error
 	}

 	for _, b := range bindings {
-		if b.ServiceID != serviceID {
-			return fmt.Errorf("binding service %q does not match rebuild service %q", b.ServiceID, serviceID)
-		}
 		key := chainKey(b.ServiceID, b.PathID)
 		m.lastBindings[key] = cloneBinding(b)
 		chain := m.buildChain(b)
@@ -217,10 +239,11 @@ func (m *Manager) Rebuild(serviceID string, bindings []PathTargetBinding) error
 		}
 		next.addChain(key, chain)
 	}
+	next.serviceRevision[serviceID] = revision

 	m.chains.Store(next)
 	m.closeChainsAsync(retired)
-	return nil
+	return revision, nil
 }

 // Invalidate drops every chain for the given service ID.
@@ -246,6 +269,7 @@ func (m *Manager) Invalidate(serviceID string) {
 			delete(m.lastBindings, k)
 		}
 	}
+	delete(next.serviceRevision, serviceID)
 	m.chains.Store(next)
 	m.closeChainsAsync(retired)
 }
@@ -359,6 +383,19 @@ func (m *Manager) ChainFor(serviceID, pathID string) *Chain {
 	return c
 }

+// ChainForRevision returns a chain only if the service revision matches.
+// A matching revision with a nil chain is a coherent chainless target. Zero never matches.
+func (m *Manager) ChainForRevision(serviceID, pathID string, revision Revision) (*Chain, bool) {
+	if revision == 0 {
+		return nil, false
+	}
+	tbl := m.chains.Load()
+	if tbl == nil || tbl.serviceRevision[serviceID] != revision {
+		return nil, false
+	}
+	return tbl.byTarget[chainKey(serviceID, pathID)], true
+}
+
 // buildChain resolves each enabled spec and returns the assembled
 // chain. Returns a nil chain when no middlewares are bound; resolver
 // errors per middleware are logged and counted but do not abort the
@@ -0,0 +1,230 @@
+package middleware
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+const revisionTestMiddlewareID = "revision_test"
+
+type revisionTestFactory struct {
+	builds atomic.Int64
+}
+
+func (f *revisionTestFactory) ID() string {
+	return revisionTestMiddlewareID
+}
+
+func (f *revisionTestFactory) New(_ []byte) (Middleware, error) {
+	f.builds.Add(1)
+	return revisionTestMiddleware{}, nil
+}
+
+type revisionTestMiddleware struct{}
+
+func (revisionTestMiddleware) ID() string                     { return revisionTestMiddlewareID }
+func (revisionTestMiddleware) Version() string                { return "test" }
+func (revisionTestMiddleware) Slot() Slot                     { return SlotOnRequest }
+func (revisionTestMiddleware) AcceptedContentTypes() []string { return nil }
+func (revisionTestMiddleware) MetadataKeys() []string         { return nil }
+func (revisionTestMiddleware) MutationsSupported() bool       { return false }
+func (revisionTestMiddleware) Invoke(context.Context, *Input) (*Output, error) {
+	return &Output{Decision: DecisionAllow}, nil
+}
+func (revisionTestMiddleware) Close() error { return nil }
+
+func newRevisionTestManager(t *testing.T) (*Manager, *revisionTestFactory) {
+	t.Helper()
+	factory := &revisionTestFactory{}
+	registry := NewRegistry()
+	require.NoError(t, registry.Register(factory))
+	manager := NewManager(1, nil, nil)
+	manager.SetResolver(NewResolver(registry))
+	return manager, factory
+}
+
+func revisionTestBinding(serviceID, pathID string) PathTargetBinding {
+	return PathTargetBinding{
+		ServiceID: serviceID,
+		PathID:    pathID,
+		Specs: []Spec{{
+			ID:      revisionTestMiddlewareID,
+			Slot:    SlotOnRequest,
+			Enabled: true,
+		}},
+	}
+}
+
+func TestManagerRebuildSnapshotRevisionLookup(t *testing.T) {
+	manager, _ := newRevisionTestManager(t)
+	binding := revisionTestBinding("service", "/")
+
+	revision1, err := manager.RebuildSnapshot("service", []PathTargetBinding{binding})
+	require.NoError(t, err)
+	require.NotZero(t, revision1)
+	chain1, match := manager.ChainForRevision("service", "/", revision1)
+	require.True(t, match)
+	require.NotNil(t, chain1)
+
+	revision2, err := manager.RebuildSnapshot("service", []PathTargetBinding{binding})
+	require.NoError(t, err)
+	assert.NotEqual(t, revision1, revision2)
+	chain2, match := manager.ChainForRevision("service", "/", revision2)
+	require.True(t, match)
+	require.NotNil(t, chain2)
+	assert.NotSame(t, chain1, chain2)
+
+	stale, match := manager.ChainForRevision("service", "/", revision1)
+	assert.False(t, match)
+	assert.Nil(t, stale)
+	zero, match := manager.ChainForRevision("service", "/", 0)
+	assert.False(t, match)
+	assert.Nil(t, zero)
+}
+
+func TestManagerRebuildSnapshotTracksChainlessTransitions(t *testing.T) {
+	manager, _ := newRevisionTestManager(t)
+
+	chainlessRevision, err := manager.RebuildSnapshot("service", nil)
+	require.NoError(t, err)
+	chain, match := manager.ChainForRevision("service", "/", chainlessRevision)
+	assert.True(t, match)
+	assert.Nil(t, chain)
+
+	chainRevision, err := manager.RebuildSnapshot("service", []PathTargetBinding{
+		revisionTestBinding("service", "/"),
+	})
+	require.NoError(t, err)
+	_, match = manager.ChainForRevision("service", "/", chainlessRevision)
+	assert.False(t, match)
+	chain, match = manager.ChainForRevision("service", "/", chainRevision)
+	assert.True(t, match)
+	assert.NotNil(t, chain)
+
+	chainlessRevision2, err := manager.RebuildSnapshot("service", nil)
+	require.NoError(t, err)
+	_, match = manager.ChainForRevision("service", "/", chainRevision)
+	assert.False(t, match)
+	chain, match = manager.ChainForRevision("service", "/", chainlessRevision2)
+	assert.True(t, match)
+	assert.Nil(t, chain)
+}
+
+func TestManagerInvalidateMiddlewarePreservesServiceRevision(t *testing.T) {
+	manager, factory := newRevisionTestManager(t)
+	revision, err := manager.RebuildSnapshot("service", []PathTargetBinding{
+		revisionTestBinding("service", "/"),
+	})
+	require.NoError(t, err)
+	chain1, match := manager.ChainForRevision("service", "/", revision)
+	require.True(t, match)
+	require.NotNil(t, chain1)
+
+	manager.InvalidateMiddleware(revisionTestMiddlewareID)
+
+	chain2, match := manager.ChainForRevision("service", "/", revision)
+	require.True(t, match)
+	require.NotNil(t, chain2)
+	assert.NotSame(t, chain1, chain2)
+	assert.EqualValues(t, 2, factory.builds.Load())
+}
+
+func TestManagerInvalidateRemovesRevisionAndReaddDoesNotReuseIt(t *testing.T) {
+	manager, _ := newRevisionTestManager(t)
+	revision1, err := manager.RebuildSnapshot("service", nil)
+	require.NoError(t, err)
+
+	manager.Invalidate("service")
+	_, match := manager.ChainForRevision("service", "/", revision1)
+	assert.False(t, match)
+
+	revision2, err := manager.RebuildSnapshot("service", nil)
+	require.NoError(t, err)
+	assert.NotEqual(t, revision1, revision2)
+	_, match = manager.ChainForRevision("service", "/", revision2)
+	assert.True(t, match)
+
+	manager.InvalidateAll()
+	_, match = manager.ChainForRevision("service", "/", revision2)
+	assert.False(t, match)
+
+	revision3, err := manager.RebuildSnapshot("service", nil)
+	require.NoError(t, err)
+	assert.NotEqual(t, revision1, revision3)
+	assert.NotEqual(t, revision2, revision3)
+}
+
+func TestManagerRebuildSnapshotRejectsMismatchedBindingBeforeMutation(t *testing.T) {
+	manager, _ := newRevisionTestManager(t)
+	binding := revisionTestBinding("service", "/")
+	revision1, err := manager.RebuildSnapshot("service", []PathTargetBinding{binding})
+	require.NoError(t, err)
+	chain1, match := manager.ChainForRevision("service", "/", revision1)
+	require.True(t, match)
+
+	invalid := revisionTestBinding("other-service", "/replacement")
+	revision, err := manager.RebuildSnapshot("service", []PathTargetBinding{invalid})
+	require.Error(t, err)
+	assert.Zero(t, revision)
+	chainAfter, match := manager.ChainForRevision("service", "/", revision1)
+	assert.True(t, match)
+	assert.Same(t, chain1, chainAfter)
+
+	revision2, err := manager.RebuildSnapshot("service", nil)
+	require.NoError(t, err)
+	assert.Equal(t, revision1+1, revision2, "a rejected rebuild must not consume a revision")
+}
+
+func TestManagerConcurrentRebuildSnapshotRevisionsAreUnique(t *testing.T) {
+	manager, _ := newRevisionTestManager(t)
+	const rebuilds = 64
+
+	start := make(chan struct{})
+	revisions := make(chan Revision, rebuilds)
+	errors := make(chan error, rebuilds)
+	var wg sync.WaitGroup
+	for range rebuilds {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-start
+			revision, err := manager.RebuildSnapshot("service", nil)
+			if err != nil {
+				errors <- err
+				return
+			}
+			revisions <- revision
+		}()
+	}
+	close(start)
+	wg.Wait()
+	close(errors)
+	close(revisions)
+
+	for err := range errors {
+		require.NoError(t, err)
+	}
+	seen := make(map[Revision]struct{}, rebuilds)
+	for revision := range revisions {
+		require.NotZero(t, revision)
+		_, duplicate := seen[revision]
+		assert.False(t, duplicate, "each successful rebuild must receive a fresh revision")
+		seen[revision] = struct{}{}
+	}
+	require.Len(t, seen, rebuilds)
+
+	matches := 0
+	for revision := range seen {
+		chain, match := manager.ChainForRevision("service", "/", revision)
+		assert.Nil(t, chain)
+		if match {
+			matches++
+		}
+	}
+	assert.Equal(t, 1, matches, "only the current immutable table revision may match")
+}