-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
144 lines (127 loc) · 3.15 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
package main
import (
"encoding/json"
"flag"
"fmt"
"log"
"strconv"
"github.com/pengubco/algorithms/maglev_hash"
"github.com/samber/lo"
)
// This program writes out four lines.
// 1. number of nodes
// 2. number of slots
// 3: number of slots assigned to each node
// 4: number of slots that need to be redistributed after removing a node.
func main() {
var nodeCnt int
var slotCnt int
flag.IntVar(&nodeCnt, "nodeCnt", 0, "number of nodes")
flag.IntVar(&slotCnt, "slotCnt", 0, "number of slots. must be a prime number and larger than nodeCnt")
flag.Parse()
e, err := NewExperiment(nodeCnt, slotCnt)
if err != nil {
log.Fatal(err)
}
fmt.Println(nodeCnt)
fmt.Println(slotCnt)
slotToNode := e.getSlotAssignment()
nodeLoads := lo.Values(calculateNodeLoads(slotToNode))
b, _ := json.Marshal(nodeLoads)
fmt.Printf("%s\n", string(b))
e.removeNode("1")
newSlotToNode := e.getSlotAssignment()
slotMoved := calculateSlotMove(slotToNode, newSlotToNode)
fmt.Printf("%d\n", slotMoved)
}
func calculateNodeLoads(slotToNode map[int]int) map[int]int {
load := make(map[int]int)
for _, v := range slotToNode {
load[v]++
}
return load
}
// calculateSlotMove returns number of slots that are redistributed.
func calculateSlotMove(a, b map[int]int) int {
result := 0
for k, _ := range a {
if a[k] != b[k] {
result++
}
}
return result
}
func sanityCheck() {
e, err := NewExperiment(3, 7)
if err != nil {
log.Fatal(err)
}
slotToNode := e.getSlotAssignment()
fmt.Printf("%+v\n", slotToNode)
e.removeNode("1")
slotToNode2 := e.getSlotAssignment()
fmt.Printf("%+v\n", slotToNode2)
}
// Experiment has the following settings.
// 1. Nodes: Nodes are number 0, 1, 2, ...
// 2. Use Atoi as the key-hash function which returns the integer form a string.
// So that we can use Node([]byte("1")) to get the node assigned for slot-1.
type Experiment struct {
nodeCnt int
slotCnt int
nodes []string
mh *maglev_hash.MaglevHash
}
func NewExperiment(nodeCnt int, slotCnt int) (*Experiment, error) {
nodes := lo.Times(nodeCnt, func(i int) string {
return fmt.Sprintf("%d", i)
})
mh, err := maglev_hash.NewMaglevWithTableSize(slotCnt, nodes, func(b []byte) uint32 {
i, _ := strconv.Atoi(string(b))
return uint32(i)
})
if err != nil {
return nil, err
}
return &Experiment{
nodeCnt: nodeCnt,
slotCnt: slotCnt,
nodes: nodes,
mh: mh,
}, nil
}
// returns slot -> node mapping.
func (e *Experiment) getSlotAssignment() map[int]int {
result := make(map[int]int)
for i := 0; i < e.slotCnt; i++ {
nodeStr := e.mh.Node([]byte(strconv.Itoa(i)))
nodeID, _ := strconv.Atoi(nodeStr)
result[i] = nodeID
}
return result
}
func (e *Experiment) removeNode(node string) error {
found := false
newNodes := make([]string, 0)
for _, v := range e.nodes {
if v != node {
newNodes = append(newNodes, v)
} else {
found = true
}
}
if !found {
return fmt.Errorf("node does not exist, %s", node)
}
e.nodeCnt = len(newNodes)
e.nodes = newNodes
mh, err := maglev_hash.NewMaglevWithTableSize(e.slotCnt, e.nodes, func(b []byte) uint32 {
i, _ := strconv.Atoi(string(b))
return uint32(i)
})
if err != nil {
return err
}
e.mh = mh
return nil
}