github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/datastore/crdb/pool/balancer_test.go (about)

     1  package pool
     2  
     3  import (
     4  	"context"
     5  	"math/rand"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  func TestNodeConnectionBalancerPrune(t *testing.T) {
    13  	tests := []struct {
    14  		name       string
    15  		maxConns   uint32
    16  		nodes      []uint32
    17  		conns      []uint32
    18  		expectedGC []uint32
    19  	}{
    20  		{
    21  			name:       "no extra, no pruning needed",
    22  			nodes:      []uint32{1, 2, 3},
    23  			maxConns:   9,
    24  			conns:      []uint32{1, 1, 1, 2, 2, 2, 3, 3, 3},
    25  			expectedGC: []uint32{},
    26  		},
    27  		{
    28  			name:       "no extra, max 1",
    29  			nodes:      []uint32{1, 2, 3},
    30  			maxConns:   1,
    31  			conns:      []uint32{1},
    32  			expectedGC: []uint32{},
    33  		},
    34  		{
    35  			name:       "prune 1, max 1",
    36  			nodes:      []uint32{1, 2, 3},
    37  			maxConns:   1,
    38  			conns:      []uint32{1, 2},
    39  			expectedGC: []uint32{2},
    40  		},
    41  		{
    42  			name:       "no extra, max 2",
    43  			nodes:      []uint32{1, 2, 3},
    44  			maxConns:   2,
    45  			conns:      []uint32{1, 2},
    46  			expectedGC: []uint32{},
    47  		},
    48  		{
    49  			name:       "prune 1, max 2",
    50  			nodes:      []uint32{1, 2, 3},
    51  			maxConns:   2,
    52  			conns:      []uint32{1, 2, 3},
    53  			expectedGC: []uint32{3},
    54  		},
    55  		{
    56  			name:       "no extra, max 1 per node",
    57  			nodes:      []uint32{1, 2, 3},
    58  			maxConns:   3,
    59  			conns:      []uint32{1, 2, 3},
    60  			expectedGC: []uint32{},
    61  		},
    62  		{
    63  			name:       "1 extra, max 1 per node",
    64  			nodes:      []uint32{1, 2, 3},
    65  			maxConns:   3,
    66  			conns:      []uint32{1, 2, 2, 3},
    67  			expectedGC: []uint32{2},
    68  		},
    69  		{
    70  			name:       "no extra, max 2 per node",
    71  			nodes:      []uint32{1, 2, 3},
    72  			maxConns:   6,
    73  			conns:      []uint32{1, 1, 2, 2, 3, 3},
    74  			expectedGC: []uint32{},
    75  		},
    76  		{
    77  			name:       "1 extra, max 2 per node",
    78  			nodes:      []uint32{1, 2, 3},
    79  			maxConns:   6,
    80  			conns:      []uint32{1, 1, 2, 2, 3, 3, 3},
    81  			expectedGC: []uint32{3},
    82  		},
    83  		{
    84  			name:       "1 extra, prune 1",
    85  			nodes:      []uint32{1, 2, 3},
    86  			maxConns:   9,
    87  			conns:      []uint32{1, 1, 1, 1, 2, 3},
    88  			expectedGC: []uint32{1},
    89  		},
    90  		{
    91  			name:       "2 extra, prune 1",
    92  			nodes:      []uint32{1, 2, 3},
    93  			maxConns:   9,
    94  			conns:      []uint32{1, 1, 1, 1, 1, 2, 3},
    95  			expectedGC: []uint32{1},
    96  		},
    97  		{
    98  			name:       "5 extra, prune 2",
    99  			nodes:      []uint32{1, 2, 3},
   100  			maxConns:   9,
   101  			conns:      []uint32{1, 1, 1, 1, 1, 1, 1, 1, 2, 3},
   102  			expectedGC: []uint32{1, 1},
   103  		},
   104  		{
   105  			name:       "7 extra, prune 3",
   106  			nodes:      []uint32{1, 2, 3},
   107  			maxConns:   9,
   108  			conns:      []uint32{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 3},
   109  			expectedGC: []uint32{1, 1, 1},
   110  		},
   111  		{
   112  			name:     "prune 2 from each node",
   113  			nodes:    []uint32{1, 2, 3},
   114  			maxConns: 9,
   115  			conns: []uint32{
   116  				1, 1, 1, 1, 1, 1, 1, 1,
   117  				2, 2, 2, 2, 2, 2, 2, 2,
   118  				3, 3, 3, 3, 3, 3, 3, 3,
   119  			},
   120  			expectedGC: []uint32{1, 1, 2, 2, 3, 3},
   121  		},
   122  		{
   123  			name:     "uneven split should fill max pool exactly",
   124  			nodes:    []uint32{1, 2, 3},
   125  			maxConns: 10,
   126  			// note that node 2 gets the extra connection due to shuffling.
   127  			// this is randomized between runs of the server but we pin
   128  			// the seed for the tests
   129  			conns:      []uint32{1, 1, 1, 2, 2, 2, 2, 3, 3, 3},
   130  			expectedGC: []uint32{},
   131  		},
   132  		{
   133  			name:       "uneven split should prune to max pool exactly",
   134  			nodes:      []uint32{1, 2, 3},
   135  			maxConns:   11,
   136  			conns:      []uint32{1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3},
   137  			expectedGC: []uint32{3},
   138  		},
   139  	}
   140  	for _, tt := range tests {
   141  		tt := tt
   142  		t.Run(tt.name, func(t *testing.T) {
   143  			tracker, err := NewNodeHealthChecker("")
   144  			require.NoError(t, err)
   145  			for _, n := range tt.nodes {
   146  				tracker.healthyNodes[n] = struct{}{}
   147  			}
   148  
   149  			pool := NewFakePool(tt.maxConns)
   150  
   151  			p := newNodeConnectionBalancer[*FakePoolConn[*FakeConn], *FakeConn](pool, tracker, 1*time.Minute)
   152  			p.seed = 0
   153  			// nolint:gosec
   154  			// G404 use of non cryptographically secure random number generator is not concern here,
   155  			// as it's used for jittering the interval for health checks.
   156  			p.rnd = rand.New(rand.NewSource(0))
   157  
   158  			for _, n := range tt.conns {
   159  				pool.nodeForConn[NewFakeConn()] = n
   160  			}
   161  
   162  			ctx, cancel := context.WithCancel(context.Background())
   163  			defer cancel()
   164  
   165  			p.mustPruneConnections(ctx)
   166  			require.Equal(t, len(tt.expectedGC), len(pool.gc))
   167  			gcFromNodes := make([]uint32, 0, len(tt.expectedGC))
   168  			for _, n := range pool.gc {
   169  				gcFromNodes = append(gcFromNodes, n)
   170  			}
   171  			require.ElementsMatch(t, tt.expectedGC, gcFromNodes)
   172  		})
   173  	}
   174  }