Browse Source

raft: add msgDenial to deny removed nodes

Yicheng Qin 11 years ago
parent
commit
f1853b4364
3 changed files with 72 additions and 0 deletions
  1. 23 0
      raft/node.go
  2. 47 0
      raft/node_test.go
  3. 2 0
      raft/raft.go

+ 23 - 0
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")
 		}

+ 47 - 0
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)

+ 2 - 0
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 {