github.com/uber/kraken@v0.1.4/lib/hrw/fixtures.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  package hrw
    15  
    16  import (
    17  	"crypto/rand"
    18  	"encoding/hex"
    19  	"strconv"
    20  )
    21  
    22  // NodeKeysTable is a Node to keys map utility type
    23  type NodeKeysTable map[string]map[string]struct{}
    24  
    25  // All functions defined in this file are intended for testing purposes only.
    26  
    27  // RendezvousHashFixture creates a weigthed rendezvous hashing object with
    28  // specified #numKey of keys, a hash function (we use sha256.New primarely for testing) and
    29  // the list of weights which define nodes in rendezvous hashing schema and assign weights to
    30  // them in their natural order. The nodes will be also initialized with seed names equal
    31  // to their corresponding indices in the array, so for
    32  // RendezvousHashFixture(10, sha256.New, 100, 200, 300) there will be a RendezvousHash object created
    33  // with 3 nodes "0": 100, "1":200, "2":300
    34  // The fixture will return RendezvousHash object and the node key buckets table
    35  func RendezvousHashFixture(numKeys int, hashFactory HashFactory, scoreFunc UIntToFloat, weights ...int) (*RendezvousHash, map[string]map[string]struct{}) {
    36  	rh := NewRendezvousHash(hashFactory, scoreFunc)
    37  
    38  	keys := NodeKeysTable{}
    39  	b := make([]byte, 64)
    40  
    41  	totalWeights := 0
    42  	for index, weight := range weights {
    43  		totalWeights += weight
    44  		rh.AddNode(strconv.Itoa(index), weight)
    45  		keys[strconv.Itoa(index)] = make(map[string]struct{})
    46  	}
    47  	// 1500 the sum of all weights
    48  	for i := 0; i < numKeys; i++ {
    49  		rand.Read(b)
    50  		key := hex.EncodeToString(b)
    51  		nodes := rh.GetOrderedNodes(key, 1)
    52  		keys[nodes[0].Label][key] = struct{}{}
    53  	}
    54  
    55  	return rh, keys
    56  }
    57  
    58  //HashKeyFixture generate #numkeys random keys according to a hash function
    59  func HashKeyFixture(numKeys int, hashFactory HashFactory) []string {
    60  	var keys []string
    61  	b := make([]byte, 64)
    62  
    63  	for i := 0; i < numKeys; i++ {
    64  		rand.Read(b)
    65  		key := hex.EncodeToString(b)
    66  		keys = append(keys, key)
    67  	}
    68  
    69  	return keys
    70  }