vip: switch VIP when the owner rejects connections

pkg/manager/health/health.go

@@ -52,9 +52,10 @@ func (m *Manager) Healthy() (bool, string) {
 }
 
 // RejectConns reports whether new connections should be rejected and returns a
-// reason string. Used by the proxy server. It returns true only on memory
-// pressure; graceful shutdown does NOT reject here, because the proxy keeps
-// accepting until its listeners are closed.
+// reason string. Used by the proxy server and the VIP manager. It returns true
+// only on memory pressure; graceful shutdown does NOT reject here, because the
+// proxy keeps accepting until its listeners are closed. The VIP release on
+// graceful shutdown is handled by PreClose.
 func (m *Manager) RejectConns() (bool, string) {
 	if m.rejectCheck != nil {
 		if reject, reason := m.rejectCheck(); reject {

pkg/manager/vip/manager.go

@@ -31,12 +31,28 @@ const (
 	garpRefreshInterval = 1 * time.Second
 )
 
+// servingCheckInterval is how often the VIP manager re-evaluates whether the
+// instance is rejecting connections. It is a variable so tests can
+// shorten it to speed up the resign/recompete transitions.
+var servingCheckInterval = 1 * time.Second
+
 type VIPManager interface {
 	Start(context.Context, *clientv3.Client) error
+	// SetConnRejecter configures the checker that decides whether the instance is
+	// rejecting new connections. When it reports a rejection, the manager resigns
+	// the VIP owner so a healthy node takes over, and re-campaigns once it stops
+	// rejecting.
+	SetConnRejecter(ConnRejecter)
 	PreClose()
 	Close()
 }
 
+// ConnRejecter reports whether the instance is currently rejecting new
+// connections.
+type ConnRejecter interface {
+	RejectConns() (bool, string)
+}
+
 var _ VIPManager = (*vipManager)(nil)
 
 type vipManager struct {
@@ -52,6 +68,17 @@ type vipManager struct {
 	cfgGetter config.ConfigGetter
 	election  elect.Election
 	lg        *zap.Logger
+	// rejecter, when set, drives the resign/recompete loop. nil keeps the
+	// historical behavior of always competing for the VIP.
+	rejecter ConnRejecter
+	// newElection creates and starts a fresh election. It is set in Start,
+	// closing over the etcd client, election parameters, and the start context
+	// so they don't need to live as separate fields. Tests override it to avoid
+	// a real etcd server.
+	newElection func() elect.Election
+	startCtx    context.Context
+	startCancel context.CancelFunc
+	watchWG     waitgroup.WaitGroup
 }
 
 func NewVIPManager(lg *zap.Logger, cfgGetter config.ConfigGetter) (*vipManager, error) {
@@ -76,6 +103,14 @@ func NewVIPManager(lg *zap.Logger, cfgGetter config.ConfigGetter) (*vipManager,
 	return vm, nil
 }
 
+// SetConnRejecter configures the checker that decides whether this instance is
+// rejecting new connections. It must be called before Start. When the checker
+// reports a rejection, the VIP manager resigns the owner so a healthy node
+// takes over, and re-campaigns once it stops rejecting.
+func (vm *vipManager) SetConnRejecter(r ConnRejecter) {
+	vm.rejecter = r
+}
+
 func (vm *vipManager) Start(ctx context.Context, etcdCli *clientv3.Client) error {
 	vm.mu.Lock()
 	defer vm.mu.Unlock()
@@ -93,11 +128,83 @@ func (vm *vipManager) Start(ctx context.Context, etcdCli *clientv3.Client) error
 	id := net.JoinHostPort(ip, port)
 	electionCfg := elect.DefaultElectionConfig(sessionTTL)
 	key := fmt.Sprintf(vipKey, vm.operation.Addr())
-	vm.election = elect.NewElection(vm.lg.Named("elect"), etcdCli, electionCfg, id, key, vm)
-	vm.election.Start(ctx)
+	vm.startCtx, vm.startCancel = context.WithCancel(ctx)
+	startCtx := vm.startCtx
+	vm.newElection = func() elect.Election {
+		e := elect.NewElection(vm.lg.Named("elect"), etcdCli, electionCfg, id, key, vm)
+		e.Start(startCtx)
+		return e
+	}
+	vm.election = vm.newElection()
+
+	// The watcher is a no-op when there is no rejecter, preserving the
+	// historical always-compete behavior.
+	if vm.rejecter != nil {
+		watchCtx := vm.startCtx
+		vm.watchWG.RunWithRecover(func() {
+			vm.watchRejecter(watchCtx)
+		}, nil, vm.lg)
+	}
 	return nil
 }
 
+// resign drops the current election so the etcd lease is revoked and another
+// node can take over. election.Close retires this member, which removes the
+// VIP locally via OnRetired.
+func (vm *vipManager) resign() {
+	vm.mu.Lock()
+	if vm.closing || vm.election == nil {
+		vm.mu.Unlock()
+		return
+	}
+	election := vm.election
+	vm.election = nil
+	vm.mu.Unlock()
+	// election.Close synchronously invokes OnRetired, which needs vm.mu, so it
+	// must run outside the lock to avoid a self-deadlock.
+	vm.lg.Info("resign VIP owner because the instance rejects new connections")
+	election.Close()
+}
+
+// recompete creates a new election and starts campaigning after recovery.
+func (vm *vipManager) recompete() {
+	vm.mu.Lock()
+	defer vm.mu.Unlock()
+	if vm.closing || vm.election != nil {
+		return
+	}
+	vm.lg.Info("recompete for VIP because the instance accepts new connections")
+	vm.election = vm.newElection()
+}
+
+// watchRejecter polls the rejecter and flips the VIP ownership on transitions.
+// Polling is intentional: the health signal is a simple boolean with no event
+// channel, and the interval is bounded by the election TTL so failover latency
+// stays comparable to an etcd session expiry. The instance never rejects at
+// startup, so the initial state is assumed to be not rejecting.
+func (vm *vipManager) watchRejecter(ctx context.Context) {
+	ticker := time.NewTicker(servingCheckInterval)
+	defer ticker.Stop()
+	wasRejecting := false
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			reject, _ := vm.rejecter.RejectConns()
+			if reject == wasRejecting {
+				continue
+			}
+			wasRejecting = reject
+			if reject {
+				vm.resign()
+			} else {
+				vm.recompete()
+			}
+		}
+	}
+}
+
 func (vm *vipManager) OnElected() {
 	vm.mu.Lock()
 	// Election.Close may race with an already in-flight OnElected callback.
@@ -208,6 +315,7 @@ func (vm *vipManager) stopARPRefresh() {
 // shutdowns do not expose the VIP on two nodes at the same time.
 func (vm *vipManager) PreClose() {
 	election := vm.prepareForClose()
+	vm.watchWG.Wait()
 	if election != nil {
 		election.Close()
 	}
@@ -216,6 +324,7 @@ func (vm *vipManager) PreClose() {
 // Close resigns the owner and makes sure the VIP is removed locally.
 func (vm *vipManager) Close() {
 	election := vm.prepareForClose()
+	vm.watchWG.Wait()
 	if election != nil {
 		election.Close()
 	}
@@ -231,5 +340,14 @@ func (vm *vipManager) prepareForClose() elect.Election {
 	vm.closing = true
 	vm.stopARPRefresh()
 	vm.delVIP(context.Background())
-	return vm.election
+	// Cancel the start context so the rejecter watcher and any in-flight election
+	// campaign loop stop. watchWG is waited outside the lock to avoid deadlocking
+	// with a watcher callback that may be waiting for vm.mu.
+	if vm.startCancel != nil {
+		vm.startCancel()
+		vm.startCancel = nil
+	}
+	election := vm.election
+	vm.election = nil
+	return election
 }

pkg/manager/vip/manager_test.go

@@ -7,12 +7,15 @@ import (
 	"context"
 	"runtime"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/pingcap/tiproxy/lib/config"
 	"github.com/pingcap/tiproxy/lib/util/logger"
 	"github.com/pingcap/tiproxy/pkg/manager/cert"
+	"github.com/pingcap/tiproxy/pkg/manager/elect"
 	"github.com/pingcap/tiproxy/pkg/util/etcd"
 	"github.com/stretchr/testify/require"
 )
@@ -295,6 +298,58 @@ func TestStartAndClose(t *testing.T) {
 	}
 }
 
+func TestRejecterWatcherResignsAndRecompetes(t *testing.T) {
+	// Shorten the polling interval so transitions happen within the test window.
+	origInterval := servingCheckInterval
+	servingCheckInterval = 10 * time.Millisecond
+	defer func() { servingCheckInterval = origInterval }()
+
+	lg, _ := logger.CreateLoggerForTest(t)
+	operation := newMockNetworkOperation()
+	operation.hasIP.Store(false)
+	rejecter := &mockConnRejecter{}
+	rejecter.reject.Store(false)
+
+	vm := &vipManager{
+		lg:        lg,
+		cfgGetter: newMockConfigGetter(newMockConfig()),
+		operation: operation,
+		rejecter:  rejecter,
+	}
+	vm.newElection = func() elect.Election {
+		e := &autoElection{member: vm}
+		e.Start(vm.startCtx)
+		return e
+	}
+	vm.startCtx, vm.startCancel = context.WithCancel(context.Background())
+	// Start competing: the auto election wins immediately and binds the VIP.
+	vm.mu.Lock()
+	vm.election = vm.newElection()
+	vm.mu.Unlock()
+	require.Eventually(t, func() bool { return operation.hasIP.Load() }, time.Second, 10*time.Millisecond)
+	require.EqualValues(t, 1, operation.addIPCnt.Load())
+
+	// Start the rejecter watcher, driven by vm.startCtx so that PreClose's
+	// cancellation of startCtx (and watchWG wait) is the shutdown path under test.
+	vm.watchWG.RunWithRecover(func() { vm.watchRejecter(vm.startCtx) }, nil, vm.lg)
+
+	// Flip to rejecting: the watcher resigns, which retires and drops the VIP.
+	rejecter.reject.Store(true)
+	require.Eventually(t, func() bool { return !operation.hasIP.Load() }, 3*time.Second, 10*time.Millisecond)
+	require.EqualValues(t, 1, operation.delIPCnt.Load())
+
+	// Flip back to not rejecting: the watcher recompetes, wins, and rebinds the VIP.
+	rejecter.reject.Store(false)
+	require.Eventually(t, func() bool { return operation.hasIP.Load() }, 3*time.Second, 10*time.Millisecond)
+	require.EqualValues(t, 2, operation.addIPCnt.Load())
+
+	// Shutdown: closing must release the VIP and stop the watcher without deadlock.
+	// PreClose cancels vm.startCtx (stopping the watcher) and waits for watchWG.
+	vm.PreClose()
+	vm.Close()
+	require.False(t, operation.hasIP.Load())
+}
+
 func TestMultiVIP(t *testing.T) {
 	if runtime.GOOS != "linux" {
 		return
@@ -356,3 +411,43 @@ func (e *closeHookElection) Close() {
 		e.closeFn()
 	}
 }
+
+var _ elect.Election = (*autoElection)(nil)
+
+// autoElection calls OnElected on Start (simulating an immediate win) and
+// OnRetired on Close (simulating losing ownership). OnElected is delivered
+// asynchronously, like the real election, so callers that start the election
+// while holding the member's lock do not deadlock.
+type autoElection struct {
+	member  elect.Member
+	started atomic.Bool
+	closed  atomic.Bool
+	wg      sync.WaitGroup
+}
+
+func (e *autoElection) Start(context.Context) {
+	if e.started.Swap(true) {
+		return
+	}
+	e.wg.Add(1)
+	go func() {
+		defer e.wg.Done()
+		e.member.OnElected()
+	}()
+}
+
+func (e *autoElection) ID() string {
+	return ""
+}
+
+func (e *autoElection) GetOwnerID(context.Context) (string, error) {
+	return "", nil
+}
+
+func (e *autoElection) Close() {
+	if !e.started.Load() || e.closed.Swap(true) {
+		return
+	}
+	e.wg.Wait()
+	e.member.OnRetired()
+}

pkg/manager/vip/mock_test.go

@@ -89,6 +89,19 @@ func (me *mockElection) Close() {
 	me.wg.Wait()
 }
 
+var _ ConnRejecter = (*mockConnRejecter)(nil)
+
+type mockConnRejecter struct {
+	reject atomic.Bool
+}
+
+func (m *mockConnRejecter) RejectConns() (bool, string) {
+	if m.reject.Load() {
+		return true, "mock reject"
+	}
+	return false, ""
+}
+
 var _ NetworkOperation = (*mockNetworkOperation)(nil)
 
 type mockNetworkOperation struct {

pkg/server/server.go

@@ -217,6 +217,9 @@ func NewServer(ctx context.Context, sctx *sctx.Context) (srv *Server, err error)
 			return
 		}
 		if srv.vipManager != nil && !reflect.ValueOf(srv.vipManager).IsNil() {
+			// Release the VIP while this instance is rejecting connections (e.g. under
+			// memory pressure) so a healthy node takes over, and re-campaign on recovery.
+			srv.vipManager.SetConnRejecter(srv.healthMgr)
 			if vipEtcdCli != nil {
 				if err = srv.vipManager.Start(ctx, vipEtcdCli); err != nil {
 					return