github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/dbnode/sharding/shardset.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 sharding
    22  
    23  import (
    24  	"errors"
    25  	"math"
    26  
    27  	"github.com/m3db/m3/src/cluster/shard"
    28  	"github.com/m3db/m3/src/x/ident"
    29  	"github.com/m3db/stackmurmur3/v2"
    30  )
    31  
    32  var (
    33  	// ErrDuplicateShards returned when shard set contains duplicate shards
    34  	ErrDuplicateShards = errors.New("duplicate shards")
    35  
    36  	// ErrInvalidShardID is returned on an invalid shard ID
    37  	ErrInvalidShardID = errors.New("no shard with given ID")
    38  )
    39  
    40  type shardSet struct {
    41  	shards   []shard.Shard
    42  	ids      []uint32
    43  	shardMap map[uint32]shard.Shard
    44  	fn       HashFn
    45  }
    46  
    47  // NewShardSet creates a new sharding scheme with a set of shards
    48  func NewShardSet(shards []shard.Shard, fn HashFn) (ShardSet, error) {
    49  	if err := validateShards(shards); err != nil {
    50  		return nil, err
    51  	}
    52  	return newValidatedShardSet(shards, fn), nil
    53  }
    54  
    55  // NewEmptyShardSet creates a new sharding scheme with an empty set of shards
    56  func NewEmptyShardSet(fn HashFn) ShardSet {
    57  	return newValidatedShardSet(nil, fn)
    58  }
    59  
    60  func newValidatedShardSet(shards []shard.Shard, fn HashFn) ShardSet {
    61  	ids := make([]uint32, len(shards))
    62  	shardMap := make(map[uint32]shard.Shard, len(shards))
    63  	for i, shard := range shards {
    64  		ids[i] = shard.ID()
    65  		shardMap[shard.ID()] = shard
    66  	}
    67  	return &shardSet{
    68  		shards:   shards,
    69  		ids:      ids,
    70  		shardMap: shardMap,
    71  		fn:       fn,
    72  	}
    73  }
    74  
    75  func (s *shardSet) Lookup(identifier ident.ID) uint32 {
    76  	return s.fn(identifier)
    77  }
    78  
    79  func (s *shardSet) LookupShard(shardID uint32) (shard.Shard, error) {
    80  	hostShard, ok := s.shardMap[shardID]
    81  	if !ok {
    82  		return nil, ErrInvalidShardID
    83  	}
    84  	return hostShard, nil
    85  }
    86  
    87  func (s *shardSet) LookupStateByID(shardID uint32) (shard.State, error) {
    88  	hostShard, ok := s.shardMap[shardID]
    89  	if !ok {
    90  		return shard.State(0), ErrInvalidShardID
    91  	}
    92  	return hostShard.State(), nil
    93  }
    94  
    95  func (s *shardSet) All() []shard.Shard {
    96  	return s.shards[:]
    97  }
    98  
    99  func (s *shardSet) AllIDs() []uint32 {
   100  	return s.ids[:]
   101  }
   102  
   103  func (s *shardSet) Min() uint32 {
   104  	min := uint32(math.MaxUint32)
   105  	for _, shard := range s.ids {
   106  		if shard < min {
   107  			min = shard
   108  		}
   109  	}
   110  	return min
   111  }
   112  
   113  func (s *shardSet) Max() uint32 {
   114  	max := uint32(0)
   115  	for _, shard := range s.ids {
   116  		if shard > max {
   117  			max = shard
   118  		}
   119  	}
   120  	return max
   121  }
   122  
   123  func (s *shardSet) HashFn() HashFn {
   124  	return s.fn
   125  }
   126  
   127  // NewShards returns a new slice of shards with a specified state
   128  func NewShards(ids []uint32, state shard.State) []shard.Shard {
   129  	shards := make([]shard.Shard, len(ids))
   130  	for i, id := range ids {
   131  		shards[i] = shard.NewShard(id).SetState(state)
   132  	}
   133  	return shards
   134  }
   135  
   136  // IDs returns a new slice of shard IDs for a set of shards
   137  func IDs(shards []shard.Shard) []uint32 {
   138  	ids := make([]uint32, len(shards))
   139  	for i := range ids {
   140  		ids[i] = shards[i].ID()
   141  	}
   142  	return ids
   143  }
   144  
   145  func validateShards(shards []shard.Shard) error {
   146  	uniqueShards := make(map[uint32]struct{}, len(shards))
   147  	for _, s := range shards {
   148  		if _, exist := uniqueShards[s.ID()]; exist {
   149  			return ErrDuplicateShards
   150  		}
   151  		uniqueShards[s.ID()] = struct{}{}
   152  	}
   153  	return nil
   154  }
   155  
   156  // DefaultHashFn generates a HashFn based on murmur32
   157  func DefaultHashFn(length int) HashFn {
   158  	return NewHashFn(length, 0)
   159  }
   160  
   161  // NewHashGenWithSeed generates a HashFnGen based on murmur32 with a given seed
   162  func NewHashGenWithSeed(seed uint32) HashGen {
   163  	return func(length int) HashFn {
   164  		return NewHashFn(length, seed)
   165  	}
   166  }
   167  
   168  // NewHashFn generates a HashFN based on murmur32 with a given seed
   169  func NewHashFn(length int, seed uint32) HashFn {
   170  	return func(id ident.ID) uint32 {
   171  		return murmur3.SeedSum32(seed, id.Bytes()) % uint32(length)
   172  	}
   173  }