04-Raft State Machine

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

	// 获取client发送的消息，交给Raft模块去处理，这里的client并不是用户的client，而是要执行命令的sql端
	Some((request, response_tx)) = client_rx.next() => {
		let id = Uuid::new_v4().as_bytes().to_vec();
		let msg = Message{
			from: Address::Client,
			to: Address::Node(node.id()),
			term: 0,
			event: Event::ClientRequest{id: id.clone(), request},
		};
		node = node.step(msg)?;
		requests.insert(id, response_tx);
	}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

/// Drives a state machine, taking operations from state_rx and sending results via node_tx.
pub struct Driver {
    node_id: NodeID,
    state_rx: UnboundedReceiverStream<Instruction>,
    node_tx: mpsc::UnboundedSender<Message>,
    /// Notify clients when their mutation is applied. <index, (client, id)>
    notify: HashMap<Index, (Address, Vec<u8>)>,
    /// Execute client queries when they receive a quorum. <index, <id, query>>
    queries: BTreeMap<Index, BTreeMap<Vec<u8>, Query>>,
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

    pub async fn drive(mut self, mut state: Box<dyn State>) -> Result<()> {
        debug!("Starting state machine driver at applied index {}", state.get_applied_index());
        while let Some(instruction) = self.state_rx.next().await {
            if let Err(error) = self.execute(instruction, &mut *state) {
                error!("Halting state machine due to error: {}", error);
                return Err(error);
            }
        }
        debug!("Stopping state machine driver");
        Ok(())
    }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

/// A Raft-managed state machine.
pub trait State: Send {
    /// Returns the last applied index from the state machine.
    fn get_applied_index(&self) -> Index;

    /// Applies a log entry to the state machine. If it returns Error::Internal,
    /// the Raft node halts. Any other error is considered applied and returned
    /// to the caller.
    ///
    /// The entry may contain a noop command, which is committed by Raft during
    /// leader changes. This still needs to be applied to the state machine to
    /// properly track the applied index, and returns an empty result.
    ///
    /// TODO: consider using runtime assertions instead of Error::Internal.
    fn apply(&mut self, entry: Entry) -> Result<Vec<u8>>;

    /// Queries the state machine. All errors are propagated to the caller.
    fn query(&self, command: Vec<u8>) -> Result<Vec<u8>>;
}

1
2
3
4
5
6
7

/// The Raft state machine for the Raft-based SQL engine, using a KV SQL engine
pub struct State<E: storage::engine::Engine> {
    /// The underlying KV SQL engine
    engine: super::KV<E>,
    /// The last applied index
    applied_index: u64,
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

pub enum Instruction {
    /// Abort all pending operations, e.g. due to leader change.
    Abort,
    /// Apply a log entry.
    Apply { entry: Entry },
    /// Notify the given address with the result of applying the entry at the given index.
    Notify { id: Vec<u8>, address: Address, index: Index },
    /// Query the state machine when the given term and index has been confirmed by vote.
    Query { id: Vec<u8>, address: Address, command: Vec<u8>, term: Term, index: Index, quorum: u64 },
    /// Extend the given server status and return it to the given address.
    Status { id: Vec<u8>, address: Address, status: Box<Status> },
    /// Votes for queries at the given term and commit index.
    Vote { term: Term, index: Index, address: Address },
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

	// Leader.step()
	pub fn step(mut self, msg: Message) -> Result<Node> {
		Event::ClientRequest { id, request: Request::Mutate(command) } => {
			let index = self.propose(Some(command))?;
			self.state_tx.send(Instruction::Notify { id, address: msg.from, index })?;
			if self.peers.is_empty() {
				self.maybe_commit()?;
			}
		}
	}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

    /// Executes a state machine instruction. 
    fn execute(&mut self, i: Instruction, state: &mut dyn State) -> Result<()> {
        match i {
            Instruction::Notify { id, address, index } => {
                if index > state.get_applied_index() {
                    self.notify.insert(index, (address, id));
                } else {
                    self.send(address, Event::ClientResponse { id, response: Err(Error::Abort) })?;
                }
            }
        }
        Ok(())
    }

1
2
3
4
5
6

	Event::ClientResponse { id, mut response } => {
		if let Ok(Response::Status(ref mut status)) = response {
			status.server = self.id;
		}
		self.send(Address::Client, Event::ClientResponse { id, response })?;
	}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

