git.frostfs.info/TrueCloudLab/frostfs-sdk-go@v0.0.0-20241022124111-5361f0ecebd3/pool/sampler_test.go (about)

     1  package pool
     2  
     3  import (
     4  	"context"
     5  	"math/rand"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/require"
     9  )
    10  
    11  func TestSamplerStability(t *testing.T) {
    12  	const COUNT = 100000
    13  
    14  	cases := []struct {
    15  		probabilities []float64
    16  		expected      []int
    17  	}{
    18  		{
    19  			probabilities: []float64{1, 0},
    20  			expected:      []int{COUNT, 0},
    21  		},
    22  		{
    23  			probabilities: []float64{0.1, 0.2, 0.7},
    24  			expected:      []int{10138, 19813, 70049},
    25  		},
    26  		{
    27  			probabilities: []float64{0.2, 0.2, 0.4, 0.1, 0.1, 0},
    28  			expected:      []int{19824, 20169, 39900, 10243, 9864, 0},
    29  		},
    30  	}
    31  
    32  	for _, tc := range cases {
    33  		sampler := newSampler(tc.probabilities, rand.NewSource(0))
    34  		res := make([]int, len(tc.probabilities))
    35  		for range COUNT {
    36  			res[sampler.Next()]++
    37  		}
    38  
    39  		require.Equal(t, tc.expected, res, "probabilities: %v", tc.probabilities)
    40  	}
    41  }
    42  
    43  func TestHealthyReweight(t *testing.T) {
    44  	var (
    45  		weights = []float64{0.9, 0.1}
    46  		names   = []string{"node0", "node1"}
    47  		buffer  = make([]float64, len(weights))
    48  	)
    49  
    50  	cache, err := newCache(0)
    51  	require.NoError(t, err)
    52  
    53  	client1 := newMockClient(names[0], *newPrivateKey(t))
    54  	client1.errOnDial()
    55  
    56  	client2 := newMockClient(names[1], *newPrivateKey(t))
    57  
    58  	inner := &innerPool{
    59  		sampler: newSampler(weights, rand.NewSource(0)),
    60  		clients: []client{client1, client2},
    61  	}
    62  	p := &Pool{
    63  		innerPools:      []*innerPool{inner},
    64  		cache:           cache,
    65  		key:             newPrivateKey(t),
    66  		rebalanceParams: rebalanceParameters{nodesParams: []*nodesParam{{weights: weights}}},
    67  	}
    68  
    69  	// check getting first node connection before rebalance happened
    70  	connection0, err := p.connection()
    71  	require.NoError(t, err)
    72  	mock0 := connection0.(*mockClient)
    73  	require.Equal(t, names[0], mock0.address())
    74  
    75  	p.updateInnerNodesHealth(context.TODO(), 0, buffer)
    76  
    77  	connection1, err := p.connection()
    78  	require.NoError(t, err)
    79  	mock1 := connection1.(*mockClient)
    80  	require.Equal(t, names[1], mock1.address())
    81  
    82  	// enabled first node again
    83  	inner.lock.Lock()
    84  	inner.clients[0] = newMockClient(names[0], *newPrivateKey(t))
    85  	inner.lock.Unlock()
    86  
    87  	p.updateInnerNodesHealth(context.TODO(), 0, buffer)
    88  	inner.sampler = newSampler(weights, rand.NewSource(0))
    89  
    90  	connection0, err = p.connection()
    91  	require.NoError(t, err)
    92  	mock0 = connection0.(*mockClient)
    93  	require.Equal(t, names[0], mock0.address())
    94  }
    95  
    96  func TestHealthyNoReweight(t *testing.T) {
    97  	var (
    98  		weights = []float64{0.9, 0.1}
    99  		names   = []string{"node0", "node1"}
   100  		buffer  = make([]float64, len(weights))
   101  	)
   102  
   103  	sampl := newSampler(weights, rand.NewSource(0))
   104  	inner := &innerPool{
   105  		sampler: sampl,
   106  		clients: []client{
   107  			newMockClient(names[0], *newPrivateKey(t)),
   108  			newMockClient(names[1], *newPrivateKey(t)),
   109  		},
   110  	}
   111  	p := &Pool{
   112  		innerPools:      []*innerPool{inner},
   113  		rebalanceParams: rebalanceParameters{nodesParams: []*nodesParam{{weights: weights}}},
   114  	}
   115  
   116  	p.updateInnerNodesHealth(context.TODO(), 0, buffer)
   117  
   118  	inner.lock.RLock()
   119  	defer inner.lock.RUnlock()
   120  	require.Equal(t, inner.sampler, sampl)
   121  }