├── README.md ├── LICENSE ├── example_test.go ├── consistent.go └── consistent_test.go /README.md: -------------------------------------------------------------------------------- 1 | consistent 2 | ========== 3 | 4 | Consistent hash package for Go. 5 | 6 | Installation 7 | ------------ 8 | 9 | go get github.com/stathat/consistent 10 | 11 | Examples 12 | -------- 13 | 14 | Look at the godoc. 15 | 16 | Status 17 | ------ 18 | 19 | This package was extracted from production code powering [StatHat](http://www.stathat.com), 20 | so clearly we feel that it is production-ready, but it should still be considered 21 | experimental as other uses of it could reveal issues we aren't experiencing. 22 | 23 | Contact us 24 | ---------- 25 | 26 | We'd love to hear from you if you are using `consistent`. 27 | Get in touch: [@stat_hat](http://twitter.com/stat_hat) or [contact us here](http://www.stathat.com/docs/contact). 28 | 29 | About 30 | ----- 31 | 32 | Written by Patrick Crosby at [StatHat](http://www.stathat.com). Twitter: [@stat_hat](http://twitter.com/stat_hat) 33 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (C) 2012 Numerotron 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 | -------------------------------------------------------------------------------- /example_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (C) 2012 Numerotron Inc. 2 | // Use of this source code is governed by an MIT-style license 3 | // that can be found in the LICENSE file. 4 | 5 | package consistent_test 6 | 7 | import ( 8 | "../consistent" 9 | "fmt" 10 | "log" 11 | ) 12 | 13 | func ExampleNew() { 14 | c := consistent.New() 15 | c.Add("cacheA") 16 | c.Add("cacheB") 17 | c.Add("cacheC") 18 | users := []string{"user_mcnulty", "user_bunk", "user_omar", "user_bunny", "user_stringer"} 19 | for _, u := range users { 20 | server, err := c.Get(u) 21 | if err != nil { 22 | log.Fatal(err) 23 | } 24 | fmt.Printf("%s => %s\n", u, server) 25 | } 26 | // Output: 27 | // user_mcnulty => cacheA 28 | // user_bunk => cacheC 29 | // user_omar => cacheC 30 | // user_bunny => cacheA 31 | // user_stringer => cacheB 32 | } 33 | 34 | func ExampleAdd() { 35 | c := consistent.New() 36 | c.Add("cacheA") 37 | c.Add("cacheB") 38 | c.Add("cacheC") 39 | users := []string{"user_mcnulty", "user_bunk", "user_omar", "user_bunny", "user_stringer"} 40 | fmt.Println("initial state [A, B, C]") 41 | for _, u := range users { 42 | server, err := c.Get(u) 43 | if err != nil { 44 | log.Fatal(err) 45 | } 46 | fmt.Printf("%s => %s\n", u, server) 47 | } 48 | c.Add("cacheD") 49 | c.Add("cacheE") 50 | fmt.Println("\nwith cacheD, cacheE [A, B, C, D, E]") 51 | for _, u := range users { 52 | server, err := c.Get(u) 53 | if err != nil { 54 | log.Fatal(err) 55 | } 56 | fmt.Printf("%s => %s\n", u, server) 57 | } 58 | // Output: 59 | // initial state [A, B, C] 60 | // user_mcnulty => cacheA 61 | // user_bunk => cacheC 62 | // user_omar => cacheC 63 | // user_bunny => cacheA 64 | // user_stringer => cacheB 65 | // 66 | // with cacheD, cacheE [A, B, C, D, E] 67 | // user_mcnulty => cacheE 68 | // user_bunk => cacheD 69 | // user_omar => cacheD 70 | // user_bunny => cacheA 71 | // user_stringer => cacheB 72 | } 73 | 74 | func ExampleRemove() { 75 | c := consistent.New() 76 | c.Add("cacheA") 77 | c.Add("cacheB") 78 | c.Add("cacheC") 79 | users := []string{"user_mcnulty", "user_bunk", "user_omar", "user_bunny", "user_stringer"} 80 | fmt.Println("initial state [A, B, C]") 81 | for _, u := range users { 82 | server, err := c.Get(u) 83 | if err != nil { 84 | log.Fatal(err) 85 | } 86 | fmt.Printf("%s => %s\n", u, server) 87 | } 88 | c.Remove("cacheC") 89 | fmt.Println("\ncacheC removed [A, B]") 90 | for _, u := range users { 91 | server, err := c.Get(u) 92 | if err != nil { 93 | log.Fatal(err) 94 | } 95 | fmt.Printf("%s => %s\n", u, server) 96 | } 97 | // Output: 98 | // initial state [A, B, C] 99 | // user_mcnulty => cacheA 100 | // user_bunk => cacheC 101 | // user_omar => cacheC 102 | // user_bunny => cacheA 103 | // user_stringer => cacheB 104 | // 105 | // cacheC removed [A, B] 106 | // user_mcnulty => cacheA 107 | // user_bunk => cacheA 108 | // user_omar => cacheA 109 | // user_bunny => cacheA 110 | // user_stringer => cacheB 111 | } 112 | -------------------------------------------------------------------------------- /consistent.go: -------------------------------------------------------------------------------- 1 | // Copyright (C) 2012 Numerotron Inc. 2 | // Use of this source code is governed by an MIT-style license 3 | // that can be found in the LICENSE file. 4 | 5 | // Package consistent provides a consistent hashing function. 6 | // 7 | // Consistent hashing is often used to distribute requests to a changing set of servers. For example, 8 | // say you have some cache servers cacheA, cacheB, and cacheC. You want to decide which cache server 9 | // to use to look up information on a user. 10 | // 11 | // You could use a typical hash table and hash the user id 12 | // to one of cacheA, cacheB, or cacheC. But with a typical hash table, if you add or remove a server, 13 | // almost all keys will get remapped to different results, which basically could bring your service 14 | // to a grinding halt while the caches get rebuilt. 15 | // 16 | // With a consistent hash, adding or removing a server drastically reduces the number of keys that 17 | // get remapped. 18 | // 19 | // Read more about consistent hashing on wikipedia: http://en.wikipedia.org/wiki/Consistent_hashing 20 | // 21 | package consistent 22 | 23 | import ( 24 | "errors" 25 | "hash/crc32" 26 | "sort" 27 | "strconv" 28 | "sync" 29 | ) 30 | 31 | type uints []uint32 32 | 33 | // Len returns the length of the uints array. 34 | func (x uints) Len() int { return len(x) } 35 | 36 | // Less returns true if element i is less than element j. 37 | func (x uints) Less(i, j int) bool { return x[i] < x[j] } 38 | 39 | // Swap exchanges elements i and j. 40 | func (x uints) Swap(i, j int) { x[i], x[j] = x[j], x[i] } 41 | 42 | // ErrEmptyCircle is the error returned when trying to get an element when nothing has been added to hash. 43 | var ErrEmptyCircle = errors.New("empty circle") 44 | 45 | // Consistent holds the information about the members of the consistent hash circle. 46 | type Consistent struct { 47 | circle map[uint32]string 48 | members map[string]bool 49 | sortedHashes uints 50 | NumberOfReplicas int 51 | count int64 52 | scratch [64]byte 53 | sync.RWMutex 54 | } 55 | 56 | // New creates a new Consistent object with a default setting of 20 replicas for each entry. 57 | // 58 | // To change the number of replicas, set NumberOfReplicas before adding entries. 59 | func New() *Consistent { 60 | c := new(Consistent) 61 | c.NumberOfReplicas = 20 62 | c.circle = make(map[uint32]string) 63 | c.members = make(map[string]bool) 64 | return c 65 | } 66 | 67 | // eltKey generates a string key for an element with an index. 68 | func (c *Consistent) eltKey(elt string, idx int) string { 69 | return elt + "|" + strconv.Itoa(idx) 70 | } 71 | 72 | // Add inserts a string element in the consistent hash. 73 | func (c *Consistent) Add(elt string) { 74 | c.Lock() 75 | defer c.Unlock() 76 | for i := 0; i < c.NumberOfReplicas; i++ { 77 | c.circle[c.hashKey(c.eltKey(elt, i))] = elt 78 | } 79 | c.members[elt] = true 80 | c.updateSortedHashes() 81 | c.count++ 82 | } 83 | 84 | // Remove removes an element from the hash. 85 | func (c *Consistent) Remove(elt string) { 86 | c.Lock() 87 | defer c.Unlock() 88 | for i := 0; i < c.NumberOfReplicas; i++ { 89 | delete(c.circle, c.hashKey(c.eltKey(elt, i))) 90 | } 91 | delete(c.members, elt) 92 | c.updateSortedHashes() 93 | c.count-- 94 | } 95 | 96 | // Set sets all the elements in the hash. If there are existing elements not present in elts, they will be removed. 97 | func (c *Consistent) Set(elts []string) { 98 | mems := c.Members() 99 | for _, k := range mems { 100 | found := false 101 | for _, v := range elts { 102 | if k == v { 103 | found = true 104 | break 105 | } 106 | } 107 | if !found { 108 | c.Remove(k) 109 | } 110 | } 111 | for _, v := range elts { 112 | c.RLock() 113 | _, exists := c.members[v] 114 | c.RUnlock() 115 | if exists { 116 | continue 117 | } 118 | c.Add(v) 119 | } 120 | } 121 | 122 | func (c *Consistent) Members() []string { 123 | c.RLock() 124 | defer c.RUnlock() 125 | var m []string 126 | for k := range c.members { 127 | m = append(m, k) 128 | } 129 | return m 130 | } 131 | 132 | // Get returns an element close to where name hashes to in the circle. 133 | func (c *Consistent) Get(name string) (string, error) { 134 | c.RLock() 135 | defer c.RUnlock() 136 | if len(c.circle) == 0 { 137 | return "", ErrEmptyCircle 138 | } 139 | key := c.hashKey(name) 140 | i := c.search(key) 141 | return c.circle[c.sortedHashes[i]], nil 142 | } 143 | 144 | func (c *Consistent) search(key uint32) (i int) { 145 | f := func(x int) bool { 146 | return c.sortedHashes[x] > key 147 | } 148 | i = sort.Search(len(c.sortedHashes), f) 149 | if i >= len(c.sortedHashes) { 150 | i = 0 151 | } 152 | return 153 | } 154 | 155 | // GetTwo returns the two closest distinct elements to the name input in the circle. 156 | func (c *Consistent) GetTwo(name string) (string, string, error) { 157 | c.RLock() 158 | defer c.RUnlock() 159 | if len(c.circle) == 0 { 160 | return "", "", ErrEmptyCircle 161 | } 162 | key := c.hashKey(name) 163 | i := c.search(key) 164 | a := c.circle[c.sortedHashes[i]] 165 | 166 | if c.count == 1 { 167 | return a, "", nil 168 | } 169 | 170 | start := i 171 | var b string 172 | for i = start + 1; i != start; i++ { 173 | if i >= len(c.sortedHashes) { 174 | i = 0 175 | } 176 | b = c.circle[c.sortedHashes[i]] 177 | if b != a { 178 | break 179 | } 180 | } 181 | return a, b, nil 182 | } 183 | 184 | // GetN returns the N closest distinct elements to the name input in the circle. 185 | func (c *Consistent) GetN(name string, n int) ([]string, error) { 186 | c.RLock() 187 | defer c.RUnlock() 188 | 189 | if len(c.circle) == 0 { 190 | return nil, ErrEmptyCircle 191 | } 192 | 193 | if c.count < int64(n) { 194 | n = int(c.count) 195 | } 196 | 197 | var ( 198 | key = c.hashKey(name) 199 | i = c.search(key) 200 | start = i 201 | res = make([]string, 0, n) 202 | elem = c.circle[c.sortedHashes[i]] 203 | ) 204 | 205 | res = append(res, elem) 206 | 207 | if len(res) == n { 208 | return res, nil 209 | } 210 | 211 | for i = start + 1; i != start; i++ { 212 | if i >= len(c.sortedHashes) { 213 | i = 0 214 | } 215 | elem = c.circle[c.sortedHashes[i]] 216 | if !slice_contains_member(res, elem) { 217 | res = append(res, elem) 218 | } 219 | if len(res) == n { 220 | break 221 | } 222 | } 223 | 224 | return res, nil 225 | } 226 | 227 | func (c *Consistent) hashKey(key string) uint32 { 228 | if len(key) < 64 { 229 | copy(c.scratch[:], key) 230 | return crc32.ChecksumIEEE(c.scratch[:len(key)]) 231 | } 232 | return crc32.ChecksumIEEE([]byte(key)) 233 | } 234 | 235 | func (c *Consistent) updateSortedHashes() { 236 | hashes := c.sortedHashes[:0] 237 | //reallocate if we're holding on to too much (1/4th) 238 | if cap(c.sortedHashes)/(c.NumberOfReplicas*4) > len(c.circle) { 239 | hashes = nil 240 | } 241 | for k := range c.circle { 242 | hashes = append(hashes, k) 243 | } 244 | sort.Sort(hashes) 245 | c.sortedHashes = hashes 246 | } 247 | 248 | func slice_contains_member(set []string, member string) bool { 249 | for _, m := range set { 250 | if m == member { 251 | return true 252 | } 253 | } 254 | return false 255 | } 256 | -------------------------------------------------------------------------------- /consistent_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (C) 2012 Numerotron Inc. 2 | // Use of this source code is governed by an MIT-style license 3 | // that can be found in the LICENSE file. 4 | 5 | package consistent 6 | 7 | import ( 8 | "runtime" 9 | "sort" 10 | "strconv" 11 | "testing" 12 | "testing/quick" 13 | ) 14 | 15 | func checkNum(num, expected int, t *testing.T) { 16 | if num != expected { 17 | t.Errorf("expected %d, got %d", expected, num) 18 | } 19 | } 20 | 21 | func TestNew(t *testing.T) { 22 | x := New() 23 | if x == nil { 24 | t.Errorf("expected obj") 25 | } 26 | checkNum(x.NumberOfReplicas, 20, t) 27 | } 28 | 29 | func TestAdd(t *testing.T) { 30 | x := New() 31 | x.Add("abcdefg") 32 | checkNum(len(x.circle), 20, t) 33 | checkNum(len(x.sortedHashes), 20, t) 34 | if sort.IsSorted(x.sortedHashes) == false { 35 | t.Errorf("expected sorted hashes to be sorted") 36 | } 37 | x.Add("qwer") 38 | checkNum(len(x.circle), 40, t) 39 | checkNum(len(x.sortedHashes), 40, t) 40 | if sort.IsSorted(x.sortedHashes) == false { 41 | t.Errorf("expected sorted hashes to be sorted") 42 | } 43 | } 44 | 45 | func TestRemove(t *testing.T) { 46 | x := New() 47 | x.Add("abcdefg") 48 | x.Remove("abcdefg") 49 | checkNum(len(x.circle), 0, t) 50 | checkNum(len(x.sortedHashes), 0, t) 51 | } 52 | 53 | func TestRemoveNonExisting(t *testing.T) { 54 | x := New() 55 | x.Add("abcdefg") 56 | x.Remove("abcdefghijk") 57 | checkNum(len(x.circle), 20, t) 58 | } 59 | 60 | func TestGetEmpty(t *testing.T) { 61 | x := New() 62 | _, err := x.Get("asdfsadfsadf") 63 | if err == nil { 64 | t.Errorf("expected error") 65 | } 66 | if err != ErrEmptyCircle { 67 | t.Errorf("expected empty circle error") 68 | } 69 | } 70 | 71 | func TestGetSingle(t *testing.T) { 72 | x := New() 73 | x.Add("abcdefg") 74 | f := func(s string) bool { 75 | y, err := x.Get(s) 76 | if err != nil { 77 | t.Logf("error: %q", err) 78 | return false 79 | } 80 | t.Logf("s = %q, y = %q", s, y) 81 | return y == "abcdefg" 82 | } 83 | if err := quick.Check(f, nil); err != nil { 84 | t.Fatal(err) 85 | } 86 | } 87 | 88 | func TestGetMultiple(t *testing.T) { 89 | x := New() 90 | x.Add("abcdefg") 91 | x.Add("hijklmn") 92 | x.Add("opqrstu") 93 | result, err := x.Get("ggg") 94 | if err != nil { 95 | t.Fatal(err) 96 | } 97 | if result != "opqrstu" { 98 | t.Errorf("expected 'opqrstu', got %q", result) 99 | } 100 | result, err = x.Get("hhh") 101 | if err != nil { 102 | t.Fatal(err) 103 | } 104 | if result != "abcdefg" { 105 | t.Errorf("expected 'abcdefg', got %q", result) 106 | } 107 | result, err = x.Get("iiiiii") 108 | if err != nil { 109 | t.Fatal(err) 110 | } 111 | if result != "hijklmn" { 112 | t.Errorf("expected 'hijklmn', got %q", result) 113 | } 114 | } 115 | 116 | func TestGetMultipleQuick(t *testing.T) { 117 | x := New() 118 | x.Add("abcdefg") 119 | x.Add("hijklmn") 120 | x.Add("opqrstu") 121 | f := func(s string) bool { 122 | y, err := x.Get(s) 123 | if err != nil { 124 | t.Logf("error: %q", err) 125 | return false 126 | } 127 | t.Logf("s = %q, y = %q", s, y) 128 | return y == "abcdefg" || y == "hijklmn" || y == "opqrstu" 129 | } 130 | if err := quick.Check(f, nil); err != nil { 131 | t.Fatal(err) 132 | } 133 | } 134 | 135 | func TestGetMultipleRemove(t *testing.T) { 136 | x := New() 137 | x.Add("abcdefg") 138 | x.Add("hijklmn") 139 | x.Add("opqrstu") 140 | result, err := x.Get("ggg") 141 | if err != nil { 142 | t.Fatal(err) 143 | } 144 | if result != "opqrstu" { 145 | t.Errorf("expected 'opqrstu', got %q", result) 146 | } 147 | result, err = x.Get("hhh") 148 | if err != nil { 149 | t.Fatal(err) 150 | } 151 | if result != "abcdefg" { 152 | t.Errorf("expected 'abcdefg', got %q", result) 153 | } 154 | result, err = x.Get("iiiiii") 155 | if err != nil { 156 | t.Fatal(err) 157 | } 158 | if result != "hijklmn" { 159 | t.Errorf("expected 'hijklmn', got %q", result) 160 | } 161 | x.Remove("hijklmn") 162 | result, err = x.Get("ggg") 163 | if err != nil { 164 | t.Fatal(err) 165 | } 166 | if result != "opqrstu" { 167 | t.Errorf("expected 'opqrstu', got %q", result) 168 | } 169 | result, err = x.Get("hhh") 170 | if err != nil { 171 | t.Fatal(err) 172 | } 173 | if result != "abcdefg" { 174 | t.Errorf("expected 'abcdefg', got %q", result) 175 | } 176 | result, err = x.Get("iiiiii") 177 | if err != nil { 178 | t.Fatal(err) 179 | } 180 | if result != "opqrstu" { 181 | t.Errorf("expected 'opqrstu', got %q", result) 182 | } 183 | } 184 | 185 | func TestGetMultipleRemoveQuick(t *testing.T) { 186 | x := New() 187 | x.Add("abcdefg") 188 | x.Add("hijklmn") 189 | x.Add("opqrstu") 190 | x.Remove("opqrstu") 191 | f := func(s string) bool { 192 | y, err := x.Get(s) 193 | if err != nil { 194 | t.Logf("error: %q", err) 195 | return false 196 | } 197 | t.Logf("s = %q, y = %q", s, y) 198 | return y == "abcdefg" || y == "hijklmn" 199 | } 200 | if err := quick.Check(f, nil); err != nil { 201 | t.Fatal(err) 202 | } 203 | } 204 | 205 | func TestGetTwo(t *testing.T) { 206 | x := New() 207 | x.Add("abcdefg") 208 | x.Add("hijklmn") 209 | x.Add("opqrstu") 210 | a, b, err := x.GetTwo("99999999") 211 | if err != nil { 212 | t.Fatal(err) 213 | } 214 | if a == b { 215 | t.Errorf("a shouldn't equal b") 216 | } 217 | if a != "abcdefg" { 218 | t.Errorf("wrong a: %q", a) 219 | } 220 | if b != "opqrstu" { 221 | t.Errorf("wrong b: %q", b) 222 | } 223 | } 224 | 225 | func TestGetTwoQuick(t *testing.T) { 226 | x := New() 227 | x.Add("abcdefg") 228 | x.Add("hijklmn") 229 | x.Add("opqrstu") 230 | f := func(s string) bool { 231 | a, b, err := x.GetTwo(s) 232 | if err != nil { 233 | t.Logf("error: %q", err) 234 | return false 235 | } 236 | if a == b { 237 | t.Logf("a == b") 238 | return false 239 | } 240 | if a != "abcdefg" && a != "hijklmn" && a != "opqrstu" { 241 | t.Logf("invalid a: %q", a) 242 | return false 243 | } 244 | 245 | if b != "abcdefg" && b != "hijklmn" && b != "opqrstu" { 246 | t.Logf("invalid b: %q", b) 247 | return false 248 | } 249 | return true 250 | } 251 | if err := quick.Check(f, nil); err != nil { 252 | t.Fatal(err) 253 | } 254 | } 255 | 256 | func TestGetTwoOnlyTwoQuick(t *testing.T) { 257 | x := New() 258 | x.Add("abcdefg") 259 | x.Add("hijklmn") 260 | f := func(s string) bool { 261 | a, b, err := x.GetTwo(s) 262 | if err != nil { 263 | t.Logf("error: %q", err) 264 | return false 265 | } 266 | if a == b { 267 | t.Logf("a == b") 268 | return false 269 | } 270 | if a != "abcdefg" && a != "hijklmn" { 271 | t.Logf("invalid a: %q", a) 272 | return false 273 | } 274 | 275 | if b != "abcdefg" && b != "hijklmn" { 276 | t.Logf("invalid b: %q", b) 277 | return false 278 | } 279 | return true 280 | } 281 | if err := quick.Check(f, nil); err != nil { 282 | t.Fatal(err) 283 | } 284 | } 285 | 286 | func TestGetTwoOnlyOneInCircle(t *testing.T) { 287 | x := New() 288 | x.Add("abcdefg") 289 | a, b, err := x.GetTwo("99999999") 290 | if err != nil { 291 | t.Fatal(err) 292 | } 293 | if a == b { 294 | t.Errorf("a shouldn't equal b") 295 | } 296 | if a != "abcdefg" { 297 | t.Errorf("wrong a: %q", a) 298 | } 299 | if b != "" { 300 | t.Errorf("wrong b: %q", b) 301 | } 302 | } 303 | 304 | func TestGetN(t *testing.T) { 305 | x := New() 306 | x.Add("abcdefg") 307 | x.Add("hijklmn") 308 | x.Add("opqrstu") 309 | members, err := x.GetN("99999999", 3) 310 | if err != nil { 311 | t.Fatal(err) 312 | } 313 | if len(members) != 3 { 314 | t.Errorf("expected 3 members instead of %d", len(members)) 315 | } 316 | if members[0] != "abcdefg" { 317 | t.Errorf("wrong members[0]: %q", members[0]) 318 | } 319 | if members[1] != "opqrstu" { 320 | t.Errorf("wrong members[1]: %q", members[1]) 321 | } 322 | if members[2] != "hijklmn" { 323 | t.Errorf("wrong members[2]: %q", members[2]) 324 | } 325 | } 326 | 327 | func TestGetNLess(t *testing.T) { 328 | x := New() 329 | x.Add("abcdefg") 330 | x.Add("hijklmn") 331 | x.Add("opqrstu") 332 | members, err := x.GetN("99999999", 2) 333 | if err != nil { 334 | t.Fatal(err) 335 | } 336 | if len(members) != 2 { 337 | t.Errorf("expected 2 members instead of %d", len(members)) 338 | } 339 | if members[0] != "abcdefg" { 340 | t.Errorf("wrong members[0]: %q", members[0]) 341 | } 342 | if members[1] != "opqrstu" { 343 | t.Errorf("wrong members[1]: %q", members[1]) 344 | } 345 | } 346 | 347 | func TestGetNMore(t *testing.T) { 348 | x := New() 349 | x.Add("abcdefg") 350 | x.Add("hijklmn") 351 | x.Add("opqrstu") 352 | members, err := x.GetN("99999999", 5) 353 | if err != nil { 354 | t.Fatal(err) 355 | } 356 | if len(members) != 3 { 357 | t.Errorf("expected 3 members instead of %d", len(members)) 358 | } 359 | if members[0] != "abcdefg" { 360 | t.Errorf("wrong members[0]: %q", members[0]) 361 | } 362 | if members[1] != "opqrstu" { 363 | t.Errorf("wrong members[1]: %q", members[1]) 364 | } 365 | if members[2] != "hijklmn" { 366 | t.Errorf("wrong members[2]: %q", members[2]) 367 | } 368 | } 369 | 370 | func TestGetNQuick(t *testing.T) { 371 | x := New() 372 | x.Add("abcdefg") 373 | x.Add("hijklmn") 374 | x.Add("opqrstu") 375 | f := func(s string) bool { 376 | members, err := x.GetN(s, 3) 377 | if err != nil { 378 | t.Logf("error: %q", err) 379 | return false 380 | } 381 | if len(members) != 3 { 382 | t.Logf("expected 3 members instead of %d", len(members)) 383 | return false 384 | } 385 | set := make(map[string]bool, 4) 386 | for _, member := range members { 387 | if set[member] { 388 | t.Logf("duplicate error") 389 | return false 390 | } 391 | set[member] = true 392 | if member != "abcdefg" && member != "hijklmn" && member != "opqrstu" { 393 | t.Logf("invalid member: %q", member) 394 | return false 395 | } 396 | } 397 | return true 398 | } 399 | if err := quick.Check(f, nil); err != nil { 400 | t.Fatal(err) 401 | } 402 | } 403 | 404 | func TestGetNLessQuick(t *testing.T) { 405 | x := New() 406 | x.Add("abcdefg") 407 | x.Add("hijklmn") 408 | x.Add("opqrstu") 409 | f := func(s string) bool { 410 | members, err := x.GetN(s, 2) 411 | if err != nil { 412 | t.Logf("error: %q", err) 413 | return false 414 | } 415 | if len(members) != 2 { 416 | t.Logf("expected 2 members instead of %d", len(members)) 417 | return false 418 | } 419 | set := make(map[string]bool, 4) 420 | for _, member := range members { 421 | if set[member] { 422 | t.Logf("duplicate error") 423 | return false 424 | } 425 | set[member] = true 426 | if member != "abcdefg" && member != "hijklmn" && member != "opqrstu" { 427 | t.Logf("invalid member: %q", member) 428 | return false 429 | } 430 | } 431 | return true 432 | } 433 | if err := quick.Check(f, nil); err != nil { 434 | t.Fatal(err) 435 | } 436 | } 437 | 438 | func TestGetNMoreQuick(t *testing.T) { 439 | x := New() 440 | x.Add("abcdefg") 441 | x.Add("hijklmn") 442 | x.Add("opqrstu") 443 | f := func(s string) bool { 444 | members, err := x.GetN(s, 5) 445 | if err != nil { 446 | t.Logf("error: %q", err) 447 | return false 448 | } 449 | if len(members) != 3 { 450 | t.Logf("expected 3 members instead of %d", len(members)) 451 | return false 452 | } 453 | set := make(map[string]bool, 4) 454 | for _, member := range members { 455 | if set[member] { 456 | t.Logf("duplicate error") 457 | return false 458 | } 459 | set[member] = true 460 | if member != "abcdefg" && member != "hijklmn" && member != "opqrstu" { 461 | t.Logf("invalid member: %q", member) 462 | return false 463 | } 464 | } 465 | return true 466 | } 467 | if err := quick.Check(f, nil); err != nil { 468 | t.Fatal(err) 469 | } 470 | } 471 | 472 | func TestSet(t *testing.T) { 473 | x := New() 474 | x.Add("abc") 475 | x.Add("def") 476 | x.Add("ghi") 477 | x.Set([]string{"jkl", "mno"}) 478 | if x.count != 2 { 479 | t.Errorf("expected 2 elts, got %d", x.count) 480 | } 481 | a, b, err := x.GetTwo("qwerqwerwqer") 482 | if err != nil { 483 | t.Fatal(err) 484 | } 485 | if a != "jkl" && a != "mno" { 486 | t.Errorf("expected jkl or mno, got %s", a) 487 | } 488 | if b != "jkl" && b != "mno" { 489 | t.Errorf("expected jkl or mno, got %s", b) 490 | } 491 | if a == b { 492 | t.Errorf("expected a != b, they were both %s", a) 493 | } 494 | x.Set([]string{"pqr", "mno"}) 495 | if x.count != 2 { 496 | t.Errorf("expected 2 elts, got %d", x.count) 497 | } 498 | a, b, err = x.GetTwo("qwerqwerwqer") 499 | if err != nil { 500 | t.Fatal(err) 501 | } 502 | if a != "pqr" && a != "mno" { 503 | t.Errorf("expected jkl or mno, got %s", a) 504 | } 505 | if b != "pqr" && b != "mno" { 506 | t.Errorf("expected jkl or mno, got %s", b) 507 | } 508 | if a == b { 509 | t.Errorf("expected a != b, they were both %s", a) 510 | } 511 | x.Set([]string{"pqr", "mno"}) 512 | if x.count != 2 { 513 | t.Errorf("expected 2 elts, got %d", x.count) 514 | } 515 | a, b, err = x.GetTwo("qwerqwerwqer") 516 | if err != nil { 517 | t.Fatal(err) 518 | } 519 | if a != "pqr" && a != "mno" { 520 | t.Errorf("expected jkl or mno, got %s", a) 521 | } 522 | if b != "pqr" && b != "mno" { 523 | t.Errorf("expected jkl or mno, got %s", b) 524 | } 525 | if a == b { 526 | t.Errorf("expected a != b, they were both %s", a) 527 | } 528 | } 529 | 530 | // allocBytes returns the number of bytes allocated by invoking f. 531 | func allocBytes(f func()) uint64 { 532 | var stats runtime.MemStats 533 | runtime.ReadMemStats(&stats) 534 | t := stats.TotalAlloc 535 | f() 536 | runtime.ReadMemStats(&stats) 537 | return stats.TotalAlloc - t 538 | } 539 | 540 | func mallocNum(f func()) uint64 { 541 | var stats runtime.MemStats 542 | runtime.ReadMemStats(&stats) 543 | t := stats.Mallocs 544 | f() 545 | runtime.ReadMemStats(&stats) 546 | return stats.Mallocs - t 547 | } 548 | 549 | func BenchmarkAllocations(b *testing.B) { 550 | x := New() 551 | x.Add("stays") 552 | b.ResetTimer() 553 | allocSize := allocBytes(func() { 554 | for i := 0; i < b.N; i++ { 555 | x.Add("Foo") 556 | x.Remove("Foo") 557 | } 558 | }) 559 | b.Logf("%d: Allocated %d bytes (%.2fx)", b.N, allocSize, float64(allocSize)/float64(b.N)) 560 | } 561 | 562 | func BenchmarkMalloc(b *testing.B) { 563 | x := New() 564 | x.Add("stays") 565 | b.ResetTimer() 566 | mallocs := mallocNum(func() { 567 | for i := 0; i < b.N; i++ { 568 | x.Add("Foo") 569 | x.Remove("Foo") 570 | } 571 | }) 572 | b.Logf("%d: Mallocd %d times (%.2fx)", b.N, mallocs, float64(mallocs)/float64(b.N)) 573 | } 574 | 575 | func BenchmarkCycle(b *testing.B) { 576 | x := New() 577 | x.Add("nothing") 578 | b.ResetTimer() 579 | for i := 0; i < b.N; i++ { 580 | x.Add("foo" + strconv.Itoa(i)) 581 | x.Remove("foo" + strconv.Itoa(i)) 582 | } 583 | } 584 | 585 | func BenchmarkCycleLarge(b *testing.B) { 586 | x := New() 587 | for i := 0; i < 10; i++ { 588 | x.Add("start" + strconv.Itoa(i)) 589 | } 590 | b.ResetTimer() 591 | for i := 0; i < b.N; i++ { 592 | x.Add("foo" + strconv.Itoa(i)) 593 | x.Remove("foo" + strconv.Itoa(i)) 594 | } 595 | } 596 | 597 | func BenchmarkGet(b *testing.B) { 598 | x := New() 599 | x.Add("nothing") 600 | b.ResetTimer() 601 | for i := 0; i < b.N; i++ { 602 | x.Get("nothing") 603 | } 604 | } 605 | 606 | func BenchmarkGetLarge(b *testing.B) { 607 | x := New() 608 | for i := 0; i < 10; i++ { 609 | x.Add("start" + strconv.Itoa(i)) 610 | } 611 | b.ResetTimer() 612 | for i := 0; i < b.N; i++ { 613 | x.Get("nothing") 614 | } 615 | } 616 | 617 | func BenchmarkGetN(b *testing.B) { 618 | x := New() 619 | x.Add("nothing") 620 | b.ResetTimer() 621 | for i := 0; i < b.N; i++ { 622 | x.GetN("nothing", 3) 623 | } 624 | } 625 | 626 | func BenchmarkGetNLarge(b *testing.B) { 627 | x := New() 628 | for i := 0; i < 10; i++ { 629 | x.Add("start" + strconv.Itoa(i)) 630 | } 631 | b.ResetTimer() 632 | for i := 0; i < b.N; i++ { 633 | x.GetN("nothing", 3) 634 | } 635 | } 636 | 637 | func BenchmarkGetTwo(b *testing.B) { 638 | x := New() 639 | x.Add("nothing") 640 | b.ResetTimer() 641 | for i := 0; i < b.N; i++ { 642 | x.GetTwo("nothing") 643 | } 644 | } 645 | 646 | func BenchmarkGetTwoLarge(b *testing.B) { 647 | x := New() 648 | for i := 0; i < 10; i++ { 649 | x.Add("start" + strconv.Itoa(i)) 650 | } 651 | b.ResetTimer() 652 | for i := 0; i < b.N; i++ { 653 | x.GetTwo("nothing") 654 | } 655 | } 656 | --------------------------------------------------------------------------------