[release-24.0] discovery: fix keyspaceState leak for keyspaces missing in localCell (#19993) (#20109)

Signed-off-by: Kyle Johnson <kyjohnson@hubspot.com>
Signed-off-by: Kyle Johnson <kylej@hubspot.com>
Co-authored-by: vitess-bot[bot] <108069721+vitess-bot[bot]@users.noreply.github.com>
Co-authored-by: Kyle Johnson <kyjohnson@hubspot.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: 3 people

go/vt/discovery/keyspace_events.go

@@ -43,6 +43,15 @@ var (
 	// waitConsistentKeyspacesCheck is the amount of time to wait for between checks to verify the keyspace is consistent.
 	waitConsistentKeyspacesCheck = 100 * time.Millisecond
 	kewHcSubscriberName          = "KeyspaceEventWatcher"
+	// missingKeyspaceTTL is how long getKeyspaceStatus caches a NoNode result
+	// for a keyspace's SrvKeyspace in the local cell. Without this cache,
+	// every healthcheck event for a tablet whose keyspace has no SrvKeyspace
+	// in the local cell would allocate a keyspaceState, register watchers,
+	// and immediately tear it down — leaking the SrvVSchema listener (whose
+	// callback always returns true) on every event. This is observable in
+	// multi-cell deployments where some keyspaces are only served from a
+	// subset of cells.
+	missingKeyspaceTTL = 30 * time.Second
 )
 
 // KeyspaceEventWatcher is an auxiliary watcher that watches all availability incidents
@@ -60,6 +69,9 @@ type KeyspaceEventWatcher struct {
 
 	mu        sync.Mutex
 	keyspaces map[string]*keyspaceState
+	// missingKeyspaces is a negative cache for keyspaces whose SrvKeyspace
+	// does not exist in localCell. Entries expire after missingKeyspaceTTL.
+	missingKeyspaces map[string]time.Time
 
 	subsMu sync.Mutex
 	subs   map[chan *KeyspaceEvent]struct{}
@@ -91,11 +103,12 @@ type ShardEvent struct {
 // will be used to detect unhealthy nodes.
 func NewKeyspaceEventWatcher(ctx context.Context, topoServer srvtopo.Server, hc HealthCheck, localCell string) *KeyspaceEventWatcher {
 	kew := &KeyspaceEventWatcher{
-		hc:        hc,
-		ts:        topoServer,
-		localCell: localCell,
-		keyspaces: make(map[string]*keyspaceState),
-		subs:      make(map[chan *KeyspaceEvent]struct{}),
+		hc:               hc,
+		ts:               topoServer,
+		localCell:        localCell,
+		keyspaces:        make(map[string]*keyspaceState),
+		missingKeyspaces: make(map[string]time.Time),
+		subs:             make(map[chan *KeyspaceEvent]struct{}),
 	}
 	kew.run(ctx)
 	log.Info(fmt.Sprintf("started watching keyspace events in %q", localCell))
@@ -126,6 +139,14 @@ func (kss *keyspaceState) isConsistent() bool {
 	return kss.consistent
 }
 
+// isDeleted returns whether the keyspace has been marked deleted by a NoNode
+// SrvKeyspace watch result.
+func (kss *keyspaceState) isDeleted() bool {
+	kss.mu.Lock()
+	defer kss.mu.Unlock()
+	return kss.deleted
+}
+
 // Format prints the internal state for this keyspace for debug purposes.
 func (kss *keyspaceState) Format(f fmt.State, verb rune) {
 	kss.mu.Lock()
@@ -618,17 +639,30 @@ func (kss *keyspaceState) isServing() bool {
 // In addition, the traffic switcher updates SrvVSchema when the DeniedTables attributes in a Shard
 // record is modified.
 func (kss *keyspaceState) onSrvVSchema(vs *vschemapb.SrvVSchema, err error) bool {
+	kss.mu.Lock()
+	defer kss.mu.Unlock()
+	// If onSrvKeyspace has already marked this keyspace deleted (NoNode in
+	// localCell), unregister this listener too — including on nil-payload
+	// updates from a server shutdown. onSrvVSchema otherwise always returns
+	// true, so the resilient SrvVSchema watcher would keep the closure —
+	// and the orphan keyspaceState it captures — in its listeners slice
+	// forever and re-run this callback's work on every update.
+	if kss.deleted {
+		return false
+	}
 	// The vschema can be nil if the server is currently shutting down.
 	if vs == nil {
 		return true
 	}
-
-	kss.mu.Lock()
-	defer kss.mu.Unlock()
-	var kerr error
-	if kss.moveTablesState, kerr = kss.getMoveTablesStatus(vs); err != nil {
+	// Use a local for the new state — getMoveTablesStatus returns (nil, err)
+	// on failure, and assigning directly into kss.moveTablesState would
+	// silently clobber the previously-tracked state on a transient topo blip.
+	newState, kerr := kss.getMoveTablesStatus(vs)
+	if kerr != nil {
 		log.Error(fmt.Sprintf("onSrvVSchema: keyspace %s failed to get move tables status: %v", kss.keyspace, kerr))
+		return true
 	}
+	kss.moveTablesState = newState
 	if kss.moveTablesState != nil && kss.moveTablesState.Typ != MoveTablesNone {
 		// Mark the keyspace as inconsistent. ensureConsistentLocked() checks if the workflow is
 		// switched, and if so, it will send an event to the buffering subscribers to indicate that
@@ -650,6 +684,14 @@ func newKeyspaceState(ctx context.Context, kew *KeyspaceEventWatcher, cell, keys
 		shards:   make(map[string]*shardState),
 	}
 	kew.ts.WatchSrvKeyspace(ctx, cell, keyspace, kss.onSrvKeyspace)
+	if kss.isDeleted() {
+		// SrvKeyspace returned NoNode synchronously via the cached value
+		// during addListener, so this keyspaceState is already discarded.
+		// Skip the SrvVSchema listener: onSrvVSchema always returns true,
+		// and the listener would never be reaped — pinning the orphan
+		// keyspaceState in the SrvVSchema watcher's listeners slice.
+		return kss
+	}
 	kew.ts.WatchSrvVSchema(ctx, cell, kss.onSrvVSchema)
 	return kss
 }
@@ -673,12 +715,22 @@ func (kew *KeyspaceEventWatcher) getKeyspaceStatus(ctx context.Context, keyspace
 	defer kew.mu.Unlock()
 	kss := kew.keyspaces[keyspace]
 	if kss == nil {
+		// Skip allocation if we recently confirmed this keyspace has no
+		// SrvKeyspace in the local cell. Healthchecks for tablets in other
+		// watched cells can otherwise drive this path >1k times/sec.
+		if t, ok := kew.missingKeyspaces[keyspace]; ok {
+			if time.Since(t) < missingKeyspaceTTL {
+				return nil
+			}
+			delete(kew.missingKeyspaces, keyspace)
+		}
 		kss = newKeyspaceState(ctx, kew, kew.localCell, keyspace)
 		kew.keyspaces[keyspace] = kss
 	}
-	if kss.deleted {
+	if kss.isDeleted() {
 		kss = nil
 		delete(kew.keyspaces, keyspace)
+		kew.missingKeyspaces[keyspace] = time.Now()
 		// Delete from the sidecar database identifier cache as well.
 		// Ignore any errors as they should all mean that the entry
 		// does not exist in the cache (which will be common).

go/vt/discovery/keyspace_events_test.go

@@ -24,6 +24,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"vitess.io/vitess/go/test/utils"
@@ -396,9 +397,10 @@ func TestWaitForConsistentKeyspaces(t *testing.T) {
 			ctx, cancel := context.WithCancel(context.Background())
 			cancel()
 			kew := KeyspaceEventWatcher{
-				keyspaces: tt.ksMap,
-				mu:        sync.Mutex{},
-				ts:        &fakeTopoServer{},
+				keyspaces:        tt.ksMap,
+				missingKeyspaces: make(map[string]time.Time),
+				mu:               sync.Mutex{},
+				ts:               &fakeTopoServer{},
 			}
 			err := kew.WaitForConsistentKeyspaces(ctx, tt.ksList)
 			if tt.errExpected != "" {
@@ -701,3 +703,118 @@ func (f *fakeTopoServer) WatchSrvVSchema(ctx context.Context, cell string, callb
 	sv, err := f.GetSrvVSchema(ctx, cell)
 	callback(sv, err)
 }
+
+// fakeMissingKeyspaceTopoServer is a fakeTopoServer whose WatchSrvKeyspace
+// returns NoNode for keyspaces in the missing set, simulating a keyspace
+// that exists in the cluster but has no SrvKeyspace in the local cell.
+// It also counts how many times each Watch is invoked so tests can assert
+// the negative cache and the SrvVSchema-skip behavior in
+// KeyspaceEventWatcher.
+type fakeMissingKeyspaceTopoServer struct {
+	fakeTopoServer
+	missing               map[string]struct{}
+	watchSrvKeyspaceCalls atomic.Int64
+	watchSrvVSchemaCalls  atomic.Int64
+}
+
+func (f *fakeMissingKeyspaceTopoServer) WatchSrvKeyspace(ctx context.Context, cell, keyspace string, callback func(*topodatapb.SrvKeyspace, error) bool) {
+	f.watchSrvKeyspaceCalls.Add(1)
+	if _, ok := f.missing[keyspace]; ok {
+		callback(nil, topo.NewError(topo.NoNode, keyspace))
+		return
+	}
+	f.fakeTopoServer.WatchSrvKeyspace(ctx, cell, keyspace, callback)
+}
+
+func (f *fakeMissingKeyspaceTopoServer) WatchSrvVSchema(ctx context.Context, cell string, callback func(*vschemapb.SrvVSchema, error) bool) {
+	f.watchSrvVSchemaCalls.Add(1)
+	f.fakeTopoServer.WatchSrvVSchema(ctx, cell, callback)
+}
+
+// TestKeyspaceEventWatcherMissingKeyspaceCache verifies that healthchecks for
+// a keyspace whose SrvKeyspace does not exist in the local cell don't allocate
+// a new keyspaceState (and don't register a SrvVSchema listener) on every
+// event. Without the negative cache and the SrvVSchema-skip, this path used
+// to fire on every healthcheck event and pin orphan keyspaceStates in the
+// SrvVSchema watcher's listeners slice (onSrvVSchema always returns true,
+// so the listener was never reaped).
+func TestKeyspaceEventWatcherMissingKeyspaceCache(t *testing.T) {
+	cell := "cell1"
+	missing := "missing-keyspace"
+
+	sts := &fakeMissingKeyspaceTopoServer{
+		missing: map[string]struct{}{missing: {}},
+	}
+	hc := NewFakeHealthCheck(make(chan *TabletHealth))
+	t.Cleanup(func() { hc.Close() })
+
+	kew := &KeyspaceEventWatcher{
+		hc:               hc,
+		ts:               sts,
+		localCell:        cell,
+		keyspaces:        make(map[string]*keyspaceState),
+		missingKeyspaces: make(map[string]time.Time),
+		subs:             make(map[chan *KeyspaceEvent]struct{}),
+	}
+
+	const lookups = 100
+	for range lookups {
+		require.Nil(t, kew.getKeyspaceStatus(t.Context(), missing),
+			"getKeyspaceStatus must return nil for a keyspace missing from localCell")
+	}
+
+	assert.Equal(t, int64(1), sts.watchSrvKeyspaceCalls.Load(),
+		"WatchSrvKeyspace should be called at most once within missingKeyspaceTTL — the negative cache should short-circuit subsequent lookups")
+	assert.Equal(t, int64(0), sts.watchSrvVSchemaCalls.Load(),
+		"WatchSrvVSchema must never be called when SrvKeyspace returns NoNode synchronously — otherwise onSrvVSchema (always returning true) pins an orphan keyspaceState in the listeners slice")
+
+	// Force expiry of the negative cache and confirm exactly one re-allocation.
+	kew.mu.Lock()
+	for k := range kew.missingKeyspaces {
+		kew.missingKeyspaces[k] = time.Now().Add(-2 * missingKeyspaceTTL)
+	}
+	kew.mu.Unlock()
+
+	require.Nil(t, kew.getKeyspaceStatus(t.Context(), missing))
+	assert.Equal(t, int64(2), sts.watchSrvKeyspaceCalls.Load(),
+		"after the negative-cache TTL elapses, the next lookup should re-probe SrvKeyspace exactly once")
+	assert.Equal(t, int64(0), sts.watchSrvVSchemaCalls.Load(),
+		"WatchSrvVSchema must still not be called after re-probing a missing keyspace")
+}
+
+// TestOnSrvVSchemaUnregistersAfterDelete covers the async-deletion path: a
+// keyspace exists in localCell when newKeyspaceState registers onSrvVSchema,
+// then later disappears. onSrvKeyspace marks the state deleted and returns
+// false (reaping itself), but onSrvVSchema must also return false on its next
+// invocation — for every payload shape, including the nil "server shutting
+// down" case. Otherwise the resilient SrvVSchema watcher keeps the closure
+// (and the orphan keyspaceState it captures) in its listeners slice forever
+// and re-runs the callback's work on every SrvVSchema update.
+func TestOnSrvVSchemaUnregistersAfterDelete(t *testing.T) {
+	cases := []struct {
+		name string
+		vs   *vschemapb.SrvVSchema
+		err  error
+	}{
+		{"non-nil payload", &vschemapb.SrvVSchema{}, nil},
+		{"nil payload (server shutdown)", nil, nil},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			kss := &keyspaceState{
+				keyspace: "ks1",
+				shards:   make(map[string]*shardState),
+			}
+
+			// Simulate the async deletion: SrvKeyspace returns NoNode for localCell.
+			require.False(t, kss.onSrvKeyspace(nil, topo.NewError(topo.NoNode, "ks1")),
+				"onSrvKeyspace must return false on NoNode so the SrvKeyspace watcher reaps it")
+			require.True(t, kss.isDeleted())
+
+			// The next SrvVSchema update must reap the orphan listener
+			// regardless of payload shape.
+			require.False(t, kss.onSrvVSchema(tc.vs, tc.err),
+				"onSrvVSchema must return false once the keyspace is deleted, otherwise the listener pins an orphan keyspaceState")
+		})
+	}
+}