vip: support multiple VIP (#705) (#825)

Co-authored-by: djshow832 <[email protected]>

Author: ti-chi-bot

pkg/balance/metricsreader/backend_reader_test.go

@@ -1297,8 +1297,7 @@ func TestElection(t *testing.T) {
 	// test owner
 	suite.delKV(ownerKey)
 	require.Eventually(t, func() bool {
-		kvs := suite.getKV(ownerKey)
-		return len(kvs) == 1 && strings.HasSuffix(string(kvs[0].Value), "3080")
+		return br.isOwner.Load()
 	}, 3*time.Second, 10*time.Millisecond)
 	err = br.ReadMetrics(context.Background())
 	require.NoError(t, err)

pkg/manager/vip/manager.go

@@ -5,6 +5,7 @@ package vip
 
 import (
 	"context"
+	"fmt"
 	"net"
 	"sync/atomic"
 
@@ -15,8 +16,9 @@ import (
 )
 
 const (
-	// vipKey is the key in etcd for VIP election.
-	vipKey = "/tiproxy/vip/owner"
+	// vipKey is the key in etcd for VIP election. The key contains the VIP address so that
+	// multiple VIP can coexist.
+	vipKey = "/tiproxy/vip/%s/owner"
 	// sessionTTL is the session's TTL in seconds for VIP election.
 	// The etcd client keeps alive every TTL/3 seconds.
 	// The TTL determines the failover time so it should be short.
@@ -74,7 +76,8 @@ func (vm *vipManager) Start(ctx context.Context, etcdCli *clientv3.Client) error
 
 	id := net.JoinHostPort(ip, port)
 	electionCfg := elect.DefaultElectionConfig(sessionTTL)
-	vm.election = elect.NewElection(vm.lg.Named("elect"), etcdCli, electionCfg, id, vipKey, vm)
+	key := fmt.Sprintf(vipKey, vm.operation.Addr())
+	vm.election = elect.NewElection(vm.lg.Named("elect"), etcdCli, electionCfg, id, key, vm)
 	vm.election.Start(ctx)
 	return nil
 }

pkg/manager/vip/manager_test.go

@@ -123,7 +123,7 @@ func TestNetworkOperation(t *testing.T) {
 	operation := newMockNetworkOperation()
 	vm := &vipManager{
 		lg:        lg,
-		cfgGetter: newMockConfigGetter(&config.Config{HA: config.HA{VirtualIP: "10.10.10.10/24", Interface: "eth0"}}),
+		cfgGetter: newMockConfigGetter(newMockConfig()),
 		operation: operation,
 	}
 	vm.delOnRetire.Store(true)
@@ -149,11 +149,7 @@ func TestNetworkOperation(t *testing.T) {
 
 func TestStartAndClose(t *testing.T) {
 	lg, _ := logger.CreateLoggerForTest(t)
-	vm, err := NewVIPManager(lg, newMockConfigGetter(&config.Config{
-		Proxy: config.ProxyServer{Addr: "0.0.0.0:6000"},
-		API:   config.API{Addr: "0.0.0.0:3080"},
-		HA:    config.HA{VirtualIP: "127.0.0.2/24", Interface: "lo"},
-	}))
+	vm, err := NewVIPManager(lg, newMockConfigGetter(newMockConfig()))
 	if runtime.GOOS != "linux" {
 		require.Error(t, err)
 	} else {
@@ -178,3 +174,44 @@ func TestStartAndClose(t *testing.T) {
 		vm.Close()
 	}
 }
+
+func TestMultiVIP(t *testing.T) {
+	if runtime.GOOS != "linux" {
+		return
+	}
+	config1, config2 := newMockConfig(), newMockConfig()
+	config1.HA.VirtualIP = "127.0.0.2/24"
+	config2.HA.VirtualIP = "127.0.0.3/24"
+	lg, text := logger.CreateLoggerForTest(t)
+	vm1, err := NewVIPManager(lg, newMockConfigGetter(config1))
+	// Maybe interface doesn't exist or lack of permission.
+	if err != nil {
+		return
+	}
+	vm2, err := NewVIPManager(lg, newMockConfigGetter(config2))
+	require.NoError(t, err)
+
+	server, err := etcd.CreateEtcdServer("0.0.0.0:0", t.TempDir(), lg)
+	require.NoError(t, err)
+	endpoint := server.Clients[0].Addr().String()
+	cfg := etcd.ConfigForEtcdTest(endpoint)
+
+	certMgr := cert.NewCertManager()
+	err = certMgr.Init(cfg, lg, nil)
+	require.NoError(t, err)
+	client, err := etcd.InitEtcdClient(lg, cfg, certMgr)
+	require.NoError(t, err)
+
+	err = vm1.Start(context.Background(), client)
+	require.NoError(t, err)
+	err = vm2.Start(context.Background(), client)
+	require.NoError(t, err)
+	require.Eventually(t, func() bool {
+		return strings.Count(text.String(), "adding VIP success") >= 2 ||
+			strings.Count(text.String(), "ip: command not found") >= 2
+	}, 3*time.Second, 10*time.Millisecond)
+	vm1.PreClose()
+	vm2.PreClose()
+	vm1.Close()
+	vm2.Close()
+}

pkg/manager/vip/mock_test.go

@@ -121,3 +121,15 @@ func (mno *mockNetworkOperation) SendARP() error {
 	}
 	return nil
 }
+
+func (mno *mockNetworkOperation) Addr() string {
+	return ""
+}
+
+func newMockConfig() *config.Config {
+	return &config.Config{
+		Proxy: config.ProxyServer{Addr: "0.0.0.0:6000"},
+		API:   config.API{Addr: "0.0.0.0:3080"},
+		HA:    config.HA{VirtualIP: "127.0.0.2/24", Interface: "lo"},
+	}
+}

pkg/manager/vip/network.go

@@ -20,6 +20,7 @@ type NetworkOperation interface {
 	AddIP() error
 	DeleteIP() error
 	SendARP() error
+	Addr() string
 }
 
 var _ NetworkOperation = (*networkOperation)(nil)
@@ -93,3 +94,10 @@ func (no *networkOperation) SendARP() error {
 	}
 	return errors.WithStack(err)
 }
+
+func (no *networkOperation) Addr() string {
+	if no.address == nil {
+		return ""
+	}
+	return no.address.IP.String()
+}

pkg/manager/vip/network_test.go

@@ -21,18 +21,19 @@ func TestAddDelIP(t *testing.T) {
 	tests := []struct {
 		virtualIP string
 		link      string
+		addr      string
 		initErr   string
-		addErr    string
 		delErr    string
-		sendErr   string
 	}{
 		{
 			virtualIP: "127.0.0.2/24",
+			addr:      "127.0.0.2",
 			link:      "lo",
 		},
 		{
 			virtualIP: "0.0.0.0/24",
 			link:      "lo",
+			addr:      "0.0.0.0",
 			delErr:    "cannot assign requested address",
 		},
 		{
@@ -55,27 +56,18 @@ func TestAddDelIP(t *testing.T) {
 		}
 		require.NoError(t, err, "case %d", i)
 		require.NotNil(t, operation, "case %d", i)
+		require.Equal(t, test.addr, operation.Addr(), "case %d", i)
 
 		err = operation.AddIP()
 		// Maybe the command is not installed.
-		if err != nil && isOtherErr(err) {
+		if err != nil {
+			require.True(t, isOtherErr(err))
 			continue
 		}
-		if test.addErr != "" {
-			require.Error(t, err, "case %d", i)
-			require.Contains(t, err.Error(), test.addErr, "case %d", i)
-		} else {
-			require.NoError(t, err, "case %d", i)
-		}
 
 		err = operation.SendARP()
-		if err == nil || !isOtherErr(err) {
-			if test.sendErr != "" {
-				require.Error(t, err, "case %d", i)
-				require.Contains(t, err.Error(), test.sendErr, "case %d", i)
-			} else {
-				require.NoError(t, err, "case %d", i)
-			}
+		if err != nil {
+			require.True(t, isOtherErr(err))
 		}
 
 		if err := operation.DeleteIP(); test.delErr != "" {