elect: retire the owner after the etcd server recovers (#627)

Author: djshow832

pkg/balance/metricsreader/backend_reader.go

@@ -119,9 +119,8 @@ func (br *BackendReader) initElection(ctx context.Context, cfg *config.Config) e
 	} else {
 		key = fmt.Sprintf("%s/%s", readerOwnerKeyPrefix, readerOwnerKeySuffix)
 	}
-	election := elect.NewElection(br.lg, br.etcdCli, br.electionCfg, id, key, br)
-	br.election = election
-	election.Start(ctx)
+	br.election = elect.NewElection(br.lg.Named("elect"), br.etcdCli, br.electionCfg, id, key, br)
+	br.election.Start(ctx)
 	return nil
 }
 

pkg/manager/elect/election.go

@@ -6,7 +6,6 @@ package elect
 import (
 	"context"
 	"strings"
-	"sync/atomic"
 	"time"
 
 	"github.com/pingcap/tiproxy/lib/util/errors"
@@ -39,8 +38,6 @@ type Election interface {
 	Start(context.Context)
 	// ID returns the member ID.
 	ID() string
-	// IsOwner returns whether the member is the owner.
-	IsOwner() bool
 	// GetOwnerID gets the owner ID.
 	GetOwnerID(ctx context.Context) (string, error)
 	// Close stops compaining the owner.
@@ -75,7 +72,6 @@ type election struct {
 	trimedKey string
 	lg        *zap.Logger
 	etcdCli   *clientv3.Client
-	elec      atomic.Pointer[concurrency.Election]
 	wg        waitgroup.WaitGroup
 	cancel    context.CancelFunc
 	member    Member
@@ -133,19 +129,17 @@ func (m *election) initSession(ctx context.Context) (*concurrency.Session, error
 	return session, err
 }
 
-func (m *election) IsOwner() bool {
-	ownerID, err := m.GetOwnerID(context.Background())
-	if err != nil {
-		return false
-	}
-	return ownerID == m.id
-}
-
 func (m *election) campaignLoop(ctx context.Context) {
 	session, err := m.initSession(ctx)
 	if err != nil {
 		return
 	}
+	isOwner := false
+	defer func() {
+		if isOwner {
+			m.onRetired()
+		}
+	}()
 	for {
 		select {
 		case <-session.Done():
@@ -172,34 +166,44 @@ func (m *election) campaignLoop(ctx context.Context) {
 			continue
 		}
 
+		// Retire after the etcd server can be connected so that there will always be an owner.
+		// It's allowed if multiple members act as the owner but it's not allowed if no member acts as the owner.
+		// E.g. at least one member needs to bind the VIP.
+		if isOwner {
+			m.onRetired()
+			isOwner = false
+		}
+
 		elec := concurrency.NewElection(session, m.key)
 		if err = elec.Campaign(ctx, m.id); err != nil {
 			m.lg.Info("failed to campaign", zap.Error(err))
 			continue
 		}
 
-		ownerID, err := m.GetOwnerID(ctx)
-		if err != nil || ownerID != m.id {
+		kv, err := m.getOwnerInfo(ctx)
+		if err != nil {
+			m.lg.Warn("failed to get owner info", zap.Error(err))
+			continue
+		}
+		if hack.String(kv.Value) != m.id {
+			m.lg.Warn("owner id mismatches", zap.String("owner", hack.String(kv.Value)))
 			continue
 		}
 
-		m.onElected(elec)
-		// NOTICE: watchOwner won't revoke the lease.
-		m.watchOwner(ctx, session, ownerID)
-		m.onRetired()
+		m.onElected()
+		isOwner = true
+		m.watchOwner(ctx, session, hack.String(kv.Key))
 	}
 }
 
-func (m *election) onElected(elec *concurrency.Election) {
+func (m *election) onElected() {
 	m.member.OnElected()
-	m.elec.Store(elec)
 	metrics.OwnerGauge.WithLabelValues(m.trimedKey).Set(1)
 	m.lg.Info("elected as the owner")
 }
 
 func (m *election) onRetired() {
 	m.member.OnRetired()
-	m.elec.Store(nil)
 	// Delete the metric so that it doesn't show on Grafana.
 	metrics.OwnerGauge.MetricVec.DeletePartialMatch(map[string]string{metrics.LblType: m.trimedKey})
 	m.lg.Info("the owner retires")
@@ -218,17 +222,25 @@ func (m *election) revokeLease(leaseID clientv3.LeaseID) {
 
 // GetOwnerID is similar to concurrency.Election.Leader() but it doesn't need an concurrency.Election.
 func (m *election) GetOwnerID(ctx context.Context) (string, error) {
+	kv, err := m.getOwnerInfo(ctx)
+	if err != nil {
+		return "", err
+	}
+	return hack.String(kv.Value), nil
+}
+
+func (m *election) getOwnerInfo(ctx context.Context) (*mvccpb.KeyValue, error) {
 	if m.etcdCli == nil {
-		return "", concurrency.ErrElectionNoLeader
+		return nil, concurrency.ErrElectionNoLeader
 	}
 	kvs, err := etcd.GetKVs(ctx, m.etcdCli, m.key, clientv3.WithFirstCreate(), m.cfg.Timeout, m.cfg.RetryIntvl, m.cfg.RetryCnt)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 	if len(kvs) == 0 {
-		return "", concurrency.ErrElectionNoLeader
+		return nil, concurrency.ErrElectionNoLeader
 	}
-	return hack.String(kvs[0].Value), nil
+	return kvs[0], nil
 }
 
 func (m *election) watchOwner(ctx context.Context, session *concurrency.Session, key string) {

pkg/manager/elect/election_test.go

@@ -75,6 +75,8 @@ func TestEtcdServerDown(t *testing.T) {
 	addr := ts.shutdownServer()
 	_, err := elec1.GetOwnerID(context.Background())
 	require.Error(t, err)
+	// the owner should not retire before the server is up again
+	ts.expectNoEvent("1")
 	ts.startServer(addr)
 	// the previous owner only retires when the new one is elected
 	ts.expectEvent("1", eventTypeRetired, eventTypeElected)
@@ -133,15 +135,15 @@ func TestOwnerMetric(t *testing.T) {
 	}
 
 	elec1 := NewElection(lg, nil, electionConfigForTest(1), "1", ownerKeyPrefix+"key"+ownerKeySuffix, newMockMember())
-	elec1.onElected(nil)
+	elec1.onElected()
 	checkMetric("key", true)
 
 	elec2 := NewElection(lg, nil, electionConfigForTest(1), "1", "key2/1", newMockMember())
-	elec2.onElected(nil)
+	elec2.onElected()
 	checkMetric("key2/1", true)
 
 	elec3 := NewElection(lg, nil, electionConfigForTest(1), "1", ownerKeyPrefix+"key3/1", newMockMember())
-	elec3.onElected(nil)
+	elec3.onElected()
 	checkMetric("key3/1", true)
 
 	elec1.onRetired()

pkg/manager/elect/mock_test.go

@@ -51,6 +51,15 @@ func (mo *mockMember) expectEvent(t *testing.T, expected ...int) {
 	}
 }
 
+func (mo *mockMember) expectNoEvent(t *testing.T) {
+	select {
+	case <-time.After(100 * time.Millisecond):
+		return
+	case event := <-mo.ch:
+		require.Fail(t, "unexpected event", event)
+	}
+}
+
 func (mo *mockMember) hang(hang bool) {
 	contn := true
 	for contn {
@@ -144,7 +153,6 @@ func (ts *etcdTestSuite) getOwnerID() string {
 		} else {
 			require.Equal(ts.t, ownerID, id)
 		}
-		require.Equal(ts.t, elec.id == ownerID, elec.IsOwner())
 	}
 	return ownerID
 }
@@ -154,6 +162,11 @@ func (ts *etcdTestSuite) expectEvent(id string, event ...int) {
 	elec.member.(*mockMember).expectEvent(ts.t, event...)
 }
 
+func (ts *etcdTestSuite) expectNoEvent(id string) {
+	elec := ts.getElection(id)
+	elec.member.(*mockMember).expectNoEvent(ts.t)
+}
+
 func (ts *etcdTestSuite) hang(id string, hang bool) {
 	elec := ts.getElection(id)
 	elec.member.(*mockMember).hang(hang)

pkg/manager/vip/manager.go

@@ -71,15 +71,7 @@ func (vm *vipManager) Start(ctx context.Context, etcdCli *clientv3.Client) error
 
 	id := net.JoinHostPort(ip, port)
 	electionCfg := elect.DefaultElectionConfig(sessionTTL)
-	election := elect.NewElection(vm.lg, etcdCli, electionCfg, id, vipKey, vm)
-	vm.election = election
-	// Check the ownership at startup just in case the node is just down and restarted.
-	// Before it was down, it may be either the owner or not.
-	if election.IsOwner() {
-		vm.OnElected()
-	} else {
-		vm.OnRetired()
-	}
+	vm.election = elect.NewElection(vm.lg.Named("elect"), etcdCli, electionCfg, id, vipKey, vm)
 	vm.election.Start(ctx)
 	return nil
 }

pkg/manager/vip/mock_test.go

@@ -57,10 +57,6 @@ func (me *mockElection) GetOwnerID(ctx context.Context) (string, error) {
 	return "", nil
 }
 
-func (me *mockElection) IsOwner() bool {
-	return true
-}
-
 func (me *mockElection) Start(ctx context.Context) {
 	me.wg.Run(func() {
 		for {