diff --git a/raft/node.go b/raft/node.go index e9d78b13d..62dff5be4 100644 --- a/raft/node.go +++ b/raft/node.go @@ -25,6 +25,10 @@ type Node struct { elapsed tick election tick heartbeat tick + + // TODO: it needs garbage collection later + rmNodes map[int64]struct{} + removed bool } func New(id int64, heartbeat, election tick) *Node { @@ -36,6 +40,7 @@ func New(id int64, heartbeat, election tick) *Node { heartbeat: heartbeat, election: election, sm: newStateMachine(id, []int64{id}), + rmNodes: make(map[int64]struct{}), } return n @@ -57,6 +62,8 @@ func (n *Node) IsLeader() bool { return n.Leader() == n.Id() } func (n *Node) Leader() int64 { return n.sm.lead.Get() } +func (n *Node) IsRemoved() bool { return n.removed } + // Propose asynchronously proposes data be applied to the underlying state machine. func (n *Node) Propose(data []byte) { n.propose(Normal, data) } @@ -75,6 +82,17 @@ func (n *Node) Remove(id int64) { n.updateConf(RemoveNode, &Config{NodeId: id}) func (n *Node) Msgs() []Message { return n.sm.Msgs() } func (n *Node) Step(m Message) bool { + if m.Type == msgDenied { + n.removed = true + return false + } + if m.Term != 0 { + if _, ok := n.rmNodes[m.From]; ok { + n.sm.send(Message{To: m.From, Type: msgDenied}) + return true + } + } + l := len(n.sm.msgs) if !n.sm.Step(m) { return false @@ -107,6 +125,7 @@ func (n *Node) Next() []Entry { continue } n.sm.addNode(c.NodeId) + delete(n.rmNodes, c.NodeId) case RemoveNode: c := new(Config) if err := json.Unmarshal(ents[i].Data, c); err != nil { @@ -114,6 +133,10 @@ func (n *Node) Next() []Entry { continue } n.sm.removeNode(c.NodeId) + n.rmNodes[c.NodeId] = struct{}{} + if c.NodeId == n.sm.id { + n.removed = true + } default: panic("unexpected entry type") } diff --git a/raft/node_test.go b/raft/node_test.go index 0d27cc561..d5a9f3136 100644 --- a/raft/node_test.go +++ b/raft/node_test.go @@ -1,6 +1,7 @@ package raft import ( + "reflect" "testing" ) @@ -141,6 +142,52 @@ func TestRemove(t *testing.T) { } } +func TestDenial(t *testing.T) { + logents := []Entry{ + {Type: AddNode, Term: 1, Data: []byte(`{"NodeId":1}`)}, + {Type: AddNode, Term: 1, Data: []byte(`{"NodeId":2}`)}, + {Type: RemoveNode, Term: 1, Data: []byte(`{"NodeId":2}`)}, + } + + tests := []struct { + ent Entry + wdenied map[int64]bool + }{ + { + Entry{Type: AddNode, Term: 1, Data: []byte(`{"NodeId":2}`)}, + map[int64]bool{0: false, 1: false, 2: false}, + }, + { + Entry{Type: RemoveNode, Term: 1, Data: []byte(`{"NodeId":1}`)}, + map[int64]bool{0: false, 1: true, 2: true}, + }, + { + Entry{Type: RemoveNode, Term: 1, Data: []byte(`{"NodeId":0}`)}, + map[int64]bool{0: true, 1: false, 2: true}, + }, + } + + for i, tt := range tests { + n := dictate(New(0, defaultHeartbeat, defaultElection)) + n.Next() + n.Msgs() + n.sm.log.append(n.sm.log.committed, append(logents, tt.ent)...) + n.sm.log.committed += int64(len(logents) + 1) + n.Next() + + for id, denied := range tt.wdenied { + n.Step(Message{From: id, To: 0, Type: msgApp, Term: 1}) + w := []Message{} + if denied { + w = []Message{{From: 0, To: id, Term: 1, Type: msgDenied}} + } + if g := n.Msgs(); !reflect.DeepEqual(g, w) { + t.Errorf("#%d: msgs for %d = %+v, want %+v", i, id, g, w) + } + } + } +} + func dictate(n *Node) *Node { n.Step(Message{Type: msgHup}) n.Add(n.Id(), "", nil) diff --git a/raft/raft.go b/raft/raft.go index 3f536449e..a6bae130d 100644 --- a/raft/raft.go +++ b/raft/raft.go @@ -19,6 +19,7 @@ const ( msgVote msgVoteResp msgSnap + msgDenied ) var mtmap = [...]string{ @@ -30,6 +31,7 @@ var mtmap = [...]string{ msgVote: "msgVote", msgVoteResp: "msgVoteResp", msgSnap: "msgSnap", + msgDenied: "msgDenied", } func (mt messageType) String() string {