    pub fn step(mut self, msg: Message) -> Result<Node> {
		Event::ClientRequest { id, request: Request::Query(command) } => {
			let (commit_index, _) = self.log.get_commit_index();
			self.state_tx.send(Instruction::Query {
				id,
				address: msg.from,
				command,
				term: self.term,
				index: commit_index,
				quorum: self.quorum(),
			})?;
			self.state_tx.send(Instruction::Vote {
				term: self.term,
				index: commit_index,
				address: Address::Node(self.id),
			})?;
			self.heartbeat()?;
		}
	}

1
2
3
4
5
6
7

    // Follower.step()
    pub fn step(mut self, msg: Message) -> Result<Node> {
		// 先处理commit和apply
		// ....
		// 确认Leader身份
		self.send(msg.from, Event::ConfirmLeader { commit_index, has_committed })?;
	}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

    pub fn step(mut self, msg: Message) -> Result<Node> {	
		// A follower received one of our heartbeats and confirms that we
		// are its leader. If it doesn't have the commit index in its local
		// log, replicate the log to it.
		Event::ConfirmLeader { commit_index, has_committed } => {
			let from = msg.from.unwrap();
			// 与上层状态机进行交互
			self.state_tx.send(Instruction::Vote {
				term: msg.term,
				index: commit_index,
				address: msg.from,
			})?;
			if !has_committed {
				self.send_log(from)?;
			}
		}
	}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

struct Query {
    id: Vec<u8>,
    term: Term,
    address: Address,
    command: Vec<u8>,
    quorum: u64,
    votes: HashSet<Address>,
}

pub struct Driver {
	// ...
    /// Execute client queries when they receive a quorum. <index, <id, query>>
    queries: BTreeMap<Index, BTreeMap<Vec<u8>, Query>>,
}

    fn execute(&mut self, i: Instruction, state: &mut dyn State) -> Result<()> {

		Instruction::Query { id, address, command, index, term, quorum } => {
			self.queries.entry(index).or_default().insert(
				id.clone(),
				Query { id, term, address, command, quorum, votes: HashSet::new() },
			);
		}
	}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

    /// Executes a state machine instruction. 
    fn execute(&mut self, i: Instruction, state: &mut dyn State) -> Result<()> {
        debug!("Executing {:?}", i);
        match i {
            Instruction::Vote { term, index, address } => {
                self.query_vote(term, index, address);
                self.query_execute(state)?;
            }
        }
        Ok(())
    }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

    /// Votes for queries up to and including a given commit index for a term by an address.
    fn query_vote(&mut self, term: Term, commit_index: Index, address: Address) {
        for (_, queries) in self.queries.range_mut(..=commit_index) {
            for (_, query) in queries.iter_mut() {
                if term >= query.term {
                // 使用HashSet记录保证同一个节点(使用Address表示)，只能对同一个commit_index投票一次
                    query.votes.insert(address);
                }
            }
        }
    }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

    /// Executes any queries that are ready.
    fn query_execute(&mut self, state: &mut dyn State) -> Result<()> {
	    // self.query_ready()会获取出所有的达到quorum并且index <= applied_index的Query，并且从self.queries当中移除
        for query in self.query_ready(state.get_applied_index()) {
            debug!("Executing query {:?}", query.command);
            let result = state.query(query.command);
            if let Err(error @ Error::Internal(_)) = result {
                return Err(error);
            }
            self.send(
                query.address,
                Event::ClientResponse { id: query.id, response: result.map(Response::Query) },
            )?
        }
        Ok(())
    }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

		// If we receive a message for a future term, become a leaderless
	// follower in it and step the message. If the message is a Heartbeat or
	// AppendEntries from the leader, stepping it will follow the leader.
	if msg.term > self.term {
		return self.become_follower(msg.term)?.step(msg);
	}

    /// Transforms the leader into a follower. This can only happen if we find a
    /// new term, so we become a leaderless follower.
    fn become_follower(mut self, term: Term) -> Result<RoleNode<Follower>> {
        assert!(term >= self.term, "Term regression {} -> {}", self.term, term);
        assert!(term > self.term, "Can only become follower in later term");

        info!("Discovered new term {}", term);
        self.term = term;
        self.log.set_term(term, None)?;
        self.state_tx.send(Instruction::Abort)?;
        Ok(self.become_role(Follower::new(None, None)))
    }

 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

    fn execute(&mut self, i: Instruction, state: &mut dyn State) -> Result<()> {
        debug!("Executing {:?}", i);
        match i {
            Instruction::Abort => {
                self.notify_abort()?;
                self.query_abort()?;
            }
            // ...
            // ...
        }
    }
    
    
    
