github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/pkg/client/internal/leader_resolver_test.go (about)

     1  // Copyright 2022 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package internal
    15  
    16  import (
    17  	"reflect"
    18  	"testing"
    19  	"time"
    20  
    21  	"github.com/stretchr/testify/mock"
    22  	"github.com/stretchr/testify/require"
    23  	"google.golang.org/grpc/resolver"
    24  	"google.golang.org/grpc/serviceconfig"
    25  )
    26  
    27  type mockResolverClientConn struct {
    28  	resolver.ClientConn
    29  	mock.Mock
    30  }
    31  
    32  func (cc *mockResolverClientConn) UpdateState(state resolver.State) error {
    33  	args := cc.Called(state)
    34  	return args.Error(0)
    35  }
    36  
    37  func (cc *mockResolverClientConn) ReportError(err error) {
    38  	cc.Called(err)
    39  }
    40  
    41  func (cc *mockResolverClientConn) ParseServiceConfig(serviceConfigJSON string) *serviceconfig.ParseResult {
    42  	args := cc.Called(serviceConfigJSON)
    43  	return args.Get(0).(*serviceconfig.ParseResult)
    44  }
    45  
    46  func TestLeaderResolver(t *testing.T) {
    47  	t.Parallel()
    48  
    49  	mockConn := &mockResolverClientConn{}
    50  	leaderResolver := newLeaderResolverWithFollowerSorter(map[string]bool{
    51  		"leader":     true,
    52  		"follower-1": false,
    53  		"follower-2": false,
    54  	}, orderedFollowerSorter)
    55  	defer leaderResolver.Close()
    56  
    57  	mockConn.On("UpdateState", mock.MatchedBy(func(state resolver.State) bool {
    58  		return reflect.DeepEqual(state.Addresses, []resolver.Address{
    59  			{Addr: "leader", ServerName: "leader"},
    60  			{Addr: "follower-1", ServerName: "follower-1"},
    61  			{Addr: "follower-2", ServerName: "follower-2"},
    62  		})
    63  	})).Return(nil).Once()
    64  	mockConn.On("ParseServiceConfig", `{"loadBalancingPolicy": "pick_first"}`).
    65  		Return(&serviceconfig.ParseResult{}).Once()
    66  
    67  	res, err := leaderResolver.Build(resolver.Target{}, mockConn, resolver.BuildOptions{})
    68  	require.NoError(t, err)
    69  
    70  	res.ResolveNow(struct{}{})
    71  
    72  	mockConn.On("UpdateState", mock.MatchedBy(func(state resolver.State) bool {
    73  		return reflect.DeepEqual(state.Addresses, []resolver.Address{
    74  			{Addr: "leader", ServerName: "leader"},
    75  			{Addr: "follower-2", ServerName: "follower-2"},
    76  			{Addr: "follower-3", ServerName: "follower-3"},
    77  		})
    78  	})).Return(nil).Once()
    79  	leaderResolver.UpdateServerList(map[string]bool{
    80  		"leader":     true,
    81  		"follower-2": false,
    82  		"follower-3": false,
    83  	})
    84  	time.Sleep(100 * time.Millisecond)
    85  
    86  	mockConn.AssertExpectations(t)
    87  }
    88  
    89  func TestLeaderResolverFrequentUpdate(t *testing.T) {
    90  	t.Parallel()
    91  
    92  	mockConn := &mockResolverClientConn{}
    93  	leaderResolver := newLeaderResolverWithFollowerSorter(map[string]bool{
    94  		"leader":     true,
    95  		"follower-1": false,
    96  		"follower-2": false,
    97  	}, orderedFollowerSorter)
    98  	defer leaderResolver.Close()
    99  
   100  	mockConn.On("UpdateState", mock.MatchedBy(func(state resolver.State) bool {
   101  		return reflect.DeepEqual(state.Addresses, []resolver.Address{
   102  			{Addr: "leader", ServerName: "leader"},
   103  			{Addr: "follower-1", ServerName: "follower-1"},
   104  			{Addr: "follower-2", ServerName: "follower-2"},
   105  		})
   106  	})).Return(nil)
   107  	mockConn.On("ParseServiceConfig", `{"loadBalancingPolicy": "pick_first"}`).
   108  		Return(&serviceconfig.ParseResult{}).Once()
   109  
   110  	_, err := leaderResolver.Build(resolver.Target{}, mockConn, resolver.BuildOptions{})
   111  	require.NoError(t, err)
   112  
   113  	for i := 0; i < 1000; i++ {
   114  		leaderResolver.UpdateServerList(map[string]bool{
   115  			"leader":     true,
   116  			"follower-1": false,
   117  			"follower-2": false,
   118  		})
   119  	}
   120  	time.Sleep(100 * time.Millisecond)
   121  
   122  	mockConn.AssertExpectations(t)
   123  }
   124  
   125  func TestLeaderResolverUpdateAfterClose(t *testing.T) {
   126  	t.Parallel()
   127  
   128  	mockConn := &mockResolverClientConn{}
   129  	leaderResolver := newLeaderResolverWithFollowerSorter(map[string]bool{
   130  		"leader":     true,
   131  		"follower-1": false,
   132  		"follower-2": false,
   133  	}, orderedFollowerSorter)
   134  
   135  	mockConn.On("UpdateState", mock.MatchedBy(func(state resolver.State) bool {
   136  		return reflect.DeepEqual(state.Addresses, []resolver.Address{
   137  			{Addr: "leader", ServerName: "leader"},
   138  			{Addr: "follower-1", ServerName: "follower-1"},
   139  			{Addr: "follower-2", ServerName: "follower-2"},
   140  		})
   141  	})).Return(nil)
   142  	mockConn.On("ParseServiceConfig", `{"loadBalancingPolicy": "pick_first"}`).
   143  		Return(&serviceconfig.ParseResult{}).Once()
   144  
   145  	_, err := leaderResolver.Build(resolver.Target{}, mockConn, resolver.BuildOptions{})
   146  	require.NoError(t, err)
   147  
   148  	// Close the resolver now.
   149  	leaderResolver.Close()
   150  
   151  	// Updating the server list after the resolver is closed
   152  	// should not panic or deadlock.
   153  	for i := 0; i < 1000; i++ {
   154  		leaderResolver.UpdateServerList(map[string]bool{
   155  			"leader":     true,
   156  			"follower-1": false,
   157  			"follower-2": false,
   158  		})
   159  	}
   160  }