raft: break Step into pieces

raft/raft.go

@@ -49,6 +49,12 @@ var stmap = [...]string{
 	stateLeader:    "stateLeader",
 }
 
+var stepmap = [...]stepFunc{
+	stateFollower:  stepFollower,
+	stateCandidate: stepCandidate,
+	stateLeader:    stepLeader,
+}
+
 func (st stateType) String() string {
 	return stmap[int(st)]
 }
@@ -241,8 +247,7 @@ func (sm *stateMachine) Msgs() []Message {
 }
 
 func (sm *stateMachine) Step(m Message) (ok bool) {
-	switch m.Type {
+	if m.Type == msgHup {
-	case msgHup:
 		sm.becomeCandidate()
 		if sm.q() == sm.poll(sm.id, true) {
 			sm.becomeLeader()
@@ -256,19 +261,49 @@ func (sm *stateMachine) Step(m Message) (ok bool) {
 			sm.send(Message{To: i, Type: msgVote, Index: lasti, LogTerm: sm.log.term(lasti)})
 		}
 		return true
-	case msgBeat:
+	}
-		if sm.state != stateLeader {
+
+	switch {
+	case m.Term == 0:
+		// local message
+	case m.Term > sm.term:
+		sm.becomeFollower(m.Term, m.From)
+	case m.Term < sm.term:
+		// ignore
 		return true
 	}
+
+	return stepmap[sm.state](sm, m)
+}
+
+func (sm *stateMachine) handleAppendEntries(m Message) {
+	if sm.log.maybeAppend(m.Index, m.LogTerm, m.Commit, m.Entries...) {
+		sm.send(Message{To: m.From, Type: msgAppResp, Index: sm.log.lastIndex()})
+	} else {
+		sm.send(Message{To: m.From, Type: msgAppResp, Index: -1})
+	}
+}
+
+func (sm *stateMachine) addNode(id int) {
+	sm.ins[id] = &index{next: sm.log.lastIndex() + 1}
+	sm.pendingConf = false
+}
+
+func (sm *stateMachine) removeNode(id int) {
+	delete(sm.ins, id)
+	sm.pendingConf = false
+}
+
+type stepFunc func(sm *stateMachine, m Message) bool
+
+func stepLeader(sm *stateMachine, m Message) bool {
+	switch m.Type {
+	case msgBeat:
 		sm.bcastAppend()
-		return
 	case msgProp:
 		if len(m.Entries) != 1 {
 			panic("unexpected length(entries) of a msgProp")
 		}
-
-		switch sm.lead {
-		case sm.id:
 		e := m.Entries[0]
 		if e.isConfig() {
 			if sm.pendingConf {
@@ -282,35 +317,6 @@ func (sm *stateMachine) Step(m Message) (ok bool) {
 		sm.ins[sm.id].update(sm.log.lastIndex())
 		sm.maybeCommit()
 		sm.bcastAppend()
-		case none:
-			// msgProp given without leader
-			return false
-		default:
-			m.To = sm.lead
-			sm.send(m)
-		}
-		return true
-	}
-
-	switch {
-	case m.Term > sm.term:
-		sm.becomeFollower(m.Term, m.From)
-	case m.Term < sm.term:
-		// ignore
-		return true
-	}
-
-	handleAppendEntries := func() {
-		if sm.log.maybeAppend(m.Index, m.LogTerm, m.Commit, m.Entries...) {
-			sm.send(Message{To: m.From, Type: msgAppResp, Index: sm.log.lastIndex()})
-		} else {
-			sm.send(Message{To: m.From, Type: msgAppResp, Index: -1})
-		}
-	}
-
-	switch sm.state {
-	case stateLeader:
-		switch m.Type {
 	case msgAppResp:
 		if m.Index < 0 {
 			sm.ins[m.From].decr()
@@ -324,11 +330,16 @@ func (sm *stateMachine) Step(m Message) (ok bool) {
 	case msgVote:
 		sm.send(Message{To: m.From, Type: msgVoteResp, Index: -1})
 	}
-	case stateCandidate:
+	return true
+}
+
+func stepCandidate(sm *stateMachine, m Message) bool {
 	switch m.Type {
+	case msgProp:
+		return false
 	case msgApp:
 		sm.becomeFollower(sm.term, m.From)
-			handleAppendEntries()
+		sm.handleAppendEntries(m)
 	case msgVote:
 		sm.send(Message{To: m.From, Type: msgVoteResp, Index: -1})
 	case msgVoteResp:
@@ -341,10 +352,19 @@ func (sm *stateMachine) Step(m Message) (ok bool) {
 			sm.becomeFollower(sm.term, none)
 		}
 	}
-	case stateFollower:
+	return true
+}
+
+func stepFollower(sm *stateMachine, m Message) bool {
 	switch m.Type {
+	case msgProp:
+		if sm.lead == none {
+			return false
+		}
+		m.To = sm.lead
+		sm.send(m)
 	case msgApp:
-			handleAppendEntries()
+		sm.handleAppendEntries(m)
 	case msgVote:
 		if (sm.vote == none || sm.vote == m.From) && sm.log.isUpToDate(m.Index, m.LogTerm) {
 			sm.vote = m.From
@@ -353,16 +373,5 @@ func (sm *stateMachine) Step(m Message) (ok bool) {
 			sm.send(Message{To: m.From, Type: msgVoteResp, Index: -1})
 		}
 	}
-	}
 	return true
 }
-
-func (sm *stateMachine) addNode(id int) {
-	sm.ins[id] = &index{next: sm.log.lastIndex() + 1}
-	sm.pendingConf = false
-}
-
-func (sm *stateMachine) removeNode(id int) {
-	delete(sm.ins, id)
-	sm.pendingConf = false
-}