    /// Aborts all pending notifications.
    fn notify_abort(&mut self) -> Result<()> {
        for (_, (address, id)) in std::mem::take(&mut self.notify) {
            self.send(address, Event::ClientResponse { id, response: Err(Error::Abort) })?;
        }
        Ok(())
    }
    
    /// Aborts all pending queries.
    fn query_abort(&mut self) -> Result<()> {
        for (_, queries) in std::mem::take(&mut self.queries) {
            for (id, query) in queries {
                self.send(
                    query.address,
                    Event::ClientResponse { id, response: Err(Error::Abort) },
                )?;
            }
        }
        Ok(())
    }

 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

	// Leader
	
	if commit_index > prev_commit_index {
		self.log.commit(commit_index)?;
		// TODO: Move application elsewhere, but needs access to applied index.
		let mut scan = self.log.scan((prev_commit_index + 1)..=commit_index)?;
		while let Some(entry) = scan.next().transpose()? {
			self.state_tx.send(Instruction::Apply { entry })?;
		}
	}
	
	// Follower
	Event::Heartbeat { commit_index, commit_term } => {
		// ....


		// Advance commit index and apply entries if possible.
		let has_committed = self.log.has(commit_index, commit_term)?;
		let (old_commit_index, _) = self.log.get_commit_index();
		if has_committed && commit_index > old_commit_index {
			self.log.commit(commit_index)?;
			let mut scan = self.log.scan((old_commit_index + 1)..=commit_index)?;
			while let Some(entry) = scan.next().transpose()? {
				self.state_tx.send(Instruction::Apply { entry })?;
			}
		}
		self.send(msg.from, Event::ConfirmLeader { commit_index, has_committed })?;
	}

 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

    fn execute(&mut self, i: Instruction, state: &mut dyn State) -> Result<()> {
        match i {
        // ...
            Instruction::Apply { entry } => {
                self.apply(state, entry)?;
            }
        /// ...
        }
    }

    /// Applies an entry to the state machine.
    pub fn apply(&mut self, state: &mut dyn State, entry: Entry) -> Result<Index> {
        // Apply the command.
        debug!("Applying {:?}", entry);
        match state.apply(entry) {
            Err(error @ Error::Internal(_)) => return Err(error),
            result => self.notify_applied(state.get_applied_index(), result)?,
        };
        // Try to execute any pending queries, since they may have been submitted for a
        // commit_index which hadn't been applied yet.
        self.query_execute(state)?;
        Ok(state.get_applied_index())
    }


    /// Notifies a client about an applied log entry, if any.
    fn notify_applied(&mut self, index: Index, result: Result<Vec<u8>>) -> Result<()> {
        if let Some((to, id)) = self.notify.remove(&index) {
            self.send(to, Event::ClientResponse { id, response: result.map(Response::Mutate) })?;
        }
        Ok(())
    }

 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

impl Node {
    /// Creates a new Raft node, starting as a follower, or leader if no peers.
    pub async fn new(
        id: NodeID,
        peers: HashSet<NodeID>,
        mut log: Log,
        mut state: Box<dyn State>,
        node_tx: mpsc::UnboundedSender<Message>,
    ) -> Result<Self> {
        let (state_tx, state_rx) = mpsc::unbounded_channel();
        // Driver用于和Raft状态机进行交互，启动状态机运行
        let mut driver = Driver::new(id, state_rx, node_tx.clone());
        driver.apply_log(&mut *state, &mut log)?;
        tokio::spawn(driver.drive(state));

        let (term, voted_for) = log.get_term()?;
        let mut node = RoleNode {
            id,
            peers,
            term,
            log,
            node_tx,
            state_tx,
            role: Follower::new(None, voted_for),
        };
        if node.peers.is_empty() {
            info!("No peers specified, starting as leader");
            // If we didn't vote for ourself in the persisted term, bump the
            // term and vote for ourself to ensure we have a valid leader term.
            if voted_for != Some(id) {
                node.term += 1;
                node.log.set_term(node.term, Some(id))?;
            }
            let (last_index, _) = node.log.get_last_index();
            Ok(node.become_role(Leader::new(HashSet::new(), last_index)).into())
        } else {
            Ok(node.into())
        }
    }
}

1
2
3
4
5
6
7
8

    /// Sends a message.
    /// 状态机使用node_tx进行消息发送，那么对应的接收端就是node_rx
    fn send(&self, to: Address, event: Event) -> Result<()> {
        // TODO: This needs to use the correct term.
        let msg = Message { from: Address::Node(self.node_id), to, term: 0, event };
        debug!("Sending {:?}", msg);
        Ok(self.node_tx.send(msg)?)
    }