github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/balancer/cluster/cluster_test.go (about)

     1  package cluster
     2  
     3  import (
     4  	"context"
     5  	"math"
     6  	"strconv"
     7  	"sync"
     8  	"testing"
     9  
    10  	"github.com/stretchr/testify/require"
    11  
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
    13  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/mock"
    14  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    15  )
    16  
    17  func TestCluster(t *testing.T) {
    18  	ctx := xtest.Context(t)
    19  
    20  	t.Run("Nil", func(t *testing.T) {
    21  		var s *Cluster
    22  
    23  		require.Empty(t, s.All())
    24  
    25  		e, err := s.Next(ctx)
    26  		require.ErrorIs(t, err, ErrNilPtr)
    27  		require.Nil(t, e)
    28  	})
    29  
    30  	t.Run("Empty", func(t *testing.T) {
    31  		s := New(nil)
    32  
    33  		require.Empty(t, s.All())
    34  
    35  		e, err := s.Next(ctx)
    36  		require.ErrorIs(t, err, ErrNoEndpoints)
    37  		require.Nil(t, e)
    38  	})
    39  
    40  	t.Run("One", func(t *testing.T) {
    41  		s := New([]endpoint.Endpoint{&mock.Endpoint{
    42  			AddrField:   "1",
    43  			NodeIDField: 1,
    44  		}})
    45  
    46  		require.Len(t, s.All(), 1)
    47  
    48  		e, err := s.Next(ctx)
    49  		require.NoError(t, err)
    50  		require.NotNil(t, e)
    51  		require.Equal(t, "1", e.Address())
    52  		require.Equal(t, uint32(1), e.NodeID())
    53  	})
    54  
    55  	t.Run("ContextError", func(t *testing.T) {
    56  		ctxWithCancel, cancel := context.WithCancel(ctx)
    57  		cancel()
    58  
    59  		s := New([]endpoint.Endpoint{&mock.Endpoint{
    60  			AddrField:   "1",
    61  			NodeIDField: 1,
    62  		}})
    63  
    64  		require.Len(t, s.All(), 1)
    65  
    66  		e, err := s.Next(ctxWithCancel)
    67  		require.ErrorIs(t, err, context.Canceled)
    68  		require.Nil(t, e)
    69  	})
    70  
    71  	t.Run("Without", func(t *testing.T) {
    72  		s := New([]endpoint.Endpoint{
    73  			&mock.Endpoint{
    74  				AddrField:   "1",
    75  				NodeIDField: 1,
    76  			},
    77  			&mock.Endpoint{
    78  				AddrField:   "2",
    79  				NodeIDField: 2,
    80  			},
    81  			&mock.Endpoint{
    82  				AddrField:   "3",
    83  				NodeIDField: 3,
    84  			},
    85  			&mock.Endpoint{
    86  				AddrField:   "4",
    87  				NodeIDField: 4,
    88  			},
    89  			&mock.Endpoint{
    90  				AddrField:   "5",
    91  				NodeIDField: 5,
    92  			},
    93  		})
    94  
    95  		{ // initial state
    96  			require.Len(t, s.All(), 5)
    97  			require.Len(t, s.index, 5)
    98  			require.Len(t, s.prefer, 5)
    99  		}
   100  
   101  		{ // without first endpoint
   102  			e, err := s.Next(ctx)
   103  			require.NoError(t, err)
   104  			require.NotNil(t, e)
   105  			s = Without(s, e)
   106  			require.Len(t, s.All(), 5)
   107  			require.Len(t, s.index, 5)
   108  			require.Len(t, s.prefer, 4)
   109  			require.Len(t, s.fallback, 1)
   110  		}
   111  
   112  		{ // without second endpoint
   113  			e, err := s.Next(ctx)
   114  			require.NoError(t, err)
   115  			require.NotNil(t, e)
   116  			s = Without(s, e)
   117  			require.Len(t, s.All(), 5)
   118  			require.Len(t, s.index, 5)
   119  			require.Len(t, s.prefer, 3)
   120  			require.Len(t, s.fallback, 2)
   121  		}
   122  
   123  		{ // without third endpoint
   124  			e, err := s.Next(ctx)
   125  			require.NoError(t, err)
   126  			require.NotNil(t, e)
   127  			s = Without(s, e)
   128  			require.Len(t, s.All(), 5)
   129  			require.Len(t, s.index, 5)
   130  			require.Len(t, s.prefer, 2)
   131  			require.Len(t, s.fallback, 3)
   132  		}
   133  
   134  		{ // without fourth endpoint
   135  			e, err := s.Next(ctx)
   136  			require.NoError(t, err)
   137  			require.NotNil(t, e)
   138  			s = Without(s, e)
   139  			require.Len(t, s.All(), 5)
   140  			require.Len(t, s.index, 5)
   141  			require.Len(t, s.prefer, 1)
   142  			require.Len(t, s.fallback, 4)
   143  		}
   144  
   145  		{ // without fifth endpoint
   146  			e, err := s.Next(ctx)
   147  			require.NoError(t, err)
   148  			require.NotNil(t, e)
   149  			s = Without(s, e)
   150  			require.Len(t, s.All(), 5)
   151  			require.Len(t, s.index, 5)
   152  			require.Empty(t, s.prefer)
   153  			require.Len(t, s.fallback, 5)
   154  		}
   155  
   156  		{ // next from fallback is ok
   157  			e, err := s.Next(ctx)
   158  			require.NoError(t, err)
   159  			require.NotNil(t, e)
   160  		}
   161  	})
   162  
   163  	t.Run("WithFilter", func(t *testing.T) {
   164  		s := New([]endpoint.Endpoint{
   165  			&mock.Endpoint{
   166  				AddrField:   "1",
   167  				NodeIDField: 1,
   168  			},
   169  			&mock.Endpoint{
   170  				AddrField:   "2",
   171  				NodeIDField: 2,
   172  			},
   173  			&mock.Endpoint{
   174  				AddrField:   "3",
   175  				NodeIDField: 3,
   176  			},
   177  			&mock.Endpoint{
   178  				AddrField:   "4",
   179  				NodeIDField: 4,
   180  			},
   181  		}, WithFilter(func(e endpoint.Info) bool {
   182  			return e.NodeID()%2 == 0
   183  		}))
   184  
   185  		require.Len(t, s.index, 2)
   186  		require.Len(t, s.All(), 2)
   187  		require.Len(t, s.prefer, 2)
   188  		require.Empty(t, s.fallback)
   189  	})
   190  
   191  	t.Run("WithFallback", func(t *testing.T) {
   192  		t.Run("SplittedPreferAndFallback", func(t *testing.T) {
   193  			s := New([]endpoint.Endpoint{
   194  				&mock.Endpoint{
   195  					AddrField:   "1",
   196  					NodeIDField: 1,
   197  				},
   198  				&mock.Endpoint{
   199  					AddrField:   "2",
   200  					NodeIDField: 2,
   201  				},
   202  				&mock.Endpoint{
   203  					AddrField:   "3",
   204  					NodeIDField: 3,
   205  				},
   206  				&mock.Endpoint{
   207  					AddrField:   "4",
   208  					NodeIDField: 4,
   209  				},
   210  			}, WithFilter(func(e endpoint.Info) bool {
   211  				return e.NodeID()%2 == 0
   212  			}), WithFallback(true))
   213  
   214  			require.Len(t, s.index, 4)
   215  			require.Len(t, s.All(), 4)
   216  			require.Len(t, s.prefer, 2)
   217  			require.Len(t, s.fallback, 2)
   218  		})
   219  
   220  		t.Run("OnlyFallback", func(t *testing.T) {
   221  			s := New([]endpoint.Endpoint{
   222  				&mock.Endpoint{
   223  					AddrField:   "1",
   224  					NodeIDField: 1,
   225  				},
   226  			}, WithFilter(func(e endpoint.Info) bool {
   227  				return false
   228  			}), WithFallback(true))
   229  
   230  			require.Len(t, s.index, 1)
   231  			require.Len(t, s.All(), 1)
   232  			require.Empty(t, s.prefer)
   233  			require.Len(t, s.fallback, 1)
   234  
   235  			e, err := s.Next(ctx)
   236  			require.NoError(t, err)
   237  			require.Equal(t, "1", e.Address())
   238  			require.Equal(t, uint32(1), e.NodeID())
   239  		})
   240  	})
   241  
   242  	t.Run("PreferByNodeID", func(t *testing.T) {
   243  		s := New([]endpoint.Endpoint{
   244  			&mock.Endpoint{
   245  				AddrField:   "1",
   246  				NodeIDField: 1,
   247  			},
   248  			&mock.Endpoint{
   249  				AddrField:   "2",
   250  				NodeIDField: 2,
   251  			},
   252  			&mock.Endpoint{
   253  				AddrField:   "3",
   254  				NodeIDField: 3,
   255  			},
   256  			&mock.Endpoint{
   257  				AddrField:   "4",
   258  				NodeIDField: 4,
   259  			},
   260  		})
   261  
   262  		xtest.TestManyTimes(t, func(t testing.TB) {
   263  			e, err := s.Next(endpoint.WithNodeID(ctx, 3))
   264  			require.NoError(t, err)
   265  			require.NotNil(t, e)
   266  			require.Equal(t, "3", e.Address())
   267  			require.Equal(t, uint32(3), e.NodeID())
   268  		})
   269  	})
   270  
   271  	t.Run("NormalDistribution", func(t *testing.T) {
   272  		const (
   273  			buckets = 10
   274  			total   = 1000000
   275  			epsilon = int(float64(total) / float64(buckets) * 0.015)
   276  		)
   277  		endpoints := make([]endpoint.Endpoint, buckets)
   278  
   279  		for i := 0; i < buckets; i++ {
   280  			endpoints[i] = &mock.Endpoint{
   281  				AddrField:   strconv.Itoa(i),
   282  				NodeIDField: uint32(i),
   283  			}
   284  		}
   285  
   286  		s := New(endpoints)
   287  
   288  		distribution := make([]int, len(endpoints))
   289  		for i := 0; i < total; i++ {
   290  			e, err := s.Next(ctx)
   291  			require.NoError(t, err)
   292  			require.NotNil(t, e)
   293  			distribution[e.NodeID()]++
   294  		}
   295  
   296  		for i := range distribution {
   297  			if distribution[i] < total/buckets-epsilon || distribution[i] > total/buckets+epsilon {
   298  				t.Errorf("unexpected distribuition[%d] = %0.1f%%", i,
   299  					math.Abs(float64(distribution[i]-total/buckets)/float64(total/buckets)*100),
   300  				)
   301  			}
   302  		}
   303  	})
   304  }
   305  
   306  func BenchmarkNext1(b *testing.B) {
   307  	benchmarkNextParallel(b, 1)
   308  }
   309  
   310  func BenchmarkNext4(b *testing.B) {
   311  	benchmarkNextParallel(b, 4)
   312  }
   313  
   314  func BenchmarkNext16(b *testing.B) {
   315  	benchmarkNextParallel(b, 16)
   316  }
   317  
   318  func BenchmarkNext32(b *testing.B) {
   319  	benchmarkNextParallel(b, 32)
   320  }
   321  
   322  func BenchmarkNext64(b *testing.B) {
   323  	benchmarkNextParallel(b, 64)
   324  }
   325  
   326  func BenchmarkNext128(b *testing.B) {
   327  	benchmarkNextParallel(b, 128)
   328  }
   329  
   330  func BenchmarkNext256(b *testing.B) {
   331  	benchmarkNextParallel(b, 256)
   332  }
   333  
   334  func BenchmarkNext512(b *testing.B) {
   335  	benchmarkNextParallel(b, 512)
   336  }
   337  
   338  func benchmarkNextParallel(b *testing.B, parallelism int) {
   339  	ctx := xtest.Context(b)
   340  
   341  	s := New([]endpoint.Endpoint{
   342  		&mock.Endpoint{
   343  			AddrField:   "1",
   344  			NodeIDField: 1,
   345  		},
   346  		&mock.Endpoint{
   347  			AddrField:   "2",
   348  			NodeIDField: 2,
   349  		},
   350  		&mock.Endpoint{
   351  			AddrField:   "3",
   352  			NodeIDField: 3,
   353  		},
   354  		&mock.Endpoint{
   355  			AddrField:   "4",
   356  			NodeIDField: 4,
   357  		},
   358  	}, WithFilter(func(e endpoint.Info) bool {
   359  		return e.NodeID()%2 == 0
   360  	}), WithFallback(true))
   361  
   362  	b.ReportAllocs()
   363  
   364  	b.ResetTimer()
   365  
   366  	var wg sync.WaitGroup
   367  	wg.Add(parallelism)
   368  	for range make([]struct{}, parallelism) {
   369  		go func() {
   370  			defer wg.Done()
   371  
   372  			for i := 0; i < b.N/parallelism; i++ {
   373  				e, err := s.Next(ctx)
   374  				require.NoError(b, err)
   375  				require.NotNil(b, e)
   376  			}
   377  		}()
   378  	}
   379  	wg.Wait()
   380  }