github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/dbnode/topology/map_test.go (about)

     1  // Copyright (c) 2016 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package topology
    22  
    23  import (
    24  	"testing"
    25  
    26  	"github.com/m3db/m3/src/cluster/shard"
    27  	"github.com/m3db/m3/src/dbnode/sharding"
    28  	"github.com/m3db/m3/src/x/ident"
    29  
    30  	"github.com/stretchr/testify/assert"
    31  	"github.com/stretchr/testify/require"
    32  )
    33  
    34  func newTestShardSet(
    35  	t *testing.T,
    36  	shards []testShard,
    37  	hashFn sharding.HashFn,
    38  ) sharding.ShardSet {
    39  	var values []shard.Shard
    40  	for _, elem := range shards {
    41  		value := shard.NewShard(elem.id).SetState(elem.state)
    42  		values = append(values, value)
    43  	}
    44  	shardSet, err := sharding.NewShardSet(values, hashFn)
    45  	require.NoError(t, err)
    46  	return shardSet
    47  }
    48  
    49  type testShard struct {
    50  	id    uint32
    51  	state shard.State
    52  }
    53  
    54  type testShards []testShard
    55  
    56  func (s testShards) IDs() []uint32 {
    57  	var ids []uint32
    58  	for _, elem := range s {
    59  		ids = append(ids, elem.id)
    60  	}
    61  	return ids
    62  }
    63  
    64  func TestStaticMap(t *testing.T) {
    65  	hashFn := func(id ident.ID) uint32 {
    66  		switch id.String() {
    67  		case "foo":
    68  			return 0
    69  		case "bar":
    70  			return 1
    71  		case "unowned":
    72  			return 999
    73  		default:
    74  			return 2
    75  		}
    76  	}
    77  
    78  	hosts := []struct {
    79  		id     string
    80  		addr   string
    81  		shards testShards
    82  	}{
    83  		{"h1", "h1:9000", []testShard{{id: 0, state: shard.Available}}},
    84  		{"h2", "h2:9000", []testShard{{id: 1, state: shard.Available}}},
    85  		{"h3", "h3:9000", []testShard{{id: 0, state: shard.Available}}},
    86  		{"h4", "h4:9000", []testShard{{id: 1, state: shard.Initializing}}},
    87  	}
    88  
    89  	var hostShardSets []HostShardSet
    90  	for _, h := range hosts {
    91  		hostShardSets = append(hostShardSets,
    92  			NewHostShardSet(
    93  				NewHost(h.id, h.addr),
    94  				newTestShardSet(t, h.shards, hashFn)))
    95  	}
    96  
    97  	seedShardSet := newTestShardSet(t, []testShard{
    98  		{id: 0, state: shard.Available},
    99  		{id: 1, state: shard.Available},
   100  	}, hashFn)
   101  	opts := NewStaticOptions().
   102  		SetShardSet(seedShardSet).
   103  		SetReplicas(2).
   104  		SetHostShardSets(hostShardSets)
   105  
   106  	m := NewStaticMap(opts)
   107  	require.Equal(t, 4, len(m.Hosts()))
   108  	require.Equal(t, 4, m.HostsLen())
   109  	for i, h := range hosts {
   110  		assert.Equal(t, h.id, m.Hosts()[i].ID())
   111  		assert.Equal(t, h.addr, m.Hosts()[i].Address())
   112  	}
   113  
   114  	require.Equal(t, 4, len(m.HostShardSets()))
   115  	for i, h := range hosts {
   116  		assert.Equal(t, h.id, m.HostShardSets()[i].Host().ID())
   117  		assert.Equal(t, h.addr, m.HostShardSets()[i].Host().Address())
   118  		assert.Equal(t, h.shards.IDs(), m.HostShardSets()[i].ShardSet().AllIDs())
   119  	}
   120  
   121  	targetShard, targetHosts, err := m.Route(ident.StringID("foo"))
   122  	require.NoError(t, err)
   123  	assert.Equal(t, uint32(0), targetShard)
   124  	require.Equal(t, 2, len(targetHosts))
   125  	assert.Equal(t, "h1", targetHosts[0].ID())
   126  	assert.Equal(t, "h3", targetHosts[1].ID())
   127  
   128  	_, _, err = m.Route(ident.StringID("unowned"))
   129  	require.Error(t, err)
   130  	assert.Equal(t, errUnownedShard, err)
   131  
   132  	targetHosts, err = m.RouteShard(1)
   133  	require.NoError(t, err)
   134  	require.Equal(t, 2, len(targetHosts))
   135  	assert.Equal(t, "h2", targetHosts[0].ID())
   136  	assert.Equal(t, "h4", targetHosts[1].ID())
   137  
   138  	_, err = m.RouteShard(999)
   139  	require.Error(t, err)
   140  	assert.Equal(t, errUnownedShard, err)
   141  
   142  	err = m.RouteForEach(ident.StringID("bar"), func(
   143  		idx int,
   144  		s shard.Shard,
   145  		h Host,
   146  	) {
   147  		switch idx {
   148  		case 1:
   149  			assert.Equal(t, "h2", h.ID())
   150  			assert.Equal(t, shard.Available, s.State())
   151  		case 3:
   152  			assert.Equal(t, "h4", h.ID())
   153  			assert.Equal(t, shard.Initializing, s.State())
   154  		default:
   155  			assert.Fail(t, "routed to wrong host")
   156  		}
   157  	})
   158  	assert.NoError(t, err)
   159  
   160  	err = m.RouteForEach(ident.StringID("unowned"), func(_ int, _ shard.Shard, _ Host) {})
   161  	require.Error(t, err)
   162  	assert.Equal(t, errUnownedShard, err)
   163  
   164  	assert.Equal(t, 2, m.Replicas())
   165  	assert.Equal(t, 2, m.MajorityReplicas())
   166  }