diff --git a/pkg/manager/health/health.go b/pkg/manager/health/health.go index 5b5a1200b..27e6ed33e 100644 --- a/pkg/manager/health/health.go +++ b/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 { diff --git a/pkg/manager/vip/manager.go b/pkg/manager/vip/manager.go index 953e17909..4e5e16de2 100644 --- a/pkg/manager/vip/manager.go +++ b/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 } diff --git a/pkg/manager/vip/manager_test.go b/pkg/manager/vip/manager_test.go index d42ac8a4a..852352b35 100644 --- a/pkg/manager/vip/manager_test.go +++ b/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() +} diff --git a/pkg/manager/vip/mock_test.go b/pkg/manager/vip/mock_test.go index 17285b150..c81e0d8b7 100644 --- a/pkg/manager/vip/mock_test.go +++ b/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 { diff --git a/pkg/server/server.go b/pkg/server/server.go index b14e75293..6385ec20d 100644 --- a/pkg/server/server.go +++ b/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